From 9179bcd4166bd03f4903aafed39d2c284eeb8c91 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 21 Sep 2026 15:28:35 -0500 Subject: [PATCH] [Bugfix][MoE] Reject hash routing for unsupported monolithic backends (#57867) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/kernels/moe/conftest.py | 2 +- tests/kernels/moe/test_topk_softplus_sqrt.py | 59 +++++++++++++++++++ .../model_executor/layers/fused_moe/config.py | 1 + vllm/model_executor/layers/fused_moe/layer.py | 1 + .../layers/fused_moe/modular_kernel.py | 2 + 5 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/conftest.py b/tests/kernels/moe/conftest.py index fdb3088087..af99dc756b 100644 --- a/tests/kernels/moe/conftest.py +++ b/tests/kernels/moe/conftest.py @@ -58,6 +58,6 @@ def pytest_runtest_setup(item): # Not tryfirst: the skipping plugin evaluates skip marks in a tryfirst # hook, so tests about to be skipped never reach this, while fixture setup # still runs afterwards. - if item.path.name in NEEDS_CLEAN_ENTRY: + if item.path.name in NEEDS_CLEAN_ENTRY or "dist_init" in item.fixturenames: reset_workspace_manager() cleanup_dist_env_and_memory() diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index ab9072576d..bc438c6693 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -81,6 +81,65 @@ def _torch_topk_softplus_sqrt( return topk_weights.to(torch.float32), topk_ids.to(torch.int32) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection") +@pytest.mark.parametrize("has_hash_routing", [False, True]) +@pytest.mark.parametrize( + "backend,expected_experts", + [ + ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), + ("aiter", "AiterExperts"), + ("triton_unfused", "UnfusedOAITritonExperts"), + ], +) +def test_hash_routing_backend_selection( + dist_init, + default_vllm_config, + monkeypatch, + has_hash_routing, + backend, + expected_experts, +): + """Hash tables must reach backend selection; modular AITER remains eligible.""" + from dataclasses import replace + + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + select_deepseek_v4_mxfp4_moe_backend, + ) + from vllm.platforms import rocm + + # Exercise both native MX and monolithic AITER selection on any ROCm device. + monkeypatch.setattr(rocm, "on_gfx950", lambda: True) + monkeypatch.setattr(rocm, "on_gfx1250", lambda: False) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: True) + default_vllm_config.kernel_config.moe_backend = "triton" + layer = FusedMoEFactory( + num_experts=8, + top_k=2, + hidden_size=256, + intermediate_size=256, + params_dtype=torch.bfloat16, + scoring_func="sqrtsoftplus", + e_score_correction_bias=torch.zeros(8), + hash_indices_table=( + torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) + if has_hash_routing + else None + ), + prefix="hash_routing_selection", + ) + config = replace(layer.routed_experts.moe_config, moe_backend=backend) + + if has_hash_routing and backend == "aiter_triton_mxfp4_bf16": + with pytest.raises(ValueError, match="hash routing"): + select_deepseek_v4_mxfp4_moe_backend(config) + else: + _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) + assert experts_cls.__name__ == expected_experts + + def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id(): gating_output = torch.tensor([[2.0, 1.0, 1.0, 0.0]]) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 62ab2bf88a..0c82d28096 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1282,6 +1282,7 @@ class FusedMoEConfig: elastic_ep_max_dp_size: int | None = None has_bias: bool = False is_lora_enabled: bool = False + has_hash_routing: bool = False # When True, the MoE skips its final cross-rank all-reduce (and the separate # shared-expert reduce), returning the partial per-rank sum. The caller is diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 49557bb3f9..ff620622c1 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -363,6 +363,7 @@ def FusedMoEFactory( activation=moe_activation, device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal + has_hash_routing=hash_indices_table is not None, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d91743dd95..fba16f67b0 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -544,6 +544,8 @@ class FusedMoEExperts(ABC): return False, _make_reason( f"parallel config {moe_config.moe_parallel_config}" ) + elif moe_config.has_hash_routing and cls.is_monolithic(): + return False, _make_reason("hash routing") elif not cls._supports_routing_method( moe_config.routing_method, weight_key, activation_key ):