diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index 212f3a9a25..1155662436 100644 --- a/tests/config/test_model_arch_config.py +++ b/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": { diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index f43312a372..21b7ba4239 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/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 = [] diff --git a/vllm/config/model.py b/vllm/config/model.py index e6c3dd944d..d45f775217 100644 --- a/vllm/config/model.py +++ b/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 diff --git a/vllm/config/model_arch.py b/vllm/config/model_arch.py index 0b4744de48..2264ba2727 100644 --- a/vllm/config/model_arch.py +++ b/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.""" diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index df6805eb1f..618e82bea9 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/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, ) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 90c026cf9b..3425648df1 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/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), diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index a1dceeab46..66a600af14 100644 --- a/vllm/v1/metrics/perf.py +++ b/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 )