[R3] Unify routed expert shape configuration (#50940)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: yihengz <yihengz@nvidia.com>

tests/config/test_model_arch_config.py

@@ -4,6 +4,8 @@
 
 import json
 from pathlib import Path
+from types import SimpleNamespace
+from typing import cast
 
 import pytest
 from transformers import PretrainedConfig
@@ -131,6 +133,38 @@ def test_head_size_falls_back_when_head_dim_is_zero():
     assert convertor.get_head_size() == 128
 
 
+@pytest.mark.parametrize(
+    "attribute",
+    [
+        "num_experts_per_tok",
+        "num_experts_per_token",
+        "top_k_experts",
+        "moe_topk",
+        "moe_top_k",
+    ],
+)
+def test_num_experts_per_tok_aliases(attribute: str):
+    hf_config = PretrainedConfig(**{attribute: 4})
+    convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
+    model_config = cast(
+        ModelConfig,
+        SimpleNamespace(
+            model_arch_config=SimpleNamespace(
+                num_experts_per_token=convertor.get_num_experts_per_token()
+            )
+        ),
+    )
+
+    assert ModelConfig.get_num_experts_per_tok(model_config) == 4
+
+
+def test_num_experts_per_tok_none_is_normalized():
+    hf_config = PretrainedConfig(top_k_experts=None)
+    convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
+
+    assert convertor.get_num_experts_per_token() == 0
+
+
 def test_legacy_modelopt_config_without_producer_is_normalized():
     quantization_config = {
         "quantization": {

tests/model_executor/test_routed_experts_capture.py

@@ -7,16 +7,25 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm.config import VllmConfig
+from vllm.config import ModelConfig, VllmConfig
 from vllm.config.compilation import CompilationMode
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
     RoutedExpertsCapturer,
+    RoutedExpertsManager,
     bind_routed_experts_capturer,
     get_routed_experts_attn_gid,
 )
 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+from vllm.transformers_utils.model_arch_config_convertor import (
+    ModelArchConfigConvertorBase,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+)
 
 pytestmark = pytest.mark.cpu_test
 
@@ -76,6 +85,72 @@ def _make_modular_routed_experts():
     )
 
 
+def _make_model_config(hf_config):
+    num_experts_per_token = ModelArchConfigConvertorBase(
+        hf_config, hf_config
+    ).get_num_experts_per_token()
+    model_config = SimpleNamespace(
+        hf_text_config=hf_config,
+        model_arch_config=SimpleNamespace(
+            num_experts_per_token=num_experts_per_token,
+        ),
+    )
+    model_config.get_num_experts = lambda: hf_config.num_experts
+    model_config.get_num_experts_per_tok = lambda: (
+        ModelConfig.get_num_experts_per_tok(model_config)
+    )
+    model_config.get_total_num_hidden_layers = lambda: hf_config.num_hidden_layers
+    return model_config
+
+
+def test_routed_experts_manager_uses_gemma4_top_k_experts():
+    hf_config = SimpleNamespace(
+        num_experts=8,
+        top_k_experts=2,
+        num_hidden_layers=3,
+    )
+    vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
+    kv_cache_spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[],
+        kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
+    )
+
+    manager = RoutedExpertsManager(vllm_config, kv_cache_config)
+
+    assert manager.routed_experts_by_slot.shape == (8, 3, 2)
+
+
+def test_routed_experts_manager_uses_kimi_k3_experts_per_token():
+    hf_config = SimpleNamespace(
+        num_experts=8,
+        num_experts_per_token=2,
+        num_hidden_layers=3,
+    )
+    vllm_config = SimpleNamespace(model_config=_make_model_config(hf_config))
+    kv_cache_spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[],
+        kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)],
+    )
+
+    manager = RoutedExpertsManager(vllm_config, kv_cache_config)
+
+    assert manager.routed_experts_by_slot.shape == (8, 3, 2)
+
+
 def test_base_router_capture_pre_eplb_mapping():
     router = _make_router()
     captured = []

vllm/config/model.py

@@ -1436,6 +1436,9 @@ class ModelConfig:
     def get_num_experts(self) -> int:
         return self.model_arch_config.num_experts
 
+    def get_num_experts_per_tok(self) -> int:
+        return self.model_arch_config.num_experts_per_token
+
     def get_total_num_hidden_layers(self) -> int:
         return self.model_arch_config.total_num_hidden_layers
 

vllm/config/model_arch.py

@@ -47,6 +47,9 @@ class ModelArchitectureConfig:
     num_experts: int
     """Number of experts in the model."""
 
+    num_experts_per_token: int
+    """Number of routed experts selected per token."""
+
     quantization_config: dict[str, Any] | None
     """Quantization configuration dictionary containing quantization parameters."""
 

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

@@ -20,40 +20,18 @@ from vllm.v1.outputs import RoutedExpertsTensors
 logger = logging.getLogger(__name__)
 
 
-def _get_num_experts_per_tok(hf_config) -> int:
-    """Resolve the per-token expert count from the HF config.
-
-    Different model families store this under different attribute names
-    (e.g. ``num_experts_per_tok`` for DeepSeek, ``top_k_experts`` for Gemma 4).
-    """
-    val = getattr(hf_config, "num_experts_per_tok", None)
-    if val is None:
-        val = getattr(hf_config, "top_k_experts", None)
-    if val is None:
+def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]:
+    model_config = vllm_config.model_config
+    num_layers = model_config.get_total_num_hidden_layers()
+    num_experts = model_config.get_num_experts()
+    num_experts_per_tok = model_config.get_num_experts_per_tok()
+    if num_layers <= 0 or num_experts <= 0 or num_experts_per_tok <= 0:
         raise ValueError(
-            "Cannot determine num_experts_per_tok: HF config has neither "
-            "'num_experts_per_tok' nor 'top_k_experts'"
+            "Routed-experts capture requires positive layer, expert, and "
+            "experts-per-token counts, got "
+            f"{num_layers=}, {num_experts=}, {num_experts_per_tok=}."
         )
-    return val
-
-
-def get_num_experts(hf_config) -> int:
-    """Resolve ``num_experts`` across HuggingFace config naming conventions.
-
-    Different MoE model families expose this under different keys:
-      - ``num_experts``: Mixtral, Qwen2-MoE, Qwen3-MoE
-      - ``n_routed_experts``: DeepSeek-V2/V3
-      - ``num_local_experts``: Mixtral (older exports)
-    """
-    for key in ("num_experts", "n_routed_experts", "num_local_experts"):
-        val = getattr(hf_config, key, None)
-        if val is not None:
-            return val
-    raise ValueError(
-        "Could not resolve num_experts from model config. "
-        "Expected one of 'num_experts', 'n_routed_experts', "
-        "or 'num_local_experts'."
-    )
+    return num_layers, num_experts, num_experts_per_tok
 
 
 class RoutedExpertsCapturer:
@@ -89,9 +67,7 @@ class RoutedExpertsCapturer:
         vllm_config: VllmConfig,
         kv_cache_config: KVCacheConfig,
     ) -> None:
-        hf_config = vllm_config.model_config.hf_text_config
-        num_experts_per_tok = _get_num_experts_per_tok(hf_config)
-        num_layers = hf_config.num_hidden_layers
+        num_layers, _, num_experts_per_tok = _get_routed_experts_shape(vllm_config)
         logger.info(
             "RoutedExpertsCapturer: allocating buffer with "
             "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
@@ -99,7 +75,7 @@ class RoutedExpertsCapturer:
             max_num_batched_tokens,
             num_layers,
             num_experts_per_tok,
-            getattr(hf_config, "model_type", "unknown"),
+            vllm_config.model_config.hf_text_config.model_type,
         )
         self.device_buffer = torch.zeros(
             (
@@ -344,9 +320,9 @@ class RoutedExpertsManager:
         # block IDs span [0, num_blocks) regardless of how many groups
         # exist. Sizing to the full pool avoids index-out-of-range
         # when different groups happen to land on the same block.
-        hf_config = vllm_config.model_config.hf_text_config
-        num_experts = get_num_experts(hf_config)
-        num_experts_per_tok = _get_num_experts_per_tok(hf_config)
+        num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
+            vllm_config
+        )
         max_num_slots = kv_cache_config.num_blocks * self.block_size
         # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct
         # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping
@@ -356,7 +332,7 @@ class RoutedExpertsManager:
         self.routed_experts_by_slot = np.zeros(
             (
                 max_num_slots,
-                hf_config.num_hidden_layers,
+                num_layers,
                 num_experts_per_tok,
             ),
             dtype=expert_id_dtype,
@@ -366,8 +342,8 @@ class RoutedExpertsManager:
             "(slots=%d, layers=%d, top_k=%d, dtype=%s)",
             self.routed_experts_by_slot.nbytes / 1e9,
             max_num_slots,
-            hf_config.num_hidden_layers,
-            hf_config.num_experts_per_tok,
+            num_layers,
+            num_experts_per_tok,
             self.routed_experts_by_slot.dtype.name,
         )
 

vllm/transformers_utils/model_arch_config_convertor.py

@@ -164,6 +164,16 @@ class ModelArchConfigConvertorBase:
             num_experts = self.get_num_experts_from_block_configs()
         return num_experts
 
+    def get_num_experts_per_token(self) -> int:
+        names = [
+            "num_experts_per_tok",
+            "num_experts_per_token",
+            "top_k_experts",
+            "moe_topk",
+            "moe_top_k",
+        ]
+        return getattr_iter(self.hf_text_config, names, 0) or 0
+
     @final
     @classmethod
     def get_torch_dtype(
@@ -379,6 +389,7 @@ class ModelArchConfigConvertorBase:
             vocab_size=self.get_vocab_size(),
             total_num_kv_heads=self.get_total_num_kv_heads(),
             num_experts=self.get_num_experts(),
+            num_experts_per_token=self.get_num_experts_per_token(),
             quantization_config=self.get_quantization_config(),
             is_deepseek_mla=self.is_deepseek_mla(),
             is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal),

vllm/v1/metrics/perf.py

@@ -828,9 +828,7 @@ class BaseFfnConfigParser(Parser):
 
         # Try different naming conventions.
         args.num_experts = vllm_config.model_config.get_num_experts()
-        args.num_experts_per_tok = getattr_from_list(
-            cfg, ["num_experts_per_tok", "moe_topk"], 0
-        )
+        args.num_experts_per_tok = vllm_config.model_config.get_num_experts_per_tok()
         args.moe_intermediate_size = getattr_from_list(
             cfg, ["moe_intermediate_size", "intermediate_size"], 0
         )