[None][perf] Enable FLUX2 VisualGen fused NVFP4 SwiGLU path

[pst2154]

Summary

• Enable the CuTe DSL block-scaling MLP path for FLUX.2 dual-stream FFNs.
• Add a guarded FLUX.2 single-stream NVFP4 SwiGLU producer path that emits FP4 directly into the downstream projection.
• Keep gate/up interleaving limited to the single-stream MLP projection and NVFP4 SM100/SM103; QKV and SM120 retain their existing layouts.
• Add a correct split fallback that reorders FLUX's [gate | up] layout to the quantizer's [up | gate] contract.
• Keep fast paths limited to static NVFP4 Blackwell cases with compatible packed widths and no pre-quant/dynamic activation quantization requirements.
• Inherit merged [None][fix] visual_gen FLUX: enable fused DiT QK-norm + RoPE by default #15676 from main, which independently restores fused FLUX QK-norm + RoPE at TP=1. That path is complementary to this PR's TP>1 FLUX.2 NVFP4 producer optimization.

Checks

• File-scoped pre-commit, Ruff/format, byte-compilation, and diff checks pass.
• Rebased cleanly onto current upstream/main, including merged [None][fix] visual_gen FLUX: enable fused DiT QK-norm + RoPE by default #15676.
• DCO sign-off is present.
• GitHub PR Checks, LLM API Compatibility, and Release Checks pass on rebased commit f5c3d58d6d.

B300 validation

Using nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc18 on SM103:

• Split SwiGLU + FP4: 99.9909% packed-value match and 100% scale-factor match versus direct SwiGLU then FP4 quantization.
• Producer-fused BF16: cosine 0.99999446 and normalized MAE 0.00192503 versus ordinary NVFP4 GEMM + SwiGLU.
• Producer FP4 versus quantized fused-BF16 reference: 99.6643% packed-value match and 99.1699% scale-factor match.
• Shared NVFP4 down projection output: cosine 0.99961090.
• FLUX.2-dev TP8-local dimensions (M=4096, K=6144, N=2304): 0.1028 ms baseline versus 0.0658 ms producer-fused (1.56x faster).

Scope note

Validation used focused real operators and model construction, not a full FLUX.2 checkpoint generation run. #15676 separately includes TP1 end-to-end FLUX.2-dev BF16 validation.

Summary by CodeRabbit

• New Features

  • Added hardware- and quantization-aware acceleration for FLUX.2 attention and MLP projections.
  • Enabled optimized FP4 and CUDA paths when supported by the GPU and model configuration.
  • Improved handling of combined attention and MLP outputs across supported execution paths.

• Tests

  • Added coverage for FP4 compatibility, GPU capability checks, projection routing, and SwiGLU behavior.
  • Verified optimized execution is enabled or disabled appropriately based on hardware and quantization settings.

[pst2154]

/bot run

[pst2154]

B300 follow-up validation for commit f5c3d58d6d (rebased; NVFP4 code change equivalent to previously validated 3c289d1d8c) using nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc18 on SM103.

The branch now includes merged #15676 from main: FLUX TP=1 uses its fused QK-norm + RoPE path, while this PR's producer-fused NVFP4 path applies to the split TP>1 FLUX.2 projection layout. The two optimizations are complementary.

Fixes added after GPU validation

• Keep CuTe SwiGLU weight interleaving limited to NVFP4 on SM100/SM103, so SM120 retains the normal weight layout.
• Mark only the single-stream MLP projection for gate/up interleaving; QKV remains unchanged.
• Reorder the split fallback from FLUX's [gate | up] layout to the MoE quantizer's [up | gate] contract.
• Route interleaved single-stream weights through the producer-fused CuTe kernel, using BF16 output below the FP4-output token threshold.

Focused real-op correctness

• Split SwiGLU+FP4 quantization: 99.9909% packed-value match and 100% scale-factor match versus direct SwiGLU then FP4 quantization.
• Producer-fused BF16: cosine 0.99999446 and normalized MAE 0.00192503 versus ordinary NVFP4 GEMM + SwiGLU.
• Producer FP4 output versus quantized fused-BF16 reference: 99.6643% packed-value match and 99.1699% scale-factor match (the existing kernel tests use a 95% threshold).
• Feeding fused and ordinary FP4 outputs through the same NVFP4 down projection: cosine 0.99961090.

Focused performance

FLUX.2-dev TP8-local dimensions (M=4096, K=6144, local_intermediate=2304):

• NVFP4 GEMM + SwiGLU + FP4 quantization: 0.1028 ms
• Producer-fused GEMM + SwiGLU + FP4 output: 0.0658 ms (1.56x faster)

The repaired split fallback measured 0.1505 ms because the layout reorder requires a copy; the producer-fused path is the expected static-NVFP4 route.

Local pre-commit, Ruff, formatting, byte-compilation, and diff checks pass after rebasing onto #15676. This was focused real-op/model-construction validation, not a full FLUX.2 checkpoint generation run.

[coderabbitai]

📝 Walkthrough

Walkthrough

FLUX.2 adds hardware- and quantization-gated CUTEDSL/NVFP4 MLP projection paths, propagates the configuration through joint projections and transformer FFNs, and adds unit tests for eligibility guards, tensor ordering, and SM100f behavior.

Changes

FLUX.2 FP4 projection acceleration

Layer / File(s)	Summary
Hardware-aware SwiGLU capability wiring tensorrt_llm/_torch/visual_gen/models/flux/attention.py, tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py, tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py	SM100f/NVFP4 checks select CUTEDSL blockscaling and pass the option through joint QKV/MLP projections and both transformer FFNs.
FP4 and CUTEDSL attention projection paths tensorrt_llm/_torch/visual_gen/models/flux/attention.py	Attention adds eligibility predicates, FP4 quantization and tiling, split-output combination, and early forward routes for accelerated MLP projections.
Capability and ordering tests tests/unittest/_torch/visual_gen/test_flux_transformer.py	Tests cover FP4-out and CUTEDSL guards, SwiGLU gate/up reordering, and FFN enablement on SM100f with NVFP4.

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

Suggested labels: Community want to contribute

Suggested reviewers: chang-l, qijune, shixiaowei02

Sequence Diagram(s)

sequenceDiagram
  participant Flux2ParallelSelfAttention
  participant CUTEDSL_NVFP4_Kernels
  participant ToOutProjection
  Flux2ParallelSelfAttention->>Flux2ParallelSelfAttention: evaluate FP4/CUTEDSL eligibility
  Flux2ParallelSelfAttention->>CUTEDSL_NVFP4_Kernels: quantize and project MLP states
  CUTEDSL_NVFP4_Kernels-->>Flux2ParallelSelfAttention: return MLP projection
  Flux2ParallelSelfAttention->>ToOutProjection: combine attention and MLP outputs
  ToOutProjection-->>Flux2ParallelSelfAttention: return projected output

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change and follows the requested [ticket][type] format.
Description check	✅ Passed	The description covers the summary, validation, and checklist context, with only minor template sections missing.
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

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch codex/flux2-nvfp4-swiglu-clean

---

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/visual_gen/models/flux/attention.py (1)

607-627: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Gate weight interleaving with the same CuteDSL eligibility. use_cute_dsl_blockscaling_mm still interleaves gate_up_proj weights during load, but forward() only uses the CuteDSL path when _can_project_hidden_mlp_with_cute_dsl() passes. On tp_size == 1, that guard is false, so the split fallback reads an already-interleaved [up | gate] layout as contiguous [gate | up] and returns wrong activations.

🤖 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 `@tensorrt_llm/_torch/visual_gen/models/flux/attention.py` around lines 607 -
627, The weight interleaving controlled by use_cute_dsl_blockscaling_mm must use
the same eligibility condition as the forward CuteDSL path. Update the
gate_up_proj loading/interleaving logic and its related guard to rely on
_can_project_hidden_mlp_with_cute_dsl(), so tp_size == 1 uses the
non-interleaved layout expected by the split fallback while eligible CuteDSL
execution retains interleaving.

🤖 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 `@tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py`:
- Around line 253-257: Update the use_cute_dsl_swiglu initialization to avoid
unconditionally calling config.get_quant_config() when config may be None. Reuse
the already-resolved self.quant_config as in attention.py::_init_qkv_proj, or
otherwise guard the config access while preserving the existing optional-config
behavior.

---

Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/flux/attention.py`:
- Around line 607-627: The weight interleaving controlled by
use_cute_dsl_blockscaling_mm must use the same eligibility condition as the
forward CuteDSL path. Update the gate_up_proj loading/interleaving logic and its
related guard to rely on _can_project_hidden_mlp_with_cute_dsl(), so tp_size ==
1 uses the non-interleaved layout expected by the split fallback while eligible
CuteDSL execution retains interleaving.

🪄 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: aebde702-b46b-4a28-87e7-365188413298

📥 Commits

Reviewing files that changed from the base of the PR and between fab39a7 and f5c3d58.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/visual_gen/models/flux/attention.py
• tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py
• tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py
• tests/unittest/_torch/visual_gen/test_flux_transformer.py

[karljang]

@pst2154, could you please resolve any conflicts? Also, I’m curious if you had a chance to verify the e2e LPIPS results as well?

[pst2154]

/bot run --disable-fail-fast

[pst2154]

Resolved the conflict against current main in 87c6b42f51, preserving the new uneven-TP QKV sharding and the producer-fused NVFP4 path. I also fixed the optional-config guard and aligned static weight interleaving with runtime eligibility: SM100 or SM103 NVFP4, TP greater than 1, and the exact packed local layout. On B300, focused tests pass: 3 passed, 13 deselected. The exact FLUX.2 TP2 NVFP4 e2e prompt completed at 256x256 for 4 steps, with 48 eligible fused modules per rank and 192 fused calls per rank. Fused versus the same TP2 NVFP4 unfused path measured LPIPS 0.02360791, below 0.05. The checked-in BF16 golden is not directly comparable to this quantized TP2 mode: fused and unfused both differ from it by about 0.10 LPIPS, while the fused versus unfused A/B remains below the quality gate.

[luyiyun1021]

Fused SwiGLU paths have no numerical coverage — only guard tests

The new unit tests exercise only the dispatch guards; every one of them mocks out the actual fused op, so nothing validates that the fused math matches the eager reference. E.g. test_flux2_moe_swiglu_reorders_gate_up monkeypatches torch.ops.trtllm.moe_swiglu_nvfp4_quantize, so it only proves "we reorder to [up|gate]" — not that [up|gate] + the kernel equals the reference silu(gate)*up. The cute_dsl bf16/fp4-out paths are similarly only reached through SimpleNamespace fakes.

That equivalence is load-bearing and subtle: the [gate|up] → [up|gate] swap at attention.py#L600-L601 is only correct because moe_swiglu_nvfp4_quantize applies SiLU to the second half (silu(2nd)*1st), the opposite convention from silu_and_mul (silu(1st)*2nd). Right now the only thing verifying this end-to-end is the manual B300 run in the PR description. If the kernel's gate/up convention ever changes, this path silently produces wrong output and CI stays green.

Suggestion: add an SM100/103-gated parity test that runs the real ops and asserts each fast path matches the eager swiglu(gate_up) + down_proj reference — essentially promoting the B300 numbers already in the description into CI. Gate it the same way the op's own tests in test_cute_dsl_moe.py do:

from tensorrt_llm._utils import get_sm_version

@pytest.mark.skipif(
    get_sm_version() not in (100, 103),
    reason="cute_dsl NVFP4 SwiGLU is only supported on SM 100/103 GPUs",
)
def test_flux2_fused_swiglu_matches_eager():
    # build a small static-NVFP4 Flux2ParallelSelfAttention (tp>1),
    # run the fp4out / bf16 cute_dsl paths and the moe_swiglu split fallback,
    # assert cosine / packed-value match vs eager swiglu + down_proj
    # (reuse the thresholds from the PR's B300 validation)

This keeps the opposite-convention invariant from regressing unnoticed.

[pst2154]

@luyiyun1021 @zhenhuaw-me The requested numerical coverage is now in d0241b0b33. It runs real SM100/103 kernels through actual TP2 NVFP4 FLUX.2 modules for CuTe BF16 output, CuTe FP4 output, and the MoE split fallback, each against eager SwiGLU plus the same down projection. All three passed on B300 SM103 (3 passed, 16 deselected). This coverage also found and fixed the packed FP4 storage dtype at the MoE fallback boundary.

[pst2154]

/bot run

[karljang]

LGTM

[chang-l]

/bot run --disable-fail-fast

[chang-l]

FYI @pst2154 — the three earlier /bot run comments never started a pipeline. .github/workflows/blossom-ci.yml gates the Authorization job on contains(fromJson('[...]'), github.actor), and your handle is not in that allowlist, so the comments were silently skipped (no bot ack, no Jenkins run — which is why this PR has zero CI history).

I have triggered it from an allowlisted account; the pipeline is running now. Worth opening a small PR to add yourself to the allowlist in .github/workflows/blossom-ci.yml so you can self-trigger going forward.

[tensorrt-cicd]

PR_Github #61968 [ run ] triggered by Bot. Commit: d0241b0 Link to invocation

[tensorrt-cicd]

PR_Github #61968 [ run ] completed with state FAILURE. Commit: d0241b0
/LLM/main/L0_MergeRequest_PR pipeline #50159 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]

/bot run --disable-fail-fast

[chang-l]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62012 [ run ] triggered by Bot. Commit: d0241b0 Link to invocation

[tensorrt-cicd]

PR_Github #62012 [ run ] completed with state SUCCESS. Commit: d0241b0
/LLM/main/L0_MergeRequest_PR pipeline #50202 completed with status: 'SUCCESS'

CI Report

Link to invocation