iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI] Update DeepSeek V4.1 MegaMoE routing test (#56599)

Browse Source 
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Kevin H. Luu
2026-09-12 11:31:13 -07:00
committed by GitHub
co-authored by OpenAI Codex
parent fd7cd4b883
commit 1ee4be4dbc
1 changed files with 3 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/test_deepseek_v4_mega_moe.py
+3 -6
View File
@@ -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)