anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
jiangmengyu18betta18Hongsheng Liu
 authored and GitHub committed 2026-07-27 22:41:34 +08:00
1 parent 82172b112d
commit e108cbb510
12 files changed
+379 -318
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/distributed/test_expert_parallel_layout.py
+4 -2
View File
@@ -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
+215
View File
@@ -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
-197
View File
@@ -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,
)