"tests/tokenizers_/test_mistral.py" did not exist on "2c0f58203c111bcc331f931664400acfc94cb9bc"
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"),
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
    "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"),
136
    "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"),
137
138
    "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
    "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
139
    "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
140
    "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
141
142
    "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),
    "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
143
    "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
144
    "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
145
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
146
    "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),  # noqa: E501
147
    "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
Paul Pak's avatar
Paul Pak committed
148
    "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
149
    "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
150
    "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
151
152
    # For decapoda-research/llama-*
    "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
XuruiYang's avatar
XuruiYang committed
153
    "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
154
    "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
155
    "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
156
157
    "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
    "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
158
159
160
    "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
    "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
    "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
161
    "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
162
    "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
163
    "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
164
165
166
167
    "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
    # transformers's mpt class has lower case
    "MptForCausalLM": ("mpt", "MPTForCausalLM"),
    "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
168
    "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
169
    "MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),
170
    "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
Luis Vega's avatar
Luis Vega committed
171
    "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
172
    "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
173
    "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
174
    "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
175
    "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
176
    "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),
177
178
179
    "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
    "OPTForCausalLM": ("opt", "OPTForCausalLM"),
    "OrionForCausalLM": ("orion", "OrionForCausalLM"),
180
    "OuroForCausalLM": ("ouro", "OuroForCausalLM"),
181
    "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
182
    "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
183
    "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
184
185
186
187
    "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
    "PhiForCausalLM": ("phi", "PhiForCausalLM"),
    "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
    "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
188
    "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
189
    "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
190
    "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
191
192
    "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
193
194
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
195
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
196
197
    "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
    "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
198
    "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
Li Xie's avatar
Li Xie committed
199
    "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
Song's avatar
Song committed
200
    "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
csy0225's avatar
csy0225 committed
201
    "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),
202
203
204
205
    "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
    "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
    "SolarForCausalLM": ("solar", "SolarForCausalLM"),
206
    "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
207
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
208
    "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
209
    "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
210
    "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
211
212
213
}

_EMBEDDING_MODELS = {
214
    # [Text-only]
215
    "BertModel": ("bert", "BertEmbeddingModel"),
216
    "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
217
    "ErnieModel": ("ernie", "ErnieEmbeddingModel"),
218
    "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
219
    "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
220
    "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
221
    "Gemma3TextModel": ("gemma3", "Gemma3Model"),
222
    "GlmForCausalLM": ("glm", "GlmForCausalLM"),
223
    "GritLM": ("gritlm", "GritLM"),
224
225
    "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
    "GteNewModel": ("bert_with_rope", "GteNewModel"),
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
241
    "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
    "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
242
    "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
243
    "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
chengchengpei's avatar
chengchengpei committed
244
245
246
247
    "VoyageQwen3BidirectionalEmbedModel": (
        "voyage",
        "VoyageQwen3BidirectionalEmbedModel",
    ),
248
    "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
249
    # [Multimodal]
250
    "CLIPModel": ("clip", "CLIPEmbeddingModel"),
251
    "ColPaliForRetrieval": ("colpali", "ColPaliModel"),
252
253
254
    "LlavaNextForConditionalGeneration": (
        "llava_next",
        "LlavaNextForConditionalGeneration",
255
    ),
256
    "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
257
    "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"),  # noqa: E501
258
    "SiglipModel": ("siglip", "SiglipEmbeddingModel"),
259
260
261
262
    "LlamaNemotronVLModel": (
        "nemotron_vl",
        "LlamaNemotronVLForEmbedding",
    ),
263
264
    # Technically Terratorch models work on images, both in
    # input and output. I am adding it here because it piggy-backs on embedding
265
    # models for the time being.
266
267
    "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"),
    "Terratorch": ("terratorch", "Terratorch"),
268
269
}

270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
_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 = {
289
    "BertForTokenClassification": ("bert", "BertForTokenClassification"),
290
    "ErnieForTokenClassification": ("ernie", "ErnieForTokenClassification"),
291
292
293
294
295
296
297
298
299
    "ModernBertForTokenClassification": (
        "modernbert",
        "ModernBertForTokenClassification",
    ),
}

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

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

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

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

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

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

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

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

663

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

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


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

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


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

    interfaces: _ModelInfo
734
    model_cls: type[nn.Module]
735
736

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

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

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


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

756
757
758
    module_name: str
    class_name: str

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

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

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

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

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

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

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

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

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

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

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

        return mi
853

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

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


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


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

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

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

902
        `model_cls` can be either:
903

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

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

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

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

938
        self.models[model_arch] = model
939

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

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

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

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

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

969
        return _try_load_model_cls(model_arch, self.models[model_arch])
970

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

975
976
977
978
979
980
        return _try_inspect_model_cls(model_arch, self.models[model_arch])

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

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

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

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

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

1043
        return model_config._get_transformers_backend_cls()
1044

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

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

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

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

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

1121
        return self._raise_for_unsupported(architectures)
1122

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

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

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

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

1175
        return self._raise_for_unsupported(architectures)
1176

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

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

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

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

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

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

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

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

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

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

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

1265

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

_T = TypeVar("_T")


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

1285
        # `cloudpickle` allows pickling lambda functions directly
1286
        import cloudpickle
1287

1288
        input_bytes = cloudpickle.dumps((fn, output_filepath))
1289
1290
1291

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

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

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


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

1313
1314
1315
1316
1317
    load_general_plugins()

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

    result = fn()
1318
1319
1320

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


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