[None][perf] Add Qwen Image VisualGen perf fastpaths

[pst2154]

Summary

• Add a Qwen-Image prompt-mask fastpath so all-valid prompt batches avoid materialized attention masks.
• Route Qwen joint masks through key_padding_mask for VANILLA attention, including VANILLA-backed Ulysses; reject unsupported masked sequence-parallel combinations instead of silently running rank-local SDPA.
• Keep FA4 on the explicit mask fallback because its key-padding path assumes valid tokens form one prefix, while Qwen's joint mask is [valid text | padded text | valid image].
• Enable guarded fused QK-norm + RoPE preparation for Qwen Image on CUDA BF16 TP=1 and hoist complex-to-real RoPE conversion from every block to once per model forward.
• Make Qwen TP>1 shard both image and text Q/K/V streams consistently, row-shard/reduce the text output projection, and reduce row-parallel MLP output before residual updates.
• Integrate SequenceSharder for Qwen sequence parallelism: pad and shard image/text tokens and RoPE independently, translate the key-padding mask into Ulysses rank-interleaved order, then gather and unpad image output.
• Make Qwen adaLN helpers compile-friendly and route gelu-approximate through shared gelu_tanh for the fused NVFP4 producer path.
• Preserve the positive noise prediction across the negative CUDA Graph replay during true CFG, with regression coverage for graph-output buffer aliasing.

Checks

• pre-commit run --files tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
• ruff check and ruff format --check pass for the changed pipeline and regression test.
• python3 -m py_compile ...
• git diff --check
• Rebased cleanly onto current upstream/main.
• DCO sign-off is present.
• The prior rebased head passed GitHub PR Checks, LLM API Compatibility, and Release Checks; CI is being rerun for c46072b716.

B300 validation

Using Qwen-Image-2512 weights with current PR Python over nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc18 on SM103:

• Full 60-block BF16 TP=1 transformer smoke with an interior padded-text mask: finite (1, 256, 64) output, 4.31 s forward, 38.20 GiB peak, and exactly one complex RoPE conversion.
• Full 60-block BF16 TP=2 transformer smoke: rank outputs matched exactly; 3.38-3.45 s forward per rank.
• TP=2 unfused output versus TP=1 fused output: cosine 0.99982893, max absolute difference 0.13671875, mean absolute difference 0.01553461.
• Full 60-block BF16 Ulysses=2 run with odd image/text lengths, forced padding, and an interior text mask: 3.49-3.52 s forward; rank outputs matched exactly; versus TP=1 cosine 0.99974769, max absolute difference 0.166015625, mean absolute difference 0.01943409.
• The same full forced-padding case on Ulysses=4: 5.37-5.40 s forward; all four rank outputs matched exactly; versus TP=1 cosine 0.99986815, max absolute difference 0.1328125, mean absolute difference 0.01413129.
• A smaller forced-padding Ulysses=2 parity case matched the unpadded TP=1 reference exactly.
• Focused fused versus unfused Q/K preparation: cosine 0.999996 for both Q and K.

Scope note

This is real-checkpoint transformer validation at a reduced token count, not a full text-encoder/VAE image-generation quality run. Local targeted pytest collection remains blocked by workstation dependencies (mpi4py through the normal harness; nvtx and a usable PyTorch install when bypassing the parent conftest). Repository hooks and focused rc18/B300 validation pass.

Summary by CodeRabbit

• Performance

  • Improved Qwen-Image transformer execution for tensor-parallel and sequence-parallel workloads.
  • Added optimized attention and rotary-embedding paths where supported.
  • Improved compatibility across attention backends and CPU fallback execution.

• Reliability

  • Improved handling of padded prompts and unsupported masked sequence-parallel configurations.
  • Prevented prediction corruption during CUDA graph-based guidance execution.

• Tests

  • Expanded coverage for prompt masking, attention behavior, parallel execution, rotary embeddings, and transformer operations.

[pst2154]

/bot run

[pst2154]

B300 follow-up validation for commit 75bd9f2978 (rebased; code change equivalent to previously validated 4cd65d50fb) using current PR Python over nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc18 on SM103.

Additional issues found and fixed

• Text-stream Q/K/V and to_add_out were not tensor-parallel sharded/reduced alongside the image stream.
• Qwen's MLP explicitly disabled its row-parallel output reduction, so each rank added a different partial MLP result to a replicated residual.
• The Ulysses wrapper was selected, but Qwen never sequence-sharded its image/text tokens or gathered the image output.
• Padded sequence-parallel masks needed translation to Ulysses' rank-interleaved key order.
• Complex-to-real RoPE conversion was repeated in every transformer block.

Real Qwen-Image-2512 BF16 checkpoint

TP=1, full 60-block transformer, 256 image tokens + 64 text tokens with an interior padded-text mask:

• finite (1, 256, 64) BF16 output
• 4.31 s forward
• 38.20 GiB peak allocation
• complex RoPE conversion called exactly once

TP=2, same checkpoint and broadcast inputs:

• both rank outputs matched exactly (max abs = 0.0)
• 3.38-3.45 s forward per rank
• TP=2 unfused versus TP=1 fused: cosine 0.99982893, max abs 0.13671875, mean abs 0.01553461

Ulysses=2, full 60-block transformer with 255 image tokens, 63 text tokens, forced padding, and an interior text mask:

• both rank outputs matched exactly (max abs = 0.0)
• 3.49-3.52 s forward per rank
• complex RoPE conversion called exactly once per rank
• versus fresh TP=1 reference: cosine 0.99974769, max abs 0.166015625, mean abs 0.01943409

The same full forced-padding case on Ulysses=4:

• all four rank outputs matched exactly (max abs = 0.0)
• 5.37-5.40 s forward per rank
• complex RoPE conversion called exactly once per rank
• versus fresh TP=1 reference: cosine 0.99986815, max abs 0.1328125, mean abs 0.01413129

A smaller odd-length forced-padding Ulysses=2 case matched its unpadded TP=1 reference exactly. Focused fused/unfused Q/K preparation measured cosine 0.999996 for both Q and K.

File-scoped pre-commit, Ruff/format, byte-compilation, diff checks, and DCO pass after the rebase. Local pytest collection remains blocked by the workstation environment (mpi4py in the normal harness and a missing usable torch.nn when bypassing the parent conftest). These were real-checkpoint reduced-token transformer tests; they did not include text encoder/VAE decoding or an NVFP4 checkpoint.

[pst2154]

CUDA Graph + true-CFG follow-up in c46072b716:

• Confirmed the output-lifetime bug. CUDAGraphRunner keeps graph outputs as non-owning tensor views, and graph keys share a memory pool. The negative-prompt replay can therefore overwrite the positive noise_pred before the CFG combination consumes it.
• The pipeline now clones the positive prediction immediately before the negative replay, only when CUDA Graph runners are active. Eager CFG does not take the extra copy.
• Added a CPU regression test whose transformer deliberately returns the same backing output buffer on both calls, reproducing the graph aliasing behavior.
• File-scoped pre-commit, Ruff/format, byte-compilation, diff checks, and DCO pass. Targeted pytest collection is blocked on this workstation by missing runtime dependencies (mpi4py in the normal harness; nvtx/PyTorch when bypassing it), so CI is being rerun.

[pst2154]

/bot run

[pst2154]

/bot run

[pst2154]

Cache-DiT compatibility follow-up for b238aff0cc combined with #16339 at cd065a7165:

• Validated on 1x NVIDIA B300 SXM6 AC using the current TensorRT-LLM main staging image (reports TensorRT-LLM 1.3.0rc21, PyTorch 2.12.0a0, and cache-dit 1.3.13.dev9).
• The focused wrapper regression passed: test_qwen_sequence_sharding_rejects_padding_without_key_mask_support.
• The real Qwen-Image Cache-DiT integration passed: TestCacheDiTRealPipelineForward::test_qwen_image_cache_dit_skips_blocks_after_forward.
• The real test loaded qwen-image, enabled Cache-DiT, ran 512x512 true CFG for 16 denoise steps, and verified accumulated cached steps were greater than zero. Result: 1 passed in 137.58 seconds.

The new _first_block_attn() helper successfully resolves the real first Qwen block through the Cache-DiT wrapper at both the fused-RoPE and sequence-padding capability probes; no AttributeError occurred.

A merge-tree audit against #16339 finds one expected textual conflict in the true-CFG denoise loop. The combined resolution preserves both behaviors: clone the positive CUDA Graph output before the negative replay, then switch Cache-DiT to the uncond branch before computing the negative prediction.

[coderabbitai]

📝 Walkthrough

Walkthrough

Qwen-Image pipeline prompt masks are now optional, CUDA graph true-CFG execution protects positive predictions from buffer aliasing, and the transformer adds fused RoPE/QK-norm paths, tensor-parallel attention handling, backend-aware masking, sequence sharding, and expanded tests.

Changes

Qwen-Image runtime and validation

Layer / File(s)	Summary
Prompt masks and CFG replay safety tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py, tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py, tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py	Prompt encoding returns None for fully valid masks, and true-CFG CUDA graph execution clones positive predictions before negative replay. Tests cover optional masks and aliased transformer outputs.
Fused attention and mask-aware QKV execution tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py, tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py	Joint attention adds shared RoPE conversion, fused QK normalization gating, tensor-parallel projections and local heads, backend-specific mask dispatch, and unsupported masked sequence-parallel validation.
Transformer blocks and sequence sharding tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py, tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py	Transformer blocks use compiled modulation and residual helpers; model execution shards and gathers sequences, reuses fused rotary inputs, and supports configurable feed-forward activations and tensor-parallel reductions.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant QwenImageTransformer2DModel
  participant SequenceSharder
  participant QwenImageTransformerBlock
  participant QwenJointAttention
  QwenImageTransformer2DModel->>SequenceSharder: shard sequences and masks
  SequenceSharder->>QwenImageTransformer2DModel: return sharded inputs
  QwenImageTransformer2DModel->>QwenImageTransformerBlock: execute block with rotary inputs
  QwenImageTransformerBlock->>QwenJointAttention: run masked joint attention
  QwenImageTransformer2DModel->>SequenceSharder: gather and unpad output

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16391: Updates the same true-CFG CUDA graph path to protect positive noise_pred values from replay aliasing.

Suggested reviewers: yibinl-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 24.14% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and matches the PR’s main goal of adding Qwen Image performance fastpaths.
Description check	✅ Passed	The description is detailed and covers summary, checks, and validation, though it doesn’t mirror all template headings exactly.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py (1)

550-625: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cover final sequence gather and unpadding.

This test exercises _shard_sequences, but would still pass if model forward() omitted the final gather or supplied the wrong unpad_to. Coverage is insufficient for that end-to-end contract.

Add test_qwen_transformer_forward_gathers_unpadded_image_tokens to tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py, asserting the original image sequence length is restored.

As per path instructions, test feedback should state coverage sufficiency and provide a concrete test file and follow-up.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py` around
lines 550 - 625, Extend the Qwen transformer tests with
test_qwen_transformer_forward_gathers_unpadded_image_tokens, exercising
model.forward() rather than only _shard_sequences. Configure the sharder to
observe the final gather and unpad_to value, then assert the returned image
tokens regain the original sequence length and that the gather uses that length,
covering the end-to-end unpadding contract.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py`:
- Around line 493-506: Update
test_qwen_joint_attention_fused_rope_requires_qk_norm to construct
QwenJointAttention with a fused-kernel-supported attention head dimension (64,
128, or 256), adjusting dim consistently with the head configuration while
preserving the qk_norm assertions.

---

Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py`:
- Around line 550-625: Extend the Qwen transformer tests with
test_qwen_transformer_forward_gathers_unpadded_image_tokens, exercising
model.forward() rather than only _shard_sequences. Configure the sharder to
observe the final gather and unpad_to value, then assert the returned image
tokens regain the original sequence length and that the gather uses that length,
covering the end-to-end unpadding contract.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f958c480-1b68-4663-b16a-ad0f0aba08c4

📥 Commits

Reviewing files that changed from the base of the PR and between 287c12a and b238aff.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
• tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
• tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
• tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py

[pst2154]

/bot run

[karljang]

Additionally, please check following issues that my agent found:

1. Hybrid CP x Ulysses is incorrect. Qwen shards across the full CP x Ulysses group, but SEPARATE_QKV disables Ulysses when CP is also active. Each rank attends to an incomplete sequence before outputs are gathered. Ring configurations are also rejected during construction.
2. A newly added prompt-mask test cannot initialize the pipeline. It supplies only torch_dtype, while BasePipeline immediately requires primary_pretrained_config and other configuration fields.
   Thank you!

[yibinl-nvidia]

/bot run

[yibinl-nvidia]

LGTM

[tensorrt-cicd]

PR_Github #61620 [ run ] triggered by Bot. Commit: 9723ac8 Link to invocation

[tensorrt-cicd]

PR_Github #61620 [ run ] completed with state FAILURE. Commit: 9723ac8
/LLM/main/L0_MergeRequest_PR pipeline #49828 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #61673 [ run ] triggered by Bot. Commit: 9723ac8 Link to invocation

[tensorrt-cicd]

PR_Github #61673 [ run ] completed with state FAILURE. Commit: 9723ac8
/LLM/main/L0_MergeRequest_PR pipeline #49879 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #61793 [ run ] triggered by Bot. Commit: 9723ac8 Link to invocation

[tensorrt-cicd]

PR_Github #61793 [ run ] completed with state SUCCESS. Commit: 9723ac8
/LLM/main/L0_MergeRequest_PR pipeline #49993 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chang-l]

/bot run

[tensorrt-cicd]

PR_Github #62278 [ run ] triggered by Bot. Commit: aabf891 Link to invocation

[tensorrt-cicd]

PR_Github #62278 [ run ] completed with state SUCCESS. Commit: aabf891
/LLM/main/L0_MergeRequest_PR pipeline #50449 completed with status: 'SUCCESS'

CI Report

Link to invocation

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #62570 [ run ] triggered by Bot. Commit: dc70e7c Link to invocation

[tensorrt-cicd]

PR_Github #62570 [ run ] completed with state FAILURE. Commit: dc70e7c
/LLM/main/L0_MergeRequest_PR pipeline #50718 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #62617 [ run ] triggered by Bot. Commit: dc70e7c Link to invocation

[tensorrt-cicd]

PR_Github #62617 [ run ] completed with state SUCCESS. Commit: dc70e7c
/LLM/main/L0_MergeRequest_PR pipeline #50759 completed with status: 'SUCCESS'

CI Report

Link to invocation