Skip to content

[https://nvbugs/6572800][fix] Added _use_fused_ln(fp4_scale) and gated each of the three norm sites on its… #17509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
trtllm-agent wants to merge 1 commit into from
Closed

[https://nvbugs/6572800][fix] Added _use_fused_ln(fp4_scale) and gated each of the three norm sites on its… #17509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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

@brnguyen2 brnguyen2 Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
):
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 0 additions & 4 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@brnguyen2 brnguyen2 Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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)
Expand Down
37 changes: 34 additions & 3 deletions tests/unittest/_torch/visual_gen/test_wan_transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
)


Expand Down Expand Up @@ -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)])

@brnguyen2 brnguyen2 Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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
Expand Down
Loading