[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>
2026-09-25 18:05:54 +00:00
parent 1b922ffcbe
commit e3dd5b5a75
1 changed files with 126 additions and 0 deletions
@@ -16,6 +16,132 @@ pytestmark = pytest.mark.skipif(
 )


+@pytest.mark.parametrize("hash_routing", [False, True], ids=["regular", "hash"])
+@pytest.mark.parametrize("vision", [False, True], ids=["text", "vision"])
+@torch.inference_mode()
+def test_rocm_moe_routing_and_shared_experts_match_reference(
+    hash_routing, vision, monkeypatch, default_vllm_config, dist_init
+) -> None:
+    """AMD V4 routing, clamped experts and shared output agree with PyTorch."""
+    from vllm.forward_context import set_forward_context
+    from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE
+    from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
+    from vllm.utils.torch_utils import set_default_torch_dtype
+    from vllm.v1.worker.workspace import (
+        init_workspace_manager,
+        is_workspace_manager_initialized,
+    )
+
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")
+    default_vllm_config.kernel_config.moe_backend = "triton"
+    config = SimpleNamespace(
+        hidden_size=512,
+        n_routed_experts=8,
+        num_experts_per_tok=2,
+        moe_intermediate_size=512,
+        swiglu_limit=0.05,
+        norm_topk_prob=True,
+        scoring_func="sqrtsoftplus",
+        routed_scaling_factor=1.5,
+        num_hash_layers=int(hash_routing),
+        vocab_size=IMAGE_SENTINEL_BASE_ID + 8,
+        topk_method="noaux_tc",
+        vision_n_layers=int(vision),
+        n_shared_experts=1,
+        hidden_act="silu",
+    )
+    torch.manual_seed(42)
+    with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16):
+        moe = DeepseekV4MoE(
+            SimpleNamespace(
+                model_config=SimpleNamespace(hf_config=config),
+                quant_config=None,
+                parallel_config=ParallelConfig(),
+            ),
+            prefix="model.layers.0.ffn",
+        )
+        for parameter in moe.parameters():
+            if parameter.is_floating_point():
+                parameter.normal_(std=0.04)
+        hidden = torch.randn(17, config.hidden_size) * 0.1
+        input_ids = torch.arange(17, dtype=torch.int32)
+        if vision:
+            # Include both ends of the five-token sentinel interval and the
+            # first ordinary token above it.
+            input_ids[1:4] = torch.tensor(
+                [
+                    IMAGE_SENTINEL_BASE_ID,
+                    IMAGE_SENTINEL_BASE_ID + 4,
+                    IMAGE_SENTINEL_BASE_ID + 5,
+                ],
+                dtype=torch.int32,
+            )
+        if moe.gate.tid2eid is not None:
+            rows = torch.arange(config.vocab_size)
+            moe.gate.tid2eid.copy_(torch.stack((rows % 8, (rows + 3) % 8), dim=1))
+        if moe.gate.e_score_correction_bias is not None:
+            moe.gate.e_score_correction_bias.copy_(torch.arange(8) * 0.1)
+        if moe.gate.bias_vl is not None:
+            moe.gate.bias_vl.copy_(torch.arange(8) * -0.1)
+
+    scores = torch.nn.functional.softplus(
+        hidden.float() @ moe.gate.weight.float().T
+    ).sqrt()
+    bias = moe.gate.e_score_correction_bias
+    selection_scores = scores if bias is None else scores + bias
+    expected_ids = selection_scores.topk(2, dim=-1).indices
+    if hash_routing:
+        expected_ids = moe.gate.tid2eid[input_ids.long()].long()
+    if vision:
+        image_mask = (input_ids >= IMAGE_SENTINEL_BASE_ID) & (
+            input_ids < IMAGE_SENTINEL_BASE_ID + 5
+        )
+        expected_ids[image_mask] = (
+            (scores + moe.gate.bias_vl).topk(2, dim=-1).indices[image_mask]
+        )
+    expected_weights = scores.gather(1, expected_ids)
+    expected_weights *= config.routed_scaling_factor / expected_weights.sum(
+        dim=-1, keepdim=True
+    )
+    router_logits, _ = moe.gate(hidden)
+    weights, ids = moe.experts.router.select_experts(
+        hidden, router_logits, topk_indices_dtype=torch.int32, input_ids=input_ids
+    )
+    torch.testing.assert_close(ids.long(), expected_ids)
+    torch.testing.assert_close(weights, expected_weights, rtol=1e-5, atol=1e-6)
+
+    def expert_reference(x, gate_up_weight, down_weight):
+        gate, up = (
+            torch.nn.functional.linear(x, gate_up_weight).float().chunk(2, dim=-1)
+        )
+        gate = gate.clamp(max=config.swiglu_limit)
+        up = up.clamp(-config.swiglu_limit, config.swiglu_limit)
+        activated = (torch.nn.functional.silu(gate) * up).to(x.dtype)
+        return torch.nn.functional.linear(activated, down_weight).float()
+
+    routed = moe.experts.routed_experts
+    expected = torch.zeros_like(hidden, dtype=torch.float32)
+    for expert_id in range(config.n_routed_experts):
+        token, slot = torch.where(expected_ids == expert_id)
+        expected[token] += (
+            expert_reference(
+                hidden[token], routed.w13_weight[expert_id], routed.w2_weight[expert_id]
+            )
+            * expected_weights[token, slot, None]
+        )
+    shared = moe.shared_experts
+    expected += expert_reference(
+        hidden, shared.gate_up_proj.weight, shared.down_proj.weight
+    )
+    moe.experts._quant_method.process_weights_after_loading(routed)
+    if not is_workspace_manager_initialized():
+        init_workspace_manager(torch.device("cuda:0"))
+    with set_forward_context(None, default_vllm_config, num_tokens=hidden.shape[0]):
+        actual = moe(hidden, input_ids)
+    torch.testing.assert_close(actual.float(), expected, rtol=0.02, atol=5e-5)
+
+
 def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None:
    from vllm.config import CacheConfig
    from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype