[R3] Unify routed expert shape configuration (#50940)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: yihengz <yihengz@nvidia.com>
aoshen02 committed
Aug 5, 2026 at 11:20 UTC
6153dbe036abe11628ab5f7e120707e4f7331652
7 files changed
+146
-46
tests/config/test_model_arch_config.py
+34
@@ -4,6 +4,8 @@
4
5
import json
6
from pathlib import Path
7
+from types import SimpleNamespace
8
+from typing import cast
9
10
import pytest
11
from transformers import PretrainedConfig
@@ -131,6 +133,38 @@ def test_head_size_falls_back_when_head_dim_is_zero():
133
assert convertor.get_head_size() == 128
134
135
136
+@pytest.mark.parametrize(
137
+ "attribute",
138
+ [
139
+ "num_experts_per_tok",
140
+ "num_experts_per_token",
141
+ "top_k_experts",
142
+ "moe_topk",
143
+ "moe_top_k",
144
+ ],
145
+)
146
+def test_num_experts_per_tok_aliases(attribute: str):
147
+ hf_config = PretrainedConfig(**{attribute: 4})
148
+ convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
149
+ model_config = cast(
150
+ ModelConfig,
151
+ SimpleNamespace(
152
+ model_arch_config=SimpleNamespace(
153
+ num_experts_per_token=convertor.get_num_experts_per_token()
154
+ )
155
+ ),
156
+ )
157
+
158
+ assert ModelConfig.get_num_experts_per_tok(model_config) == 4
159
+
160
+
161
+def test_num_experts_per_tok_none_is_normalized():
162
+ hf_config = PretrainedConfig(top_k_experts=None)
163
+ convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
164
+
165
+ assert convertor.get_num_experts_per_token() == 0
166
+
167
+
168
def test_legacy_modelopt_config_without_producer_is_normalized():
169
quantization_config = {
170
"quantization": {
tests/model_executor/test_routed_experts_capture.py
+76
-1
@@ -7,16 +7,25 @@ from unittest.mock import Mock, patch
7
import pytest
8
import torch
9
10
-from vllm.config import VllmConfig
10
+from vllm.config import ModelConfig, VllmConfig
11
from vllm.config.compilation import CompilationMode
12
from vllm.distributed.eplb.eplb_state import EplbLayerState
13
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
14
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
15
RoutedExpertsCapturer,
16
+ RoutedExpertsManager,
17
bind_routed_experts_capturer,
18
get_routed_experts_attn_gid,
19
)
20
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
21
+from vllm.transformers_utils.model_arch_config_convertor import (
22
+ ModelArchConfigConvertorBase,
23
+)
24
+from vllm.v1.kv_cache_interface import (
25
+ FullAttentionSpec,
26
+ KVCacheConfig,
27
+ KVCacheGroupSpec,
28
+)
29
30
pytestmark = pytest.mark.cpu_test
31
@@ -76,6 +85,72 @@ def _make_modular_routed_experts():
85
)
86
87
88
+def _make_model_config(hf_config):
89
+ num_experts_per_token = ModelArchConfigConvertorBase(
90
+ hf_config, hf_config
91
+ ).get_num_experts_per_token()
92
+ model_config = SimpleNamespace(
93
+ hf_text_config=hf_config,
94
+ model_arch_config=SimpleNamespace(
95
+ num_experts_per_token=num_experts_per_token,
96
+ ),
97
+ )
98
+ model_config.get_num_experts = lambda: hf_config.num_experts
99
+ model_config.get_num_experts_per_tok = lambda: (
100
+ ModelConfig.get_num_experts_per_tok(model_config)
101
+ )
102
+ model_config.get_total_num_hidden_layers = lambda: hf_config.num_hidden_layers
103
+ return model_config
104
+
105
+
106
+def test_routed_experts_manager_uses_gemma4_top_k_experts():
107
+ hf_config = SimpleNamespace(
108
+ num_experts=8,
109
+ top_k_experts=2,
110
+ num_hidden_layers=3,
111
+ )
112
+ vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
113
+ kv_cache_spec = FullAttentionSpec(
114
+ block_size=4,
115
+ num_kv_heads=1,
116
+ head_size=1,
117
+ dtype=torch.float32,
118
+ )
119
+ kv_cache_config = KVCacheConfig(
120
+ num_blocks=2,
121
+ kv_cache_tensors=[],
122
+ kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
123
+ )
124
+
125
+ manager = RoutedExpertsManager(vllm_config, kv_cache_config)
126
+
127
+ assert manager.routed_experts_by_slot.shape == (8, 3, 2)
128
+
129
+
130
+def test_routed_experts_manager_uses_kimi_k3_experts_per_token():
131
+ hf_config = SimpleNamespace(
132
+ num_experts=8,
133
+ num_experts_per_token=2,
134
+ num_hidden_layers=3,
135
+ )
136
+ vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
137
+ kv_cache_spec = FullAttentionSpec(
138
+ block_size=4,
139
+ num_kv_heads=1,
140
+ head_size=1,
141
+ dtype=torch.float32,
142
+ )
143
+ kv_cache_config = KVCacheConfig(
144
+ num_blocks=2,
145
+ kv_cache_tensors=[],
146
+ kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
147
+ )
148
+
149
+ manager = RoutedExpertsManager(vllm_config, kv_cache_config)
150
+
151
+ assert manager.routed_experts_by_slot.shape == (8, 3, 2)
152
+
153
+
154
def test_base_router_capture_pre_eplb_mapping():
155
router = _make_router()
156
captured = []
vllm/config/model.py
+3
@@ -1436,6 +1436,9 @@ class ModelConfig:
1436
def get_num_experts(self) -> int:
1437
return self.model_arch_config.num_experts
1438
1439
+ def get_num_experts_per_tok(self) -> int:
1440
+ return self.model_arch_config.num_experts_per_token
1441
+
1442
def get_total_num_hidden_layers(self) -> int:
1443
return self.model_arch_config.total_num_hidden_layers
1444
vllm/config/model_arch.py
+3
@@ -47,6 +47,9 @@ class ModelArchitectureConfig:
47
num_experts: int
48
"""Number of experts in the model."""
49
50
+ num_experts_per_token: int
51
+ """Number of routed experts selected per token."""
52
+
53
quantization_config: dict[str, Any] | None
54
"""Quantization configuration dictionary containing quantization parameters."""
55
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+18
-42
@@ -20,40 +20,18 @@ from vllm.v1.outputs import RoutedExpertsTensors
20
logger = logging.getLogger(__name__)
21
22
23
-def _get_num_experts_per_tok(hf_config) -> int:
24
- """Resolve the per-token expert count from the HF config.
25
-
26
- Different model families store this under different attribute names
27
- (e.g. ``num_experts_per_tok`` for DeepSeek, ``top_k_experts`` for Gemma 4).
28
- """
29
- val = getattr(hf_config, "num_experts_per_tok", None)
30
- if val is None:
31
- val = getattr(hf_config, "top_k_experts", None)
32
- if val is None:
23
+def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]:
24
+ model_config = vllm_config.model_config
25
+ num_layers = model_config.get_total_num_hidden_layers()
26
+ num_experts = model_config.get_num_experts()
27
+ num_experts_per_tok = model_config.get_num_experts_per_tok()
28
+ if num_layers <= 0 or num_experts <= 0 or num_experts_per_tok <= 0:
29
raise ValueError(
34
- "Cannot determine num_experts_per_tok: HF config has neither "
35
- "'num_experts_per_tok' nor 'top_k_experts'"
30
+ "Routed-experts capture requires positive layer, expert, and "
31
+ "experts-per-token counts, got "
32
+ f"{num_layers=}, {num_experts=}, {num_experts_per_tok=}."
33
)
37
- return val
38
-
39
-
40
-def get_num_experts(hf_config) -> int:
41
- """Resolve ``num_experts`` across HuggingFace config naming conventions.
42
-
43
- Different MoE model families expose this under different keys:
44
- - ``num_experts``: Mixtral, Qwen2-MoE, Qwen3-MoE
45
- - ``n_routed_experts``: DeepSeek-V2/V3
46
- - ``num_local_experts``: Mixtral (older exports)
47
- """
48
- for key in ("num_experts", "n_routed_experts", "num_local_experts"):
49
- val = getattr(hf_config, key, None)
50
- if val is not None:
51
- return val
52
- raise ValueError(
53
- "Could not resolve num_experts from model config. "
54
- "Expected one of 'num_experts', 'n_routed_experts', "
55
- "or 'num_local_experts'."
56
- )
34
+ return num_layers, num_experts, num_experts_per_tok
35
36
37
class RoutedExpertsCapturer:
@@ -89,9 +67,7 @@ class RoutedExpertsCapturer:
67
vllm_config: VllmConfig,
68
kv_cache_config: KVCacheConfig,
69
) -> None:
92
- hf_config = vllm_config.model_config.hf_text_config
93
- num_experts_per_tok = _get_num_experts_per_tok(hf_config)
94
- num_layers = hf_config.num_hidden_layers
70
+ num_layers, _, num_experts_per_tok = _get_routed_experts_shape(vllm_config)
71
logger.info(
72
"RoutedExpertsCapturer: allocating buffer with "
73
"max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
@@ -99,7 +75,7 @@ class RoutedExpertsCapturer:
75
max_num_batched_tokens,
76
num_layers,
77
num_experts_per_tok,
102
- getattr(hf_config, "model_type", "unknown"),
78
+ vllm_config.model_config.hf_text_config.model_type,
79
)
80
self.device_buffer = torch.zeros(
81
(
@@ -344,9 +320,9 @@ class RoutedExpertsManager:
320
# block IDs span [0, num_blocks) regardless of how many groups
321
# exist. Sizing to the full pool avoids index-out-of-range
322
# when different groups happen to land on the same block.
347
- hf_config = vllm_config.model_config.hf_text_config
348
- num_experts = get_num_experts(hf_config)
349
- num_experts_per_tok = _get_num_experts_per_tok(hf_config)
323
+ num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
324
+ vllm_config
325
+ )
326
max_num_slots = kv_cache_config.num_blocks * self.block_size
327
# Expert IDs are 0..num_experts-1; uint8 fits 256 distinct
328
# values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping
@@ -356,7 +332,7 @@ class RoutedExpertsManager:
332
self.routed_experts_by_slot = np.zeros(
333
(
334
max_num_slots,
359
- hf_config.num_hidden_layers,
335
+ num_layers,
336
num_experts_per_tok,
337
),
338
dtype=expert_id_dtype,
@@ -366,8 +342,8 @@ class RoutedExpertsManager:
342
"(slots=%d, layers=%d, top_k=%d, dtype=%s)",
343
self.routed_experts_by_slot.nbytes / 1e9,
344
max_num_slots,
369
- hf_config.num_hidden_layers,
370
- hf_config.num_experts_per_tok,
345
+ num_layers,
346
+ num_experts_per_tok,
347
self.routed_experts_by_slot.dtype.name,
348
)
349
vllm/transformers_utils/model_arch_config_convertor.py
+11
@@ -164,6 +164,16 @@ class ModelArchConfigConvertorBase:
164
num_experts = self.get_num_experts_from_block_configs()
165
return num_experts
166
167
+ def get_num_experts_per_token(self) -> int:
168
+ names = [
169
+ "num_experts_per_tok",
170
+ "num_experts_per_token",
171
+ "top_k_experts",
172
+ "moe_topk",
173
+ "moe_top_k",
174
+ ]
175
+ return getattr_iter(self.hf_text_config, names, 0) or 0
176
+
177
@final
178
@classmethod
179
def get_torch_dtype(
@@ -379,6 +389,7 @@ class ModelArchConfigConvertorBase:
389
vocab_size=self.get_vocab_size(),
390
total_num_kv_heads=self.get_total_num_kv_heads(),
391
num_experts=self.get_num_experts(),
392
+ num_experts_per_token=self.get_num_experts_per_token(),
393
quantization_config=self.get_quantization_config(),
394
is_deepseek_mla=self.is_deepseek_mla(),
395
is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal),
vllm/v1/metrics/perf.py
+1
-3
@@ -828,9 +828,7 @@ class BaseFfnConfigParser(Parser):
828
829
# Try different naming conventions.
830
args.num_experts = vllm_config.model_config.get_num_experts()
831
- args.num_experts_per_tok = getattr_from_list(
832
- cfg, ["num_experts_per_tok", "moe_topk"], 0
833
- )
831
+ args.num_experts_per_tok = vllm_config.model_config.get_num_experts_per_tok()
832
args.moe_intermediate_size = getattr_from_list(
833
cfg, ["moe_intermediate_size", "intermediate_size"], 0
834
)