From 1d8497d653a8a4d357eb496d3069aed05d13b90c Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:04:35 -0700 Subject: [PATCH] [nvbugs/6572800][fix] Only fuse Wan LayerNorm where its quantize exists The Wan 2.2 LayerNorm+shiftscale+quant fusion (#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> --- .../visual_gen/models/wan/transformer_wan.py | 21 +++++++++-- tests/integration/test_lists/waives.txt | 4 -- .../_torch/visual_gen/test_wan_transformer.py | 37 +++++++++++++++++-- 3 files changed, 51 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index 057e088138cd..87b0ed0dfb56 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -370,7 +370,7 @@ def __init__( self._norm1_fp4_scale: Optional[torch.Tensor] = None self._norm2_fp4_scale: Optional[torch.Tensor] = None self._norm3_fp4_scale: Optional[torch.Tensor] = None - self._fused_ln_supported = hidden_size == 5120 + self._fused_ln_shape_supported = hidden_size == 5120 self.ffn = MLP( hidden_size=hidden_size, @@ -467,6 +467,19 @@ def __init__( torch.empty(1, 6, hidden_size).normal_(std=hidden_size**-0.5) ) + def _use_fused_ln(self, fp4_scale: Optional[torch.Tensor]) -> bool: + """Whether a norm site should take the fused LayerNorm kernel. + + The kernel's value is folding the site's downstream NVFP4 quantize into + the normalization epilogue, so only take it where that quantize exists: + fp4_scale is its input scale, None for an unquantized checkpoint or a + quant-excluded layer. The op computes the LayerNorm statistics itself + and is not bit-exact with F.layer_norm, so with no quantize to fold in + 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 + def _fused_adaln_quant(self, x, scale_msa, shift_msa, temb, fp4_scale, eps): """Shared norm1/norm3 path: flatten x to 2D, build the per-token or per-batch modulation rows, and run the fused LayerNorm+AdaLN+NVFP4 op. @@ -536,7 +549,7 @@ def forward( self.scale_shift_table.float() + temb.float() ).chunk(6, dim=1) - if self._fused_ln_supported: + if self._use_fused_ln(self._norm1_fp4_scale): # x is [B, S, D]; flatten to 2D for the fused op, reshape output back. normed = self._fused_adaln_quant( x, scale_msa, shift_msa, temb, self._norm1_fp4_scale, self.norm1.variance_epsilon @@ -566,7 +579,7 @@ def forward( x = (x.float() + attn1_out.float() * gate_msa).to(x.dtype) if ( - self._fused_ln_supported + self._use_fused_ln(self._norm2_fp4_scale) and isinstance(self.norm2, LayerNorm) and self.norm2.weight is not None ): @@ -625,7 +638,7 @@ def forward( # 3. Feed-forward. Mirrors norm1: fused LN+AdaLN (with optional NVFP4 # quant) reshaped back to [B, S, D]; self.ffn consumes it. - if self._fused_ln_supported: + if self._use_fused_ln(self._norm3_fp4_scale): normed = self._fused_adaln_quant( x, c_scale_msa, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 61b3e57b020f..89478ce2bb15 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -122,10 +122,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) -examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan21-nvfp4] SKIP (https://nvbugs/6572800) -examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-cuda-graph] SKIP (https://nvbugs/6572800) -examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-fp8-blockwise] SKIP (https://nvbugs/6572800) -examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_golden[wan22-nvfp4] SKIP (https://nvbugs/6572800) full:A100/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856) full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) diff --git a/tests/unittest/_torch/visual_gen/test_wan_transformer.py b/tests/unittest/_torch/visual_gen/test_wan_transformer.py index b3132939f3d9..83e6ff093859 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_wan_transformer.py @@ -42,7 +42,10 @@ DiffusionPipelineConfig, VisualGenArgs, ) -from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanTransformer3DModel +from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import ( + WanBlock, + WanTransformer3DModel, +) from tensorrt_llm.models.modeling_utils import QuantConfig @@ -163,14 +166,16 @@ def _load_models(checkpoint_dir: str): } -def _make_model_config(config_dict: dict) -> DiffusionModelConfig: +def _make_model_config( + config_dict: dict, *, skip_create_weights: bool = False +) -> DiffusionModelConfig: return DiffusionModelConfig( pretrained_config=SimpleNamespace(**config_dict), quant_config=QuantConfig(), quant_config_dict=None, dynamic_weight_quant=False, force_dynamic_quantization=False, - skip_create_weights_in_init=False, + skip_create_weights_in_init=skip_create_weights, ) @@ -341,6 +346,32 @@ def test_allclose_to_hf(self): 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)]) + def test_fused_layernorm_requires_its_quantize(self, num_heads, shape_supported): + """The fused LN kernel is only taken where a quantize exists to fold in. + + The kernel derives the LayerNorm statistics itself and is not bit-exact + with F.layer_norm, so taking it with no downstream NVFP4 quantize costs + accuracy for nothing (nvbugs/6535765, nvbugs/6572800). Off the supported + hidden size (40 heads x 128 = 5120) it is never taken, scale or not. + """ + cfg = { + **WAN_1_3B_CONFIG, + "num_layers": 1, + "num_attention_heads": num_heads, + "hidden_size": num_heads * WAN_1_3B_CONFIG["attention_head_dim"], + } + block = WanBlock( + model_config=_make_model_config(cfg, skip_create_weights=True), _layer_idx=0 + ) + assert block._fused_ln_shape_supported is shape_supported + + # An unquantized checkpoint leaves every norm's fp4 scale unset. + assert block._norm1_fp4_scale is None + assert not block._use_fused_ln(block._norm1_fp4_scale) + # With a scale present, the fusion is taken iff the shape supports it. + assert block._use_fused_ln(torch.empty(1)) is shape_supported + # ============================================================================ # T2V correctness test — Wan2.1-T2V-1.3B