registry.py 50.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
"""
Whenever you add an architecture to this page, please also update
`tests/models/registry.py` with example HuggingFace models for it.
"""
7

8
import importlib
9
import json
10
import os
11
import pickle
12
13
import subprocess
import sys
14
import tempfile
15
from abc import ABC, abstractmethod
16
from collections.abc import Callable, Set
17
from dataclasses import asdict, dataclass, field
18
from functools import lru_cache
19
from pathlib import Path
20
from typing import TYPE_CHECKING, Any, TypeVar
21
22

import torch.nn as nn
23
import transformers
24

25
from vllm import envs
26
27
28
29
30
from vllm.config import (
    ModelConfig,
    iter_architecture_defaults,
    try_match_architecture_defaults,
)
31
from vllm.logger import init_logger
32
from vllm.logging_utils import logtime
33
from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module
34
from vllm.utils.hashing import safe_hash
35

36
37
if TYPE_CHECKING:
    from vllm.config.model import AttnTypeStr
38
    from vllm.config.pooler import SequencePoolingType, TokenPoolingType
39
40
else:
    AttnTypeStr = Any
41
42
    SequencePoolingType = Any
    TokenPoolingType = Any
43
44


45
46
47
48
49
from .interfaces import (
    has_inner_state,
    has_noops,
    is_attention_free,
    is_hybrid,
Patrick von Platen's avatar
Patrick von Platen committed
50
    requires_raw_input_tokens,
51
    supports_cross_encoding,
52
    supports_late_interaction,
53
    supports_mamba_prefix_caching,
54
55
56
57
58
59
60
    supports_multimodal,
    supports_multimodal_encoder_tp_data,
    supports_multimodal_raw_input_only,
    supports_pp,
    supports_transcription,
)
from .interfaces_base import (
61
    get_attn_type,
62
63
    get_default_seq_pooling_type,
    get_default_tok_pooling_type,
64
65
66
    is_pooling_model,
    is_text_generation_model,
)
67
68
69

logger = init_logger(__name__)

70
71
_TEXT_GENERATION_MODELS = {
    # [Decoder-only]
72
    "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"),
73
    "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"),
74
75
    "AquilaModel": ("llama", "LlamaForCausalLM"),
    "AquilaForCausalLM": ("llama", "LlamaForCausalLM"),  # AquilaChat2
Raghav Ravishankar's avatar
Raghav Ravishankar committed
76
    "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),
77
    "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
78
    "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"),
79
80
81
82
    # baichuan-7b, upper case 'C' in the class name
    "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
    # baichuan-13b, lower case 'c' in the class name
    "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
83
    "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
ant-yy's avatar
ant-yy committed
84
    "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"),
Jiangyun Zhu's avatar
Jiangyun Zhu committed
85
    "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"),
Yu Chin Fabian Lim's avatar
Yu Chin Fabian Lim committed
86
    "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
87
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
88
    "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
89
    "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
90
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
91
    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
92
    "CwmForCausalLM": ("llama", "LlamaForCausalLM"),
93
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
94
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
95
    "DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"),
96
    "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
97
    "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
98
    "DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
99
    "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"),
100
    "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),
101
    "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
102
    "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
103
    "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"),
Kyungmin Lee's avatar
Kyungmin Lee committed
104
    "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"),
105
    "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
106
107
108
    "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
    "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
    "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"),
109
    "FlexOlmoForCausalLM": ("flex_olmo", "FlexOlmoForCausalLM"),
110
111
    "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
    "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
112
    "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
113
    "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
114
    "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
115
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
116
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
117
    "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
118
    "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
Jee Jee Li's avatar
Jee Jee Li committed
119
    "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
120
    "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
121
122
123
124
125
126
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
127
128
    "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),  # noqa: E501
    "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),  # noqa: E501
129
    "GritLM": ("gritlm", "GritLM"),
Bijaya Dangol's avatar
Bijaya Dangol committed
130
131
    "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),
    "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),
132
133
    "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
    "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
134
    "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),
135
136
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
137
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
138
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
139
140
    "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),
    "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
141
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
142
    "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
143
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
144
    "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),  # noqa: E501
145
    "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
Paul Pak's avatar
Paul Pak committed
146
    "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
147
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
148
    "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
149
150
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
XuruiYang's avatar
XuruiYang committed
151
    "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
152
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
153
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
154
155
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
156
157
158
    "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
    "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
    "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
159
    "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
160
    "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
161
    "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
162
163
164
165
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
166
    "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
167
    "MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),
168
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
Luis Vega's avatar
Luis Vega committed
169
    "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
170
    "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
171
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
172
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
173
    "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
174
    "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),
175
176
177
    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
178
    "OuroForCausalLM": ("ouro", "OuroForCausalLM"),
179
    "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
180
    "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
181
    "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
182
183
184
185
    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
186
    "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
187
    "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
188
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
189
190
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
191
192
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
193
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
194
    "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
Li Xie's avatar
Li Xie committed
195
    "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
Song's avatar
Song committed
196
    "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
csy0225's avatar
csy0225 committed
197
    "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),
198
199
200
201
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
202
    "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
203
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
204
    "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
205
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
206
    "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
207
208
209
}

_EMBEDDING_MODELS = {
210
    # [Text-only]
211
    "BertModel": ("bert", "BertEmbeddingModel"),
212
    "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
213
    "HF_ColBERT": ("colbert", "ColBERTModel"),
214
215
    "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
    "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
216
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
217
    "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
218
    "Gemma3TextModel": ("gemma3", "Gemma3Model"),
219
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
220
    "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
221
    "GritLM": ("gritlm", "GritLM"),
222
223
    "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
    "GteNewModel": ("bert_with_rope", "GteNewModel"),
224
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
225
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
226
    "LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"),
227
    "LlamaModel": ("llama", "LlamaForCausalLM"),
228
229
    **{
        # Multiple models share the same architecture, so we include them all
230
231
        k: (mod, arch)
        for k, (mod, arch) in _TEXT_GENERATION_MODELS.items()
232
233
        if arch == "LlamaForCausalLM"
    },
234
    "MistralModel": ("llama", "LlamaForCausalLM"),
235
    "ModernBertModel": ("modernbert", "ModernBertModel"),
236
    "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
237
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
238
    "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
239
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
240
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
241
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
242
243
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
244
    "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
245
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
chengchengpei's avatar
chengchengpei committed
246
247
248
249
    "VoyageQwen3BidirectionalEmbedModel": (
        "voyage",
        "VoyageQwen3BidirectionalEmbedModel",
    ),
250
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
251
    "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
252
    # [Multimodal]
253
    "CLIPModel": ("clip", "CLIPEmbeddingModel"),
254
    "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
255
256
257
    "LlavaNextForConditionalGeneration": (
        "llava_next",
        "LlavaNextForConditionalGeneration",
258
    ),
259
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
260
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
261
262
    "ColQwen3": ("colqwen3", "ColQwen3Model"),
    "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
263
    "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"),
264
    "SiglipModel": ("siglip", "SiglipEmbeddingModel"),
265
266
267
268
    "LlamaNemotronVLModel": (
        "nemotron_vl",
        "LlamaNemotronVLForEmbedding",
    ),
269
270
    # Technically Terratorch models work on images, both in
    # input and output. I am adding it here because it piggy-backs on embedding
271
    # models for the time being.
272
273
    "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"),
    "Terratorch": ("terratorch", "Terratorch"),
274
275
}

276
277
_CROSS_ENCODER_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
278
    "BertForTokenClassification": ("bert", "BertForTokenClassification"),
279
280
281
282
    "GteNewForSequenceClassification": (
        "bert_with_rope",
        "GteNewForSequenceClassification",
    ),
283
284
285
286
    "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"),
    "LlamaBidirectionalForSequenceClassification": (
        "llama",
        "LlamaBidirectionalForSequenceClassification",
287
288
289
290
    ),
    "LlamaNemotronVLForSequenceClassification": (
        "nemotron_vl",
        "LlamaNemotronVLForSequenceClassification",
291
    ),
292
293
294
295
    "ModernBertForSequenceClassification": (
        "modernbert",
        "ModernBertForSequenceClassification",
    ),
296
297
298
299
    "ModernBertForTokenClassification": (
        "modernbert",
        "ModernBertForTokenClassification",
    ),
300
301
302
303
304
    "RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"),
    "XLMRobertaForSequenceClassification": (
        "roberta",
        "RobertaForSequenceClassification",
    ),
305
306
}

307
_MULTIMODAL_MODELS = {
308
    # [Decoder-only]
309
    "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"),
310
311
312
313
    "AudioFlamingo3ForConditionalGeneration": (
        "audioflamingo3",
        "AudioFlamingo3ForConditionalGeneration",
    ),
314
315
316
317
    "MusicFlamingoForConditionalGeneration": (
        "musicflamingo",
        "MusicFlamingoForConditionalGeneration",
    ),
318
319
320
    "AyaVisionForConditionalGeneration": (
        "aya_vision",
        "AyaVisionForConditionalGeneration",
321
    ),
322
    "BagelForConditionalGeneration": ("bagel", "BagelForConditionalGeneration"),
323
    "BeeForConditionalGeneration": ("bee", "BeeForConditionalGeneration"),
324
    "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
325
326
327
    "ChameleonForConditionalGeneration": (
        "chameleon",
        "ChameleonForConditionalGeneration",
328
    ),
329
330
331
    "Cohere2VisionForConditionalGeneration": (
        "cohere2_vision",
        "Cohere2VisionForConditionalGeneration",
332
    ),
333
    "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
334
    "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
RED's avatar
RED committed
335
    "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
Roger Wang's avatar
Roger Wang committed
336
    "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
337
338
339
340
    "Eagle2_5_VLForConditionalGeneration": (
        "eagle2_5_vl",
        "Eagle2_5_VLForConditionalGeneration",
    ),
341
342
343
    "Ernie4_5_VLMoeForConditionalGeneration": (
        "ernie45_vl",
        "Ernie4_5_VLMoeForConditionalGeneration",
344
    ),
345
346
347
348
    "FireRedASR2ForConditionalGeneration": (
        "fireredasr2",
        "FireRedASR2ForConditionalGeneration",
    ),
349
    "FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"),  # noqa: E501
350
351
352
353
    "FunAudioChatForConditionalGeneration": (
        "funaudiochat",
        "FunAudioChatForConditionalGeneration",
    ),
354
    "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
355
    "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"),  # noqa: E501
356
357
358
    "Gemma3nForConditionalGeneration": (
        "gemma3n_mm",
        "Gemma3nForConditionalGeneration",
359
    ),
360
    "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"),
361
    "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
362
363
364
    "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"),
    "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"),
    "GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"),  # noqa: E501
365
366
367
    "GraniteSpeechForConditionalGeneration": (
        "granite_speech",
        "GraniteSpeechForConditionalGeneration",
368
    ),
369
    "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
370
371
372
373
    "HunYuanVLForConditionalGeneration": (
        "hunyuan_vision",
        "HunYuanVLForConditionalGeneration",
    ),
ltd0924's avatar
ltd0924 committed
374
    "StepVLForConditionalGeneration": ("step_vl", "StepVLForConditionalGeneration"),
375
    "InternVLChatModel": ("internvl", "InternVLChatModel"),
376
    "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
Zero's avatar
Zero committed
377
378
379
380
    "OpenCUAForConditionalGeneration": (
        "opencua",
        "OpenCUAForConditionalGeneration",
    ),
381
382
383
    "InternS1ForConditionalGeneration": (
        "interns1",
        "InternS1ForConditionalGeneration",
384
    ),
385
386
387
    "InternVLForConditionalGeneration": (
        "interns1",
        "InternS1ForConditionalGeneration",
388
    ),
zxy's avatar
zxy committed
389
390
391
392
    "InternS1ProForConditionalGeneration": (
        "interns1_pro",
        "InternS1ProForConditionalGeneration",
    ),
393
394
395
396
    "Idefics3ForConditionalGeneration": (
        "idefics3",
        "Idefics3ForConditionalGeneration",
    ),
oscardev256's avatar
oscardev256 committed
397
    "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"),
398
    "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"),  # noqa: E501
399
    "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"),
400
    "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),
401
402
403
    "KeyeVL1_5ForConditionalGeneration": (
        "keye_vl1_5",
        "KeyeVL1_5ForConditionalGeneration",
404
    ),
405
    "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),
406
    "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),  # noqa: E501
Roger Wang's avatar
Roger Wang committed
407
    "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"),  # noqa: E501
408
409
410
411
    "LightOnOCRForConditionalGeneration": (
        "lightonocr",
        "LightOnOCRForConditionalGeneration",
    ),
412
    "Lfm2VlForConditionalGeneration": ("lfm2_vl", "Lfm2VLForConditionalGeneration"),
413
    "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"),
414
    "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"),  # noqa: E501
415
    "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
416
417
418
    "LlavaNextForConditionalGeneration": (
        "llava_next",
        "LlavaNextForConditionalGeneration",
419
    ),
420
421
422
    "LlavaNextVideoForConditionalGeneration": (
        "llava_next_video",
        "LlavaNextVideoForConditionalGeneration",
423
    ),
424
425
426
    "LlavaOnevisionForConditionalGeneration": (
        "llava_onevision",
        "LlavaOnevisionForConditionalGeneration",
427
    ),
428
    "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),  # noqa: E501
429
    "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),
430
431
432
    "MiniMaxVL01ForConditionalGeneration": (
        "minimax_vl_01",
        "MiniMaxVL01ForConditionalGeneration",
433
    ),
434
    "MiniCPMO": ("minicpmo", "MiniCPMO"),
435
    "MiniCPMV": ("minicpmv", "MiniCPMV"),
436
437
438
    "Mistral3ForConditionalGeneration": (
        "mistral3",
        "Mistral3ForConditionalGeneration",
439
    ),
440
    "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
441
    "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),
442
    "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
443
444
445
446
    "OpenPanguVLForConditionalGeneration": (
        "openpangu_vl",
        "OpenPanguVLForConditionalGeneration",
    ),
447
    "Ovis": ("ovis", "Ovis"),
448
    "Ovis2_5": ("ovis2_5", "Ovis2_5"),
449
450
    "Ovis2_6ForCausalLM": ("ovis2_5", "Ovis2_5"),
    "Ovis2_6_MoeForCausalLM": ("ovis2_5", "Ovis2_5"),
451
452
453
454
    "PaddleOCRVLForConditionalGeneration": (
        "paddleocr_vl",
        "PaddleOCRVLForConditionalGeneration",
    ),
455
456
457
458
    "PaliGemmaForConditionalGeneration": (
        "paligemma",
        "PaliGemmaForConditionalGeneration",
    ),
459
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
460
    "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),
461
    "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"),  # noqa: E501
462
    "QwenVLForConditionalGeneration": ("qwen_vl", "QwenVLForConditionalGeneration"),  # noqa: E501
463
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
464
465
466
    "Qwen2_5_VLForConditionalGeneration": (
        "qwen2_5_vl",
        "Qwen2_5_VLForConditionalGeneration",
467
    ),
468
469
470
    "Qwen2AudioForConditionalGeneration": (
        "qwen2_audio",
        "Qwen2AudioForConditionalGeneration",
471
    ),
472
473
474
    "Qwen2_5OmniModel": (
        "qwen2_5_omni_thinker",
        "Qwen2_5OmniThinkerForConditionalGeneration",
475
    ),
476
477
478
    "Qwen2_5OmniForConditionalGeneration": (
        "qwen2_5_omni_thinker",
        "Qwen2_5OmniThinkerForConditionalGeneration",
479
    ),
480
481
482
483
    "Qwen3OmniMoeForConditionalGeneration": (
        "qwen3_omni_moe_thinker",
        "Qwen3OmniMoeThinkerForConditionalGeneration",
    ),
Roger Wang's avatar
Roger Wang committed
484
485
486
487
    "Qwen3ASRForConditionalGeneration": (
        "qwen3_asr",
        "Qwen3ASRForConditionalGeneration",
    ),
488
489
490
491
    "Qwen3ASRRealtimeGeneration": (
        "qwen3_asr_realtime",
        "Qwen3ASRRealtimeGeneration",
    ),
492
    "Qwen3VLForConditionalGeneration": ("qwen3_vl", "Qwen3VLForConditionalGeneration"),  # noqa: E501
493
494
495
    "Qwen3VLMoeForConditionalGeneration": (
        "qwen3_vl_moe",
        "Qwen3VLMoeForConditionalGeneration",
496
    ),
497
498
499
500
501
502
503
504
    "Qwen3_5ForConditionalGeneration": (
        "qwen3_5",
        "Qwen3_5ForConditionalGeneration",
    ),
    "Qwen3_5MoeForConditionalGeneration": (
        "qwen3_5",
        "Qwen3_5MoeForConditionalGeneration",
    ),
505
    "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"),
Song's avatar
Song committed
506
    "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"),  # noqa: E501
汪志鹏's avatar
汪志鹏 committed
507
    "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"),  # noqa: E501
508
509
510
    "Tarsier2ForConditionalGeneration": (
        "qwen2_vl",
        "Tarsier2ForConditionalGeneration",
511
    ),
512
    "UltravoxModel": ("ultravox", "UltravoxModel"),
Patrick von Platen's avatar
Patrick von Platen committed
513
    "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"),  # noqa: E501
514
    "VoxtralRealtimeGeneration": ("voxtral_realtime", "VoxtralRealtimeGeneration"),  # noqa: E501
515
    # [Encoder-decoder]
516
517
518
519
    "NemotronParseForConditionalGeneration": (
        "nemotron_parse",
        "NemotronParseForConditionalGeneration",
    ),
520
    "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"),  # noqa: E501
521
}
522
523

_SPECULATIVE_DECODING_MODELS = {
524
    "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"),
525
    "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
526
    "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
zhiweiz's avatar
zhiweiz committed
527
    "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
528
    "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
529
    "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
530
    "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
531
    "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
532
    "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
533
534
535
536
    "EagleMistralLarge3ForCausalLM": (
        "mistral_large_3_eagle",
        "EagleMistralLarge3ForCausalLM",
    ),
537
    "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
538
    "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
539
    "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
Kyungmin Lee's avatar
Kyungmin Lee committed
540
    "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
541
    "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"),
XuruiYang's avatar
XuruiYang committed
542
    "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
543
    "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
544
    "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
545
    "GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"),
546
    "MedusaModel": ("medusa", "Medusa"),
547
    "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
548
    "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),
csy0225's avatar
csy0225 committed
549
    "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),
550
551
    "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
    "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
552
553
554
    # Temporarily disabled.
    # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
    # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
555
}
556

557
_TRANSFORMERS_SUPPORTED_MODELS = {
558
559
560
    # Text generation models
    "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
    # Multimodal models
561
562
563
564
    "Emu3ForConditionalGeneration": (
        "transformers",
        "TransformersMultiModalForCausalLM",
    ),
565
566
567
}

_TRANSFORMERS_BACKEND_MODELS = {
568
    # Text generation models
569
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
    "TransformersMoEForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
    # Multimodal models
    "TransformersMultiModalForCausalLM": (
        "transformers",
        "TransformersMultiModalForCausalLM",
    ),
    "TransformersMultiModalMoEForCausalLM": (
        "transformers",
        "TransformersMultiModalMoEForCausalLM",
    ),
    # Embedding models
    "TransformersEmbeddingModel": ("transformers", "TransformersEmbeddingModel"),
    "TransformersMoEEmbeddingModel": ("transformers", "TransformersMoEEmbeddingModel"),
    "TransformersMultiModalEmbeddingModel": (
        "transformers",
        "TransformersMultiModalEmbeddingModel",
    ),
    # Sequence classification models
588
    "TransformersForSequenceClassification": (
589
        "transformers",
590
        "TransformersForSequenceClassification",
591
    ),
592
    "TransformersMoEForSequenceClassification": (
593
        "transformers",
594
        "TransformersMoEForSequenceClassification",
595
    ),
596
597
598
    "TransformersMultiModalForSequenceClassification": (
        "transformers",
        "TransformersMultiModalForSequenceClassification",
599
    ),
600
}
601

602
_VLLM_MODELS = {
603
    **_TEXT_GENERATION_MODELS,
604
    **_EMBEDDING_MODELS,
605
    **_CROSS_ENCODER_MODELS,
606
    **_MULTIMODAL_MODELS,
607
    **_SPECULATIVE_DECODING_MODELS,
608
609
    **_TRANSFORMERS_SUPPORTED_MODELS,
    **_TRANSFORMERS_BACKEND_MODELS,
610
611
}

612
613
614
615
# This variable is used as the args for subprocess.run(). We
# can modify  this variable to alter the args if needed. e.g.
# when we use par format to pack things together, sys.executable
# might not be the target we want to run.
616
_SUBPROCESS_COMMAND = [sys.executable, "-m", "vllm.model_executor.models.registry"]
617

618
_PREVIOUSLY_SUPPORTED_MODELS = {
619
    "MotifForCausalLM": "0.10.2",
620
    "Phi3SmallForCausalLM": "0.9.2",
621
    "Phi4FlashForCausalLM": "0.10.2",
622
    "Phi4MultimodalForCausalLM": "0.12.0",
623
624
625
626
627
628
629
630
631
    # encoder-decoder models except whisper
    # have been removed for V0 deprecation.
    "BartModel": "0.10.2",
    "BartForConditionalGeneration": "0.10.2",
    "DonutForConditionalGeneration": "0.10.2",
    "Florence2ForConditionalGeneration": "0.10.2",
    "MBartForConditionalGeneration": "0.10.2",
    "MllamaForConditionalGeneration": "0.10.2",
}
632

633

634
635
@dataclass(frozen=True)
class _ModelInfo:
636
    architecture: str
637
    is_text_generation_model: bool
638
    is_pooling_model: bool
639
    attn_type: AttnTypeStr
640
641
    default_seq_pooling_type: SequencePoolingType
    default_tok_pooling_type: TokenPoolingType
642
    supports_cross_encoding: bool
643
    supports_late_interaction: bool
644
    supports_multimodal: bool
645
    supports_multimodal_raw_input_only: bool
Patrick von Platen's avatar
Patrick von Platen committed
646
    requires_raw_input_tokens: bool
647
    supports_multimodal_encoder_tp_data: bool
648
    supports_pp: bool
649
650
    has_inner_state: bool
    is_attention_free: bool
651
    is_hybrid: bool
652
    has_noops: bool
653
    supports_mamba_prefix_caching: bool
654
    supports_transcription: bool
655
    supports_transcription_only: bool
656
657

    @staticmethod
658
    def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
659
        return _ModelInfo(
660
            architecture=model.__name__,
661
            is_text_generation_model=is_text_generation_model(model),
662
            is_pooling_model=is_pooling_model(model),
663
664
            default_seq_pooling_type=get_default_seq_pooling_type(model),
            default_tok_pooling_type=get_default_tok_pooling_type(model),
665
            attn_type=get_attn_type(model),
666
            supports_cross_encoding=supports_cross_encoding(model),
667
            supports_late_interaction=supports_late_interaction(model),
668
            supports_multimodal=supports_multimodal(model),
669
670
671
            supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
                model
            ),
Patrick von Platen's avatar
Patrick von Platen committed
672
            requires_raw_input_tokens=requires_raw_input_tokens(model),
673
674
675
            supports_multimodal_encoder_tp_data=supports_multimodal_encoder_tp_data(
                model
            ),
676
            supports_pp=supports_pp(model),
677
678
            has_inner_state=has_inner_state(model),
            is_attention_free=is_attention_free(model),
679
            is_hybrid=is_hybrid(model),
680
            supports_mamba_prefix_caching=supports_mamba_prefix_caching(model),
681
            supports_transcription=supports_transcription(model),
682
683
684
            supports_transcription_only=(
                supports_transcription(model) and model.supports_transcription_only
            ),
685
            has_noops=has_noops(model),
686
        )
687
688


689
690
691
692
class _BaseRegisteredModel(ABC):
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
693

694
    @abstractmethod
695
    def load_model_cls(self) -> type[nn.Module]:
696
        raise NotImplementedError
697
698


699
700
701
702
703
704
705
@dataclass(frozen=True)
class _RegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has already been imported in the main process.
    """

    interfaces: _ModelInfo
706
    model_cls: type[nn.Module]
707
708

    @staticmethod
709
    def from_model_cls(model_cls: type[nn.Module]):
710
711
712
713
714
715
716
717
        return _RegisteredModel(
            interfaces=_ModelInfo.from_model_cls(model_cls),
            model_cls=model_cls,
        )

    def inspect_model_cls(self) -> _ModelInfo:
        return self.interfaces

718
    def load_model_cls(self) -> type[nn.Module]:
719
720
721
722
723
724
725
726
        return self.model_cls


@dataclass(frozen=True)
class _LazyRegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has not been imported in the main process.
    """
727

728
729
730
    module_name: str
    class_name: str

731
732
733
734
735
736
737
738
    @staticmethod
    def _get_cache_dir() -> Path:
        return Path(envs.VLLM_CACHE_ROOT) / "modelinfos"

    def _get_cache_filename(self) -> str:
        cls_name = f"{self.module_name}-{self.class_name}".replace(".", "-")
        return f"{cls_name}.json"

739
    def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None:
740
741
        try:
            try:
742
                modelinfo_path = self._get_cache_dir() / self._get_cache_filename()
743
744
745
                with open(modelinfo_path, encoding="utf-8") as file:
                    mi_dict = json.load(file)
            except FileNotFoundError:
746
                logger.debug(
747
                    "Cached model info file for class %s.%s not found",
748
749
750
                    self.module_name,
                    self.class_name,
                )
751
752
753
                return None

            if mi_dict["hash"] != module_hash:
754
                logger.debug(
755
                    "Cached model info file for class %s.%s is stale",
756
757
758
                    self.module_name,
                    self.class_name,
                )
759
760
761
762
763
                return None

            # file not changed, use cached _ModelInfo properties
            return _ModelInfo(**mi_dict["modelinfo"])
        except Exception:
764
            logger.debug(
765
                "Cached model info for class %s.%s error. ",
766
767
768
                self.module_name,
                self.class_name,
            )
769
770
            return None

771
    def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None:
772
773
        """save dictionary json file to cache"""
        from vllm.model_executor.model_loader.weight_utils import atomic_writer
774

775
776
777
778
779
780
781
782
        try:
            modelinfo_dict = {
                "hash": module_hash,
                "modelinfo": asdict(mi),
            }
            cache_dir = self._get_cache_dir()
            cache_dir.mkdir(parents=True, exist_ok=True)
            modelinfo_path = cache_dir / self._get_cache_filename()
783
            with atomic_writer(modelinfo_path, encoding="utf-8") as f:
784
785
786
787
788
                json.dump(modelinfo_dict, f, indent=2)
        except Exception:
            logger.exception("Error saving model info cache.")

    @logtime(logger=logger, msg="Registry inspect model class")
789
    def inspect_model_cls(self) -> _ModelInfo:
790
        model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py"
791
        module_hash = None
792

793
794
        if model_path.exists():
            with open(model_path, "rb") as f:
795
                module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest()
796
797
798

            mi = self._load_modelinfo_from_cache(module_hash)
            if mi is not None:
799
                logger.debug(
800
                    "Loaded model info for class %s.%s from cache",
801
802
803
                    self.module_name,
                    self.class_name,
                )
804
805
                return mi
            else:
806
                logger.debug(
807
                    "Cache model info for class %s.%s miss. Loading model instead.",
808
809
810
                    self.module_name,
                    self.class_name,
                )
811
812
813

        # Performed in another process to avoid initializing CUDA
        mi = _run_in_subprocess(
814
815
816
817
818
            lambda: _ModelInfo.from_model_cls(self.load_model_cls())
        )
        logger.debug(
            "Loaded model info for class %s.%s", self.module_name, self.class_name
        )
819
820

        # save cache file
821
822
        if module_hash is not None:
            self._save_modelinfo_to_cache(mi, module_hash)
823
824

        return mi
825

826
    def load_model_cls(self) -> type[nn.Module]:
827
828
829
830
831
832
833
834
        mod = importlib.import_module(self.module_name)
        return getattr(mod, self.class_name)


@lru_cache(maxsize=128)
def _try_load_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
835
) -> type[nn.Module] | None:
836
    from vllm.platforms import current_platform
837

838
    current_platform.verify_model_arch(model_arch)
839
840
841
    try:
        return model.load_model_cls()
    except Exception:
842
        logger.exception("Error in loading model architecture '%s'", model_arch)
843
        return None
844
845


846
847
848
849
@lru_cache(maxsize=128)
def _try_inspect_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
850
) -> _ModelInfo | None:
851
852
853
    try:
        return model.inspect_model_cls()
    except Exception:
854
        logger.exception("Error in inspecting model architecture '%s'", model_arch)
855
        return None
856
857


858
859
860
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
861
    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
862

863
    def get_supported_archs(self) -> Set[str]:
864
        return self.models.keys()
865

866
867
868
    def register_model(
        self,
        model_arch: str,
869
        model_cls: type[nn.Module] | str,
870
    ) -> None:
871
872
873
        """
        Register an external model to be used in vLLM.

874
        `model_cls` can be either:
875

876
        - A [`torch.nn.Module`][] class directly referencing the model.
877
        - A string in the format `<module>:<class>` which can be used to
878
879
          lazily import the model. This is useful to avoid initializing CUDA
          when importing the model and thus the related error
880
          `RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
881
        """
882
883
884
885
        if not isinstance(model_arch, str):
            msg = f"`model_arch` should be a string, not a {type(model_arch)}"
            raise TypeError(msg)

886
        if model_arch in self.models:
887
888
            logger.warning(
                "Model architecture %s is already registered, and will be "
889
890
891
892
                "overwritten by the new model class %s.",
                model_arch,
                model_cls,
            )
893
894
895
896
897
898

        if isinstance(model_cls, str):
            split_str = model_cls.split(":")
            if len(split_str) != 2:
                msg = "Expected a string in the format `<module>:<class>`"
                raise ValueError(msg)
899

900
            model = _LazyRegisteredModel(*split_str)
901
        elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module):
902
            model = _RegisteredModel.from_model_cls(model_cls)
903
        else:
904
905
906
907
            msg = (
                "`model_cls` should be a string or PyTorch model class, "
                f"not a {type(model_arch)}"
            )
908
            raise TypeError(msg)
909

910
        self.models[model_arch] = model
911

912
    def _raise_for_unsupported(self, architectures: list[str]):
913
        all_supported_archs = self.get_supported_archs()
914

915
916
917
        if any(arch in all_supported_archs for arch in architectures):
            raise ValueError(
                f"Model architectures {architectures} failed "
918
919
                "to be inspected. Please check the logs for more details."
            )
920

921
922
923
924
925
926
927
928
        for arch in architectures:
            if arch in _PREVIOUSLY_SUPPORTED_MODELS:
                previous_version = _PREVIOUSLY_SUPPORTED_MODELS[arch]

                raise ValueError(
                    f"Model architecture {arch} was supported in vLLM until "
                    f"v{previous_version}, and is not supported anymore. "
                    "Please use an older version of vLLM if you want to "
929
930
                    "use this model architecture."
                )
931

932
933
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
934
935
            f"Supported architectures: {all_supported_archs}"
        )
936

937
    def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:
938
939
        if model_arch not in self.models:
            return None
940

941
        return _try_load_model_cls(model_arch, self.models[model_arch])
942

943
    def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None:
944
945
        if model_arch not in self.models:
            return None
946

947
948
949
950
951
952
        return _try_inspect_model_cls(model_arch, self.models[model_arch])

    def _try_resolve_transformers(
        self,
        architecture: str,
        model_config: ModelConfig,
953
    ) -> str | None:
954
955
956
        if architecture in _TRANSFORMERS_BACKEND_MODELS:
            return architecture

957
958
959
        auto_map: dict[str, str] = (
            getattr(model_config.hf_config, "auto_map", None) or dict()
        )
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975

        # Make sure that config class is always initialized before model class,
        # otherwise the model class won't be able to access the config class,
        # the expected auto_map should have correct order like:
        # "auto_map": {
        #     "AutoConfig": "<your-repo-name>--<config-name>",
        #     "AutoModel": "<your-repo-name>--<config-name>",
        #     "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
        # },
        for prefix in ("AutoConfig", "AutoModel"):
            for name, module in auto_map.items():
                if name.startswith(prefix):
                    try_get_class_from_dynamic_module(
                        module,
                        model_config.model,
                        revision=model_config.revision,
976
                        trust_remote_code=model_config.trust_remote_code,
977
978
979
980
981
982
983
984
985
986
987
988
                        warn_on_fail=False,
                    )

        model_module = getattr(transformers, architecture, None)

        if model_module is None:
            for name, module in auto_map.items():
                if name.startswith("AutoModel"):
                    model_module = try_get_class_from_dynamic_module(
                        module,
                        model_config.model,
                        revision=model_config.revision,
989
                        trust_remote_code=model_config.trust_remote_code,
990
991
992
993
994
                        warn_on_fail=True,
                    )
                    if model_module is not None:
                        break
            else:
995
                if model_config.model_impl != "transformers":
996
997
998
999
1000
1001
1002
                    return None

                raise ValueError(
                    f"Cannot find model module. {architecture!r} is not a "
                    "registered model in the Transformers library (only "
                    "relevant if the model is meant to be in Transformers) "
                    "and 'AutoModel' is not present in the model config's "
1003
1004
                    "'auto_map' (relevant if the model is custom)."
                )
1005
1006

        if not model_module.is_backend_compatible():
1007
            if model_config.model_impl != "transformers":
1008
                return None
1009

1010
1011
            raise ValueError(
                f"The Transformers implementation of {architecture!r} "
1012
1013
                "is not compatible with vLLM."
            )
1014

1015
        return model_config._get_transformers_backend_cls()
1016

1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
    def _normalize_arch(
        self,
        architecture: str,
        model_config: ModelConfig,
    ) -> str:
        if architecture in self.models:
            return architecture

        # This may be called in order to resolve runner_type and convert_type
        # in the first place, in which case we consider the default match
        match = try_match_architecture_defaults(
            architecture,
            runner_type=getattr(model_config, "runner_type", None),
            convert_type=getattr(model_config, "convert_type", None),
        )
        if match:
            suffix, _ = match

            # Get the name of the base model to convert
            for repl_suffix, _ in iter_architecture_defaults():
                base_arch = architecture.replace(suffix, repl_suffix)
                if base_arch in self.models:
                    return base_arch

        return architecture
1042

1043
1044
    def inspect_model_cls(
        self,
1045
        architectures: str | list[str],
1046
        model_config: ModelConfig,
1047
    ) -> tuple[_ModelInfo, str]:
1048
1049
        if isinstance(architectures, str):
            architectures = [architectures]
1050
1051
        if not architectures:
            raise ValueError("No model architectures are specified")
1052
1053

        # Require transformers impl
1054
        if model_config.model_impl == "transformers":
1055
            arch = self._try_resolve_transformers(architectures[0], model_config)
1056
1057
1058
1059
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)
1060
        elif model_config.model_impl == "terratorch":
1061
1062
            model_info = self._try_inspect_model_cls("Terratorch")
            return (model_info, "Terratorch")
1063

1064
        # Fallback to transformers impl (after resolving convert_type)
1065
1066
1067
1068
1069
1070
        if (
            all(arch not in self.models for arch in architectures)
            and model_config.model_impl == "auto"
            and getattr(model_config, "convert_type", "none") == "none"
        ):
            arch = self._try_resolve_transformers(architectures[0], model_config)
1071
1072
1073
1074
1075
1076
1077
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)

        for arch in architectures:
            normalized_arch = self._normalize_arch(arch, model_config)
1078
            model_info = self._try_inspect_model_cls(normalized_arch)
1079
            if model_info is not None:
1080
                return (model_info, arch)
1081

1082
        # Fallback to transformers impl (before resolving runner_type)
1083
1084
1085
1086
1087
        if (
            all(arch not in self.models for arch in architectures)
            and model_config.model_impl == "auto"
        ):
            arch = self._try_resolve_transformers(architectures[0], model_config)
1088
1089
1090
1091
1092
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)

1093
        return self._raise_for_unsupported(architectures)
1094

1095
1096
    def resolve_model_cls(
        self,
1097
        architectures: str | list[str],
1098
        model_config: ModelConfig,
1099
    ) -> tuple[type[nn.Module], str]:
1100
1101
        if isinstance(architectures, str):
            architectures = [architectures]
1102
1103
        if not architectures:
            raise ValueError("No model architectures are specified")
1104
1105

        # Require transformers impl
1106
        if model_config.model_impl == "transformers":
1107
            arch = self._try_resolve_transformers(architectures[0], model_config)
1108
1109
1110
1111
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)
1112
        elif model_config.model_impl == "terratorch":
1113
1114
1115
1116
            arch = "Terratorch"
            model_cls = self._try_load_model_cls(arch)
            if model_cls is not None:
                return (model_cls, arch)
1117

1118
        # Fallback to transformers impl (after resolving convert_type)
1119
1120
1121
1122
1123
1124
        if (
            all(arch not in self.models for arch in architectures)
            and model_config.model_impl == "auto"
            and getattr(model_config, "convert_type", "none") == "none"
        ):
            arch = self._try_resolve_transformers(architectures[0], model_config)
1125
1126
1127
1128
1129
1130
1131
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)

        for arch in architectures:
            normalized_arch = self._normalize_arch(arch, model_config)
1132
            model_cls = self._try_load_model_cls(normalized_arch)
1133
1134
            if model_cls is not None:
                return (model_cls, arch)
1135

1136
        # Fallback to transformers impl (before resolving runner_type)
1137
1138
1139
1140
1141
        if (
            all(arch not in self.models for arch in architectures)
            and model_config.model_impl == "auto"
        ):
            arch = self._try_resolve_transformers(architectures[0], model_config)
1142
1143
1144
1145
1146
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)

1147
        return self._raise_for_unsupported(architectures)
1148

1149
1150
    def is_text_generation_model(
        self,
1151
        architectures: str | list[str],
1152
        model_config: ModelConfig,
1153
    ) -> bool:
1154
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1155
        return model_cls.is_text_generation_model
1156

1157
    def is_pooling_model(
1158
        self,
1159
        architectures: str | list[str],
1160
        model_config: ModelConfig,
1161
    ) -> bool:
1162
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1163
        return model_cls.is_pooling_model
1164

1165
1166
    def is_cross_encoder_model(
        self,
1167
        architectures: str | list[str],
1168
        model_config: ModelConfig,
1169
    ) -> bool:
1170
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1171
        return model_cls.supports_cross_encoding
1172

1173
1174
    def is_multimodal_model(
        self,
1175
        architectures: str | list[str],
1176
        model_config: ModelConfig,
1177
    ) -> bool:
1178
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1179
        return model_cls.supports_multimodal
1180

1181
    def is_multimodal_raw_input_only_model(
1182
        self,
1183
        architectures: str | list[str],
1184
        model_config: ModelConfig,
1185
    ) -> bool:
1186
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1187
        return model_cls.supports_multimodal_raw_input_only
1188

1189
1190
    def is_pp_supported_model(
        self,
1191
        architectures: str | list[str],
1192
        model_config: ModelConfig,
1193
    ) -> bool:
1194
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1195
        return model_cls.supports_pp
1196

1197
1198
    def model_has_inner_state(
        self,
1199
        architectures: str | list[str],
1200
        model_config: ModelConfig,
1201
    ) -> bool:
1202
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1203
        return model_cls.has_inner_state
1204

1205
1206
    def is_attention_free_model(
        self,
1207
        architectures: str | list[str],
1208
        model_config: ModelConfig,
1209
    ) -> bool:
1210
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1211
        return model_cls.is_attention_free
1212

1213
1214
    def is_hybrid_model(
        self,
1215
        architectures: str | list[str],
1216
        model_config: ModelConfig,
1217
    ) -> bool:
1218
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1219
1220
        return model_cls.is_hybrid

1221
1222
    def is_noops_model(
        self,
1223
        architectures: str | list[str],
1224
        model_config: ModelConfig,
1225
    ) -> bool:
1226
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1227
1228
        return model_cls.has_noops

1229
1230
    def is_transcription_model(
        self,
1231
        architectures: str | list[str],
1232
        model_config: ModelConfig,
1233
    ) -> bool:
1234
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1235
1236
        return model_cls.supports_transcription

1237
1238
    def is_transcription_only_model(
        self,
1239
        architectures: str | list[str],
1240
        model_config: ModelConfig,
1241
    ) -> bool:
1242
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1243
1244
        return model_cls.supports_transcription_only

1245

1246
1247
1248
1249
1250
1251
1252
1253
1254
ModelRegistry = _ModelRegistry(
    {
        model_arch: _LazyRegisteredModel(
            module_name=f"vllm.model_executor.models.{mod_relname}",
            class_name=cls_name,
        )
        for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
    }
)
1255
1256
1257
1258
1259

_T = TypeVar("_T")


def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
1260
1261
1262
1263
1264
    # NOTE: We use a temporary directory instead of a temporary file to avoid
    # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
    with tempfile.TemporaryDirectory() as tempdir:
        output_filepath = os.path.join(tempdir, "registry_output.tmp")

1265
        # `cloudpickle` allows pickling lambda functions directly
1266
        import cloudpickle
1267

1268
        input_bytes = cloudpickle.dumps((fn, output_filepath))
1269
1270
1271

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
1272
1273
1274
        returned = subprocess.run(
            _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True
        )
1275
1276
1277
1278
1279
1280

        # check if the subprocess is successful
        try:
            returned.check_returncode()
        except Exception as e:
            # wrap raised exception to provide more information
1281
1282
1283
            raise RuntimeError(
                f"Error raised in subprocess:\n{returned.stderr.decode()}"
            ) from e
1284

1285
        with open(output_filepath, "rb") as f:
1286
1287
1288
1289
1290
1291
            return pickle.load(f)


def _run() -> None:
    # Setup plugins
    from vllm.plugins import load_general_plugins
1292

1293
1294
1295
1296
1297
    load_general_plugins()

    fn, output_file = pickle.loads(sys.stdin.buffer.read())

    result = fn()
1298
1299
1300

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
1301
1302
1303


if __name__ == "__main__":
1304
    _run()