registry.py 51.7 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.tasks import ScoreType
34
from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module
35
from vllm.utils.hashing import safe_hash
36

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


46
47
48
49
50
from .interfaces import (
    has_inner_state,
    has_noops,
    is_attention_free,
    is_hybrid,
Patrick von Platen's avatar
Patrick von Platen committed
51
    requires_raw_input_tokens,
52
    supports_mamba_prefix_caching,
53
54
55
56
57
58
59
    supports_multimodal,
    supports_multimodal_encoder_tp_data,
    supports_multimodal_raw_input_only,
    supports_pp,
    supports_transcription,
)
from .interfaces_base import (
60
    get_attn_type,
61
62
    get_default_seq_pooling_type,
    get_default_tok_pooling_type,
63
    get_score_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"),
luopl's avatar
luopl committed
114
    "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"),
115
    "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
116
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
117
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
118
    "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
119
    "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
Jee Jee Li's avatar
Jee Jee Li committed
120
    "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
121
    "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
122
123
124
125
126
127
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
    "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
    "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
    "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
    "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
128
129
    "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),  # noqa: E501
    "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),  # noqa: E501
130
    "GritLM": ("gritlm", "GritLM"),
Bijaya Dangol's avatar
Bijaya Dangol committed
131
132
    "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),
    "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),
133
134
    "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
    "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
135
    "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),
136
    "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"),
137
    "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"),
138
139
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
140
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
141
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
142
143
    "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),
    "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
144
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
145
    "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
146
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
147
    "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),  # noqa: E501
148
    "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
Paul Pak's avatar
Paul Pak committed
149
    "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
150
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
151
    "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
152
153
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
XuruiYang's avatar
XuruiYang committed
154
    "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
155
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
156
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
157
158
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
159
160
161
    "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
    "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
    "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
162
    "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
163
    "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
164
    "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
165
166
167
168
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
169
    "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
170
    "MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),
171
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
Luis Vega's avatar
Luis Vega committed
172
    "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
173
    "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
174
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
175
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
176
    "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
177
    "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),
178
179
180
    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
181
    "OuroForCausalLM": ("ouro", "OuroForCausalLM"),
182
    "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
183
    "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
184
    "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
185
186
187
188
    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
189
    "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
190
    "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
191
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
192
193
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
194
195
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
196
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
197
198
    "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
    "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
199
    "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
Li Xie's avatar
Li Xie committed
200
    "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
Song's avatar
Song committed
201
    "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
csy0225's avatar
csy0225 committed
202
    "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),
203
204
205
206
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
207
    "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
208
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
209
    "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
210
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
211
    "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
212
213
214
}

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

271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
_LATE_INTERACTION_MODELS = {
    # [Text-only]
    "HF_ColBERT": ("colbert", "ColBERTModel"),
    "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
    "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
    # [Multimodal]
    "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
    "ColQwen3": ("colqwen3", "ColQwen3Model"),
    "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
    "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"),
}

_REWARD_MODELS = {
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
}

_TOKEN_CLASSIFICATION_MODELS = {
290
    "BertForTokenClassification": ("bert", "BertForTokenClassification"),
291
    "ErnieForTokenClassification": ("ernie", "ErnieForTokenClassification"),
292
293
294
295
296
297
298
299
300
    "ModernBertForTokenClassification": (
        "modernbert",
        "ModernBertForTokenClassification",
    ),
}

_SEQUENCE_CLASSIFICATION_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
    "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
301
    "ErnieForSequenceClassification": ("ernie", "ErnieForSequenceClassification"),
302
303
304
305
    "GteNewForSequenceClassification": (
        "bert_with_rope",
        "GteNewForSequenceClassification",
    ),
306
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),  # noqa: E501
307
308
309
    "LlamaBidirectionalForSequenceClassification": (
        "llama",
        "LlamaBidirectionalForSequenceClassification",
310
    ),
311
312
313
314
315
316
317
318
319
    "ModernBertForSequenceClassification": (
        "modernbert",
        "ModernBertForSequenceClassification",
    ),
    "RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"),
    "XLMRobertaForSequenceClassification": (
        "roberta",
        "RobertaForSequenceClassification",
    ),
320
321
322
323
324
325
    # [Multimodal]
    "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"),
    "LlamaNemotronVLForSequenceClassification": (
        "nemotron_vl",
        "LlamaNemotronVLForSequenceClassification",
    ),
326
327
}

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

_SPECULATIVE_DECODING_MODELS = {
547
    "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"),
548
    "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
549
    "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
zhiweiz's avatar
zhiweiz committed
550
    "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
551
    "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
552
    "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
553
    "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
554
    "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
555
    "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
556
557
558
559
    "EagleMistralLarge3ForCausalLM": (
        "mistral_large_3_eagle",
        "EagleMistralLarge3ForCausalLM",
    ),
560
561
    "Eagle3DeepseekV2ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
    "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
562
    "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
563
    "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
564
    "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
Kyungmin Lee's avatar
Kyungmin Lee committed
565
    "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
566
    "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"),
XuruiYang's avatar
XuruiYang committed
567
    "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
568
    "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
569
    "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
570
    "GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"),
571
    "MedusaModel": ("medusa", "Medusa"),
572
    "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
573
    "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),
csy0225's avatar
csy0225 committed
574
    "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),
575
576
    "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
    "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
577
578
579
    # Temporarily disabled.
    # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
    # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
580
}
581

582
_TRANSFORMERS_SUPPORTED_MODELS = {
583
584
585
    # Text generation models
    "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
    # Multimodal models
586
587
588
589
    "Emu3ForConditionalGeneration": (
        "transformers",
        "TransformersMultiModalForCausalLM",
    ),
590
591
592
}

_TRANSFORMERS_BACKEND_MODELS = {
593
    # Text generation models
594
    "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
    "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
613
    "TransformersForSequenceClassification": (
614
        "transformers",
615
        "TransformersForSequenceClassification",
616
    ),
617
    "TransformersMoEForSequenceClassification": (
618
        "transformers",
619
        "TransformersMoEForSequenceClassification",
620
    ),
621
622
623
    "TransformersMultiModalForSequenceClassification": (
        "transformers",
        "TransformersMultiModalForSequenceClassification",
624
    ),
625
}
626

627
_VLLM_MODELS = {
628
    **_TEXT_GENERATION_MODELS,
629
    **_EMBEDDING_MODELS,
630
631
632
633
    **_LATE_INTERACTION_MODELS,
    **_REWARD_MODELS,
    **_TOKEN_CLASSIFICATION_MODELS,
    **_SEQUENCE_CLASSIFICATION_MODELS,
634
    **_MULTIMODAL_MODELS,
635
    **_SPECULATIVE_DECODING_MODELS,
636
637
    **_TRANSFORMERS_SUPPORTED_MODELS,
    **_TRANSFORMERS_BACKEND_MODELS,
638
639
}

640
641
642
643
# 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.
644
_SUBPROCESS_COMMAND = [sys.executable, "-m", "vllm.model_executor.models.registry"]
645

646
_PREVIOUSLY_SUPPORTED_MODELS = {
647
    "MotifForCausalLM": "0.10.2",
648
    "Phi3SmallForCausalLM": "0.9.2",
649
    "Phi4FlashForCausalLM": "0.10.2",
650
    "Phi4MultimodalForCausalLM": "0.12.0",
651
652
653
654
655
656
657
658
659
    # 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",
}
660

661

662
663
@dataclass(frozen=True)
class _ModelInfo:
664
    architecture: str
665
    is_text_generation_model: bool
666
    is_pooling_model: bool
667
    attn_type: AttnTypeStr
668
669
    default_seq_pooling_type: SequencePoolingType
    default_tok_pooling_type: TokenPoolingType
670
    score_type: ScoreType
671
    supports_multimodal: bool
672
    supports_multimodal_raw_input_only: bool
Patrick von Platen's avatar
Patrick von Platen committed
673
    requires_raw_input_tokens: bool
674
    supports_multimodal_encoder_tp_data: bool
675
    supports_pp: bool
676
677
    has_inner_state: bool
    is_attention_free: bool
678
    is_hybrid: bool
679
    has_noops: bool
680
    supports_mamba_prefix_caching: bool
681
    supports_transcription: bool
682
    supports_transcription_only: bool
683
684

    @staticmethod
685
    def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
686
        return _ModelInfo(
687
            architecture=model.__name__,
688
            is_text_generation_model=is_text_generation_model(model),
689
            is_pooling_model=is_pooling_model(model),
690
691
            default_seq_pooling_type=get_default_seq_pooling_type(model),
            default_tok_pooling_type=get_default_tok_pooling_type(model),
692
            attn_type=get_attn_type(model),
693
            score_type=get_score_type(model),
694
            supports_multimodal=supports_multimodal(model),
695
696
697
            supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
                model
            ),
Patrick von Platen's avatar
Patrick von Platen committed
698
            requires_raw_input_tokens=requires_raw_input_tokens(model),
699
700
701
            supports_multimodal_encoder_tp_data=supports_multimodal_encoder_tp_data(
                model
            ),
702
            supports_pp=supports_pp(model),
703
704
            has_inner_state=has_inner_state(model),
            is_attention_free=is_attention_free(model),
705
            is_hybrid=is_hybrid(model),
706
            supports_mamba_prefix_caching=supports_mamba_prefix_caching(model),
707
            supports_transcription=supports_transcription(model),
708
709
710
            supports_transcription_only=(
                supports_transcription(model) and model.supports_transcription_only
            ),
711
            has_noops=has_noops(model),
712
        )
713
714


715
716
717
718
class _BaseRegisteredModel(ABC):
    @abstractmethod
    def inspect_model_cls(self) -> _ModelInfo:
        raise NotImplementedError
719

720
    @abstractmethod
721
    def load_model_cls(self) -> type[nn.Module]:
722
        raise NotImplementedError
723
724


725
726
727
728
729
730
731
@dataclass(frozen=True)
class _RegisteredModel(_BaseRegisteredModel):
    """
    Represents a model that has already been imported in the main process.
    """

    interfaces: _ModelInfo
732
    model_cls: type[nn.Module]
733
734

    @staticmethod
735
    def from_model_cls(model_cls: type[nn.Module]):
736
737
738
739
740
741
742
743
        return _RegisteredModel(
            interfaces=_ModelInfo.from_model_cls(model_cls),
            model_cls=model_cls,
        )

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

744
    def load_model_cls(self) -> type[nn.Module]:
745
746
747
748
749
750
751
752
        return self.model_cls


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

754
755
756
    module_name: str
    class_name: str

757
758
759
760
761
762
763
764
    @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"

765
    def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None:
766
767
        try:
            try:
768
                modelinfo_path = self._get_cache_dir() / self._get_cache_filename()
769
770
771
                with open(modelinfo_path, encoding="utf-8") as file:
                    mi_dict = json.load(file)
            except FileNotFoundError:
772
                logger.debug(
773
                    "Cached model info file for class %s.%s not found",
774
775
776
                    self.module_name,
                    self.class_name,
                )
777
778
779
                return None

            if mi_dict["hash"] != module_hash:
780
                logger.debug(
781
                    "Cached model info file for class %s.%s is stale",
782
783
784
                    self.module_name,
                    self.class_name,
                )
785
786
787
788
789
                return None

            # file not changed, use cached _ModelInfo properties
            return _ModelInfo(**mi_dict["modelinfo"])
        except Exception:
790
            logger.debug(
791
                "Cached model info for class %s.%s error. ",
792
793
794
                self.module_name,
                self.class_name,
            )
795
796
            return None

797
    def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None:
798
799
        """save dictionary json file to cache"""
        from vllm.model_executor.model_loader.weight_utils import atomic_writer
800

801
802
803
804
805
806
807
808
        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()
809
            with atomic_writer(modelinfo_path, encoding="utf-8") as f:
810
811
812
813
814
                json.dump(modelinfo_dict, f, indent=2)
        except Exception:
            logger.exception("Error saving model info cache.")

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

819
820
        if model_path.exists():
            with open(model_path, "rb") as f:
821
                module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest()
822
823
824

            mi = self._load_modelinfo_from_cache(module_hash)
            if mi is not None:
825
                logger.debug(
826
                    "Loaded model info for class %s.%s from cache",
827
828
829
                    self.module_name,
                    self.class_name,
                )
830
831
                return mi
            else:
832
                logger.debug(
833
                    "Cache model info for class %s.%s miss. Loading model instead.",
834
835
836
                    self.module_name,
                    self.class_name,
                )
837
838
839

        # Performed in another process to avoid initializing CUDA
        mi = _run_in_subprocess(
840
841
842
843
844
            lambda: _ModelInfo.from_model_cls(self.load_model_cls())
        )
        logger.debug(
            "Loaded model info for class %s.%s", self.module_name, self.class_name
        )
845
846

        # save cache file
847
848
        if module_hash is not None:
            self._save_modelinfo_to_cache(mi, module_hash)
849
850

        return mi
851

852
    def load_model_cls(self) -> type[nn.Module]:
853
854
855
856
857
858
859
860
        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,
861
) -> type[nn.Module] | None:
862
    from vllm.platforms import current_platform
863

864
    current_platform.verify_model_arch(model_arch)
865
866
867
    try:
        return model.load_model_cls()
    except Exception:
868
        logger.exception("Error in loading model architecture '%s'", model_arch)
869
        return None
870
871


872
873
874
875
@lru_cache(maxsize=128)
def _try_inspect_model_cls(
    model_arch: str,
    model: _BaseRegisteredModel,
876
) -> _ModelInfo | None:
877
878
879
    try:
        return model.inspect_model_cls()
    except Exception:
880
        logger.exception("Error in inspecting model architecture '%s'", model_arch)
881
        return None
882
883


884
885
886
@dataclass
class _ModelRegistry:
    # Keyed by model_arch
887
    models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
888

889
    def get_supported_archs(self) -> Set[str]:
890
        return self.models.keys()
891

892
893
894
    def register_model(
        self,
        model_arch: str,
895
        model_cls: type[nn.Module] | str,
896
    ) -> None:
897
898
899
        """
        Register an external model to be used in vLLM.

900
        `model_cls` can be either:
901

902
        - A [`torch.nn.Module`][] class directly referencing the model.
903
        - A string in the format `<module>:<class>` which can be used to
904
905
          lazily import the model. This is useful to avoid initializing CUDA
          when importing the model and thus the related error
906
          `RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
907
        """
908
909
910
911
        if not isinstance(model_arch, str):
            msg = f"`model_arch` should be a string, not a {type(model_arch)}"
            raise TypeError(msg)

912
        if model_arch in self.models:
913
914
            logger.warning(
                "Model architecture %s is already registered, and will be "
915
916
917
918
                "overwritten by the new model class %s.",
                model_arch,
                model_cls,
            )
919
920
921
922
923
924

        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)
925

926
            model = _LazyRegisteredModel(*split_str)
927
        elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module):
928
            model = _RegisteredModel.from_model_cls(model_cls)
929
        else:
930
931
932
933
            msg = (
                "`model_cls` should be a string or PyTorch model class, "
                f"not a {type(model_arch)}"
            )
934
            raise TypeError(msg)
935

936
        self.models[model_arch] = model
937

938
    def _raise_for_unsupported(self, architectures: list[str]):
939
        all_supported_archs = self.get_supported_archs()
940

941
942
943
        if any(arch in all_supported_archs for arch in architectures):
            raise ValueError(
                f"Model architectures {architectures} failed "
944
945
                "to be inspected. Please check the logs for more details."
            )
946

947
948
949
950
951
952
953
954
        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 "
955
956
                    "use this model architecture."
                )
957

958
959
        raise ValueError(
            f"Model architectures {architectures} are not supported for now. "
960
961
            f"Supported architectures: {all_supported_archs}"
        )
962

963
    def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:
964
965
        if model_arch not in self.models:
            return None
966

967
        return _try_load_model_cls(model_arch, self.models[model_arch])
968

969
    def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None:
970
971
        if model_arch not in self.models:
            return None
972

973
974
975
976
977
978
        return _try_inspect_model_cls(model_arch, self.models[model_arch])

    def _try_resolve_transformers(
        self,
        architecture: str,
        model_config: ModelConfig,
979
    ) -> str | None:
980
981
982
        if architecture in _TRANSFORMERS_BACKEND_MODELS:
            return architecture

983
984
985
        auto_map: dict[str, str] = (
            getattr(model_config.hf_config, "auto_map", None) or dict()
        )
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001

        # 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,
1002
                        trust_remote_code=model_config.trust_remote_code,
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
                        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,
1015
                        trust_remote_code=model_config.trust_remote_code,
1016
1017
1018
1019
1020
                        warn_on_fail=True,
                    )
                    if model_module is not None:
                        break
            else:
1021
                if model_config.model_impl != "transformers":
1022
1023
1024
1025
1026
1027
1028
                    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 "
1029
1030
                    "'auto_map' (relevant if the model is custom)."
                )
1031
1032

        if not model_module.is_backend_compatible():
1033
            if model_config.model_impl != "transformers":
1034
                return None
1035

1036
1037
            raise ValueError(
                f"The Transformers implementation of {architecture!r} "
1038
1039
                "is not compatible with vLLM."
            )
1040

1041
        return model_config._get_transformers_backend_cls()
1042

1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
    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
1068

1069
1070
    def inspect_model_cls(
        self,
1071
        architectures: str | list[str],
1072
        model_config: ModelConfig,
1073
    ) -> tuple[_ModelInfo, str]:
1074
1075
        if isinstance(architectures, str):
            architectures = [architectures]
1076
1077
        if not architectures:
            raise ValueError("No model architectures are specified")
1078
1079

        # Require transformers impl
1080
        if model_config.model_impl == "transformers":
1081
            arch = self._try_resolve_transformers(architectures[0], model_config)
1082
1083
1084
1085
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)
1086
        elif model_config.model_impl == "terratorch":
1087
1088
            model_info = self._try_inspect_model_cls("Terratorch")
            return (model_info, "Terratorch")
1089

1090
        # Fallback to transformers impl (after resolving convert_type)
1091
1092
1093
1094
1095
1096
        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)
1097
1098
1099
1100
1101
1102
1103
            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)
1104
            model_info = self._try_inspect_model_cls(normalized_arch)
1105
            if model_info is not None:
1106
                return (model_info, arch)
1107

1108
        # Fallback to transformers impl (before resolving runner_type)
1109
1110
1111
1112
1113
        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)
1114
1115
1116
1117
1118
            if arch is not None:
                model_info = self._try_inspect_model_cls(arch)
                if model_info is not None:
                    return (model_info, arch)

1119
        return self._raise_for_unsupported(architectures)
1120

1121
1122
    def resolve_model_cls(
        self,
1123
        architectures: str | list[str],
1124
        model_config: ModelConfig,
1125
    ) -> tuple[type[nn.Module], str]:
1126
1127
        if isinstance(architectures, str):
            architectures = [architectures]
1128
1129
        if not architectures:
            raise ValueError("No model architectures are specified")
1130
1131

        # Require transformers impl
1132
        if model_config.model_impl == "transformers":
1133
            arch = self._try_resolve_transformers(architectures[0], model_config)
1134
1135
1136
1137
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)
1138
        elif model_config.model_impl == "terratorch":
1139
1140
1141
1142
            arch = "Terratorch"
            model_cls = self._try_load_model_cls(arch)
            if model_cls is not None:
                return (model_cls, arch)
1143

1144
        # Fallback to transformers impl (after resolving convert_type)
1145
1146
1147
1148
1149
1150
        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)
1151
1152
1153
1154
1155
1156
1157
            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)
1158
            model_cls = self._try_load_model_cls(normalized_arch)
1159
1160
            if model_cls is not None:
                return (model_cls, arch)
1161

1162
        # Fallback to transformers impl (before resolving runner_type)
1163
1164
1165
1166
1167
        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)
1168
1169
1170
1171
1172
            if arch is not None:
                model_cls = self._try_load_model_cls(arch)
                if model_cls is not None:
                    return (model_cls, arch)

1173
        return self._raise_for_unsupported(architectures)
1174

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

1183
    def is_pooling_model(
1184
        self,
1185
        architectures: str | list[str],
1186
        model_config: ModelConfig,
1187
    ) -> bool:
1188
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1189
        return model_cls.is_pooling_model
1190
1191
1192

    def is_multimodal_model(
        self,
1193
        architectures: str | list[str],
1194
        model_config: ModelConfig,
1195
    ) -> bool:
1196
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1197
        return model_cls.supports_multimodal
1198

1199
    def is_multimodal_raw_input_only_model(
1200
        self,
1201
        architectures: str | list[str],
1202
        model_config: ModelConfig,
1203
    ) -> bool:
1204
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1205
        return model_cls.supports_multimodal_raw_input_only
1206

1207
1208
    def is_pp_supported_model(
        self,
1209
        architectures: str | list[str],
1210
        model_config: ModelConfig,
1211
    ) -> bool:
1212
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1213
        return model_cls.supports_pp
1214

1215
1216
    def model_has_inner_state(
        self,
1217
        architectures: str | list[str],
1218
        model_config: ModelConfig,
1219
    ) -> bool:
1220
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1221
        return model_cls.has_inner_state
1222

1223
1224
    def is_attention_free_model(
        self,
1225
        architectures: str | list[str],
1226
        model_config: ModelConfig,
1227
    ) -> bool:
1228
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1229
        return model_cls.is_attention_free
1230

1231
1232
    def is_hybrid_model(
        self,
1233
        architectures: str | list[str],
1234
        model_config: ModelConfig,
1235
    ) -> bool:
1236
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1237
1238
        return model_cls.is_hybrid

1239
1240
    def is_noops_model(
        self,
1241
        architectures: str | list[str],
1242
        model_config: ModelConfig,
1243
    ) -> bool:
1244
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1245
1246
        return model_cls.has_noops

1247
1248
    def is_transcription_model(
        self,
1249
        architectures: str | list[str],
1250
        model_config: ModelConfig,
1251
    ) -> bool:
1252
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1253
1254
        return model_cls.supports_transcription

1255
1256
    def is_transcription_only_model(
        self,
1257
        architectures: str | list[str],
1258
        model_config: ModelConfig,
1259
    ) -> bool:
1260
        model_cls, _ = self.inspect_model_cls(architectures, model_config)
1261
1262
        return model_cls.supports_transcription_only

1263

1264
1265
1266
1267
1268
1269
1270
1271
1272
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()
    }
)
1273
1274
1275
1276
1277

_T = TypeVar("_T")


def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
1278
1279
1280
1281
1282
    # 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")

1283
        # `cloudpickle` allows pickling lambda functions directly
1284
        import cloudpickle
1285

1286
        input_bytes = cloudpickle.dumps((fn, output_filepath))
1287
1288
1289

        # cannot use `sys.executable __file__` here because the script
        # contains relative imports
1290
1291
1292
        returned = subprocess.run(
            _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True
        )
1293
1294
1295
1296
1297
1298

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

1303
        with open(output_filepath, "rb") as f:
1304
1305
1306
1307
1308
1309
            return pickle.load(f)


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

1311
1312
1313
1314
1315
    load_general_plugins()

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

    result = fn()
1316
1317
1318

    with open(output_file, "wb") as f:
        f.write(pickle.dumps(result))
1319
1320
1321


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