mtp registry

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

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
 

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