[ROCm][CI] Test AMD DeepSeek V4 MoE routing against a PyTorch reference (#58740)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 25, 2026 at 20:05 UTC
e3dd5b5a758a6ba1a56f0d79a652ffbb8560abab
1 file changed
+126
tests/models/test_deepseek_v4_vl_rocm.py
+126
@@ -16,6 +16,132 @@ pytestmark = pytest.mark.skipif(
16
)
17
18
19
+@pytest.mark.parametrize("hash_routing", [False, True], ids=["regular", "hash"])
20
+@pytest.mark.parametrize("vision", [False, True], ids=["text", "vision"])
21
+@torch.inference_mode()
22
+def test_rocm_moe_routing_and_shared_experts_match_reference(
23
+ hash_routing, vision, monkeypatch, default_vllm_config, dist_init
24
+) -> None:
25
+ """AMD V4 routing, clamped experts and shared output agree with PyTorch."""
26
+ from vllm.forward_context import set_forward_context
27
+ from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE
28
+ from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
29
+ from vllm.utils.torch_utils import set_default_torch_dtype
30
+ from vllm.v1.worker.workspace import (
31
+ init_workspace_manager,
32
+ is_workspace_manager_initialized,
33
+ )
34
+
35
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
36
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")
37
+ default_vllm_config.kernel_config.moe_backend = "triton"
38
+ config = SimpleNamespace(
39
+ hidden_size=512,
40
+ n_routed_experts=8,
41
+ num_experts_per_tok=2,
42
+ moe_intermediate_size=512,
43
+ swiglu_limit=0.05,
44
+ norm_topk_prob=True,
45
+ scoring_func="sqrtsoftplus",
46
+ routed_scaling_factor=1.5,
47
+ num_hash_layers=int(hash_routing),
48
+ vocab_size=IMAGE_SENTINEL_BASE_ID + 8,
49
+ topk_method="noaux_tc",
50
+ vision_n_layers=int(vision),
51
+ n_shared_experts=1,
52
+ hidden_act="silu",
53
+ )
54
+ torch.manual_seed(42)
55
+ with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16):
56
+ moe = DeepseekV4MoE(
57
+ SimpleNamespace(
58
+ model_config=SimpleNamespace(hf_config=config),
59
+ quant_config=None,
60
+ parallel_config=ParallelConfig(),
61
+ ),
62
+ prefix="model.layers.0.ffn",
63
+ )
64
+ for parameter in moe.parameters():
65
+ if parameter.is_floating_point():
66
+ parameter.normal_(std=0.04)
67
+ hidden = torch.randn(17, config.hidden_size) * 0.1
68
+ input_ids = torch.arange(17, dtype=torch.int32)
69
+ if vision:
70
+ # Include both ends of the five-token sentinel interval and the
71
+ # first ordinary token above it.
72
+ input_ids[1:4] = torch.tensor(
73
+ [
74
+ IMAGE_SENTINEL_BASE_ID,
75
+ IMAGE_SENTINEL_BASE_ID + 4,
76
+ IMAGE_SENTINEL_BASE_ID + 5,
77
+ ],
78
+ dtype=torch.int32,
79
+ )
80
+ if moe.gate.tid2eid is not None:
81
+ rows = torch.arange(config.vocab_size)
82
+ moe.gate.tid2eid.copy_(torch.stack((rows % 8, (rows + 3) % 8), dim=1))
83
+ if moe.gate.e_score_correction_bias is not None:
84
+ moe.gate.e_score_correction_bias.copy_(torch.arange(8) * 0.1)
85
+ if moe.gate.bias_vl is not None:
86
+ moe.gate.bias_vl.copy_(torch.arange(8) * -0.1)
87
+
88
+ scores = torch.nn.functional.softplus(
89
+ hidden.float() @ moe.gate.weight.float().T
90
+ ).sqrt()
91
+ bias = moe.gate.e_score_correction_bias
92
+ selection_scores = scores if bias is None else scores + bias
93
+ expected_ids = selection_scores.topk(2, dim=-1).indices
94
+ if hash_routing:
95
+ expected_ids = moe.gate.tid2eid[input_ids.long()].long()
96
+ if vision:
97
+ image_mask = (input_ids >= IMAGE_SENTINEL_BASE_ID) & (
98
+ input_ids < IMAGE_SENTINEL_BASE_ID + 5
99
+ )
100
+ expected_ids[image_mask] = (
101
+ (scores + moe.gate.bias_vl).topk(2, dim=-1).indices[image_mask]
102
+ )
103
+ expected_weights = scores.gather(1, expected_ids)
104
+ expected_weights *= config.routed_scaling_factor / expected_weights.sum(
105
+ dim=-1, keepdim=True
106
+ )
107
+ router_logits, _ = moe.gate(hidden)
108
+ weights, ids = moe.experts.router.select_experts(
109
+ hidden, router_logits, topk_indices_dtype=torch.int32, input_ids=input_ids
110
+ )
111
+ torch.testing.assert_close(ids.long(), expected_ids)
112
+ torch.testing.assert_close(weights, expected_weights, rtol=1e-5, atol=1e-6)
113
+
114
+ def expert_reference(x, gate_up_weight, down_weight):
115
+ gate, up = (
116
+ torch.nn.functional.linear(x, gate_up_weight).float().chunk(2, dim=-1)
117
+ )
118
+ gate = gate.clamp(max=config.swiglu_limit)
119
+ up = up.clamp(-config.swiglu_limit, config.swiglu_limit)
120
+ activated = (torch.nn.functional.silu(gate) * up).to(x.dtype)
121
+ return torch.nn.functional.linear(activated, down_weight).float()
122
+
123
+ routed = moe.experts.routed_experts
124
+ expected = torch.zeros_like(hidden, dtype=torch.float32)
125
+ for expert_id in range(config.n_routed_experts):
126
+ token, slot = torch.where(expected_ids == expert_id)
127
+ expected[token] += (
128
+ expert_reference(
129
+ hidden[token], routed.w13_weight[expert_id], routed.w2_weight[expert_id]
130
+ )
131
+ * expected_weights[token, slot, None]
132
+ )
133
+ shared = moe.shared_experts
134
+ expected += expert_reference(
135
+ hidden, shared.gate_up_proj.weight, shared.down_proj.weight
136
+ )
137
+ moe.experts._quant_method.process_weights_after_loading(routed)
138
+ if not is_workspace_manager_initialized():
139
+ init_workspace_manager(torch.device("cuda:0"))
140
+ with set_forward_context(None, default_vllm_config, num_tokens=hidden.shape[0]):
141
+ actual = moe(hidden, input_ids)
142
+ torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=5e-5)
143
+
144
+
145
def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None:
146
from vllm.config import CacheConfig
147
from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype