Zhou TaichangandClaude Opus 5
2ce15ab7a0
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
216 lines
8.5 KiB
Python
216 lines
8.5 KiB
Python
# 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(FusedMoEFactory=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(FusedMoEFactory=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(FusedMoEFactory=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,
|
|
)
|