[None][fix] visual_gen FLUX: enable fused DiT QK-norm + RoPE by default by chang-l · Pull Request #15676 · NVIDIA/TensorRT-LLM

chang-l · 2026-06-26T21:36:18Z

Description

The fused DiT QK-norm + RoPE kernel was added for FLUX in #11869 and was enabled by default there: the base Attention default auto-enabled it for per-head norm, and FLUX (qk_norm_mode="per_head") rode that auto-default.

# #11869 base default:
else:
    self.fuse_qk_norm_rope = qk_norm_mode != "full"   # FLUX (per_head) -> True

#13052 ("Fused cross-head QK Norm + RoPE kernel for WAN", merged 2026-05-13) then changed the base default to opt-in:

-        if fuse_qk_norm_rope is not None:
-            self.fuse_qk_norm_rope = fuse_qk_norm_rope
-        else:
-            self.fuse_qk_norm_rope = qk_norm_mode != "full"
+        self.fuse_qk_norm_rope = fuse_qk_norm_rope if fuse_qk_norm_rope is not None else False

That PR added WAN's explicit opt-in but did not migrate FLUX, which had been relying on the auto-default. So FLUX silently regressed to the unfused F.rms_norm + apply_rotary_emb path from #13052 onward (still off in main; reproduced on release:1.3.0rc18/rc19). The stale FluxJointAttention docstring ("fused path is the default") is the fossil.

Fix

Make FLUX opt in explicitly: pass fuse_qk_norm_rope=(tp_size == 1) in FluxJointAttention (propagates to Flux2ParallelSelfAttention, covering FLUX.1/FLUX.2 dual- and single-stream blocks).

The opt-in is gated on TP=1 because the fused op currently requires TP=1 — apply_packed_qk_norm_rope asserts tp_size == 1 — so under TP>1 FLUX falls back to the unfused path. This mirrors WAN's (tp_size == 1) gating. (FLUX's per-head/intra-head norm is itself TP-invariant; the gate reflects the op's current TP=1 implementation limit, not a math/correctness issue.)

Why explicit opt-in rather than restoring a base-class auto-default: kernel applicability depends on RoPE format/convention, not just norm-mode/TP. E.g. Qwen-Image is per-head but stays False because it uses complex-valued RoPE (freqs_cis), incompatible with the kernel's real cos/sin format (documented TODO in transformer_qwen_image.py). Only the model author knows all constraints hold, so per-model opt-in is the right design.

Verification

Static: self.fuse_qk_norm_rope resolves to True for every FLUX attention module when TP=1 (was False since [#12716][feat] Fused cross-head QK Norm + RoPE kernel for WAN #13052); False under TP>1.
Runtime (single B200, FLUX.2-dev BF16, 1024×1024, 50 steps, release:1.3.0rc19, TP=1): end-to-end generate() 12.298 s → 11.771 s = 1.045× faster (−4.3% wall time), against a −0.01% null-control noise floor (single-process order-rotated A/B, cold round discarded, n=4). torch.profiler confirms trtllm::fused_dit_qk_norm_rope (→ fusedDiTQKNormRopeKernel<128,true>, 56 modules × 50 steps = 2800 calls) is exercised only with the flag on; the off path runs separate inductor-fused Triton RMSNorm kernels. Output images are visually identical between conditions. (The +4.3% is over a torch-compiled baseline, so it sits a touch below [None][feat] Add fused DiT QK Norm + RoPE CUDA kernel for FLUX #11869's standalone-microbench numbers.)

Test

FLUX opt-in/TP gating: tests/unittest/_torch/visual_gen/test_flux_attention.py::TestFluxAttentionBackend::test_fused_qk_norm_rope_enabled_only_for_tp1
Kernel correctness/dual-stream coverage: tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_dit_qk_norm_rope.py

PR Checklist

PR description clearly explains what and why.
Follows TRT-LLM coding guidelines to the best of my knowledge.

coderabbitai · 2026-06-26T21:39:28Z

📝 Walkthrough

Walkthrough

FluxJointAttention now explicitly enables the fused QK-norm + RoPE CUDA path and updates its docstring to describe that behavior for the FLUX per-head QK-norm configuration.

Changes

Flux fused QK-norm + RoPE

Layer / File(s)	Summary
Fused QK-norm + RoPE opt-in `tensorrt_llm/_torch/visual_gen/models/flux/attention.py`	`FluxJointAttention` now passes `fuse_qk_norm_rope=True` to `Attention` and expands its docstring to describe the fused kernel path and tensor-parallel per-head QK-norm mode.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~5 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and accurately summarizes the main change: enabling fused DiT QK-norm + RoPE by default for FLUX.
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.
Description check	✅ Passed	The PR description covers the issue, fix, verification, tests, and checklist, and follows the requested structure well.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

The fused DiT QK-norm + RoPE kernel was added for FLUX in NVIDIA#11869 and was enabled by default there via an auto-default (fused on for per-head norm). NVIDIA#13052 (WAN cross-head) later changed that base default to opt-in (... if ... is not None else False), wired WAN's own opt-in, but did not migrate FLUX -- which had relied on the auto-default. FLUX has run the unfused F.rms_norm + apply_rotary_emb path ever since (main, rc18, rc19), contradicting the FluxJointAttention docstring. Make FLUX opt in explicitly in FluxJointAttention (propagates to Flux2ParallelSelfAttention, covering FLUX.1/FLUX.2 dual- and single-stream blocks). The fused op currently requires TP=1 (apply_packed_qk_norm_rope asserts tp_size == 1), so gate the opt-in on tp_size == 1 and fall back to the unfused path under tensor parallelism, mirroring WAN. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

chang-l · 2026-07-01T22:06:30Z

/bot run

tensorrt-cicd · 2026-07-01T22:13:13Z

PR_Github #57037 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-01T22:56:38Z

PR_Github #57037 [ run ] completed with state SUCCESS. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #45832 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

chang-l · 2026-07-01T23:20:42Z

/bot run --disable-fail-fast

chang-l · 2026-07-01T23:44:45Z

/bot run --disable-fail-fast

chang-l · 2026-07-02T00:52:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T00:59:03Z

PR_Github #57048 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

chang-l · 2026-07-02T03:32:22Z

/bot run

tensorrt-cicd · 2026-07-02T03:34:57Z

PR_Github #57048 [ run ] completed with state FAILURE. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #45843 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

tensorrt-cicd · 2026-07-02T03:38:29Z

PR_Github #57099 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-02T06:20:43Z

PR_Github #57099 [ run ] completed with state FAILURE. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #45887 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

chang-l · 2026-07-02T06:28:23Z

/bot run

tensorrt-cicd · 2026-07-02T06:34:31Z

PR_Github #57147 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-02T07:50:38Z

PR_Github #57147 [ run ] completed with state FAILURE. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #45928 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

chang-l · 2026-07-02T08:19:24Z

/bot run

tensorrt-cicd · 2026-07-02T13:35:53Z

PR_Github #57175 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-02T15:20:39Z

PR_Github #57175 [ run ] completed with state SUCCESS. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #45950 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

chang-l · 2026-07-02T17:04:24Z

/bot run

tensorrt-cicd · 2026-07-02T17:10:15Z

PR_Github #57226 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

yibinl-nvidia

LGTM

tensorrt-cicd · 2026-07-02T18:54:25Z

PR_Github #57226 [ run ] completed with state SUCCESS. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #45994 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

chang-l · 2026-07-02T23:14:02Z

/bot run

chang-l · 2026-07-02T23:40:29Z

/bot run

chang-l · 2026-07-03T00:07:25Z

/bot run

chang-l · 2026-07-03T00:39:16Z

/bot run

tensorrt-cicd · 2026-07-03T00:50:55Z

PR_Github #57282 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-03T00:59:59Z

PR_Github #57282 [ run ] completed with state FAILURE. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #46044 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

chang-l · 2026-07-03T01:32:00Z

/bot run

tensorrt-cicd · 2026-07-03T01:38:27Z

PR_Github #57294 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-03T07:33:33Z

PR_Github #57294 [ run ] completed with state FAILURE. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #46054 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

chang-l · 2026-07-04T03:07:18Z

/bot run

tensorrt-cicd · 2026-07-04T03:12:51Z

PR_Github #57523 [ run ] triggered by Bot. Commit: 2015dbe Link to invocation

tensorrt-cicd · 2026-07-04T04:33:04Z

PR_Github #57523 [ run ] completed with state SUCCESS. Commit: 2015dbe
/LLM/main/L0_MergeRequest_PR pipeline #46252 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

chang-l requested a review from a team as a code owner June 26, 2026 21:36

github-actions Bot assigned chang-l Jun 26, 2026

chang-l marked this pull request as draft June 26, 2026 22:29

chang-l force-pushed the dev-flux-enable-fuse-qk-norm-rope branch from 21f5e08 to 0df87f4 Compare June 26, 2026 22:49

[None][test] cover FLUX fused QK-norm RoPE TP gating

2015dbe

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

chang-l marked this pull request as ready for review July 1, 2026 22:06

yibinl-nvidia approved these changes Jul 2, 2026

View reviewed changes

Merge branch 'main' into dev-flux-enable-fuse-qk-norm-rope

2451bb3

Uh oh!

Conversation

chang-l commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Fix

Verification

Test

PR Checklist

Uh oh!

coderabbitai Bot commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Uh oh!

chang-l commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

chang-l commented Jul 1, 2026

Uh oh!

chang-l commented Jul 1, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

yibinl-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

chang-l commented Jul 2, 2026

Uh oh!

chang-l commented Jul 3, 2026

Uh oh!

chang-l commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

chang-l commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

chang-l commented Jul 4, 2026

Uh oh!

tensorrt-cicd commented Jul 4, 2026

chang-l commented Jun 26, 2026 •

edited

Loading

coderabbitai Bot commented Jun 26, 2026 •

edited

Loading