[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 <codex@openai.com>
Kevin H. Luu committed
Sep 12, 2026 at 11:31 UTC
1ee4be4dbc576f2158913a89598ccfa9dec9391e
1 file changed
+3
-6
tests/models/test_deepseek_v4_mega_moe.py
+3
-6
@@ -18,10 +18,7 @@ from vllm.models.deepseek_v4.nvidia.model import (
18
)
19
from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
20
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
21
-from vllm.models.deepseek_v4_1.common.mm_preprocess import (
22
- IMAGE_PAD_ID,
23
- IMAGE_SENTINEL_BASE_ID,
24
-)
21
+from vllm.models.deepseek_v4_1.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
22
from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE
23
from vllm.platforms import current_platform
24
from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
@@ -78,7 +75,7 @@ def test_deepseek_v41_moe_routes_without_hash_table(
75
moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn")
76
hidden_states = torch.randn(4, config.hidden_size)
77
input_ids = (
81
- torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257])
78
+ torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257])
79
if vision
80
else None
81
)
@@ -100,7 +97,7 @@ def test_deepseek_v41_moe_routes_without_hash_table(
97
).sqrt()
98
bias = moe.gate.e_score_correction_bias
99
if vision:
103
- image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID)
100
+ image_mask = input_ids == IMAGE_SENTINEL_BASE_ID
101
bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias)
102
expected_ids = (scores + bias).topk(top_k, dim=-1).indices
103
expected_weights = scores.gather(1, expected_ids)