[https://nvbugs/6572800][fix] Added _use_fused_ln(fp4_scale) and gated each of the three norm sites on its… - #17509
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughWan fused LayerNorm selection now requires supported hidden dimensions and downstream FP4 scales. Norm paths preserve standard LayerNorm when fusion does not apply. Unit tests cover the new conditions, and obsolete Wan accuracy waivers are removed. ChangesWan fused LayerNorm gating
Estimated code review effort: 3 (Moderate) | ~20 minutes Mergeability Score: ⚪ Minimal · up to This is a localized Wan normalization-path fix intended to restore the correct eager or fused behavior per quantization scale; no actionable merge-blocking risk remains beyond normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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/visual_gen/test_wan_transformer.py (1)
349-374: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the fused LayerNorm branches.
The new test covers only
_use_fused_ln(). Existing tests coverWanBlock.forward()andpost_load_weights(), but no test supplies NVFP4 scales to exercise fused Norm1, Norm2, or Norm3.Add focused tests or mocks for the fused and eager paths. The test is already listed in
tests/integration/test_lists/test-db/l0_b200.yml. Runpytest tests/unittest/.Coverage verdict: insufficient.
🤖 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_wan_transformer.py` around lines 349 - 374, Add focused tests in the WanBlock test suite that provide NVFP4 scales and exercise fused and eager LayerNorm execution for Norm1, Norm2, and Norm3 through WanBlock.forward(), including the no-scale eager path. Reuse existing block/config fixtures and mocks, and assert each norm path is selected or invoked as expected while preserving existing post_load_weights coverage.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_wan_transformer.py`:
- Around line 349-374: Add focused tests in the WanBlock test suite that provide
NVFP4 scales and exercise fused and eager LayerNorm execution for Norm1, Norm2,
and Norm3 through WanBlock.forward(), including the no-scale eager path. Reuse
existing block/config fixtures and mocks, and assert each norm path is selected
or invoked as expected while preserving existing post_load_weights coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c16ad821-56d4-4eb3-9a28-638ea7fb48f2
📒 Files selected for processing (3)
tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.pytests/integration/test_lists/waives.txttests/unittest/_torch/visual_gen/test_wan_transformer.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
c9b961d to
2136687
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The gate itself looks right, and matching get_nvfp4_input_scale's own conditions (dynamic quant / AWQ pre-quant-scale / group size) with the fusion gate is the correct fix for the wan22 profiles.
Two things before merge:
-
wan21-nvfp4waiver removal is not explained by this change.wan21isWan2.1-T2V-1.3B-Diffusers— hidden size 12×128 = 1536, so_fused_ln_shape_supportedwas alreadyFalseand all three norm sites already ran eager. This patch changes nothing for that case, so either its golden failure has a separate root cause (and the waiver should stay), or the failure was environmental. Please post the passing run forwan_feature_accuracy_against_golden[wan21-nvfp4]on the same GPU type, or restore that line. -
Nothing now asserts the fused path is still taken anywhere. After this change it requires hidden 5120 and a statically calibrated NVFP4 checkpoint; every profile in
test_wan_feature_accuracy_against_goldenquantizes at runtime, so all of them take the eager branch. A future change that makesget_nvfp4_input_scalereturnNoneunconditionally would silently disable the kernel with green CI. Consider a positive assertion on a static-NVFP4-checkpoint config (e.g. the A14B NVFP4 path).
Note the bug also covers qwen_image, cosmos3, flux2 and fastwan goldens, whose waivers stay — worth saying in the description that this fixes only the wan slice.
| @@ -139,10 +139,6 @@ examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accura | |||
| examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6572800) | |||
| examples/visual_gen/test_visual_gen_wan.py::test_fastwan_lpips_against_golden SKIP (https://nvbugs/6572800) | |||
| examples/visual_gen/test_visual_gen_wan.py::test_wan22_t2v_lpips_against_golden SKIP (https://nvbugs/6535765) | |||
There was a problem hiding this comment.
The [wan21-nvfp4] waiver was removed here, but wan21 maps to Wan2.1-T2V-1.3B-Diffusers (12 heads x 128 = hidden 1536). That never satisfied the old hidden_size == 5120 gate, so it already ran the eager F.layer_norm path and this patch does not change its numerics at all. Either its golden failure has a different cause and the waiver should stay, or please attach the passing run that justifies removing it.
| we would pay that difference for nothing -- on Wan 2.2 T2V it alone | ||
| moved LPIPS from 0.040 to 0.225 (nvbugs/6535765). | ||
| """ | ||
| return self._fused_ln_shape_supported and fp4_scale is not None |
There was a problem hiding this comment.
With this gate, the fused kernel is reached only when hidden==5120 and the checkpoint carries a static NVFP4 input_scale (get_nvfp4_input_scale returns None for dynamic quant, AWQ pre-quant-scale, and group size != 16). None of the test_wan_feature_accuracy_against_golden profiles meet that — they all quantize at runtime — so the fused path is now unexercised by the tests this PR un-waives, and a future regression that turns the gate permanently off would not be caught. Worth adding a positive check on a statically-quantized NVFP4 Wan2.2 checkpoint that the fused path is actually taken.
|
|
||
| torch.testing.assert_close(trt_out, hf_out, atol=0.4, rtol=0.4) | ||
|
|
||
| @pytest.mark.parametrize("num_heads,shape_supported", [(40, True), (12, False)]) |
There was a problem hiding this comment.
This test only exercises the predicate, not forward, yet it sits in TestWanUnit, which is marked @pytest.mark.integration/wan_t2v and builds a full WanBlock. It would give more value (and run in plain unit CI) as a check that a block with _norm2_fp4_scale = None produces output equal to F.layer_norm-based eager math — that is the property the goldens actually encode.
8d6bffe to
4826e78
Compare
0c0c347 to
cbd07e0
Compare
ae44b63 to
1911042
Compare
The Wan 2.2 LayerNorm+shiftscale+quant fusion (NVIDIA#15762) gates itself on ``_fused_ln_supported = hidden_size == 5120`` -- a pure shape predicate that never consults the fp4 input scales, despite the adjacent comment stating "fusion is skipped when None". So on an unquantized checkpoint all three norm sites take the fused op with nothing to fold in: with ``fp4_input_scale`` None, ``apply_fused_layernorm_*`` falls through to the dense bf16 path, which re-derives the LayerNorm statistics in the kernel instead of calling ATen's. That fall-through is not bit-exact with ``F.layer_norm``, and the goldens are pre-fusion self-goldens that admit almost no deviation, so wan22-fp8-blockwise scored 0.359163 and wan22-cuda-graph 0.242206 against a 0.05 threshold. Three things pin the diagnosis. The golden media in visual_gen_lpips_golden_media.zip is dated 2026-07-23 while the fusion landed 2026-07-28, so every wan22 feature golden encodes the pre-fusion eager numerics. hidden_size is num_attention_heads * attention_head_dim, so Wan2.1-T2V-1.3B is 1536 and never matched the predicate -- which is why its three cases pass and serve as the control -- while Wan2.2-T2V-A14B is exactly 5120. And all three wan22 profiles fail together rather than just a quantized one because the tests request dynamic quantization, and ``get_nvfp4_input_scale`` returns None when ``force_dynamic_quantization`` is set, so fp8-blockwise, cuda-graph and nvfp4 alike reach the fused op with no scale. Gate each site on its own fp4 scale, mirroring the sibling fusions in ``flux/attention.py``, which already require ``has_nvfp4 and input_scale is not None and pre_quant_scale is None and not force_dynamic_quantization``. Restoring the eager path is necessary rather than cosmetic: no reformulation of the kernel's arithmetic recovers the score -- fp32 params give 0.2268 and exact fp64 statistics 0.2351 -- because only ATen's own reduction is bit-exact with the reference the golden encodes. Verified on B200 (umb-b200-041): all 6 test_wan_feature_accuracy_against_golden cases pass at LPIPS 0.000000, exit 0, including the three unwaived wan21 controls (unchanged) and wan22-nvfp4, where the scales are present and the fusion stays engaged. test_wan_transformer.py TestWanUnit is 5/5 with two added tests pinning the predicate in both directions. Removes the four wan feature waivers this change is measured against. The remaining 6572800 waivers are deliberately retained: fastwan (hidden 3072, never matched the predicate) and the flux/ltx2/qwenimage/cosmos3 NVFP4 cases, which fail with "generator exited with code 1" from a spawn-child import error rather than an accuracy delta -- a separate defect that does not reproduce here. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
1911042 to
1d8497d
Compare
|
Clos this in favor of #17933 |
4 participants
Summary
hidden_size == 5120, so Wan2.2 took the fused op with no NVFP4 quantize to fold in and fell through to a dense path whose in-kernel LayerNorm statistics are not bit-exact withF.layer_norm, which the pre-fusion goldens encode._use_fused_ln(fp4_scale)and gated each of the three norm sites on its own fp4 scale, mirroring the sibling gates influx/attention.py; eager where there is nothing to fuse, fused where there is.pytest tests/integration/defs/examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden -vTest plan
Links
Dev Engineer Review
_use_fused_ln(fp4_scale)to require a supported hidden size and an available FP4 input scale.F.layer_normbehavior for unquantized layers and layers without FP4 scales.QA Engineer Review
test_fused_layernorm_requires_its_quantize._make_model_configto supportskip_create_weights.