[NPU][HunyuanImage3] Adapt diffusion FusedMoE for vLLM >= 0.24.0 (#5167)

Signed-off-by: betta18 <jiangmengyu1@huawei.com>
Signed-off-by: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com>
Co-authored-by: betta18 <jiangmengyu1@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -193,7 +193,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
 @pytest.mark.cpu
 @pytest.mark.core_model
 def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
-    """vLLM DP should keep diffusion DP when expert parallelism is not enabled."""
+    """Without EP, keep diffusion DP and retain rank metadata for MC2."""
     local_rank = 0
     world_size = 8
 
@@ -249,7 +249,9 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
     assert parallel_state.vllm_parallel_state._DP.local_group == [0, 4]
     assert parallel_state.vllm_parallel_state._PCP is None
     assert parallel_state.vllm_parallel_state._EP is None
-    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS == [
+        [0, 1, 2, 3, 4, 5, 6, 7],
+    ]
 
 
 @pytest.mark.cpu

tests/diffusion/layers/test_fused_moe.py

@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the diffusion FusedMoE adapter."""
+
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestSetForwardContextNumTokens:
+    """Test _set_forward_context_num_tokens defensive fix for vLLM 0.18.0."""
+
+    def test_sets_num_tokens_when_context_available(self, mocker):
+        """num_tokens should be set on ForwardContext when available."""
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        mock_ctx = mocker.MagicMock()
+        del mock_ctx.in_profile_run  # simulate missing attr
+        mocker.patch.object(fused_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(fused_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        fused_moe._set_forward_context_num_tokens(1024)
+
+        assert mock_ctx.num_tokens == 1024
+        assert mock_ctx.in_profile_run is False
+
+    def test_sets_in_profile_run_only_if_missing(self, mocker):
+        """in_profile_run should not be overwritten if already set."""
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.in_profile_run = True  # already set
+        mocker.patch.object(fused_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(fused_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        fused_moe._set_forward_context_num_tokens(512)
+
+        assert mock_ctx.num_tokens == 512
+        assert mock_ctx.in_profile_run is True  # not overwritten
+
+    def test_noop_when_context_unavailable(self, mocker):
+        """Should do nothing when ForwardContext is not available."""
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        mocker.patch.object(fused_moe._vllm_fc, "is_forward_context_available", return_value=False)
+        mock_get = mocker.patch.object(fused_moe._vllm_fc, "get_forward_context")
+
+        fused_moe._set_forward_context_num_tokens(256)
+
+        mock_get.assert_not_called()
+
+
+class TestSetForwardContextDPMetadata:
+    """Test diffusion MoE DP metadata setup for vLLM 0.24 dispatch."""
+
+    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = None
+        fake_group = mocker.MagicMock()
+        fake_group.world_size = 2
+        fake_group.cpu_group = object()
+
+        def _all_gather_object(output, obj, group):
+            assert group is fake_group.cpu_group
+            output[:] = [obj, obj + 1]
+
+        mocker.patch.object(fused_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(fused_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mocker.patch.object(fused_moe._vllm_ps, "_DP", fake_group, create=True)
+        mocker.patch.object(fused_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
+
+        fused_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is not None
+        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
+
+    def test_keeps_existing_dp_metadata(self, mocker):
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        existing_metadata = object()
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = existing_metadata
+        mocker.patch.object(fused_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(fused_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mock_all_gather = mocker.patch.object(fused_moe.torch.distributed, "all_gather_object")
+
+        fused_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is existing_metadata
+        mock_all_gather.assert_not_called()
+
+
+class TestFusedMoEFactory:
+    """Test FusedMoE factory __new__ and mapping delegation."""
+
+    def test_new_returns_vllm_runner_and_registers_hooks(self, mocker, monkeypatch):
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        mock_platform = mocker.MagicMock()
+        mocker.patch.object(fused_moe, "current_omni_platform", mock_platform)
+
+        mock_runner = mocker.MagicMock()
+        mock_vllm_fused_moe = mocker.MagicMock(return_value=mock_runner)
+        monkeypatch.setitem(
+            sys.modules,
+            "vllm.model_executor.layers.fused_moe",
+            SimpleNamespace(FusedMoE=mock_vllm_fused_moe),
+        )
+
+        result = fused_moe.FusedMoE(prefix="test", a=1, reduce_results=True)
+
+        assert result is mock_runner
+        mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("fused_moe")
+        mock_vllm_fused_moe.assert_called_once_with(prefix="test", a=1)
+        mock_runner.register_forward_pre_hook.assert_called_once_with(
+            mocker.ANY,
+            with_kwargs=True,
+        )
+        mock_platform.register_additional_diffusion_fused_moe_hooks.assert_called_once_with(mock_runner)
+
+    def test_num_tokens_pre_hook_uses_kwargs_hidden_states(self, mocker, monkeypatch):
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        captured_hook = None
+
+        class MockRunner:
+            def register_forward_pre_hook(self, hook, *, with_kwargs):
+                nonlocal captured_hook
+                assert with_kwargs is True
+                captured_hook = hook
+
+        mock_runner = MockRunner()
+        monkeypatch.setitem(
+            sys.modules,
+            "vllm.model_executor.layers.fused_moe",
+            SimpleNamespace(FusedMoE=mocker.MagicMock(return_value=mock_runner)),
+        )
+        mocker.patch.object(fused_moe, "current_omni_platform", mocker.MagicMock())
+        mock_set_num_tokens = mocker.patch.object(fused_moe, "_set_forward_context_num_tokens")
+        mock_set_dp_metadata = mocker.patch.object(fused_moe, "_set_forward_context_dp_metadata")
+
+        fused_moe.FusedMoE(prefix="")
+        assert captured_hook is not None
+
+        hidden_states = torch.empty(17, 32)
+        captured_hook(mock_runner, (), {"hidden_states": hidden_states})
+
+        mock_set_num_tokens.assert_called_once_with(17)
+        mock_set_dp_metadata.assert_called_once_with(17)
+
+    def test_num_tokens_pre_hook_uses_positional_hidden_states(self, mocker, monkeypatch):
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        captured_hook = None
+
+        class MockRunner:
+            def register_forward_pre_hook(self, hook, *, with_kwargs):
+                nonlocal captured_hook
+                assert with_kwargs is True
+                captured_hook = hook
+
+        mock_runner = MockRunner()
+        monkeypatch.setitem(
+            sys.modules,
+            "vllm.model_executor.layers.fused_moe",
+            SimpleNamespace(FusedMoE=mocker.MagicMock(return_value=mock_runner)),
+        )
+        mocker.patch.object(fused_moe, "current_omni_platform", mocker.MagicMock())
+        mock_set_num_tokens = mocker.patch.object(fused_moe, "_set_forward_context_num_tokens")
+        mock_set_dp_metadata = mocker.patch.object(fused_moe, "_set_forward_context_dp_metadata")
+
+        fused_moe.FusedMoE(prefix="")
+        assert captured_hook is not None
+
+        hidden_states = torch.empty(23, 32)
+        captured_hook(mock_runner, (hidden_states,), {})
+
+        mock_set_num_tokens.assert_called_once_with(23)
+        mock_set_dp_metadata.assert_called_once_with(23)
+
+    def test_make_expert_params_mapping_delegates_to_vllm_function(self, mocker, monkeypatch):
+        import vllm_omni.diffusion.layers.fused_moe as fused_moe
+
+        expected_mapping = [("a", "b", 0, "c")]
+        mock_mapping_fn = mocker.MagicMock(return_value=expected_mapping)
+        monkeypatch.setitem(
+            sys.modules,
+            "vllm.model_executor.layers.fused_moe",
+            SimpleNamespace(fused_moe_make_expert_params_mapping=mock_mapping_fn),
+        )
+
+        result = fused_moe.FusedMoE.make_expert_params_mapping(
+            model=None,
+            ckpt_gate_proj_name="gate",
+            ckpt_down_proj_name="down",
+            ckpt_up_proj_name="up",
+            num_experts=4,
+            num_redundant_experts=0,
+        )
+
+        assert result == expected_mapping
+        mock_mapping_fn.assert_called_once_with(
+            None,
+            ckpt_gate_proj_name="gate",
+            ckpt_down_proj_name="down",
+            ckpt_up_proj_name="up",
+            num_experts=4,
+            num_redundant_experts=0,
+        )

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -1,197 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HunyuanFusedMoE (Support HunyuanImage3 Diffusion Model, 5a779b4)."""
-
-import pytest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class TestSetForwardContextNumTokens:
-    """Test _set_forward_context_num_tokens defensive fix for vLLM 0.18.0."""
-
-    def test_sets_num_tokens_when_context_available(self, mocker):
-        """num_tokens should be set on ForwardContext when available."""
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        del mock_ctx.in_profile_run  # simulate missing attr
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-
-        hunyuan_moe._set_forward_context_num_tokens(1024)
-
-        assert mock_ctx.num_tokens == 1024
-        assert mock_ctx.in_profile_run is False
-
-    def test_sets_in_profile_run_only_if_missing(self, mocker):
-        """in_profile_run should not be overwritten if already set."""
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.in_profile_run = True  # already set
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-
-        hunyuan_moe._set_forward_context_num_tokens(512)
-
-        assert mock_ctx.num_tokens == 512
-        assert mock_ctx.in_profile_run is True  # not overwritten
-
-    def test_noop_when_context_unavailable(self, mocker):
-        """Should do nothing when ForwardContext is not available."""
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False)
-        mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context")
-
-        hunyuan_moe._set_forward_context_num_tokens(256)
-
-        mock_get.assert_not_called()
-
-
-class TestSetForwardContextDPMetadata:
-    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
-
-    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.dp_metadata = None
-        fake_group = mocker.MagicMock()
-        fake_group.world_size = 2
-        fake_group.cpu_group = object()
-
-        def _all_gather_object(output, obj, group):
-            assert group is fake_group.cpu_group
-            output[:] = [obj, obj + 1]
-
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
-        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
-
-        hunyuan_moe._set_forward_context_dp_metadata(1024)
-
-        assert mock_ctx.dp_metadata is not None
-        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
-
-    def test_keeps_existing_dp_metadata(self, mocker):
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        existing_metadata = object()
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.dp_metadata = existing_metadata
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
-
-        hunyuan_moe._set_forward_context_dp_metadata(1024)
-
-        assert mock_ctx.dp_metadata is existing_metadata
-        mock_all_gather.assert_not_called()
-
-
-class TestHunyuanFusedMoEPlatformDispatch:
-    """Test platform dispatch via platform qualname hooks."""
-
-    def test_default_platform_uses_default_impl_qualname(self, mocker):
-        """HunyuanFusedMoE should resolve the impl class from the platform hook."""
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_platform = mocker.MagicMock()
-        mock_platform.get_diffusion_model_impl_qualname.return_value = (
-            "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
-        )
-
-        mocker.patch.object(
-            hunyuan_moe,
-            "current_omni_platform",
-            mock_platform,
-        )
-        mock_resolve = mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname")
-        mock_impl = mocker.MagicMock()
-        mock_resolve.return_value = mock_impl
-
-        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import (
-            HunyuanFusedMoE,
-        )
-
-        HunyuanFusedMoE(prefix="")
-
-        mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe")
-        mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
-        mock_resolve.assert_called_once_with(
-            "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
-        )
-        mock_impl.assert_called_once_with(prefix="")
-
-
-class TestHunyuanFusedMoEFactory:
-    """Test HunyuanFusedMoE factory __new__ and make_expert_params_mapping delegation."""
-
-    def test_new_delegates_to_impl_class(self, mocker):
-        """HunyuanFusedMoE(prefix=..., **kwargs) should instantiate and return impl instance."""
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        class MockImpl:
-            def __init__(self, *, prefix: str = "", **kwargs):
-                self.prefix = prefix
-                self.kwargs = kwargs
-
-        mock_platform = mocker.MagicMock()
-        mock_platform.get_diffusion_model_impl_qualname.return_value = "mock.impl.Qualname"
-        mocker.patch.object(hunyuan_moe, "current_omni_platform", mock_platform)
-
-        mock_impl_class = mocker.MagicMock(return_value=MockImpl(prefix="test", a=1))
-        mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class)
-
-        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import (
-            HunyuanFusedMoE,
-        )
-
-        result = HunyuanFusedMoE(prefix="test", a=1)
-
-        assert isinstance(result, MockImpl)
-        assert result.prefix == "test"
-        assert result.kwargs == {"a": 1}
-        mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe")
-        mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
-        mock_impl_class.assert_called_once_with(prefix="test", a=1)
-
-    def test_make_expert_params_mapping_delegates_to_impl(self, mocker):
-        """make_expert_params_mapping should delegate to impl class method."""
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        expected_mapping = [("a", "b", 0, "c")]
-        mock_platform = mocker.MagicMock()
-        mock_platform.get_diffusion_model_impl_qualname.return_value = "mock.impl.Qualname"
-        mocker.patch.object(hunyuan_moe, "current_omni_platform", mock_platform)
-
-        mock_impl_class = mocker.MagicMock()
-        mock_impl_class.make_expert_params_mapping = mocker.MagicMock(return_value=expected_mapping)
-        mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class)
-
-        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import (
-            HunyuanFusedMoE,
-        )
-
-        result = HunyuanFusedMoE.make_expert_params_mapping(
-            model=None,
-            ckpt_gate_proj_name="gate",
-            ckpt_down_proj_name="down",
-            ckpt_up_proj_name="up",
-            num_experts=4,
-            num_redundant_experts=0,
-        )
-
-        assert result == expected_mapping
-        mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
-        mock_impl_class.make_expert_params_mapping.assert_called_once_with(
-            None,
-            ckpt_gate_proj_name="gate",
-            ckpt_down_proj_name="down",
-            ckpt_up_proj_name="up",
-            num_experts=4,
-            num_redundant_experts=0,
-        )

vllm_omni/diffusion/distributed/parallel_state.py

@@ -312,7 +312,6 @@ def get_allgather_parallel_rank():
 
 
 def get_expert_parallel_group_ranks() -> list[list[int]]:
-    assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized"
     assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized"
     return _EXPERT_PARALLEL_GROUP_RANKS
 
@@ -948,16 +947,14 @@ def initialize_model_parallel(
         parallel_mode="fully_shard",
     )
 
-    _EXPERT_PARALLEL_GROUP_RANKS = None
+    _EXPERT_PARALLEL_GROUP_RANKS = rank_generator.get_ranks("tp-sp-cfg-dp")
     if use_moe_parallel_mapping:
-        ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
         vllm_parallel_state._EP = init_vllm_model_parallel_group(
-            group_ranks=ep_group_ranks,
+            group_ranks=_EXPERT_PARALLEL_GROUP_RANKS,
             local_rank=get_world_group().local_rank,
             backend=backend,
             group_name="ep",
         )
-        _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 
     init_dit_group(dit_parallel_size, backend)
 

vllm_omni/diffusion/layers/fused_moe.py

@@ -6,7 +6,6 @@ from typing import Any
 import torch
 import vllm.distributed.parallel_state as _vllm_ps
 import vllm.forward_context as _vllm_fc
-from vllm.utils.import_utils import resolve_obj_by_qualname
 
 from vllm_omni.platforms import current_omni_platform
 
@@ -50,18 +49,18 @@ def _set_forward_context_dp_metadata(num_tokens: int) -> None:
     forward_context.dp_metadata = _vllm_fc.DPMetadata(torch.tensor(gathered_num_tokens, dtype=torch.int64))
 
 
-class HunyuanFusedMoEDefault:
-    """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
+class FusedMoE:
+    """Adapter that configures the upstream FusedMoE ``MoERunner`` for diffusion models.
 
     Upstream commit dc68bd8c41 refactored FusedMoE from a class (``nn.Module``)
     into a factory function that returns a ``MoERunner`` instance, whose expert
     weights live in a ``routed_experts`` submodule
     (``...experts.routed_experts.w13_weight`` / ``...w2_weight``).
 
-    This adapter builds that runner, installs the omni-specific forward-context
-    setup and one-shot kernel-initialisation hook, and returns the runner
-    *directly* from ``__new__`` so the parent MoE block registers it as a real
-    ``nn.Module`` submodule.
+    This adapter builds that runner, installs the diffusion forward-context
+    setup and additional platform hooks, and returns the runner *directly* from
+    ``__new__`` so the parent MoE block registers it as a real ``nn.Module``
+    submodule.
 
     Returning the runner (rather than wrapping it in a plain object that holds it
     in an attribute) is required for correctness: a non-Module wrapper hides the
@@ -71,9 +70,11 @@ class HunyuanFusedMoEDefault:
     """
 
     def __new__(cls, *, prefix: str = "", **kwargs: Any) -> Any:
+        current_omni_platform.prepare_diffusion_op_runtime("fused_moe")
+
         # Current vLLM FusedMoE handles output reduction internally.
         kwargs.pop("reduce_results", None)
-        # FusedMoE is now a factory function — call it to get a MoERunner.
+        # FusedMoE is now a factory function; call it to get a MoERunner.
         from vllm.model_executor.layers.fused_moe import FusedMoE as _FusedMoE
 
         moe_runner = _FusedMoE(prefix=prefix, **kwargs)
@@ -91,21 +92,7 @@ class HunyuanFusedMoEDefault:
                 _set_forward_context_dp_metadata(hidden_states.shape[0])
 
         moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
-
-        # One-shot lazy kernel initialisation on the first forward (no-op unless
-        # the runner exposes an uninitialised quant_method). Mirrors the prior
-        # wrapper behaviour exactly, just bound to the runner module.
-        init_handle: Any = None
-
-        def _kernel_init_pre_hook(module: Any, args: Any, kwargs: Any) -> None:
-            nonlocal init_handle
-            quant_method = getattr(module, "quant_method", None)
-            if quant_method is not None and getattr(quant_method, "moe_kernel", None) is None:
-                quant_method.process_weights_after_loading(module)
-            if init_handle is not None:
-                init_handle.remove()
-
-        init_handle = moe_runner.register_forward_pre_hook(_kernel_init_pre_hook, with_kwargs=True)
+        current_omni_platform.register_additional_diffusion_fused_moe_hooks(moe_runner)
 
         return moe_runner
 
@@ -138,35 +125,3 @@ class HunyuanFusedMoEDefault:
             num_experts=num_experts,
             num_redundant_experts=num_redundant_experts,
         )
-
-
-class HunyuanFusedMoE:
-    def __new__(cls, *, prefix: str = "", **kwargs: Any) -> Any:
-        op_name = "hunyuan_fused_moe"
-        current_omni_platform.prepare_diffusion_op_runtime(op_name)
-        impl = resolve_obj_by_qualname(
-            current_omni_platform.get_diffusion_model_impl_qualname(op_name),
-        )
-        return impl(prefix=prefix, **kwargs)
-
-    @classmethod
-    def make_expert_params_mapping(
-        cls,
-        model: Any,
-        ckpt_gate_proj_name: str,
-        ckpt_down_proj_name: str,
-        ckpt_up_proj_name: str,
-        num_experts: int,
-        num_redundant_experts: int = 0,
-    ) -> list[tuple[str, str, int, str]]:
-        impl = resolve_obj_by_qualname(
-            current_omni_platform.get_diffusion_model_impl_qualname("hunyuan_fused_moe"),
-        )
-        return impl.make_expert_params_mapping(
-            model,
-            ckpt_gate_proj_name=ckpt_gate_proj_name,
-            ckpt_down_proj_name=ckpt_down_proj_name,
-            ckpt_up_proj_name=ckpt_up_proj_name,
-            num_experts=num_experts,
-            num_redundant_experts=num_redundant_experts,
-        )

vllm_omni/diffusion/models/hunyuan_image3/__init__.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Hunyuan Image 3 diffusion model components."""
 
-from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
     HunyuanImage3Model,
     HunyuanImage3Text2ImagePipeline,
@@ -11,4 +10,4 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
     HunyuanImage3Pipeline,
 )
 
-__all__ = ["HunyuanImage3Pipeline", "HunyuanImage3Model", "HunyuanImage3Text2ImagePipeline", "HunyuanFusedMoE"]
+__all__ = ["HunyuanImage3Pipeline", "HunyuanImage3Model", "HunyuanImage3Text2ImagePipeline"]

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -76,11 +76,12 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.layers.fused_moe import FusedMoE
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
-from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
 from vllm_omni.diffusion.utils.kv_utils import repeat_kv
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
+from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
@@ -1620,7 +1621,7 @@ class HunYuanSparseMoeBlock(nn.Module):
             self.shared_mlp = None
 
         enable_expert_parallel = get_current_vllm_config().parallel_config.enable_expert_parallel
-        self.experts = HunyuanFusedMoE(
+        self.experts = FusedMoE(
             shared_experts=self.shared_mlp,
             num_experts=self.n_routed_experts,
             top_k=top_k,
@@ -2100,7 +2101,7 @@ class HunyuanImage3Model(nn.Module):
         if _is_moe(self.config):
             # Params for weights, fp8 weight scales, fp8 activation scales
             # (param_name, weight_name, expert_id, shard_id)
-            fused_moe_expert_mapping = HunyuanFusedMoE.make_expert_params_mapping(
+            fused_moe_expert_mapping = FusedMoE.make_expert_params_mapping(
                 self,
                 ckpt_gate_proj_name="gate_proj",
                 ckpt_down_proj_name="down_proj",
@@ -2381,6 +2382,8 @@ class HunyuanImage3Model(nn.Module):
         ar_kv_reuse_len: int = 0,
         full_attn_spans: list[list[tuple[int, int]]] | None = None,
     ) -> tuple | BaseModelOutputWithPast:
+        current_omni_platform.reset_diffusion_fused_moe_forward_context()
+
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -205,6 +205,12 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if load_format == "dummy":
             return
 
+        current_omni_platform.init_diffusion_model_runner_runtime(
+            vllm_config=self.vllm_config,
+            od_config=self.od_config,
+            device=self.device,
+        )
+
         load_device = (
             "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device)
         )

vllm_omni/platforms/interface.py

@@ -70,13 +70,31 @@ class OmniPlatform(Platform):
         raise NotImplementedError
 
     @classmethod
-    def get_diffusion_model_impl_qualname(cls, op_name: str) -> str:
-        if op_name == "hunyuan_fused_moe":
-            return "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
-        raise NotImplementedError(f"Unsupported diffusion model op: {op_name}")
+    def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None:
+        return None
 
     @classmethod
-    def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None:
+    def register_additional_diffusion_fused_moe_hooks(cls, moe_runner: Any) -> None:
+        # One-shot lazy kernel initialisation on the first forward (no-op unless
+        # the runner exposes an uninitialised quant_method). Mirrors the prior
+        # wrapper behaviour exactly, just bound to the runner module.
+        init_handle: Any = None
+
+        def _kernel_init_pre_hook(module: Any, args: Any, kwargs: Any) -> None:
+            nonlocal init_handle
+            quant_method = getattr(module, "quant_method", None)
+            if quant_method is not None and getattr(quant_method, "moe_kernel", None) is None:
+                quant_method.process_weights_after_loading(module)
+            if init_handle is not None:
+                init_handle.remove()
+
+        init_handle = moe_runner.register_forward_pre_hook(
+            _kernel_init_pre_hook,
+            with_kwargs=True,
+        )
+
+    @classmethod
+    def reset_diffusion_fused_moe_forward_context(cls) -> None:
         return None
 
     @classmethod
@@ -137,10 +155,23 @@ class OmniPlatform(Platform):
         return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"
 
     @classmethod
-    def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None:
+    def init_diffusion_worker_vllm_config(
+        cls,
+        vllm_config: Any,
+    ) -> None:
         """Initialize platform-specific state for diffusion worker VllmConfig."""
         return None
 
+    @classmethod
+    def init_diffusion_model_runner_runtime(
+        cls,
+        vllm_config: Any,
+        od_config: Any,
+        device: torch.device,
+    ) -> None:
+        """Initialize platform-specific runtime state for diffusion model runners."""
+        return None
+
     @classmethod
     def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
         raise NotImplementedError

vllm_omni/platforms/npu/layers/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/platforms/npu/layers/fused_moe.py

@@ -4,14 +4,14 @@
 from typing import Any
 
 import torch
+import vllm.distributed.parallel_state as vllm_parallel_state
 import vllm.forward_context as _vllm_fc
 from vllm.config import VllmConfig
-from vllm.distributed import get_ep_group
+from vllm.distributed import get_dp_group, get_ep_group
 from vllm.distributed.parallel_state import (
     init_model_parallel_group as vllm_init_model_parallel_group,
 )
 from vllm_ascend.ascend_forward_context import MoECommType
-from vllm_ascend.ops.fused_moe.fused_moe import AscendFusedMoE
 from vllm_ascend.ops.fused_moe.moe_comm_method import _MoECommMethods
 from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type
 
@@ -29,17 +29,6 @@ def _ensure_forward_context_attr(name: str, annotation: Any, default: Any) -> No
         setattr(_vllm_fc.ForwardContext, name, default)
 
 
-def _set_hunyuan_fused_moe_forward_context(num_tokens: int) -> None:
-    if not _vllm_fc.is_forward_context_available():
-        return
-
-    forward_context = _vllm_fc.get_forward_context()
-    forward_context.num_tokens = num_tokens
-    forward_context.moe_comm_type = _select_moe_comm_method(vllm_config=omni_get_ctx().vllm_config)
-    forward_context.moe_comm_method = _MoECommMethods.get(forward_context.moe_comm_type)
-    forward_context.flash_comm_v1_enabled = False
-
-
 def _init_mc2_group_for_diffusion(
     backend: str,
     local_rank: int,
@@ -74,16 +63,16 @@ def _select_moe_comm_method(vllm_config: VllmConfig) -> MoECommType | None:
     return moe_comm_type
 
 
-def prepare_hunyuan_fused_moe_runtime() -> None:
+def prepare_fused_moe_runtime() -> None:
     vllm_config = omni_get_ctx().vllm_config
-    if vllm_config.parallel_config.enable_expert_parallel:
-        backend = torch.distributed.get_backend(get_world_group().device_group)
-        local_rank = get_world_group().local_rank
-        _init_mc2_group_for_diffusion(
-            backend=backend,
-            local_rank=local_rank,
-            group_ranks=get_expert_parallel_group_ranks(),
-        )
+
+    backend = torch.distributed.get_backend(get_world_group().device_group)
+    local_rank = get_world_group().local_rank
+    _init_mc2_group_for_diffusion(
+        backend=backend,
+        local_rank=local_rank,
+        group_ranks=get_expert_parallel_group_ranks(),
+    )
 
     moe_comm_type = _select_moe_comm_method(vllm_config=vllm_config)
     _ensure_forward_context_attr("num_tokens", int | None, None)
@@ -91,24 +80,66 @@ def prepare_hunyuan_fused_moe_runtime() -> None:
     _ensure_forward_context_attr("moe_comm_type", MoECommType | None, moe_comm_type)
     _ensure_forward_context_attr("moe_comm_method", Any, _MoECommMethods.get(moe_comm_type))
     _ensure_forward_context_attr("flash_comm_v1_enabled", bool, False)
+    _ensure_forward_context_attr("max_tokens_across_dp", int | None, None)
+    _ensure_forward_context_attr("max_tokens_across_pcp", int | None, None)
 
 
-# NOTE: vLLM v0.20.0 folded SharedFusedMoE into FusedMoE, and vllm-ascend in turn
-# removed AscendSharedFusedMoE — the shared-experts / gate / multistream-overlap
-# paths now live directly on AscendFusedMoE and are activated by passing
-# shared_experts= as a kwarg.
-class AscendHunyuanFusedMoE(AscendFusedMoE):
-    def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
-        super().__init__(prefix=prefix, **kwargs)
-        self._prefix = prefix
+def reset_fused_moe_forward_context() -> None:
+    try:
+        forward_context = _vllm_fc.get_forward_context()
+    except AssertionError:
+        return
 
-    def forward(self, hidden_states: Any, router_logits: Any) -> Any:
-        _set_hunyuan_fused_moe_forward_context(hidden_states.shape[0])
-        return super().forward(hidden_states, router_logits)
+    # MoE input token counts can change between transformer forwards, for
+    # example when text tokens are present only in some forwards. Clear cached
+    # DP/PCP padding bounds so the first MoE layer recomputes them.
+    forward_context.max_tokens_across_dp = None
+    forward_context.max_tokens_across_pcp = None
 
-    def __del__(self):
-        import vllm_ascend.distributed.parallel_state as vllm_ascend_parallel_state
 
-        if vllm_ascend_parallel_state._MC2:
-            vllm_ascend_parallel_state._MC2.destroy()
-        vllm_ascend_parallel_state._MC2 = None
+def _set_max_tokens(forward_context: Any, hidden_states: torch.Tensor) -> None:
+    if (
+        getattr(forward_context, "max_tokens_across_dp", None) is not None
+        or getattr(forward_context, "max_tokens_across_pcp", None) is not None
+    ):
+        return
+
+    num_tokens_before_pcp = hidden_states.shape[0]
+    dp_metadata = getattr(forward_context, "dp_metadata", None)
+    if dp_metadata is not None:
+        max_tokens_across_dp = int(dp_metadata.num_tokens_across_dp_cpu.max().item())
+        forward_context.max_tokens_across_dp = max_tokens_across_dp
+        # DP prepare pads each rank to max_tokens_across_dp, then gathers
+        # all DP ranks before the PCP gather.
+        num_tokens_before_pcp = max_tokens_across_dp * get_dp_group().world_size
+
+    pcp_group = getattr(vllm_parallel_state, "_PCP", None)
+    if pcp_group is None or getattr(pcp_group, "world_size", 1) <= 1:
+        return
+
+    gathered_num_tokens: list[int | None] = [None] * pcp_group.world_size
+    torch.distributed.all_gather_object(
+        gathered_num_tokens,
+        int(num_tokens_before_pcp),
+        group=pcp_group.cpu_group,
+    )
+    if any(num_tokens is None for num_tokens in gathered_num_tokens):
+        raise RuntimeError(f"Failed to gather MoE PCP token counts: {gathered_num_tokens}")
+    forward_context.max_tokens_across_pcp = max(gathered_num_tokens)
+
+
+def fused_moe_forward_context_pre_hook(
+    module: Any,
+    args: Any,
+    kwargs: Any,
+) -> None:
+    forward_context = _vllm_fc.get_forward_context()
+
+    hidden_states = kwargs.get("hidden_states")
+    if hidden_states is None and args:
+        hidden_states = args[0]
+    if hidden_states is not None:
+        _set_max_tokens(forward_context, hidden_states)
+
+    forward_context.moe_comm_type = _select_moe_comm_method(vllm_config=omni_get_ctx().vllm_config)
+    forward_context.moe_comm_method = _MoECommMethods.get(forward_context.moe_comm_type)

vllm_omni/platforms/npu/platform.py

@@ -33,6 +33,8 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
     # conv2d convolution operator in the code2wav module of Qwen3-TTS not being able to run on Aclnn
     def __init__(self) -> None:
+        from vllm_ascend.utils import adapt_patch
+
         from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches
         from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import (
             apply_qwen3_tts_code2wav_patch,
@@ -41,6 +43,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
             apply_qwen3_tts_tokenizer_v2_patch,
         )
 
+        adapt_patch(is_global_patch=True)
         apply_qwen3_tts_code2wav_patch()
         apply_qwen3_tts_tokenizer_v2_patch()
         apply_310p_patches()
@@ -73,26 +76,40 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         init_ascend_config(vllm_config)
 
+    @classmethod
+    def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None:
+        from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity
+
+        set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1)
+        set_mc2_mask(vllm_config, device)
+
     @classmethod
     def get_default_stage_config_path(cls) -> str:
         return "vllm_omni/deploy"
 
-    @classmethod
-    def get_diffusion_model_impl_qualname(cls, op_name: str) -> str:
-        if op_name == "hunyuan_fused_moe":
-            return "vllm_omni.platforms.npu.models.hunyuan_fused_moe.AscendHunyuanFusedMoE"
-        return super().get_diffusion_model_impl_qualname(op_name)
-
     @classmethod
     def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None:
-        if op_name != "hunyuan_fused_moe":
+        if op_name != "fused_moe":
             return
 
-        from vllm_omni.platforms.npu.models.hunyuan_fused_moe import (
-            prepare_hunyuan_fused_moe_runtime,
+        from vllm_omni.platforms.npu.layers.fused_moe import prepare_fused_moe_runtime
+
+        prepare_fused_moe_runtime()
+
+    @classmethod
+    def register_additional_diffusion_fused_moe_hooks(cls, moe_runner: Any) -> None:
+        from vllm_omni.platforms.npu.layers.fused_moe import fused_moe_forward_context_pre_hook
+
+        moe_runner.register_forward_pre_hook(
+            fused_moe_forward_context_pre_hook,
+            with_kwargs=True,
         )
 
-        prepare_hunyuan_fused_moe_runtime()
+    @classmethod
+    def reset_diffusion_fused_moe_forward_context(cls) -> None:
+        from vllm_omni.platforms.npu.layers.fused_moe import reset_fused_moe_forward_context
+
+        reset_fused_moe_forward_context()
 
     @classmethod
     def get_diffusion_packed_modules_mapping(