diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 37542c41d90..84eeafc40e6 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -67,22 +67,6 @@ from .interfaces_base import ( logger = init_logger(__name__) -_SPECIALIZED_TEXT_GENERATION_MODELS = { - "DeepseekV32ForCausalLM": ( - "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4", - "DeepseekV32ForCausalLM", - ), -} -_SPECIALIZED_MTP_MODEL_ARCH_BY_BASE_ARCH = { - "DeepseekV32ForCausalLM": "_SpecializedDeepSeekV32MTPModel", -} -_SPECIALIZED_MTP_MODELS = { - "_SpecializedDeepSeekV32MTPModel": ( - "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4", - "DeepSeekMTP", - ), -} - _TEXT_GENERATION_MODELS = { # [Decoder-only] "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"), @@ -602,8 +586,11 @@ _SPECULATIVE_DECODING_MODELS = { } if envs.VLLM_USE_SPECIALIZED_MODELS: - _TEXT_GENERATION_MODELS.update(_SPECIALIZED_TEXT_GENERATION_MODELS) - _TEXT_GENERATION_MODELS.update(_SPECIALIZED_MTP_MODELS) + from vllm.model_executor.specialized_models import get_specialized_models + + _specialized = get_specialized_models() + _TEXT_GENERATION_MODELS.update(_specialized) + _SPECULATIVE_DECODING_MODELS.update(_specialized) _TRANSFORMERS_SUPPORTED_MODELS = { # Text generation models @@ -1082,16 +1069,6 @@ class _ModelRegistry: architecture: str, model_config: ModelConfig, ) -> str: - if architecture == "DeepSeekMTPModel" and envs.VLLM_USE_SPECIALIZED_MODELS: - base_arch = ( - "DeepseekV32ForCausalLM" - if hasattr(model_config.hf_config, "index_topk") - else None - ) - specialized_arch = _SPECIALIZED_MTP_MODEL_ARCH_BY_BASE_ARCH.get(base_arch) - if specialized_arch is not None: - return specialized_arch - if architecture in self.models: return architecture diff --git a/vllm/model_executor/specialized_models/__init__.py b/vllm/model_executor/specialized_models/__init__.py index e69de29bb2d..3d5860d0c77 100644 --- a/vllm/model_executor/specialized_models/__init__.py +++ b/vllm/model_executor/specialized_models/__init__.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Specialized model implementations. + +Each entry maps a vLLM architecture name to a (module_path, class_name) +tuple, exactly like the main model registry. When +``VLLM_USE_SPECIALIZED_MODELS=1`` the main registry merges these entries +so they take priority over the generic implementations. + +To add a new specialized model: + 1. Create a sub-package under this directory. + 2. Add the architecture -> (module, class) mapping to ``_MODELS`` below. +""" + +from __future__ import annotations + +# ── Model list ─────────────────────────────────────────────────────── +# Maps architecture name -> (fully-qualified module, class name). +# When the flag is enabled, these override the corresponding entries +# in the main registry. +_MODELS: dict[str, tuple[str, str]] = { + "DeepseekV32ForCausalLM": ( + "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4", + "DeepseekV32ForCausalLM", + ), + "DeepSeekMTPModel": ( + "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4", + "DeepSeekMTP", + ), +} + + +def get_specialized_models() -> dict[str, tuple[str, str]]: + """Return the specialized model registry.""" + return _MODELS