[None][fix] handle rank-3 Fp4 and non-contiguous Ulysses shard in LTX-2 NVFP4 fused paths#15718
Conversation
📝 WalkthroughWalkthrough
ChangesFP4 3D Input Flattening for Fused MLP Kernels
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py (1)
267-279: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd assertions for preserved FP4 metadata.
This regression covers flattening, but not the accompanying
scaling_factor/is_sf_swizzledpreservation. A future change could pass the rank/shape checks while still breaking the quantization metadata handed to_input_prepare.Suggested test tightening
def fake_input_prepare(module, x): captured["dim"] = x.fp4_tensor.dim() if isinstance(x, Fp4QuantizedTensor) else x.dim() captured["shape"] = ( tuple(x.fp4_tensor.shape) if isinstance(x, Fp4QuantizedTensor) else tuple(x.shape) ) + if isinstance(x, Fp4QuantizedTensor): + captured["scaling_factor"] = x.scaling_factor.clone() + captured["is_sf_swizzled"] = x.is_sf_swizzled raise _Stop # short-circuit before the real (HW-only) GEMM proj.quant_method = types.SimpleNamespace(_input_prepare=fake_input_prepare) fp4_3d = Fp4QuantizedTensor( fp4_tensor=torch.zeros(B, S, K // 2, dtype=torch.uint8), - scaling_factor=torch.zeros(B * S, K // 16, dtype=torch.uint8), + scaling_factor=torch.arange(B * S * (K // 16), dtype=torch.uint8).reshape( + B * S, K // 16 + ), + is_sf_swizzled=False, )Then assert:
assert torch.equal(captured["scaling_factor"], fp4_3d.scaling_factor) assert captured["is_sf_swizzled"] is FalseCoverage is sufficient for flattening, but needs this follow-up for metadata preservation. As per path instructions, "
tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM."🤖 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/thop/parallel/test_dense_gemm_act_fusion.py` around lines 267 - 279, The FP4 flattening test in fake_input_prepare only verifies rank/shape, but it should also validate that quantization metadata is preserved when the tensor is passed into _input_prepare. Update the test around fp4_3d and proj.quant_method._input_prepare to capture the forwarded metadata and add assertions that captured["scaling_factor"] matches fp4_3d.scaling_factor and captured["is_sf_swizzled"] is False, so the regression covers metadata integrity as well as shape handling.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/thop/parallel/test_dense_gemm_act_fusion.py`:
- Around line 267-279: The FP4 flattening test in fake_input_prepare only
verifies rank/shape, but it should also validate that quantization metadata is
preserved when the tensor is passed into _input_prepare. Update the test around
fp4_3d and proj.quant_method._input_prepare to capture the forwarded metadata
and add assertions that captured["scaling_factor"] matches fp4_3d.scaling_factor
and captured["is_sf_swizzled"] is False, so the regression covers metadata
integrity as well as shape handling.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 18914edd-3b53-4b47-a86d-e4253a518846
📒 Files selected for processing (3)
tensorrt_llm/_torch/modules/gated_mlp.pytensorrt_llm/_torch/modules/mlp.pytests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py
luyiyun1021
force-pushed
the
dev/ltx2-fp4-mlp-rank-fix
branch
2 times, most recently
from
5e3a272 to
2d38bfc
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56372 [ run ] triggered by Bot. Commit: |
luyiyun1021
force-pushed
the
dev/ltx2-fp4-mlp-rank-fix
branch
from
a34c581 to
1e2b7db
Compare
luyiyun1021
changed the title
luyiyun1021
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #56413 [ run ] triggered by Bot. Commit: |
|
PR_Github #56372 [ run ] completed with state |
|
PR_Github #56413 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #56824 [ run ] triggered by Bot. Commit: |
|
PR_Github #56824 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56872 [ run ] triggered by Bot. Commit: |
The fused NVFP4 MLP epilogue path (MLP._fused_gelu and GatedMLP._fused_gate_up_swiglu) calls the dense CuteDSL GEMM ops directly, bypassing NVFP4LinearMethod.apply. Those helpers flattened only plain rank-3 tensors to 2D and passed Fp4QuantizedTensor / tuple inputs through unchanged. When a fused rmsnorm emits a pre-quantized rank-3 Fp4QuantizedTensor ([B, S, D/2]) -- as the static-NVFP4 LTX-2 transformer does for its FFN -- the 3D fp4 activation reached the 2D-only GEMM and tripped "assert inputs[0].dim() == 2". Mirror the existing NVFP4LinearMethod.apply flatten into both fused helpers: flatten the packed fp4 data to [M, D/2] (the scale factor is already laid out for the flat M view, so it is reused as-is) and unflatten the output to the caller's rank. Also count tokens as B*S for the fp4-out vs bf16-out switch. Add a HW-independent regression test asserting both fused helpers flatten a rank-3 Fp4QuantizedTensor to 2D before the kernel. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
LTX2Attention.forward_async overrides the base Attention.forward_async (LTX-2 RoPE specialization) but dropped the base's isinstance(Fp4QuantizedTensor) short-circuit. When the fused AdaLN+NVFP4 self-attn norm pre-quantizes the input, x arrives as a rank-3 Fp4QuantizedTensor and the share-quant dedup called x.reshape() on it, which Fp4QuantizedTensor does not support (AttributeError). Restore the base's 3-tier branch: when x is already fp4, reuse it directly (0 extra quant) and let each to_q/to_k/to_v's NVFP4 Linear.apply consume the shared rank-3 fp4 (it flattens internally; _input_prepare skips re-quant). Same rank-3 Fp4QuantizedTensor bug class as the MLP fix in this PR. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
Model-level smoke for the rank-3 Fp4QuantizedTensor fixes in this PR: builds a static-NVFP4 AudioVideo transformer (fused-AdaLN hidden dim 2048, so the fused norm emits a rank-3 fp4 into the FFN and attention) with fabricated weights/inputs and asserts the forward does not raise. A forward-pre-hook asserts a rank-3 Fp4QuantizedTensor actually reaches a module, so a silent bf16 fallback fails the test rather than passing trivially. Blackwell-gated (NVFP4 fused kernels are SM100+). Verified on B200: passes with the fixes; reverting the MLP flatten makes the forward crash in _fused_gelu on the rank-3 fp4. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
SequenceSharder.shard returned a dim-1 block-slice view. When the leading (batch) dim is >1 -- batched serial CFG (cfg_size=1, guidance>1) under sequence parallelism (ulysses_size>=2) -- that slice is non-contiguous, and the fused DiT kernel (fused_dit_gate_resid_norm_shift_scale, via apply_fused_rmsnorm_shift_scale) asserts a dense row-major x. Materialize a contiguous shard; .contiguous() is a no-op (returns self, no copy) when the slice is already dense (B==1). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
|
/bot run --disable-fail-fast |
luyiyun1021
force-pushed
the
dev/ltx2-fp4-mlp-rank-fix
branch
from
7d8ad4b to
5e66140
Compare
|
PR_Github #56889 [ run ] triggered by Bot. Commit: |
|
PR_Github #56872 [ run ] completed with state |
|
PR_Github #56889 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56927 [ run ] triggered by Bot. Commit: |
|
PR_Github #56927 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56997 [ run ] triggered by Bot. Commit: |
|
PR_Github #56997 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #57065 [ run ] triggered by Bot. Commit: |
|
PR_Github #57065 [ run ] completed with state |
…-2 NVFP4 fused paths (NVIDIA#15718) Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…-2 NVFP4 fused paths (NVIDIA#15718) Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
5 participants
Summary by CodeRabbit
Description
The static-NVFP4 LTX-2 transformer fuses activation quantization into its AdaLN RMSNorm, and runs that fused norm + FFN + attention on a per-rank sequence shard under Ulysses parallelism. This PR fixes three crashes in the NVFP4 LTX-2 fused transformer paths, each surfacing only under a configuration the prior tests never exercised (the LLM path pre-flattens activations to 2D; the LTX-2 e2e / unit tests ran only BF16 / FP8 and never the
cfg_size=1+ Ulysses combination), so there was no prior coverage.(1) Fused MLP epilogue (
MLP._fused_gelu,GatedMLP._fused_gate_up_swiglu) calls the dense CuteDSL GEMM ops directly, bypassingNVFP4LinearMethod.apply. The helpers flattened only plain rank-3 tensors and passedFp4QuantizedTensorthrough unchanged, so the 3D fp4 activation[B, S, D/2]from the fused norm reached the 2D-only GEMM and trippedassert inputs[0].dim() == 2.(2) Attention
LTX2Attention.forward_async(the LTX-2 RoPE-specialized override of the baseAttention.forward_async) had dropped the base'sisinstance(input, Fp4QuantizedTensor)short-circuit. When the fused self-attn norm pre-quantized the input, the SEPARATE_QKV share-quant dedup calledx.reshape()on theFp4QuantizedTensor(no such method) and raisedAttributeError. Unlike the MLP, attention does not bypassLinear.apply, so the base class already handled this; only the override had diverged.(3) Ulysses sequence shard (
SequenceSharder.shard) returned adim=1block-slice view of the hidden state. Slicing[B, S, D]along the sequence dim to[B, S/U, D]is non-contiguous wheneverB > 1(gaps between batch rows). The fused DiT kernelfused_dit_gate_resid_norm_shift_scale(viaapply_fused_rmsnorm_shift_scale) asserts a dense row-majorx, so the first transformer block crashed during warmup. This was a latent issue inSequenceSharder.shard(introduced by the CP-unification refactor) that the fused-AdaLN feature first exposed — the kernel is the first shard consumer that both requires contiguity and runs atB >= 2.Fix (3) — trigger conditions
Three conditions must hold simultaneously; the bug lives in the gap none of the existing configs covered:
cfg_size == 1andguidance_scale > 1(single-modal: no STG / modality guidance)B = 2ulysses_size >= 2xalong the sequence dim → strided viewfuse=True)Reproduced on B200 (2-GPU
cfg_size=1/ulysses_size=2,guidance_scale=4):AssertionError: x must be contiguous (fused kernel assumes flat row layout)atapply_fused_rmsnorm_shift_scale, on both ranks, during warmup.Fix (3) — impact / blast radius
.contiguous()returns the same storage (zero copy) when the slice is already dense — verified. SoB=1,cfg_size>=2(CFG-parallel →B=1/rank), andulysses_size=1(no slicing) are all byte-identical. Confirmed by an 8-GPUcfg=2 × uly=4e2e that runs clean.SequenceSharder.shardis shared by Flux / Wan / Cosmos3 / LTX-2; since it only materializes a copy when the slice is genuinely non-contiguous, those models are unaffected unless they hit the sameB>=2+ sequence-parallel combination, where it likewise prevents the crash.B>=2. Measured at the real shape[2, 7680, 4096]bf16 → ~142 µs/step (251 MB R+W), i.e. ~5.7 ms over a 40-step generation — under ~0.05% of a denoise step, paid only by the config that previously crashed.Changes
MLP._fused_gelu/GatedMLP._fused_gate_up_swiglu: flatten a rank-3Fp4QuantizedTensorto[M, D/2]before the GEMM (its scale factor is already laid out for the flatMview, so it is reused as-is) and unflatten the output back to the caller's rank — mirroringNVFP4LinearMethod.apply._token_count/ SwiGLU fp4-out switch: count tokens asB*Sfor a rank-3 input (wasB), so the fp4-out vs bf16-out selection is correct.LTX2Attention.forward_async: restore the base's 3-tier input branch — whenxis already anFp4QuantizedTensor, reuse it directly (0 extra quant) and let eachto_{q,k,v}'sNVFP4LinearMethod.applyconsume the shared rank-3 fp4.SequenceSharder.shard: return a contiguous shard (tensor[...].contiguous()); no-op (no copy) when the slice is already dense (B==1), so only the non-contiguousB>=2case pays a copy.Test Coverage
tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py::test_fused_act_flattens_3d_fp4_input[gelu|swiglu]— HW-independent regression asserting both fused MLP helpers flatten a rank-3Fp4QuantizedTensorto 2D before the kernel (without the fix the kernel sees rank-3).tests/unittest/_torch/visual_gen/test_ltx2_transformer.py::TestLTX2NVFP4ForwardSmoke::test_audio_video_nvfp4_forward_smoke— Blackwell-gated model smoke: a static-NVFP4 AudioVideo transformer (fused-AdaLN hidden dim 2048) runs a forward with fabricated weights/inputs and must not raise. A forward-pre-hook asserts a rank-3Fp4QuantizedTensoractually reaches a module, so a silent bf16 fallback fails the test rather than passing trivially. Verified on B200 as a true regression guard: passes with the fixes, and reverting the MLP flatten makes the same forward crash in_fused_geluon the rank-3 fp4.cfg_size=1/ulysses_size=2/guidance_scale=4previously crashed atapply_fused_rmsnorm_shift_scaleduring warmup; with the fix it runs green through denoise → VAE decode → saved video. The contiguity mechanism (B=1 dense / B=2 strided) and the.contiguous()no-op-when-dense property were confirmed by standalone micro-checks. A dedicated HW-independent unit test forSequenceSharder.shardcontiguity can be added on request.PR Checklist
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.