A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 11:15 UTC
[Bugfix] Fix RoutedExpertsCapturer for Gemma 4 MoE (top_k_experts) (#41401)
Signed-off-by: Tran Le <tranle@fireworks.ai>
Tran Le committed
Apr 30, 2026 at 16:19 UTC
71725f6730ec7076ee914ba06fae482a10f2f159
1 file changed
+20
-2
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+20
-2
@@ -24,6 +24,24 @@ from vllm.platforms import current_platform
24
25
logger = logging.getLogger(__name__)
26
27
+
28
+def _get_num_experts_per_tok(hf_config) -> int:
29
+ """Resolve the per-token expert count from the HF config.
30
+
31
+ Different model families store this under different attribute names
32
+ (e.g. ``num_experts_per_tok`` for DeepSeek, ``top_k_experts`` for Gemma 4).
33
+ """
34
+ val = getattr(hf_config, "num_experts_per_tok", None)
35
+ if val is None:
36
+ val = getattr(hf_config, "top_k_experts", None)
37
+ if val is None:
38
+ raise ValueError(
39
+ "Cannot determine num_experts_per_tok: HF config has neither "
40
+ "'num_experts_per_tok' nor 'top_k_experts'"
41
+ )
42
+ return val
43
+
44
+
45
# Constants
46
_TMP_DIR = tempfile.gettempdir()
47
_LOCK_FILE_PREFIX = os.path.join(_TMP_DIR, "vllm_routed_experts")
@@ -127,7 +145,7 @@ class RoutedExpertsCapturer:
145
146
hf_config = vllm_config.model_config.hf_text_config
147
num_layers = hf_config.num_hidden_layers
130
- num_experts_per_tok = hf_config.num_experts_per_tok
148
+ num_experts_per_tok = _get_num_experts_per_tok(hf_config)
149
150
# Initialize device buffer
151
self._device_buffer = torch.zeros(
@@ -300,7 +318,7 @@ class RoutedExpertsReader:
318
shape = (
319
max_num_kv_tokens,
320
hf_config.num_hidden_layers,
303
- hf_config.num_experts_per_tok,
321
+ _get_num_experts_per_tok(hf_config),
322
)
323
324
self.dp_rank = vllm_config.parallel_config.data_parallel_rank