From dc461fbba418404c5b3af14d3bd61ee693ac3be4 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:15:01 +0800 Subject: [PATCH 1/4] [6959002][fix] Supply image inputs for Wan I2V export fusion Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> --- modelopt/torch/export/diffusers_utils.py | 23 +++++++++++++++++++--- modelopt/torch/export/unified_export_hf.py | 14 ++++++++----- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 2892e16a02e..d999189c99b 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -295,6 +295,7 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Requires: hidden_states, encoder_hidden_states, timestep in_channels = getattr(cfg, "in_channels", 16) text_dim = getattr(cfg, "text_dim", 4096) + image_dim = getattr(cfg, "image_dim", None) max_seq_len = getattr(cfg, "rope_max_seq_len", 512) patch_dtype = getattr(getattr(model, "patch_embedding", None), "weight", None) @@ -308,13 +309,14 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Wan expects num_frames = 4 * n + 1; keep n small for dummy forward num_frames = 5 - text_seq_len = min(max_seq_len, 512) + # I2V attention splits image context from the final 512 text tokens. + text_seq_len = 512 if image_dim is not None else min(max_seq_len, 512) # Keep spatial dims small and divisible by patch size (default 2x2) height = 8 width = 8 - return { + dummy_inputs = { "hidden_states": torch.randn( batch_size, in_channels, num_frames, height, width, device=device, dtype=patch_dtype ), @@ -324,6 +326,21 @@ def _wan_inputs() -> dict[str, torch.Tensor]: "timestep": torch.randint(0, 1000, (batch_size,), device=device), "return_dict": False, } + if image_dim is not None: + image_embedder = model.condition_embedder.image_embedder + # The projection bias retains its compute dtype when weights are packed. + image_dtype = image_embedder.ff.net[0].proj.bias.dtype + image_batch_size = batch_size + image_seq_len = 257 + pos_embed_seq_len = getattr(cfg, "pos_embed_seq_len", None) + if pos_embed_seq_len is not None: + # FLF2V concatenates first/last-frame features inside the image embedder. + image_batch_size *= 2 + image_seq_len = pos_embed_seq_len // 2 + dummy_inputs["encoder_hidden_states_image"] = torch.randn( + image_batch_size, image_seq_len, image_dim, device=device, dtype=image_dtype + ) + return dummy_inputs def _qwen_inputs() -> dict[str, Any]: # QwenImageTransformer2DModel does NOT take the standard diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 37187dd3397..99ffbae0ae4 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -1324,11 +1324,15 @@ def _export_diffusers_checkpoint( # TODO: Add FFN fusion for AWQ-style quantization (pre_quant_scale is # promoted to module keys at export by _promote_quantizer_tensors_to_module below) print(f" Running QKV fusion for {component_name}...") - # Qwen-Image's packed-latent forward signature is non-standard; if the - # dummy forward fails for it, fail loudly rather than silently skipping - # fusion (which would export un-unified amax values). + # Qwen-Image and image-conditioned Wan must not skip fusion + # when their probe forward fails. is_qwen_component = "qwen" in type(component).__name__.lower() - _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) + is_wan_i2v_component = ( + HAS_DIFFUSERS + and isinstance(component, getattr(diffusers, "WanTransformer3DModel", ())) + and getattr(getattr(component, "config", None), "image_dim", None) is not None + ) + _fuse_qkv_linears_diffusion(component, strict=is_qwen_component or is_wan_i2v_component) # Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) From fabe404618d1079929ef3538e567d4b9f1adac52 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:22:57 +0800 Subject: [PATCH 2/4] Keep original copyright headers in Wan export fix Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> --- modelopt/torch/export/diffusers_utils.py | 2 +- modelopt/torch/export/unified_export_hf.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index d999189c99b..de4d6966a61 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 99ffbae0ae4..e695572ba2b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); From 887a2b0dd705db692215993f24bedae56bdb5f9b Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:23:24 +0800 Subject: [PATCH 3/4] Preserve Qwen fusion rationale alongside Wan export guard Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> --- modelopt/torch/export/unified_export_hf.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e695572ba2b..4a2a0ae95e1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1324,9 +1324,11 @@ def _export_diffusers_checkpoint( # TODO: Add FFN fusion for AWQ-style quantization (pre_quant_scale is # promoted to module keys at export by _promote_quantizer_tensors_to_module below) print(f" Running QKV fusion for {component_name}...") - # Qwen-Image and image-conditioned Wan must not skip fusion - # when their probe forward fails. + # Qwen-Image's packed-latent forward signature is non-standard; if the + # dummy forward fails for it, fail loudly rather than silently skipping + # fusion (which would export un-unified amax values). is_qwen_component = "qwen" in type(component).__name__.lower() + # Image-conditioned Wan also requires a successful dummy forward for fusion. is_wan_i2v_component = ( HAS_DIFFUSERS and isinstance(component, getattr(diffusers, "WanTransformer3DModel", ())) From 47e51c15750e14f40502c1d6233f4583d343b3f8 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:48:34 +0800 Subject: [PATCH 4/4] test: cover Wan image-conditioned QKV fusion on CPU Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> --- .../torch/export/test_export_diffusers.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index eedeae677e1..ec7787687e5 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -286,6 +286,50 @@ def test_flux2_dummy_inputs_shape(): assert "guidance" in inputs +@pytest.mark.parametrize("pos_embed_seq_len", [None, 514], ids=["i2v", "flf2v"]) +@torch.no_grad() +def test_wan_dummy_inputs_support_qkv_fusion(pos_embed_seq_len): + torch.manual_seed(0) + model = get_tiny_wan22_transformer( + image_dim=32, added_kv_proj_dim=24, pos_embed_seq_len=pos_embed_seq_len + ).eval() + inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) + assert inputs is not None + image_batch_size = 1 if pos_embed_seq_len is None else 2 + assert inputs["encoder_hidden_states_image"].shape == (image_batch_size, 257, 32) + assert inputs["encoder_hidden_states"].shape == (1, 512, 32) + + # Per-tensor INT8 exercises both input and weight amax fusion on CPU. + quant_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}} + ) + mtq.quantize(model, quant_cfg, lambda m: m(**inputs)) + + groups = [] + for block in model.blocks: + groups.extend( + [ + (block.attn1.to_q, block.attn1.to_k, block.attn1.to_v), + (block.attn2.to_k, block.attn2.to_v), + (block.attn2.add_k_proj, block.attn2.add_v_proj), + ] + ) + # Distinct values ensure a skipped fusion cannot pass on already-equal scales. + for group in groups: + for index, module in enumerate(group, start=1): + module.input_quantizer.amax = torch.tensor(float(index)) + module.weight_quantizer.amax = torch.tensor(float(index)) + + unified_export_hf._fuse_qkv_linears_diffusion(model, strict=True) + + for group in groups: + expected_amax = torch.tensor(float(len(group))) + for module in group: + torch.testing.assert_close(module.input_quantizer.amax, expected_amax) + torch.testing.assert_close(module.weight_quantizer.amax, expected_amax) + + @pytest.mark.parametrize( ("model_factory", "export_attention_quantizers", "expect_attention_quantizers"), [