[Hardware][Ascend] Add Wan2.2 MXFP4 UOS, Smooth and W4A8 fallback (#7210)

Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>

tests/diffusion/models/wan2_2/test_create_transformer_quant.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for transformer quant-config auto-detection.
 
 The loader path at pipeline_wan2_2.py carries two quantization contracts:
@@ -27,6 +27,7 @@ from types import SimpleNamespace
 import pytest
 
 import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
 
     # od_config must remain unchanged — _create_transformer does not modify it.
     assert od_config.quantization_config is None
+
+
+def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
+    import json
+
+    import torch
+
+    import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
+    import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
+
+    monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
+    for component in ("transformer", "transformer_2"):
+        directory = tmp_path / component
+        directory.mkdir()
+        config = dict(_MIN_CFG)
+        expert_quant_config = (
+            disk_quant_config[0 if component == "transformer" else 1]
+            if isinstance(disk_quant_config, tuple)
+            else disk_quant_config
+        )
+        if expert_quant_config is not None:
+            config["quantization_config"] = expert_quant_config
+        (directory / "config.json").write_text(json.dumps(config))
+    (tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
+
+    class FakeComponent(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
+
+    class FakeTransformer(FakeComponent):
+        def __init__(self, quant_config=None, **kwargs):
+            super().__init__()
+            self.received_quant_config = quant_config
+
+    for module in (wan22_module, i2v_module):
+        monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
+        monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
+        monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+    monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
+    pipeline_cls = {
+        "t2v": wan22_module.Wan22Pipeline,
+        "i2v": i2v_module.Wan22I2VPipeline,
+        "vace": vace_module.Wan22VACEPipeline,
+    }[pipeline_kind]
+    monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
+    od_config = OmniDiffusionConfig(
+        model=str(tmp_path),
+        model_class_name=pipeline_cls.__name__,
+        dtype=torch.bfloat16,
+        quantization_config=active,
+        boundary_ratio=0.875,
+        flow_shift=12.0,
+        enable_diffusion_pipeline_profiler=False,
+    )
+    od_config.enrich_config()
+    pipeline = pipeline_cls(od_config=od_config)
+    return pipeline
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
+@pytest.mark.parametrize("offline", [False, True])
+def test_pipeline_constructor_routes_each_expert_quantization(
+    pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
+):
+    """Run the actual pipeline constructors, expert factories and quant-method dispatch."""
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.quantization import build_quant_config
+
+    active = build_quant_config(
+        {
+            "transformer": {
+                "method": method,
+                "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
+                "w4a8_fallback_steps": [0, 1],
+            },
+            "transformer_2": {
+                "method": method,
+                "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
+                "w4a8_fallback_steps": [],
+            },
+        }
+    )
+    disk_quant_config = None
+    if offline:
+        disk_quant_config = {
+            "quant_method": method,
+            flag: True,
+            "ignored_layers": ["blocks.12.ffn.net_2"],
+            "w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
+            "w4a8_fallback_steps": [9],
+        }
+    pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
+    for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
+        received = transformer.received_quant_config
+        assert received.get_name() == method
+        assert getattr(received, flag) == offline
+        assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
+        selected = f"blocks.{10 + index}.attn1.to_qkv"
+        other = f"blocks.{11 - index}.attn1.to_qkv"
+        assert received.w4a8_fallback_layers == [selected]
+        assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
+        assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
+        if offline:
+            assert isinstance(
+                received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
+            )
+    assert pipeline.od_config.quantization_config is active
+
+
+@pytest.mark.parametrize("explicit_none", [False, True])
+def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
+    from vllm_omni.quantization import build_quant_config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+    components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
+    active = build_quant_config(components)
+    pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
+    pipeline._create_transformer(_MIN_CFG, component="transformer_2")
+    assert captured[-1].get("quant_config") is None
+    with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
+        pipeline._create_transformer(
+            {**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
+            component="transformer_2",
+        )
+
+
+def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
+    from vllm_omni.quantization import build_quant_config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+    active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
+    pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
+    pipeline._create_transformer(_MIN_CFG, component="transformer_2")
+    assert captured[-1]["quant_config"] is active.default_config
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize(
+    "configured_expert, with_default",
+    [
+        ("transformer", False),
+        ("transformer", True),
+        ("transformer_2", False),
+        ("transformer_2", True),
+    ],
+)
+def test_pipeline_partial_expert_config_does_not_match_sibling(
+    pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
+):
+    from vllm_omni.quantization import build_quant_config
+
+    component_spec = {
+        configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
+    }
+    if with_default:
+        component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
+    active = build_quant_config(component_spec)
+    pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
+    sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
+    selected = getattr(pipeline, configured_expert).received_quant_config
+    other = getattr(pipeline, sibling).received_quant_config
+    assert selected is active.component_configs[configured_expert]
+    assert other is active.default_config
+    if with_default:
+        assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
+    else:
+        assert other is None
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize("different_ignored_layers", [False, True])
+def test_enriched_auto_quantization_preserves_each_expert_policy(
+    pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
+):
+    """Startup auto-detection must not turn the first expert into a user override."""
+    high = {
+        "quant_method": "mxfp4",
+        "is_checkpoint_mxfp4_serialized": True,
+        "ignored_layers": ["blocks.0.ffn.net_2"],
+        "w4a8_fallback_steps": [],
+        "w4a8_fallback_layers": [],
+        "mxfp4_scale_alg": 0,
+        "require_smooth_scale": False,
+    }
+    low = {
+        **high,
+        "ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
+        "w4a8_fallback_steps": [37],
+        "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
+        "mxfp4_scale_alg": 2,
+        "require_smooth_scale": True,
+    }
+    pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
+    assert pipeline.od_config.quantization_config is not None  # enrich_config really detected the first expert.
+    for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
+        config = transformer.received_quant_config
+        for policy in (
+            "ignored_layers",
+            "w4a8_fallback_steps",
+            "w4a8_fallback_layers",
+            "mxfp4_scale_alg",
+            "require_smooth_scale",
+        ):
+            assert getattr(config, policy) == expected[policy], policy

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
 @pytest.fixture(autouse=True)
-def _cpu_platform(monkeypatch):
+def _cpu_pipeline_runtime(monkeypatch):
+    # These lightweight pipelines have no accelerator or distributed groups.
     module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
     monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
+    monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
 
 
 class _StubTransformer(nn.Module):
@@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
     assert torch.equal(result, torch.full_like(latents, 10.0))
 
 
+@pytest.mark.parametrize("fail_first_request", [False, True])
+def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.forward_context import (
+        ForwardContext,
+        get_forward_context,
+        is_forward_context_available,
+        override_forward_context,
+    )
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
+
+    contexts: list[ForwardContext] = []
+    events = []
+
+    class InjectedPredictionError(Exception):
+        pass
+
+    class RecordingTransformer(_StubTransformer):
+        def __init__(self, expert):
+            super().__init__()
+            self.expert = expert
+
+        def forward(self, hidden_states, encoder_hidden_states, **kwargs):
+            context = get_forward_context()
+            assert context is contexts[-1]
+            request_index = len(contexts) - 1
+            step = context.denoise_step_idx
+            positive = bool(encoder_hidden_states[0, 0, 0] > 0)
+            events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
+            if fail_first_request and request_index == 0 and step == 1 and not positive:
+                raise InjectedPredictionError
+            return (torch.ones_like(hidden_states) * (1 if positive else -1),)
+
+    def encode_prompt(**kwargs):
+        context = get_forward_context()
+        assert context.denoise_step_idx is None
+        contexts.append(context)
+        embeds = torch.ones(1, 8, 8)
+        return embeds, -embeds
+
+    pipeline = _make_pipeline()
+    pipeline.transformer = RecordingTransformer("high")
+    pipeline.transformer_2 = RecordingTransformer("low")
+    pipeline.scheduler = _StubScheduler([900, 500, 100])
+    monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
+    monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
+
+    # Use the real runner context manager, Wan forward/diffuse, and CFG
+    # dispatch. Only model computation, scheduler math and unrelated KV I/O
+    # are replaced; no test code publishes a denoise step.
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = pipeline
+    runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
+    runner.vllm_config = None
+    runner.cache_backend = None
+    monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
+    with override_forward_context(None):
+        for request_index in range(2):
+            request = OmniDiffusionRequest(
+                prompt="a fox walks",
+                request_id=f"request-{request_index}",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
+                ),
+            )
+
+            def execute():
+                return runner._execute_request_list(
+                    [request],
+                    od_config=runner.od_config,
+                    allow_single_output=True,
+                    require_request_batch_support=False,
+                    record_name="test_wan_step_context",
+                    record_output_peak_memory=False,
+                )
+
+            if fail_first_request and request_index == 0:
+                with pytest.raises(InjectedPredictionError):
+                    execute()
+            else:
+                execute()
+            assert not is_forward_context_available()
+
+    assert len(contexts) == 2 and contexts[0] is not contexts[1]
+    expected = [
+        (request, step, "high" if step == 0 else "low", positive, step in (0, 2))
+        for request in range(2)
+        for step in range(2 if fail_first_request and request == 0 else 3)
+        for positive in (True, False)
+    ]
+    assert events == expected
+
+
 class _StubDMDScheduler:
     def __init__(self) -> None:
         self.predict_clean_calls: list[tuple[float, float, float]] = []