From e108cbb51049046275deefa5eb7a9668d22133be Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:41:34 +0800 Subject: [PATCH] [NPU][HunyuanImage3] Adapt diffusion FusedMoE for vLLM >= 0.24.0 (#5167) Signed-off-by: betta18 Signed-off-by: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com> Co-authored-by: betta18 Co-authored-by: Hongsheng Liu --- .../test_expert_parallel_layout.py | 6 +- tests/diffusion/layers/test_fused_moe.py | 215 ++++++++++++++++++ .../hunyuan_image3/test_hunyuan_fused_moe.py | 197 ---------------- .../diffusion/distributed/parallel_state.py | 7 +- .../fused_moe.py} | 65 +----- .../models/hunyuan_image3/__init__.py | 3 +- .../hunyuan_image3_transformer.py | 9 +- .../worker/diffusion_model_runner.py | 6 + vllm_omni/platforms/interface.py | 43 +++- vllm_omni/platforms/npu/layers/__init__.py | 2 + .../fused_moe.py} | 107 +++++---- vllm_omni/platforms/npu/platform.py | 37 ++- 12 files changed, 379 insertions(+), 318 deletions(-) create mode 100644 tests/diffusion/layers/test_fused_moe.py delete mode 100644 tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py rename vllm_omni/diffusion/{models/hunyuan_image3/hunyuan_fused_moe.py => layers/fused_moe.py} (66%) create mode 100644 vllm_omni/platforms/npu/layers/__init__.py rename vllm_omni/platforms/npu/{models/hunyuan_fused_moe.py => layers/fused_moe.py} (52%) diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py index c4d3c6ae7..bea7e9ccf 100644 --- a/tests/diffusion/distributed/test_expert_parallel_layout.py +++ b/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 diff --git a/tests/diffusion/layers/test_fused_moe.py b/tests/diffusion/layers/test_fused_moe.py new file mode 100644 index 000000000..983e146ef --- /dev/null +++ b/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, + ) diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py deleted file mode 100644 index e8ade2e51..000000000 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py +++ /dev/null @@ -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, - ) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 4b531fec9..20fcc9dea 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/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) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py b/vllm_omni/diffusion/layers/fused_moe.py similarity index 66% rename from vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py rename to vllm_omni/diffusion/layers/fused_moe.py index fc1317a48..c5ee02487 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py +++ b/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, - ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py index 6612bd855..625895d99 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/__init__.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 5f72a607d..f749d73d3 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index c54b721f8..be68f071f 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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) ) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 54725b553..c4e4c6e37 100644 --- a/vllm_omni/platforms/interface.py +++ b/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 diff --git a/vllm_omni/platforms/npu/layers/__init__.py b/vllm_omni/platforms/npu/layers/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py b/vllm_omni/platforms/npu/layers/fused_moe.py similarity index 52% rename from vllm_omni/platforms/npu/models/hunyuan_fused_moe.py rename to vllm_omni/platforms/npu/layers/fused_moe.py index 251eac2cf..e17610701 100644 --- a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py +++ b/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) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 5133fa98e..42ff0b418 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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(