[Bugfix][MoE] Reject hash routing for unsupported monolithic backends (#57867)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 21, 2026 at 15:28 UTC
9179bcd4166bd03f4903aafed39d2c284eeb8c91
5 files changed
+64
-1
tests/kernels/moe/conftest.py
+1
-1
@@ -58,6 +58,6 @@ def pytest_runtest_setup(item):
58
# Not tryfirst: the skipping plugin evaluates skip marks in a tryfirst
59
# hook, so tests about to be skipped never reach this, while fixture setup
60
# still runs afterwards.
61
- if item.path.name in NEEDS_CLEAN_ENTRY:
61
+ if item.path.name in NEEDS_CLEAN_ENTRY or "dist_init" in item.fixturenames:
62
reset_workspace_manager()
63
cleanup_dist_env_and_memory()
tests/kernels/moe/test_topk_softplus_sqrt.py
+59
@@ -81,6 +81,65 @@ def _torch_topk_softplus_sqrt(
81
return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
82
83
84
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection")
85
+@pytest.mark.parametrize("has_hash_routing", [False, True])
86
+@pytest.mark.parametrize(
87
+ "backend,expected_experts",
88
+ [
89
+ ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"),
90
+ ("aiter", "AiterExperts"),
91
+ ("triton_unfused", "UnfusedOAITritonExperts"),
92
+ ],
93
+)
94
+def test_hash_routing_backend_selection(
95
+ dist_init,
96
+ default_vllm_config,
97
+ monkeypatch,
98
+ has_hash_routing,
99
+ backend,
100
+ expected_experts,
101
+):
102
+ """Hash tables must reach backend selection; modular AITER remains eligible."""
103
+ from dataclasses import replace
104
+
105
+ from vllm._aiter_ops import rocm_aiter_ops
106
+ from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory
107
+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
108
+ select_deepseek_v4_mxfp4_moe_backend,
109
+ )
110
+ from vllm.platforms import rocm
111
+
112
+ # Exercise both native MX and monolithic AITER selection on any ROCm device.
113
+ monkeypatch.setattr(rocm, "on_gfx950", lambda: True)
114
+ monkeypatch.setattr(rocm, "on_gfx1250", lambda: False)
115
+ monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
116
+ monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: True)
117
+ default_vllm_config.kernel_config.moe_backend = "triton"
118
+ layer = FusedMoEFactory(
119
+ num_experts=8,
120
+ top_k=2,
121
+ hidden_size=256,
122
+ intermediate_size=256,
123
+ params_dtype=torch.bfloat16,
124
+ scoring_func="sqrtsoftplus",
125
+ e_score_correction_bias=torch.zeros(8),
126
+ hash_indices_table=(
127
+ torch.tensor([[0, 1], [2, 3]], dtype=torch.int32)
128
+ if has_hash_routing
129
+ else None
130
+ ),
131
+ prefix="hash_routing_selection",
132
+ )
133
+ config = replace(layer.routed_experts.moe_config, moe_backend=backend)
134
+
135
+ if has_hash_routing and backend == "aiter_triton_mxfp4_bf16":
136
+ with pytest.raises(ValueError, match="hash routing"):
137
+ select_deepseek_v4_mxfp4_moe_backend(config)
138
+ else:
139
+ _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config)
140
+ assert experts_cls.__name__ == expected_experts
141
+
142
+
143
def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id():
144
gating_output = torch.tensor([[2.0, 1.0, 1.0, 0.0]])
145
vllm/model_executor/layers/fused_moe/config.py
+1
@@ -1282,6 +1282,7 @@ class FusedMoEConfig:
1282
elastic_ep_max_dp_size: int | None = None
1283
has_bias: bool = False
1284
is_lora_enabled: bool = False
1285
+ has_hash_routing: bool = False
1286
1287
# When True, the MoE skips its final cross-rank all-reduce (and the separate
1288
# shared-expert reduce), returning the partial per-rank sum. The caller is
vllm/model_executor/layers/fused_moe/layer.py
+1
@@ -363,6 +363,7 @@ def FusedMoEFactory(
363
activation=moe_activation,
364
device=vllm_config.device_config.device,
365
routing_method=router.routing_method_type, # Not ideal
366
+ has_hash_routing=hash_indices_table is not None,
367
swiglu_limit=swiglu_limit,
368
swiglu_alpha=swiglu_alpha,
369
swiglu_beta=swiglu_beta,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+2
@@ -544,6 +544,8 @@ class FusedMoEExperts(ABC):
544
return False, _make_reason(
545
f"parallel config {moe_config.moe_parallel_config}"
546
)
547
+ elif moe_config.has_hash_routing and cls.is_monolithic():
548
+ return False, _make_reason("hash routing")
549
elif not cls._supports_routing_method(
550
moe_config.routing_method, weight_key, activation_key
551
):