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