From 1ee4be4dbc576f2158913a89598ccfa9dec9391e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 12 Sep 2026 11:31:13 -0700 Subject: [PATCH] [CI] Update DeepSeek V4.1 MegaMoE routing test (#56599) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/models/test_deepseek_v4_mega_moe.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 58d5239285..d3d9ccabad 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -18,10 +18,7 @@ from vllm.models.deepseek_v4.nvidia.model import ( ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs -from vllm.models.deepseek_v4_1.common.mm_preprocess import ( - IMAGE_PAD_ID, - IMAGE_SENTINEL_BASE_ID, -) +from vllm.models.deepseek_v4_1.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config @@ -78,7 +75,7 @@ def test_deepseek_v41_moe_routes_without_hash_table( moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") hidden_states = torch.randn(4, config.hidden_size) input_ids = ( - torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257]) + torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257]) if vision else None ) @@ -100,7 +97,7 @@ def test_deepseek_v41_moe_routes_without_hash_table( ).sqrt() bias = moe.gate.e_score_correction_bias if vision: - image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID) + image_mask = input_ids == IMAGE_SENTINEL_BASE_ID bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias) expected_ids = (scores + bias).topk(top_k, dim=-1).indices expected_weights = scores.gather(1, expected_ids)