From e3dd5b5a758a6ba1a56f0d79a652ffbb8560abab Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Fri, 25 Sep 2026 20:05:54 +0200 Subject: [PATCH] [ROCm][CI] Test AMD DeepSeek V4 MoE routing against a PyTorch reference (#58740) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas Co-authored-by: Andreas Karatzas --- tests/models/test_deepseek_v4_vl_rocm.py | 126 +++++++++++++++++++++++ 1 file changed, 126 insertions(+) diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 098e5185fc..036efbaa42 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -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