[None][perf] Fuse Qwen Image QK norm RoPE prep

[pst2154]

Summary

• convert Qwen-Image complex RoPE frequencies into the real cos/sin layout used by VisualGen fused DiT kernels
• enable the existing packed fused QK-norm + RoPE op for supported Qwen-Image bf16 CUDA runs
• keep the existing complex RoPE fallback for unsupported/local paths and add focused tests for conversion/fallback behavior

Tests

• python3 -m py_compile tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py
• git diff --check
• python3 -m pytest tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py -k 'qwen_complex_freqs_convert or cpu_fallback' (blocked locally: missing mpi4py from tests/unittest/conftest.py)

Notes

• GPU/container validation is still needed for the fused CUDA path.

Summary by CodeRabbit

• New Features

  • Added support for a shared rotary embedding format, improving compatibility across attention paths.
  • Enabled a fused attention path in more single-device configurations.

• Bug Fixes

  • Improved attention handling when masks are present for more consistent tensor shaping.
  • Added fallback coverage to keep image transformer outputs aligned with input shapes.

• Tests

  • Added coverage for rotary embedding equivalence.
  • Added a CPU fallback test for the Qwen image transformer.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a utility converting Qwen complex RoPE frequencies to cos/sin format, enables a fused QK-norm+RoPE attention path in QwenJointAttention conditional on tensor-parallel size, refactors QKV preparation into dispatching methods, updates forward's attention-mask handling, and adds corresponding unit tests.

Changes

Fused QK-norm+RoPE attention path

Layer / File(s)	Summary
Complex-to-cos/sin RoPE conversion tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py	Adds qwen_complex_freqs_to_cos_sin converting complex RoPE frequencies into interleaved real cos/sin tensors.
Fused QKV preparation and eligibility logic tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py	Enables fuse_qk_norm_rope when tp_size == 1; adds _use_fused_qk_norm_rope, _prepare_qkv_fused, and updated _prepare_qkv_unfused to build and process QKV tensors via fused or unfused paths.
Forward dispatch and attention-mask reshaping tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py	Adds _prepare_qkv dispatch method used by forward; reworks attention-mask branch to reshape/transpose Q/K/V for SDPA when a mask is present.
Tests for RoPE conversion and fused attention tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py	Updates imports and adds tests validating rotary embedding equivalence and transformer forward pass with fuse_qk_norm_rope assertion.

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

Sequence Diagram(s)

sequenceDiagram
  participant Forward as QwenJointAttention.forward
  participant Prepare as _prepare_qkv
  participant FusedPath as _prepare_qkv_fused
  participant Convert as qwen_complex_freqs_to_cos_sin
  participant PackedOp as apply_packed_qk_norm_rope

  Forward->>Prepare: request joint Q/K/V
  Prepare->>FusedPath: dispatch when fused eligible
  FusedPath->>Convert: convert text/image RoPE freqs
  Convert-->>FusedPath: cos, sin tensors
  FusedPath->>PackedOp: apply packed QK-norm+RoPE
  PackedOp-->>FusedPath: normalized/rotated qkv
  FusedPath-->>Prepare: split per-head Q/K/V
  Prepare-->>Forward: flattened joint Q/K/V

Loading

Suggested reviewers: rakib-hasan, aswinvisva, yechank-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the PR's main change and follows the required [None][type] Summary format.
Description check	✅ Passed	The description includes summary and test coverage, but it omits the explicit PR Checklist section from the template.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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]

🧹 Nitpick comments (1)

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

254-282: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Add explicit coverage for the supported fused path.

Coverage is sufficient for CPU fallback, but insufficient for the PR’s CUDA bf16 fused path: this test uses CPU and attention_head_dim=8, while _use_fused_qk_norm_rope only enables fusion for CUDA bf16 with head dims 64/128/256. Please add or follow up with a GPU-gated test in tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py that exercises _prepare_qkv_fused or compares fused vs. unfused outputs for a supported head dim.

As per path instructions, test feedback should state whether coverage is sufficient or needs follow-up outside the PR.

🤖 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 254 - 282, Coverage here only exercises the CPU fallback in
QwenImageTransformer2DModel and does not validate the supported fused path. Add
a GPU-gated test in test_qwen_image_pipeline_config that runs the CUDA bf16
fused configuration for a supported attention_head_dim (64, 128, or 256), and
explicitly exercise _use_fused_qk_norm_rope / _prepare_qkv_fused or compare
fused vs. unfused outputs. Keep the existing CPU test for fallback, but note in
the test coverage that fused-path validation is handled by the new GPU-specific
case.

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.

Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py`:
- Around line 254-282: Coverage here only exercises the CPU fallback in
QwenImageTransformer2DModel and does not validate the supported fused path. Add
a GPU-gated test in test_qwen_image_pipeline_config that runs the CUDA bf16
fused configuration for a supported attention_head_dim (64, 128, or 256), and
explicitly exercise _use_fused_qk_norm_rope / _prepare_qkv_fused or compare
fused vs. unfused outputs. Keep the existing CPU test for fallback, but note in
the test coverage that fused-path validation is handled by the new GPU-specific
case.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d32a5e7a-2427-435c-97f5-dfe4d0a53941

📥 Commits

Reviewing files that changed from the base of the PR and between 5042300 and 1dc1b5f.

📒 Files selected for processing (2)

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

[pst2154]

Closing this in favor of the combined clean Qwen VisualGen perf PR: #16142. That PR includes this QK-norm/RoPE work plus the related mask, adaLN, and NVFP4 GELU fixes in one reviewable branch.