hancheolp commited on
Commit
812ca4f
·
verified ·
1 Parent(s): 5969c57

Add patches

Browse files
Files changed (2) hide show
  1. patches/registry.py +1304 -0
  2. patches/solar_open.py +693 -0
patches/registry.py ADDED
@@ -0,0 +1,1304 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: Apache-2.0
2
+ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+ """
4
+ Whenever you add an architecture to this page, please also update
5
+ `tests/models/registry.py` with example HuggingFace models for it.
6
+ """
7
+
8
+ import importlib
9
+ import json
10
+ import os
11
+ import pickle
12
+ import subprocess
13
+ 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
+ from vllm.config import (
27
+ ModelConfig,
28
+ iter_architecture_defaults,
29
+ try_match_architecture_defaults,
30
+ )
31
+ from vllm.logger import init_logger
32
+ from vllm.logging_utils import logtime
33
+ from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module
34
+ from vllm.utils.hashing import safe_hash
35
+
36
+ if TYPE_CHECKING:
37
+ from vllm.config.model import AttnTypeStr
38
+ from vllm.config.pooler import SequencePoolingType, TokenPoolingType
39
+ else:
40
+ AttnTypeStr = Any
41
+ SequencePoolingType = Any
42
+ TokenPoolingType = Any
43
+
44
+
45
+ from .interfaces import (
46
+ has_inner_state,
47
+ has_noops,
48
+ is_attention_free,
49
+ is_hybrid,
50
+ requires_raw_input_tokens,
51
+ supports_cross_encoding,
52
+ supports_late_interaction,
53
+ supports_mamba_prefix_caching,
54
+ supports_multimodal,
55
+ supports_multimodal_encoder_tp_data,
56
+ supports_multimodal_raw_input_only,
57
+ supports_pp,
58
+ supports_transcription,
59
+ )
60
+ from .interfaces_base import (
61
+ get_attn_type,
62
+ get_default_seq_pooling_type,
63
+ get_default_tok_pooling_type,
64
+ is_pooling_model,
65
+ is_text_generation_model,
66
+ )
67
+
68
+ logger = init_logger(__name__)
69
+
70
+ _TEXT_GENERATION_MODELS = {
71
+ # [Decoder-only]
72
+ "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"),
73
+ "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"),
74
+ "AquilaModel": ("llama", "LlamaForCausalLM"),
75
+ "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2
76
+ "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),
77
+ "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),
78
+ "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"),
79
+ # baichuan-7b, upper case 'C' in the class name
80
+ "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
81
+ # baichuan-13b, lower case 'c' in the class name
82
+ "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
83
+ "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
84
+ "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"),
85
+ "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"),
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"),
104
+ "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"),
105
+ "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
106
+ "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
107
+ "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),
108
+ "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"),
109
+ "FlexOlmoForCausalLM": ("flex_olmo", "FlexOlmoForCausalLM"),
110
+ "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
111
+ "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
112
+ "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
113
+ "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
114
+ "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
115
+ "GlmForCausalLM": ("glm", "GlmForCausalLM"),
116
+ "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
117
+ "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
118
+ "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
119
+ "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
120
+ "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
121
+ "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
122
+ "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
123
+ "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
124
+ "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
125
+ "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
126
+ "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
127
+ "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), # noqa: E501
128
+ "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), # noqa: E501
129
+ "GritLM": ("gritlm", "GritLM"),
130
+ "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"),
131
+ "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"),
132
+ "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),
133
+ "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),
134
+ "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"),
135
+ "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),
136
+ "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
137
+ "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
138
+ "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
139
+ "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),
140
+ "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
141
+ "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
142
+ "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
143
+ "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
144
+ "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), # noqa: E501
145
+ "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
146
+ "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
147
+ "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
148
+ "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"),
149
+ # For decapoda-research/llama-*
150
+ "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),
151
+ "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
152
+ "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
153
+ "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
154
+ "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
155
+ "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
156
+ "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
157
+ "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
158
+ "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
159
+ "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
160
+ "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
161
+ "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
162
+ "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
163
+ # transformers's mpt class has lower case
164
+ "MptForCausalLM": ("mpt", "MPTForCausalLM"),
165
+ "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
166
+ "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
167
+ "MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),
168
+ "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
169
+ "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
170
+ "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
171
+ "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
172
+ "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
173
+ "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
174
+ "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
175
+ "OPTForCausalLM": ("opt", "OPTForCausalLM"),
176
+ "OrionForCausalLM": ("orion", "OrionForCausalLM"),
177
+ "OuroForCausalLM": ("ouro", "OuroForCausalLM"),
178
+ "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
179
+ "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
180
+ "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
181
+ "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
182
+ "PhiForCausalLM": ("phi", "PhiForCausalLM"),
183
+ "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
184
+ "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
185
+ "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
186
+ "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
187
+ "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
188
+ "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
189
+ "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
190
+ "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
191
+ "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
192
+ "RWForCausalLM": ("falcon", "FalconForCausalLM"),
193
+ "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
194
+ "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
195
+ "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
196
+ "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),
197
+ "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
198
+ "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
199
+ "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
200
+ "SolarForCausalLM": ("solar", "SolarForCausalLM"),
201
+ "SolarOpenForCausalLM": ("solar_open", "SolarOpenForCausalLM"),
202
+ "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
203
+ "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
204
+ "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"),
205
+ "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
206
+ "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
207
+ }
208
+
209
+ _EMBEDDING_MODELS = {
210
+ # [Text-only]
211
+ "BertModel": ("bert", "BertEmbeddingModel"),
212
+ "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
213
+ "HF_ColBERT": ("colbert", "ColBERTModel"),
214
+ "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
215
+ "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
216
+ "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
217
+ "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
218
+ "Gemma3TextModel": ("gemma3", "Gemma3Model"),
219
+ "GlmForCausalLM": ("glm", "GlmForCausalLM"),
220
+ "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
221
+ "GritLM": ("gritlm", "GritLM"),
222
+ "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
223
+ "GteNewModel": ("bert_with_rope", "GteNewModel"),
224
+ "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
225
+ "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
226
+ "LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"),
227
+ "LlamaModel": ("llama", "LlamaForCausalLM"),
228
+ **{
229
+ # Multiple models share the same architecture, so we include them all
230
+ k: (mod, arch)
231
+ for k, (mod, arch) in _TEXT_GENERATION_MODELS.items()
232
+ if arch == "LlamaForCausalLM"
233
+ },
234
+ "MistralModel": ("llama", "LlamaForCausalLM"),
235
+ "ModernBertModel": ("modernbert", "ModernBertModel"),
236
+ "NomicBertModel": ("bert_with_rope", "NomicBertModel"),
237
+ "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
238
+ "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
239
+ "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
240
+ "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
241
+ "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
242
+ "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
243
+ "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
244
+ "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
245
+ "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
246
+ "VoyageQwen3BidirectionalEmbedModel": (
247
+ "voyage",
248
+ "VoyageQwen3BidirectionalEmbedModel",
249
+ ),
250
+ "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
251
+ "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
252
+ # [Multimodal]
253
+ "CLIPModel": ("clip", "CLIPEmbeddingModel"),
254
+ "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
255
+ "LlavaNextForConditionalGeneration": (
256
+ "llava_next",
257
+ "LlavaNextForConditionalGeneration",
258
+ ),
259
+ "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
260
+ "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), # noqa: E501
261
+ "ColQwen3": ("colqwen3", "ColQwen3Model"),
262
+ "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
263
+ "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"),
264
+ "SiglipModel": ("siglip", "SiglipEmbeddingModel"),
265
+ "LlamaNemotronVLModel": (
266
+ "nemotron_vl",
267
+ "LlamaNemotronVLForEmbedding",
268
+ ),
269
+ # Technically Terratorch models work on images, both in
270
+ # input and output. I am adding it here because it piggy-backs on embedding
271
+ # models for the time being.
272
+ "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"),
273
+ "Terratorch": ("terratorch", "Terratorch"),
274
+ }
275
+
276
+ _CROSS_ENCODER_MODELS = {
277
+ "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
278
+ "BertForTokenClassification": ("bert", "BertForTokenClassification"),
279
+ "GteNewForSequenceClassification": (
280
+ "bert_with_rope",
281
+ "GteNewForSequenceClassification",
282
+ ),
283
+ "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"),
284
+ "LlamaBidirectionalForSequenceClassification": (
285
+ "llama",
286
+ "LlamaBidirectionalForSequenceClassification",
287
+ ),
288
+ "LlamaNemotronVLForSequenceClassification": (
289
+ "nemotron_vl",
290
+ "LlamaNemotronVLForSequenceClassification",
291
+ ),
292
+ "ModernBertForSequenceClassification": (
293
+ "modernbert",
294
+ "ModernBertForSequenceClassification",
295
+ ),
296
+ "ModernBertForTokenClassification": (
297
+ "modernbert",
298
+ "ModernBertForTokenClassification",
299
+ ),
300
+ "RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"),
301
+ "XLMRobertaForSequenceClassification": (
302
+ "roberta",
303
+ "RobertaForSequenceClassification",
304
+ ),
305
+ }
306
+
307
+ _MULTIMODAL_MODELS = {
308
+ # [Decoder-only]
309
+ "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"),
310
+ "AudioFlamingo3ForConditionalGeneration": (
311
+ "audioflamingo3",
312
+ "AudioFlamingo3ForConditionalGeneration",
313
+ ),
314
+ "MusicFlamingoForConditionalGeneration": (
315
+ "musicflamingo",
316
+ "MusicFlamingoForConditionalGeneration",
317
+ ),
318
+ "AyaVisionForConditionalGeneration": (
319
+ "aya_vision",
320
+ "AyaVisionForConditionalGeneration",
321
+ ),
322
+ "BagelForConditionalGeneration": ("bagel", "BagelForConditionalGeneration"),
323
+ "BeeForConditionalGeneration": ("bee", "BeeForConditionalGeneration"),
324
+ "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
325
+ "ChameleonForConditionalGeneration": (
326
+ "chameleon",
327
+ "ChameleonForConditionalGeneration",
328
+ ),
329
+ "Cohere2VisionForConditionalGeneration": (
330
+ "cohere2_vision",
331
+ "Cohere2VisionForConditionalGeneration",
332
+ ),
333
+ "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
334
+ "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
335
+ "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
336
+ "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
337
+ "Eagle2_5_VLForConditionalGeneration": (
338
+ "eagle2_5_vl",
339
+ "Eagle2_5_VLForConditionalGeneration",
340
+ ),
341
+ "Ernie4_5_VLMoeForConditionalGeneration": (
342
+ "ernie45_vl",
343
+ "Ernie4_5_VLMoeForConditionalGeneration",
344
+ ),
345
+ "FireRedASR2ForConditionalGeneration": (
346
+ "fireredasr2",
347
+ "FireRedASR2ForConditionalGeneration",
348
+ ),
349
+ "FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"), # noqa: E501
350
+ "FunAudioChatForConditionalGeneration": (
351
+ "funaudiochat",
352
+ "FunAudioChatForConditionalGeneration",
353
+ ),
354
+ "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
355
+ "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), # noqa: E501
356
+ "Gemma3nForConditionalGeneration": (
357
+ "gemma3n_mm",
358
+ "Gemma3nForConditionalGeneration",
359
+ ),
360
+ "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"),
361
+ "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
362
+ "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"),
363
+ "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"),
364
+ "GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"), # noqa: E501
365
+ "GraniteSpeechForConditionalGeneration": (
366
+ "granite_speech",
367
+ "GraniteSpeechForConditionalGeneration",
368
+ ),
369
+ "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
370
+ "HunYuanVLForConditionalGeneration": (
371
+ "hunyuan_vision",
372
+ "HunYuanVLForConditionalGeneration",
373
+ ),
374
+ "StepVLForConditionalGeneration": ("step_vl", "StepVLForConditionalGeneration"),
375
+ "InternVLChatModel": ("internvl", "InternVLChatModel"),
376
+ "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
377
+ "OpenCUAForConditionalGeneration": (
378
+ "opencua",
379
+ "OpenCUAForConditionalGeneration",
380
+ ),
381
+ "InternS1ForConditionalGeneration": (
382
+ "interns1",
383
+ "InternS1ForConditionalGeneration",
384
+ ),
385
+ "InternVLForConditionalGeneration": (
386
+ "interns1",
387
+ "InternS1ForConditionalGeneration",
388
+ ),
389
+ "InternS1ProForConditionalGeneration": (
390
+ "interns1_pro",
391
+ "InternS1ProForConditionalGeneration",
392
+ ),
393
+ "Idefics3ForConditionalGeneration": (
394
+ "idefics3",
395
+ "Idefics3ForConditionalGeneration",
396
+ ),
397
+ "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"),
398
+ "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), # noqa: E501
399
+ "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"),
400
+ "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),
401
+ "KeyeVL1_5ForConditionalGeneration": (
402
+ "keye_vl1_5",
403
+ "KeyeVL1_5ForConditionalGeneration",
404
+ ),
405
+ "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),
406
+ "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), # noqa: E501
407
+ "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), # noqa: E501
408
+ "LightOnOCRForConditionalGeneration": (
409
+ "lightonocr",
410
+ "LightOnOCRForConditionalGeneration",
411
+ ),
412
+ "Lfm2VlForConditionalGeneration": ("lfm2_vl", "Lfm2VLForConditionalGeneration"),
413
+ "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"),
414
+ "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"), # noqa: E501
415
+ "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
416
+ "LlavaNextForConditionalGeneration": (
417
+ "llava_next",
418
+ "LlavaNextForConditionalGeneration",
419
+ ),
420
+ "LlavaNextVideoForConditionalGeneration": (
421
+ "llava_next_video",
422
+ "LlavaNextVideoForConditionalGeneration",
423
+ ),
424
+ "LlavaOnevisionForConditionalGeneration": (
425
+ "llava_onevision",
426
+ "LlavaOnevisionForConditionalGeneration",
427
+ ),
428
+ "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"), # noqa: E501
429
+ "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),
430
+ "MiniMaxVL01ForConditionalGeneration": (
431
+ "minimax_vl_01",
432
+ "MiniMaxVL01ForConditionalGeneration",
433
+ ),
434
+ "MiniCPMO": ("minicpmo", "MiniCPMO"),
435
+ "MiniCPMV": ("minicpmv", "MiniCPMV"),
436
+ "Mistral3ForConditionalGeneration": (
437
+ "mistral3",
438
+ "Mistral3ForConditionalGeneration",
439
+ ),
440
+ "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
441
+ "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),
442
+ "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
443
+ "OpenPanguVLForConditionalGeneration": (
444
+ "openpangu_vl",
445
+ "OpenPanguVLForConditionalGeneration",
446
+ ),
447
+ "Ovis": ("ovis", "Ovis"),
448
+ "Ovis2_5": ("ovis2_5", "Ovis2_5"),
449
+ "Ovis2_6ForCausalLM": ("ovis2_5", "Ovis2_5"),
450
+ "Ovis2_6_MoeForCausalLM": ("ovis2_5", "Ovis2_5"),
451
+ "PaddleOCRVLForConditionalGeneration": (
452
+ "paddleocr_vl",
453
+ "PaddleOCRVLForConditionalGeneration",
454
+ ),
455
+ "PaliGemmaForConditionalGeneration": (
456
+ "paligemma",
457
+ "PaliGemmaForConditionalGeneration",
458
+ ),
459
+ "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
460
+ "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),
461
+ "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"), # noqa: E501
462
+ "QwenVLForConditionalGeneration": ("qwen_vl", "QwenVLForConditionalGeneration"), # noqa: E501
463
+ "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), # noqa: E501
464
+ "Qwen2_5_VLForConditionalGeneration": (
465
+ "qwen2_5_vl",
466
+ "Qwen2_5_VLForConditionalGeneration",
467
+ ),
468
+ "Qwen2AudioForConditionalGeneration": (
469
+ "qwen2_audio",
470
+ "Qwen2AudioForConditionalGeneration",
471
+ ),
472
+ "Qwen2_5OmniModel": (
473
+ "qwen2_5_omni_thinker",
474
+ "Qwen2_5OmniThinkerForConditionalGeneration",
475
+ ),
476
+ "Qwen2_5OmniForConditionalGeneration": (
477
+ "qwen2_5_omni_thinker",
478
+ "Qwen2_5OmniThinkerForConditionalGeneration",
479
+ ),
480
+ "Qwen3OmniMoeForConditionalGeneration": (
481
+ "qwen3_omni_moe_thinker",
482
+ "Qwen3OmniMoeThinkerForConditionalGeneration",
483
+ ),
484
+ "Qwen3ASRForConditionalGeneration": (
485
+ "qwen3_asr",
486
+ "Qwen3ASRForConditionalGeneration",
487
+ ),
488
+ "Qwen3ASRRealtimeGeneration": (
489
+ "qwen3_asr_realtime",
490
+ "Qwen3ASRRealtimeGeneration",
491
+ ),
492
+ "Qwen3VLForConditionalGeneration": ("qwen3_vl", "Qwen3VLForConditionalGeneration"), # noqa: E501
493
+ "Qwen3VLMoeForConditionalGeneration": (
494
+ "qwen3_vl_moe",
495
+ "Qwen3VLMoeForConditionalGeneration",
496
+ ),
497
+ "Qwen3_5ForConditionalGeneration": (
498
+ "qwen3_5",
499
+ "Qwen3_5ForConditionalGeneration",
500
+ ),
501
+ "Qwen3_5MoeForConditionalGeneration": (
502
+ "qwen3_5",
503
+ "Qwen3_5MoeForConditionalGeneration",
504
+ ),
505
+ "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"),
506
+ "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"), # noqa: E501
507
+ "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), # noqa: E501
508
+ "Tarsier2ForConditionalGeneration": (
509
+ "qwen2_vl",
510
+ "Tarsier2ForConditionalGeneration",
511
+ ),
512
+ "UltravoxModel": ("ultravox", "UltravoxModel"),
513
+ "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), # noqa: E501
514
+ "VoxtralRealtimeGeneration": ("voxtral_realtime", "VoxtralRealtimeGeneration"), # noqa: E501
515
+ # [Encoder-decoder]
516
+ "NemotronParseForConditionalGeneration": (
517
+ "nemotron_parse",
518
+ "NemotronParseForConditionalGeneration",
519
+ ),
520
+ "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"), # noqa: E501
521
+ }
522
+
523
+ _SPECULATIVE_DECODING_MODELS = {
524
+ "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"),
525
+ "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
526
+ "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
527
+ "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
528
+ "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
529
+ "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
530
+ "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
531
+ "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
532
+ "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
533
+ "EagleMistralLarge3ForCausalLM": (
534
+ "mistral_large_3_eagle",
535
+ "EagleMistralLarge3ForCausalLM",
536
+ ),
537
+ "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
538
+ "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
539
+ "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
540
+ "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
541
+ "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"),
542
+ "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),
543
+ "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
544
+ "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
545
+ "GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"),
546
+ "MedusaModel": ("medusa", "Medusa"),
547
+ "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
548
+ "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),
549
+ "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),
550
+ "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
551
+ "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
552
+ # Temporarily disabled.
553
+ # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
554
+ # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
555
+ }
556
+
557
+ _TRANSFORMERS_SUPPORTED_MODELS = {
558
+ # Text generation models
559
+ "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
560
+ # Multimodal models
561
+ "Emu3ForConditionalGeneration": (
562
+ "transformers",
563
+ "TransformersMultiModalForCausalLM",
564
+ ),
565
+ }
566
+
567
+ _TRANSFORMERS_BACKEND_MODELS = {
568
+ # Text generation models
569
+ "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
570
+ "TransformersMoEForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
571
+ # Multimodal models
572
+ "TransformersMultiModalForCausalLM": (
573
+ "transformers",
574
+ "TransformersMultiModalForCausalLM",
575
+ ),
576
+ "TransformersMultiModalMoEForCausalLM": (
577
+ "transformers",
578
+ "TransformersMultiModalMoEForCausalLM",
579
+ ),
580
+ # Embedding models
581
+ "TransformersEmbeddingModel": ("transformers", "TransformersEmbeddingModel"),
582
+ "TransformersMoEEmbeddingModel": ("transformers", "TransformersMoEEmbeddingModel"),
583
+ "TransformersMultiModalEmbeddingModel": (
584
+ "transformers",
585
+ "TransformersMultiModalEmbeddingModel",
586
+ ),
587
+ # Sequence classification models
588
+ "TransformersForSequenceClassification": (
589
+ "transformers",
590
+ "TransformersForSequenceClassification",
591
+ ),
592
+ "TransformersMoEForSequenceClassification": (
593
+ "transformers",
594
+ "TransformersMoEForSequenceClassification",
595
+ ),
596
+ "TransformersMultiModalForSequenceClassification": (
597
+ "transformers",
598
+ "TransformersMultiModalForSequenceClassification",
599
+ ),
600
+ }
601
+
602
+ _VLLM_MODELS = {
603
+ **_TEXT_GENERATION_MODELS,
604
+ **_EMBEDDING_MODELS,
605
+ **_CROSS_ENCODER_MODELS,
606
+ **_MULTIMODAL_MODELS,
607
+ **_SPECULATIVE_DECODING_MODELS,
608
+ **_TRANSFORMERS_SUPPORTED_MODELS,
609
+ **_TRANSFORMERS_BACKEND_MODELS,
610
+ }
611
+
612
+ # This variable is used as the args for subprocess.run(). We
613
+ # can modify this variable to alter the args if needed. e.g.
614
+ # when we use par format to pack things together, sys.executable
615
+ # might not be the target we want to run.
616
+ _SUBPROCESS_COMMAND = [sys.executable, "-m", "vllm.model_executor.models.registry"]
617
+
618
+ _PREVIOUSLY_SUPPORTED_MODELS = {
619
+ "MotifForCausalLM": "0.10.2",
620
+ "Phi3SmallForCausalLM": "0.9.2",
621
+ "Phi4FlashForCausalLM": "0.10.2",
622
+ "Phi4MultimodalForCausalLM": "0.12.0",
623
+ # encoder-decoder models except whisper
624
+ # have been removed for V0 deprecation.
625
+ "BartModel": "0.10.2",
626
+ "BartForConditionalGeneration": "0.10.2",
627
+ "DonutForConditionalGeneration": "0.10.2",
628
+ "Florence2ForConditionalGeneration": "0.10.2",
629
+ "MBartForConditionalGeneration": "0.10.2",
630
+ "MllamaForConditionalGeneration": "0.10.2",
631
+ }
632
+
633
+
634
+ @dataclass(frozen=True)
635
+ class _ModelInfo:
636
+ architecture: str
637
+ is_text_generation_model: bool
638
+ is_pooling_model: bool
639
+ attn_type: AttnTypeStr
640
+ default_seq_pooling_type: SequencePoolingType
641
+ default_tok_pooling_type: TokenPoolingType
642
+ supports_cross_encoding: bool
643
+ supports_late_interaction: bool
644
+ supports_multimodal: bool
645
+ supports_multimodal_raw_input_only: bool
646
+ requires_raw_input_tokens: bool
647
+ supports_multimodal_encoder_tp_data: bool
648
+ supports_pp: bool
649
+ has_inner_state: bool
650
+ is_attention_free: bool
651
+ is_hybrid: bool
652
+ has_noops: bool
653
+ supports_mamba_prefix_caching: bool
654
+ supports_transcription: bool
655
+ supports_transcription_only: bool
656
+
657
+ @staticmethod
658
+ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo":
659
+ return _ModelInfo(
660
+ architecture=model.__name__,
661
+ is_text_generation_model=is_text_generation_model(model),
662
+ is_pooling_model=is_pooling_model(model),
663
+ default_seq_pooling_type=get_default_seq_pooling_type(model),
664
+ default_tok_pooling_type=get_default_tok_pooling_type(model),
665
+ attn_type=get_attn_type(model),
666
+ supports_cross_encoding=supports_cross_encoding(model),
667
+ supports_late_interaction=supports_late_interaction(model),
668
+ supports_multimodal=supports_multimodal(model),
669
+ supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
670
+ model
671
+ ),
672
+ requires_raw_input_tokens=requires_raw_input_tokens(model),
673
+ supports_multimodal_encoder_tp_data=supports_multimodal_encoder_tp_data(
674
+ model
675
+ ),
676
+ supports_pp=supports_pp(model),
677
+ has_inner_state=has_inner_state(model),
678
+ is_attention_free=is_attention_free(model),
679
+ is_hybrid=is_hybrid(model),
680
+ supports_mamba_prefix_caching=supports_mamba_prefix_caching(model),
681
+ supports_transcription=supports_transcription(model),
682
+ supports_transcription_only=(
683
+ supports_transcription(model) and model.supports_transcription_only
684
+ ),
685
+ has_noops=has_noops(model),
686
+ )
687
+
688
+
689
+ class _BaseRegisteredModel(ABC):
690
+ @abstractmethod
691
+ def inspect_model_cls(self) -> _ModelInfo:
692
+ raise NotImplementedError
693
+
694
+ @abstractmethod
695
+ def load_model_cls(self) -> type[nn.Module]:
696
+ raise NotImplementedError
697
+
698
+
699
+ @dataclass(frozen=True)
700
+ class _RegisteredModel(_BaseRegisteredModel):
701
+ """
702
+ Represents a model that has already been imported in the main process.
703
+ """
704
+
705
+ interfaces: _ModelInfo
706
+ model_cls: type[nn.Module]
707
+
708
+ @staticmethod
709
+ def from_model_cls(model_cls: type[nn.Module]):
710
+ return _RegisteredModel(
711
+ interfaces=_ModelInfo.from_model_cls(model_cls),
712
+ model_cls=model_cls,
713
+ )
714
+
715
+ def inspect_model_cls(self) -> _ModelInfo:
716
+ return self.interfaces
717
+
718
+ def load_model_cls(self) -> type[nn.Module]:
719
+ return self.model_cls
720
+
721
+
722
+ @dataclass(frozen=True)
723
+ class _LazyRegisteredModel(_BaseRegisteredModel):
724
+ """
725
+ Represents a model that has not been imported in the main process.
726
+ """
727
+
728
+ module_name: str
729
+ class_name: str
730
+
731
+ @staticmethod
732
+ def _get_cache_dir() -> Path:
733
+ return Path(envs.VLLM_CACHE_ROOT) / "modelinfos"
734
+
735
+ def _get_cache_filename(self) -> str:
736
+ cls_name = f"{self.module_name}-{self.class_name}".replace(".", "-")
737
+ return f"{cls_name}.json"
738
+
739
+ def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None:
740
+ try:
741
+ try:
742
+ modelinfo_path = self._get_cache_dir() / self._get_cache_filename()
743
+ with open(modelinfo_path, encoding="utf-8") as file:
744
+ mi_dict = json.load(file)
745
+ except FileNotFoundError:
746
+ logger.debug(
747
+ "Cached model info file for class %s.%s not found",
748
+ self.module_name,
749
+ self.class_name,
750
+ )
751
+ return None
752
+
753
+ if mi_dict["hash"] != module_hash:
754
+ logger.debug(
755
+ "Cached model info file for class %s.%s is stale",
756
+ self.module_name,
757
+ self.class_name,
758
+ )
759
+ return None
760
+
761
+ # file not changed, use cached _ModelInfo properties
762
+ return _ModelInfo(**mi_dict["modelinfo"])
763
+ except Exception:
764
+ logger.debug(
765
+ "Cached model info for class %s.%s error. ",
766
+ self.module_name,
767
+ self.class_name,
768
+ )
769
+ return None
770
+
771
+ def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None:
772
+ """save dictionary json file to cache"""
773
+ from vllm.model_executor.model_loader.weight_utils import atomic_writer
774
+
775
+ try:
776
+ modelinfo_dict = {
777
+ "hash": module_hash,
778
+ "modelinfo": asdict(mi),
779
+ }
780
+ cache_dir = self._get_cache_dir()
781
+ cache_dir.mkdir(parents=True, exist_ok=True)
782
+ modelinfo_path = cache_dir / self._get_cache_filename()
783
+ with atomic_writer(modelinfo_path, encoding="utf-8") as f:
784
+ json.dump(modelinfo_dict, f, indent=2)
785
+ except Exception:
786
+ logger.exception("Error saving model info cache.")
787
+
788
+ @logtime(logger=logger, msg="Registry inspect model class")
789
+ def inspect_model_cls(self) -> _ModelInfo:
790
+ model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py"
791
+ module_hash = None
792
+
793
+ if model_path.exists():
794
+ with open(model_path, "rb") as f:
795
+ module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest()
796
+
797
+ mi = self._load_modelinfo_from_cache(module_hash)
798
+ if mi is not None:
799
+ logger.debug(
800
+ "Loaded model info for class %s.%s from cache",
801
+ self.module_name,
802
+ self.class_name,
803
+ )
804
+ return mi
805
+ else:
806
+ logger.debug(
807
+ "Cache model info for class %s.%s miss. Loading model instead.",
808
+ self.module_name,
809
+ self.class_name,
810
+ )
811
+
812
+ # Performed in another process to avoid initializing CUDA
813
+ mi = _run_in_subprocess(
814
+ lambda: _ModelInfo.from_model_cls(self.load_model_cls())
815
+ )
816
+ logger.debug(
817
+ "Loaded model info for class %s.%s", self.module_name, self.class_name
818
+ )
819
+
820
+ # save cache file
821
+ if module_hash is not None:
822
+ self._save_modelinfo_to_cache(mi, module_hash)
823
+
824
+ return mi
825
+
826
+ def load_model_cls(self) -> type[nn.Module]:
827
+ mod = importlib.import_module(self.module_name)
828
+ return getattr(mod, self.class_name)
829
+
830
+
831
+ @lru_cache(maxsize=128)
832
+ def _try_load_model_cls(
833
+ model_arch: str,
834
+ model: _BaseRegisteredModel,
835
+ ) -> type[nn.Module] | None:
836
+ from vllm.platforms import current_platform
837
+
838
+ current_platform.verify_model_arch(model_arch)
839
+ try:
840
+ return model.load_model_cls()
841
+ except Exception:
842
+ logger.exception("Error in loading model architecture '%s'", model_arch)
843
+ return None
844
+
845
+
846
+ @lru_cache(maxsize=128)
847
+ def _try_inspect_model_cls(
848
+ model_arch: str,
849
+ model: _BaseRegisteredModel,
850
+ ) -> _ModelInfo | None:
851
+ try:
852
+ return model.inspect_model_cls()
853
+ except Exception:
854
+ logger.exception("Error in inspecting model architecture '%s'", model_arch)
855
+ return None
856
+
857
+
858
+ @dataclass
859
+ class _ModelRegistry:
860
+ # Keyed by model_arch
861
+ models: dict[str, _BaseRegisteredModel] = field(default_factory=dict)
862
+
863
+ def get_supported_archs(self) -> Set[str]:
864
+ return self.models.keys()
865
+
866
+ def register_model(
867
+ self,
868
+ model_arch: str,
869
+ model_cls: type[nn.Module] | str,
870
+ ) -> None:
871
+ """
872
+ Register an external model to be used in vLLM.
873
+
874
+ `model_cls` can be either:
875
+
876
+ - A [`torch.nn.Module`][] class directly referencing the model.
877
+ - A string in the format `<module>:<class>` which can be used to
878
+ lazily import the model. This is useful to avoid initializing CUDA
879
+ when importing the model and thus the related error
880
+ `RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
881
+ """
882
+ if not isinstance(model_arch, str):
883
+ msg = f"`model_arch` should be a string, not a {type(model_arch)}"
884
+ raise TypeError(msg)
885
+
886
+ if model_arch in self.models:
887
+ logger.warning(
888
+ "Model architecture %s is already registered, and will be "
889
+ "overwritten by the new model class %s.",
890
+ model_arch,
891
+ model_cls,
892
+ )
893
+
894
+ if isinstance(model_cls, str):
895
+ split_str = model_cls.split(":")
896
+ if len(split_str) != 2:
897
+ msg = "Expected a string in the format `<module>:<class>`"
898
+ raise ValueError(msg)
899
+
900
+ model = _LazyRegisteredModel(*split_str)
901
+ elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module):
902
+ model = _RegisteredModel.from_model_cls(model_cls)
903
+ else:
904
+ msg = (
905
+ "`model_cls` should be a string or PyTorch model class, "
906
+ f"not a {type(model_arch)}"
907
+ )
908
+ raise TypeError(msg)
909
+
910
+ self.models[model_arch] = model
911
+
912
+ def _raise_for_unsupported(self, architectures: list[str]):
913
+ all_supported_archs = self.get_supported_archs()
914
+
915
+ if any(arch in all_supported_archs for arch in architectures):
916
+ raise ValueError(
917
+ f"Model architectures {architectures} failed "
918
+ "to be inspected. Please check the logs for more details."
919
+ )
920
+
921
+ for arch in architectures:
922
+ if arch in _PREVIOUSLY_SUPPORTED_MODELS:
923
+ previous_version = _PREVIOUSLY_SUPPORTED_MODELS[arch]
924
+
925
+ raise ValueError(
926
+ f"Model architecture {arch} was supported in vLLM until "
927
+ f"v{previous_version}, and is not supported anymore. "
928
+ "Please use an older version of vLLM if you want to "
929
+ "use this model architecture."
930
+ )
931
+
932
+ raise ValueError(
933
+ f"Model architectures {architectures} are not supported for now. "
934
+ f"Supported architectures: {all_supported_archs}"
935
+ )
936
+
937
+ def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:
938
+ if model_arch not in self.models:
939
+ return None
940
+
941
+ return _try_load_model_cls(model_arch, self.models[model_arch])
942
+
943
+ def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None:
944
+ if model_arch not in self.models:
945
+ return None
946
+
947
+ return _try_inspect_model_cls(model_arch, self.models[model_arch])
948
+
949
+ def _try_resolve_transformers(
950
+ self,
951
+ architecture: str,
952
+ model_config: ModelConfig,
953
+ ) -> str | None:
954
+ if architecture in _TRANSFORMERS_BACKEND_MODELS:
955
+ return architecture
956
+
957
+ auto_map: dict[str, str] = (
958
+ getattr(model_config.hf_config, "auto_map", None) or dict()
959
+ )
960
+
961
+ # Make sure that config class is always initialized before model class,
962
+ # otherwise the model class won't be able to access the config class,
963
+ # the expected auto_map should have correct order like:
964
+ # "auto_map": {
965
+ # "AutoConfig": "<your-repo-name>--<config-name>",
966
+ # "AutoModel": "<your-repo-name>--<config-name>",
967
+ # "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
968
+ # },
969
+ for prefix in ("AutoConfig", "AutoModel"):
970
+ for name, module in auto_map.items():
971
+ if name.startswith(prefix):
972
+ try_get_class_from_dynamic_module(
973
+ module,
974
+ model_config.model,
975
+ revision=model_config.revision,
976
+ trust_remote_code=model_config.trust_remote_code,
977
+ warn_on_fail=False,
978
+ )
979
+
980
+ model_module = getattr(transformers, architecture, None)
981
+
982
+ if model_module is None:
983
+ for name, module in auto_map.items():
984
+ if name.startswith("AutoModel"):
985
+ model_module = try_get_class_from_dynamic_module(
986
+ module,
987
+ model_config.model,
988
+ revision=model_config.revision,
989
+ trust_remote_code=model_config.trust_remote_code,
990
+ warn_on_fail=True,
991
+ )
992
+ if model_module is not None:
993
+ break
994
+ else:
995
+ if model_config.model_impl != "transformers":
996
+ return None
997
+
998
+ raise ValueError(
999
+ f"Cannot find model module. {architecture!r} is not a "
1000
+ "registered model in the Transformers library (only "
1001
+ "relevant if the model is meant to be in Transformers) "
1002
+ "and 'AutoModel' is not present in the model config's "
1003
+ "'auto_map' (relevant if the model is custom)."
1004
+ )
1005
+
1006
+ if not model_module.is_backend_compatible():
1007
+ if model_config.model_impl != "transformers":
1008
+ return None
1009
+
1010
+ raise ValueError(
1011
+ f"The Transformers implementation of {architecture!r} "
1012
+ "is not compatible with vLLM."
1013
+ )
1014
+
1015
+ return model_config._get_transformers_backend_cls()
1016
+
1017
+ def _normalize_arch(
1018
+ self,
1019
+ architecture: str,
1020
+ model_config: ModelConfig,
1021
+ ) -> str:
1022
+ if architecture in self.models:
1023
+ return architecture
1024
+
1025
+ # This may be called in order to resolve runner_type and convert_type
1026
+ # in the first place, in which case we consider the default match
1027
+ match = try_match_architecture_defaults(
1028
+ architecture,
1029
+ runner_type=getattr(model_config, "runner_type", None),
1030
+ convert_type=getattr(model_config, "convert_type", None),
1031
+ )
1032
+ if match:
1033
+ suffix, _ = match
1034
+
1035
+ # Get the name of the base model to convert
1036
+ for repl_suffix, _ in iter_architecture_defaults():
1037
+ base_arch = architecture.replace(suffix, repl_suffix)
1038
+ if base_arch in self.models:
1039
+ return base_arch
1040
+
1041
+ return architecture
1042
+
1043
+ def inspect_model_cls(
1044
+ self,
1045
+ architectures: str | list[str],
1046
+ model_config: ModelConfig,
1047
+ ) -> tuple[_ModelInfo, str]:
1048
+ if isinstance(architectures, str):
1049
+ architectures = [architectures]
1050
+ if not architectures:
1051
+ raise ValueError("No model architectures are specified")
1052
+
1053
+ # Require transformers impl
1054
+ if model_config.model_impl == "transformers":
1055
+ arch = self._try_resolve_transformers(architectures[0], model_config)
1056
+ if arch is not None:
1057
+ model_info = self._try_inspect_model_cls(arch)
1058
+ if model_info is not None:
1059
+ return (model_info, arch)
1060
+ elif model_config.model_impl == "terratorch":
1061
+ model_info = self._try_inspect_model_cls("Terratorch")
1062
+ return (model_info, "Terratorch")
1063
+
1064
+ # Fallback to transformers impl (after resolving convert_type)
1065
+ if (
1066
+ all(arch not in self.models for arch in architectures)
1067
+ and model_config.model_impl == "auto"
1068
+ and getattr(model_config, "convert_type", "none") == "none"
1069
+ ):
1070
+ arch = self._try_resolve_transformers(architectures[0], model_config)
1071
+ if arch is not None:
1072
+ model_info = self._try_inspect_model_cls(arch)
1073
+ if model_info is not None:
1074
+ return (model_info, arch)
1075
+
1076
+ for arch in architectures:
1077
+ normalized_arch = self._normalize_arch(arch, model_config)
1078
+ model_info = self._try_inspect_model_cls(normalized_arch)
1079
+ if model_info is not None:
1080
+ return (model_info, arch)
1081
+
1082
+ # Fallback to transformers impl (before resolving runner_type)
1083
+ if (
1084
+ all(arch not in self.models for arch in architectures)
1085
+ and model_config.model_impl == "auto"
1086
+ ):
1087
+ arch = self._try_resolve_transformers(architectures[0], model_config)
1088
+ if arch is not None:
1089
+ model_info = self._try_inspect_model_cls(arch)
1090
+ if model_info is not None:
1091
+ return (model_info, arch)
1092
+
1093
+ return self._raise_for_unsupported(architectures)
1094
+
1095
+ def resolve_model_cls(
1096
+ self,
1097
+ architectures: str | list[str],
1098
+ model_config: ModelConfig,
1099
+ ) -> tuple[type[nn.Module], str]:
1100
+ if isinstance(architectures, str):
1101
+ architectures = [architectures]
1102
+ if not architectures:
1103
+ raise ValueError("No model architectures are specified")
1104
+
1105
+ # Require transformers impl
1106
+ if model_config.model_impl == "transformers":
1107
+ arch = self._try_resolve_transformers(architectures[0], model_config)
1108
+ if arch is not None:
1109
+ model_cls = self._try_load_model_cls(arch)
1110
+ if model_cls is not None:
1111
+ return (model_cls, arch)
1112
+ elif model_config.model_impl == "terratorch":
1113
+ arch = "Terratorch"
1114
+ model_cls = self._try_load_model_cls(arch)
1115
+ if model_cls is not None:
1116
+ return (model_cls, arch)
1117
+
1118
+ # Fallback to transformers impl (after resolving convert_type)
1119
+ if (
1120
+ all(arch not in self.models for arch in architectures)
1121
+ and model_config.model_impl == "auto"
1122
+ and getattr(model_config, "convert_type", "none") == "none"
1123
+ ):
1124
+ arch = self._try_resolve_transformers(architectures[0], model_config)
1125
+ if arch is not None:
1126
+ model_cls = self._try_load_model_cls(arch)
1127
+ if model_cls is not None:
1128
+ return (model_cls, arch)
1129
+
1130
+ for arch in architectures:
1131
+ normalized_arch = self._normalize_arch(arch, model_config)
1132
+ model_cls = self._try_load_model_cls(normalized_arch)
1133
+ if model_cls is not None:
1134
+ return (model_cls, arch)
1135
+
1136
+ # Fallback to transformers impl (before resolving runner_type)
1137
+ if (
1138
+ all(arch not in self.models for arch in architectures)
1139
+ and model_config.model_impl == "auto"
1140
+ ):
1141
+ arch = self._try_resolve_transformers(architectures[0], model_config)
1142
+ if arch is not None:
1143
+ model_cls = self._try_load_model_cls(arch)
1144
+ if model_cls is not None:
1145
+ return (model_cls, arch)
1146
+
1147
+ return self._raise_for_unsupported(architectures)
1148
+
1149
+ def is_text_generation_model(
1150
+ self,
1151
+ architectures: str | list[str],
1152
+ model_config: ModelConfig,
1153
+ ) -> bool:
1154
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1155
+ return model_cls.is_text_generation_model
1156
+
1157
+ def is_pooling_model(
1158
+ self,
1159
+ architectures: str | list[str],
1160
+ model_config: ModelConfig,
1161
+ ) -> bool:
1162
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1163
+ return model_cls.is_pooling_model
1164
+
1165
+ def is_cross_encoder_model(
1166
+ self,
1167
+ architectures: str | list[str],
1168
+ model_config: ModelConfig,
1169
+ ) -> bool:
1170
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1171
+ return model_cls.supports_cross_encoding
1172
+
1173
+ def is_multimodal_model(
1174
+ self,
1175
+ architectures: str | list[str],
1176
+ model_config: ModelConfig,
1177
+ ) -> bool:
1178
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1179
+ return model_cls.supports_multimodal
1180
+
1181
+ def is_multimodal_raw_input_only_model(
1182
+ self,
1183
+ architectures: str | list[str],
1184
+ model_config: ModelConfig,
1185
+ ) -> bool:
1186
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1187
+ return model_cls.supports_multimodal_raw_input_only
1188
+
1189
+ def is_pp_supported_model(
1190
+ self,
1191
+ architectures: str | list[str],
1192
+ model_config: ModelConfig,
1193
+ ) -> bool:
1194
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1195
+ return model_cls.supports_pp
1196
+
1197
+ def model_has_inner_state(
1198
+ self,
1199
+ architectures: str | list[str],
1200
+ model_config: ModelConfig,
1201
+ ) -> bool:
1202
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1203
+ return model_cls.has_inner_state
1204
+
1205
+ def is_attention_free_model(
1206
+ self,
1207
+ architectures: str | list[str],
1208
+ model_config: ModelConfig,
1209
+ ) -> bool:
1210
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1211
+ return model_cls.is_attention_free
1212
+
1213
+ def is_hybrid_model(
1214
+ self,
1215
+ architectures: str | list[str],
1216
+ model_config: ModelConfig,
1217
+ ) -> bool:
1218
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1219
+ return model_cls.is_hybrid
1220
+
1221
+ def is_noops_model(
1222
+ self,
1223
+ architectures: str | list[str],
1224
+ model_config: ModelConfig,
1225
+ ) -> bool:
1226
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1227
+ return model_cls.has_noops
1228
+
1229
+ def is_transcription_model(
1230
+ self,
1231
+ architectures: str | list[str],
1232
+ model_config: ModelConfig,
1233
+ ) -> bool:
1234
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1235
+ return model_cls.supports_transcription
1236
+
1237
+ def is_transcription_only_model(
1238
+ self,
1239
+ architectures: str | list[str],
1240
+ model_config: ModelConfig,
1241
+ ) -> bool:
1242
+ model_cls, _ = self.inspect_model_cls(architectures, model_config)
1243
+ return model_cls.supports_transcription_only
1244
+
1245
+
1246
+ ModelRegistry = _ModelRegistry(
1247
+ {
1248
+ model_arch: _LazyRegisteredModel(
1249
+ module_name=f"vllm.model_executor.models.{mod_relname}",
1250
+ class_name=cls_name,
1251
+ )
1252
+ for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
1253
+ }
1254
+ )
1255
+
1256
+ _T = TypeVar("_T")
1257
+
1258
+
1259
+ def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
1260
+ # NOTE: We use a temporary directory instead of a temporary file to avoid
1261
+ # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
1262
+ with tempfile.TemporaryDirectory() as tempdir:
1263
+ output_filepath = os.path.join(tempdir, "registry_output.tmp")
1264
+
1265
+ # `cloudpickle` allows pickling lambda functions directly
1266
+ import cloudpickle
1267
+
1268
+ input_bytes = cloudpickle.dumps((fn, output_filepath))
1269
+
1270
+ # cannot use `sys.executable __file__` here because the script
1271
+ # contains relative imports
1272
+ returned = subprocess.run(
1273
+ _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True
1274
+ )
1275
+
1276
+ # check if the subprocess is successful
1277
+ try:
1278
+ returned.check_returncode()
1279
+ except Exception as e:
1280
+ # wrap raised exception to provide more information
1281
+ raise RuntimeError(
1282
+ f"Error raised in subprocess:\n{returned.stderr.decode()}"
1283
+ ) from e
1284
+
1285
+ with open(output_filepath, "rb") as f:
1286
+ return pickle.load(f)
1287
+
1288
+
1289
+ def _run() -> None:
1290
+ # Setup plugins
1291
+ from vllm.plugins import load_general_plugins
1292
+
1293
+ load_general_plugins()
1294
+
1295
+ fn, output_file = pickle.loads(sys.stdin.buffer.read())
1296
+
1297
+ result = fn()
1298
+
1299
+ with open(output_file, "wb") as f:
1300
+ f.write(pickle.dumps(result))
1301
+
1302
+
1303
+ if __name__ == "__main__":
1304
+ _run()
patches/solar_open.py ADDED
@@ -0,0 +1,693 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: Apache-2.0
2
+ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+ #
4
+ # Adopted from
5
+ # https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/models/solar_open/modeling_solar_open.py
6
+ #
7
+ #
8
+ # Copyright 2026 Upstage Inc.
9
+ # Copyright 2026 The vLLM team.
10
+ #
11
+ #
12
+ # Licensed under the Apache License, Version 2.0 (the "License");
13
+ # you may not use this file except in compliance with the License.
14
+ # You may obtain a copy of the License at
15
+ #
16
+ # http://www.apache.org/licenses/LICENSE-2.0
17
+ #
18
+ # Unless required by applicable law or agreed to in writing, software
19
+ # distributed under the License is distributed on an "AS IS" BASIS,
20
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
21
+ # See the License for the specific language governing permissions and
22
+ # limitations under the License.
23
+ """Inference-only Solar Open model compatible with HuggingFace weights."""
24
+
25
+ import typing
26
+ from collections.abc import Callable, Iterable
27
+ from itertools import islice
28
+
29
+ import torch
30
+ from torch import nn
31
+
32
+ from vllm.compilation.decorators import support_torch_compile
33
+ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
34
+ from vllm.distributed import (
35
+ get_ep_group,
36
+ get_pp_group,
37
+ get_tensor_model_parallel_world_size,
38
+ )
39
+ from vllm.logger import init_logger
40
+ from vllm.model_executor.layers.activation import SiluAndMul
41
+ from vllm.model_executor.layers.attention.attention import Attention
42
+ from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE
43
+ from vllm.model_executor.layers.layernorm import RMSNorm
44
+ from vllm.model_executor.layers.linear import (
45
+ MergedColumnParallelLinear,
46
+ QKVParallelLinear,
47
+ RowParallelLinear,
48
+ )
49
+ from vllm.model_executor.layers.logits_processor import LogitsProcessor
50
+ from vllm.model_executor.layers.quantization import QuantizationConfig
51
+ from vllm.model_executor.layers.rotary_embedding import get_rope
52
+ from vllm.model_executor.layers.vocab_parallel_embedding import (
53
+ ParallelLMHead,
54
+ VocabParallelEmbedding,
55
+ )
56
+ from vllm.model_executor.model_loader.weight_utils import (
57
+ default_weight_loader,
58
+ maybe_remap_kv_scale_name,
59
+ )
60
+ from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
61
+ from vllm.model_executor.models.utils import (
62
+ AutoWeightsLoader,
63
+ PPMissingLayer,
64
+ is_pp_missing_parameter,
65
+ make_empty_intermediate_tensors_factory,
66
+ make_layers,
67
+ maybe_prefix,
68
+ )
69
+ from vllm.sequence import IntermediateTensors
70
+
71
+ logger = init_logger(__name__)
72
+
73
+
74
+ class SolarOpenDecoderLayer(nn.Module):
75
+ def __init__(
76
+ self,
77
+ hf_config,
78
+ cache_config: CacheConfig | None = None,
79
+ quant_config: QuantizationConfig | None = None,
80
+ prefix: str = "",
81
+ enable_eplb: bool = False,
82
+ ) -> None:
83
+ super().__init__()
84
+ self.hidden_size = hf_config.hidden_size
85
+ rope_theta = getattr(hf_config, "rope_theta", 1000000)
86
+ max_position_embeddings = getattr(hf_config, "max_position_embeddings", 131072)
87
+ # DecoderLayers are created with `make_layers` which passes the prefix
88
+ # with the layer's index.
89
+ layer_idx = int(prefix.split(sep=".")[-1])
90
+ self.layer_idx = layer_idx
91
+
92
+ self.self_attn = SolarOpenAttention(
93
+ hf_config=hf_config,
94
+ hidden_size=self.hidden_size,
95
+ num_heads=hf_config.num_attention_heads,
96
+ num_kv_heads=hf_config.num_key_value_heads,
97
+ rope_theta=rope_theta,
98
+ max_position_embeddings=max_position_embeddings,
99
+ head_dim=hf_config.head_dim,
100
+ qkv_bias=hf_config.attention_bias,
101
+ cache_config=cache_config,
102
+ quant_config=quant_config,
103
+ prefix=f"{prefix}.self_attn",
104
+ )
105
+
106
+ self.mlp = SolarOpenMoE(
107
+ hf_config=hf_config,
108
+ quant_config=quant_config,
109
+ prefix=f"{prefix}.mlp",
110
+ enable_eplb=enable_eplb,
111
+ )
112
+
113
+ self.input_layernorm = RMSNorm(
114
+ hf_config.hidden_size, eps=hf_config.rms_norm_eps
115
+ )
116
+ self.post_attention_layernorm = RMSNorm(
117
+ hf_config.hidden_size, eps=hf_config.rms_norm_eps
118
+ )
119
+ self.routed_scaling_factor = hf_config.routed_scaling_factor
120
+
121
+ def forward(
122
+ self,
123
+ positions: torch.Tensor,
124
+ hidden_states: torch.Tensor,
125
+ residual: torch.Tensor | None,
126
+ ) -> tuple[torch.Tensor, torch.Tensor]:
127
+ if residual is None:
128
+ residual = hidden_states
129
+ hidden_states = self.input_layernorm(hidden_states)
130
+ else:
131
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
132
+ hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
133
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
134
+ hidden_states = self.mlp(hidden_states)
135
+ return hidden_states, residual
136
+
137
+
138
+ class SolarOpenMLP(nn.Module):
139
+ def __init__(
140
+ self,
141
+ hidden_size: int,
142
+ intermediate_size: int,
143
+ hidden_act: str,
144
+ quant_config: QuantizationConfig | None = None,
145
+ reduce_results: bool = True,
146
+ prefix: str = "",
147
+ ) -> None:
148
+ super().__init__()
149
+ self.gate_up_proj = MergedColumnParallelLinear(
150
+ hidden_size,
151
+ [intermediate_size] * 2,
152
+ bias=False,
153
+ quant_config=quant_config,
154
+ prefix=f"{prefix}.gate_up_proj",
155
+ )
156
+ self.down_proj = RowParallelLinear(
157
+ intermediate_size,
158
+ hidden_size,
159
+ bias=False,
160
+ quant_config=quant_config,
161
+ reduce_results=reduce_results,
162
+ prefix=f"{prefix}.down_proj",
163
+ )
164
+ if hidden_act != "silu":
165
+ raise ValueError(
166
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
167
+ )
168
+ self.act_fn = SiluAndMul()
169
+
170
+ def forward(self, x):
171
+ gate_up, _ = self.gate_up_proj(x)
172
+ x = self.act_fn(gate_up)
173
+ x, _ = self.down_proj(x)
174
+ return x
175
+
176
+
177
+ class SolarOpenMoE(nn.Module):
178
+ def __init__(
179
+ self,
180
+ hf_config,
181
+ quant_config: QuantizationConfig | None = None,
182
+ prefix: str = "",
183
+ enable_eplb: bool = False,
184
+ ):
185
+ super().__init__()
186
+ self.tp_size = get_tensor_model_parallel_world_size()
187
+ self.routed_scaling_factor = hf_config.routed_scaling_factor
188
+
189
+ self.ep_group = get_ep_group().device_group
190
+ self.ep_rank = self.ep_group.rank()
191
+ self.ep_size = self.ep_group.size()
192
+ self.n_routed_experts: int = hf_config.n_routed_experts
193
+ self.n_shared_experts: int = hf_config.n_shared_experts
194
+
195
+ if hf_config.hidden_act != "silu":
196
+ raise ValueError(
197
+ f"Unsupported activation: {hf_config.hidden_act}. "
198
+ "Only silu is supported for now."
199
+ )
200
+
201
+ self.gate = nn.Linear(
202
+ hf_config.hidden_size,
203
+ hf_config.n_routed_experts,
204
+ bias=False,
205
+ dtype=torch.float32,
206
+ )
207
+ self.gate.e_score_correction_bias = nn.Parameter(
208
+ torch.empty(hf_config.n_routed_experts, dtype=torch.float32)
209
+ )
210
+
211
+ # Load balancing settings.
212
+ vllm_config = get_current_vllm_config()
213
+ eplb_config = vllm_config.parallel_config.eplb_config
214
+ self.enable_eplb = enable_eplb
215
+
216
+ self.n_redundant_experts = eplb_config.num_redundant_experts
217
+ self.n_logical_experts = self.n_routed_experts
218
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
219
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
220
+
221
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
222
+ self.physical_expert_end = (
223
+ self.physical_expert_start + self.n_local_physical_experts
224
+ )
225
+
226
+ if hf_config.n_shared_experts is not None:
227
+ intermediate_size = (
228
+ hf_config.moe_intermediate_size * hf_config.n_shared_experts
229
+ )
230
+ self.shared_experts = SolarOpenMLP(
231
+ hidden_size=hf_config.hidden_size,
232
+ intermediate_size=intermediate_size,
233
+ hidden_act=hf_config.hidden_act,
234
+ quant_config=quant_config,
235
+ reduce_results=False,
236
+ prefix=f"{prefix}.shared_experts",
237
+ )
238
+ self.experts = SharedFusedMoE(
239
+ shared_experts=self.shared_experts,
240
+ num_experts=hf_config.n_routed_experts,
241
+ top_k=hf_config.num_experts_per_tok,
242
+ hidden_size=hf_config.hidden_size,
243
+ intermediate_size=hf_config.moe_intermediate_size,
244
+ reduce_results=False,
245
+ renormalize=hf_config.norm_topk_prob,
246
+ quant_config=quant_config,
247
+ use_grouped_topk=True,
248
+ num_expert_group=hf_config.n_group,
249
+ topk_group=hf_config.topk_group,
250
+ prefix=f"{prefix}.experts",
251
+ scoring_func="sigmoid",
252
+ # we do scaling outside, set factor to 1.0 to avoid double mul
253
+ routed_scaling_factor=1.0,
254
+ e_score_correction_bias=self.gate.e_score_correction_bias,
255
+ enable_eplb=self.enable_eplb,
256
+ num_redundant_experts=self.n_redundant_experts,
257
+ )
258
+ else:
259
+ self.experts = FusedMoE(
260
+ num_experts=hf_config.n_routed_experts,
261
+ top_k=hf_config.num_experts_per_tok,
262
+ hidden_size=hf_config.hidden_size,
263
+ intermediate_size=hf_config.moe_intermediate_size,
264
+ reduce_results=False,
265
+ renormalize=hf_config.norm_topk_prob,
266
+ quant_config=quant_config,
267
+ use_grouped_topk=True,
268
+ num_expert_group=hf_config.n_group,
269
+ topk_group=hf_config.topk_group,
270
+ prefix=f"{prefix}.experts",
271
+ scoring_func="sigmoid",
272
+ # we do scaling outside, set factor to 1.0 to avoid double mul
273
+ routed_scaling_factor=1.0,
274
+ e_score_correction_bias=self.gate.e_score_correction_bias,
275
+ enable_eplb=self.enable_eplb,
276
+ num_redundant_experts=self.n_redundant_experts,
277
+ )
278
+
279
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
280
+ num_tokens, hidden_dim = hidden_states.shape
281
+ hidden_states = hidden_states.view(-1, hidden_dim)
282
+
283
+ # router_logits: (num_tokens, n_experts)
284
+ router_logits = self.gate(hidden_states.to(dtype=torch.float32))
285
+
286
+ fused_moe_out = self.experts(
287
+ hidden_states=hidden_states, router_logits=router_logits
288
+ )
289
+
290
+ if self.shared_experts is not None:
291
+ shared_output, final_hidden_states = fused_moe_out
292
+ assert shared_output is not None
293
+ final_hidden_states = (
294
+ final_hidden_states * self.routed_scaling_factor + shared_output
295
+ )
296
+ else:
297
+ final_hidden_states = fused_moe_out * self.routed_scaling_factor
298
+
299
+ if self.tp_size > 1:
300
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
301
+ final_hidden_states
302
+ )
303
+ return final_hidden_states.view(num_tokens, hidden_dim)
304
+
305
+
306
+ class SolarOpenAttention(nn.Module):
307
+ def __init__(
308
+ self,
309
+ hf_config,
310
+ hidden_size: int,
311
+ num_heads: int,
312
+ num_kv_heads: int,
313
+ rope_theta: float = 1000000,
314
+ max_position_embeddings: int = 131072,
315
+ head_dim: int | None = 128,
316
+ qkv_bias: bool = False,
317
+ cache_config: CacheConfig | None = None,
318
+ quant_config: QuantizationConfig | None = None,
319
+ prefix: str = "",
320
+ ) -> None:
321
+ super().__init__()
322
+ self.hidden_size = hidden_size
323
+ tp_size = get_tensor_model_parallel_world_size()
324
+ self.total_num_heads = num_heads
325
+ assert self.total_num_heads % tp_size == 0
326
+ self.num_heads = self.total_num_heads // tp_size
327
+ self.total_num_kv_heads = num_kv_heads
328
+ if self.total_num_kv_heads >= tp_size:
329
+ # Number of KV heads is greater than TP size, so we partition
330
+ # the KV heads across multiple tensor parallel GPUs.
331
+ assert self.total_num_kv_heads % tp_size == 0
332
+ else:
333
+ # Number of KV heads is less than TP size, so we replicate
334
+ # the KV heads across multiple tensor parallel GPUs.
335
+ assert tp_size % self.total_num_kv_heads == 0
336
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
337
+ self.head_dim = head_dim or (hidden_size // self.total_num_heads)
338
+ self.q_size = self.num_heads * self.head_dim
339
+ self.kv_size = self.num_kv_heads * self.head_dim
340
+ self.scaling = self.head_dim**-0.5
341
+ self.rope_theta = rope_theta
342
+ self.max_position_embeddings = max_position_embeddings
343
+
344
+ self.qkv_proj = QKVParallelLinear(
345
+ hidden_size,
346
+ self.head_dim,
347
+ self.total_num_heads,
348
+ self.total_num_kv_heads,
349
+ bias=qkv_bias,
350
+ quant_config=quant_config,
351
+ prefix=f"{prefix}.qkv_proj",
352
+ )
353
+
354
+ self.o_proj = RowParallelLinear(
355
+ self.total_num_heads * self.head_dim,
356
+ hidden_size,
357
+ bias=False,
358
+ quant_config=quant_config,
359
+ prefix=f"{prefix}.o_proj",
360
+ )
361
+
362
+ hf_config.rope_parameters.setdefault("partial_rotary_factor", 1.0)
363
+ self.rotary_emb = get_rope(
364
+ self.head_dim,
365
+ max_position=max_position_embeddings,
366
+ rope_parameters=hf_config.rope_parameters,
367
+ )
368
+ self.attn = Attention(
369
+ self.num_heads,
370
+ self.head_dim,
371
+ self.scaling,
372
+ num_kv_heads=self.num_kv_heads,
373
+ cache_config=cache_config,
374
+ quant_config=quant_config,
375
+ prefix=f"{prefix}.attn",
376
+ )
377
+
378
+ def forward(
379
+ self,
380
+ positions: torch.Tensor,
381
+ hidden_states: torch.Tensor,
382
+ ) -> torch.Tensor:
383
+ qkv, _ = self.qkv_proj(hidden_states)
384
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
385
+ q, k = self.rotary_emb(positions, q, k)
386
+ attn_output = self.attn(q, k, v)
387
+ output, _ = self.o_proj(attn_output)
388
+ return output
389
+
390
+
391
+ @support_torch_compile(
392
+ dynamic_arg_dims={
393
+ "input_ids": 0,
394
+ "positions": -1,
395
+ "intermediate_tensors": 0,
396
+ "inputs_embeds": 0,
397
+ }
398
+ )
399
+ class SolarOpenModel(nn.Module):
400
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
401
+ super().__init__()
402
+
403
+ hf_config = vllm_config.model_config.hf_config
404
+ cache_config = vllm_config.cache_config
405
+ quant_config = vllm_config.quant_config
406
+ enable_eplb = vllm_config.parallel_config.enable_eplb
407
+ self.hf_config = hf_config
408
+
409
+ self.vocab_size = hf_config.vocab_size
410
+
411
+ if get_pp_group().is_first_rank:
412
+ self.embed_tokens = VocabParallelEmbedding(
413
+ hf_config.vocab_size,
414
+ hf_config.hidden_size,
415
+ prefix=f"{prefix}.embed_tokens",
416
+ )
417
+ else:
418
+ self.embed_tokens = PPMissingLayer()
419
+
420
+ self.start_layer, self.end_layer, self.layers = make_layers(
421
+ hf_config.num_hidden_layers,
422
+ lambda prefix: SolarOpenDecoderLayer(
423
+ hf_config=hf_config,
424
+ cache_config=cache_config,
425
+ quant_config=quant_config,
426
+ prefix=prefix,
427
+ enable_eplb=enable_eplb,
428
+ ),
429
+ prefix=f"{prefix}.layers",
430
+ )
431
+
432
+ if get_pp_group().is_last_rank:
433
+ self.norm = RMSNorm(hf_config.hidden_size, eps=hf_config.rms_norm_eps)
434
+ else:
435
+ self.norm = PPMissingLayer()
436
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
437
+ ["hidden_states", "residual"], hf_config.hidden_size
438
+ )
439
+
440
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
441
+ return self.embed_tokens(input_ids)
442
+
443
+ def forward(
444
+ self,
445
+ input_ids: torch.Tensor | None,
446
+ positions: torch.Tensor,
447
+ intermediate_tensors: IntermediateTensors | None = None,
448
+ inputs_embeds: torch.Tensor | None = None,
449
+ ) -> torch.Tensor | IntermediateTensors:
450
+ if get_pp_group().is_first_rank:
451
+ if inputs_embeds is not None:
452
+ hidden_states = inputs_embeds
453
+ else:
454
+ hidden_states = self.embed_input_ids(input_ids)
455
+ residual = None
456
+ else:
457
+ assert intermediate_tensors is not None
458
+ hidden_states = intermediate_tensors["hidden_states"]
459
+ residual = intermediate_tensors["residual"]
460
+
461
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
462
+ hidden_states, residual = layer(positions, hidden_states, residual)
463
+
464
+ if not get_pp_group().is_last_rank:
465
+ return IntermediateTensors(
466
+ {"hidden_states": hidden_states, "residual": residual}
467
+ )
468
+
469
+ hidden_states, _ = self.norm(hidden_states, residual)
470
+ return hidden_states
471
+
472
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
473
+ # Params for weights, fp8 weight scales, fp8 activation scales
474
+ # (param_name, weight_name, expert_id, shard_id)
475
+ return SharedFusedMoE.make_expert_params_mapping(
476
+ self,
477
+ ckpt_gate_proj_name="gate_proj",
478
+ ckpt_down_proj_name="down_proj",
479
+ ckpt_up_proj_name="up_proj",
480
+ num_experts=self.hf_config.n_routed_experts,
481
+ )
482
+
483
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
484
+ stacked_params_mapping = [
485
+ # (param_name, shard_name, shard_id)
486
+ ("qkv_proj", "q_proj", "q"),
487
+ ("qkv_proj", "k_proj", "k"),
488
+ ("qkv_proj", "v_proj", "v"),
489
+ ("gate_up_proj", "gate_proj", 0),
490
+ ("gate_up_proj", "up_proj", 1),
491
+ ]
492
+
493
+ params_dict = dict(self.named_parameters())
494
+ loaded_params: set[str] = set()
495
+ expert_params_mapping = self.get_expert_mapping()
496
+ for name, loaded_weight in weights:
497
+ for param_name, weight_name, shard_id in stacked_params_mapping:
498
+ # Skip non-stacked layers and experts (experts handled below).
499
+ if weight_name not in name:
500
+ continue
501
+ # We have mlp.experts[0].gate_proj in the checkpoint.
502
+ # Since we handle the experts below in expert_params_mapping,
503
+ # we need to skip here BEFORE we update the name, otherwise
504
+ # name will be updated to mlp.experts[0].gate_up_proj, which
505
+ # will then be updated below in expert_params_mapping
506
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
507
+ if ("mlp.experts." in name) and name not in params_dict:
508
+ continue
509
+ name = name.replace(weight_name, param_name)
510
+ # Skip loading extra bias for GPTQ models.
511
+ if name.endswith(".bias") and name not in params_dict:
512
+ continue
513
+ if is_pp_missing_parameter(name, self):
514
+ continue
515
+
516
+ param = params_dict[name]
517
+ weight_loader = param.weight_loader
518
+ weight_loader(param, loaded_weight, shard_id)
519
+ break
520
+ else:
521
+ is_expert_weight = False
522
+ for mapping in expert_params_mapping:
523
+ param_name, weight_name, expert_id, shard_id = mapping
524
+ if weight_name not in name:
525
+ continue
526
+
527
+ # Anyway, this is an expert weight and should not be
528
+ # attempted to load as other weights later
529
+ is_expert_weight = True
530
+
531
+ # Do not modify `name` since the loop may continue here
532
+ # Instead, create a new variable
533
+ name_mapped = name.replace(weight_name, param_name)
534
+
535
+ if is_pp_missing_parameter(name_mapped, self):
536
+ continue
537
+
538
+ param = params_dict[name_mapped]
539
+ # We should ask the weight loader to return success or not
540
+ # here since otherwise we may skip experts with other
541
+ # available replicas.
542
+ weight_loader = typing.cast(
543
+ Callable[..., bool], param.weight_loader
544
+ )
545
+ success = weight_loader(
546
+ param,
547
+ loaded_weight,
548
+ name_mapped,
549
+ shard_id=shard_id,
550
+ expert_id=expert_id,
551
+ return_success=True,
552
+ )
553
+ if success:
554
+ name = name_mapped
555
+ break
556
+ else:
557
+ if is_expert_weight:
558
+ # We've checked that this is an expert weight
559
+ # However it's not mapped locally to this rank
560
+ # So we simply skip it
561
+ continue
562
+
563
+ # Skip loading extra bias for GPTQ models.
564
+ if name.endswith(".bias") and name not in params_dict:
565
+ continue
566
+
567
+ # Remapping the name of FP8 kv-scale.
568
+ name = maybe_remap_kv_scale_name(name, params_dict)
569
+ if name is None:
570
+ continue
571
+
572
+ if is_pp_missing_parameter(name, self):
573
+ continue
574
+
575
+ param = params_dict[name]
576
+ weight_loader = getattr(
577
+ param, "weight_loader", default_weight_loader
578
+ )
579
+ weight_loader(param, loaded_weight)
580
+ loaded_params.add(name)
581
+
582
+ return loaded_params
583
+
584
+
585
+ class SolarOpenForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
586
+ packed_modules_mapping = {
587
+ "qkv_proj": [
588
+ "q_proj",
589
+ "k_proj",
590
+ "v_proj",
591
+ ],
592
+ "gate_up_proj": [
593
+ "gate_proj",
594
+ "up_proj",
595
+ ],
596
+ }
597
+
598
+ fall_back_to_pt_during_load = False
599
+
600
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
601
+ super().__init__()
602
+ config = vllm_config.model_config.hf_config
603
+ quant_config = vllm_config.quant_config
604
+ self.config = config
605
+ self.quant_config = quant_config
606
+ self.model = SolarOpenModel(
607
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
608
+ )
609
+ if get_pp_group().is_last_rank:
610
+ self.lm_head = ParallelLMHead(
611
+ config.vocab_size,
612
+ config.hidden_size,
613
+ quant_config=quant_config,
614
+ prefix=maybe_prefix(prefix, "lm_head"),
615
+ )
616
+ else:
617
+ self.lm_head = PPMissingLayer()
618
+ self.logits_processor = LogitsProcessor(config.vocab_size)
619
+ self.make_empty_intermediate_tensors = (
620
+ self.model.make_empty_intermediate_tensors
621
+ )
622
+ self.expert_weights = []
623
+
624
+ # Set MoE hyperparameters
625
+ self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
626
+ self.num_expert_groups = config.n_group
627
+
628
+ self.moe_layers: list[FusedMoE] = []
629
+ example_moe = None
630
+ for layer in self.model.layers:
631
+ if isinstance(layer, PPMissingLayer):
632
+ continue
633
+
634
+ assert isinstance(layer, SolarOpenDecoderLayer)
635
+ if isinstance(layer.mlp, SolarOpenMoE):
636
+ # Pick last one layer since the first ones may be dense layers.
637
+ example_moe = layer.mlp
638
+ self.moe_layers.append(layer.mlp.experts)
639
+
640
+ if example_moe is None:
641
+ raise RuntimeError("No SolarOpenMoE layer found in model.layers.")
642
+
643
+ self.num_logical_experts = example_moe.n_logical_experts
644
+ self.num_physical_experts = example_moe.n_physical_experts
645
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
646
+ self.num_routed_experts = example_moe.n_routed_experts
647
+ self.num_shared_experts = example_moe.n_shared_experts
648
+ self.num_redundant_experts = example_moe.n_redundant_experts
649
+
650
+ def set_eplb_state(
651
+ self,
652
+ expert_load_view: torch.Tensor,
653
+ logical_to_physical_map: torch.Tensor,
654
+ logical_replica_count: torch.Tensor,
655
+ ) -> None:
656
+ for layer_idx, layer in enumerate(self.moe_layers):
657
+ # Register the expert weights.
658
+ self.expert_weights.append(layer.get_expert_weights())
659
+ layer.set_eplb_state(
660
+ moe_layer_idx=layer_idx,
661
+ expert_load_view=expert_load_view,
662
+ logical_to_physical_map=logical_to_physical_map,
663
+ logical_replica_count=logical_replica_count,
664
+ )
665
+
666
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
667
+ return self.model.embed_input_ids(input_ids)
668
+
669
+ def forward(
670
+ self,
671
+ input_ids: torch.Tensor,
672
+ positions: torch.Tensor,
673
+ intermediate_tensors: IntermediateTensors | None = None,
674
+ inputs_embeds: torch.Tensor | None = None,
675
+ ) -> torch.Tensor | IntermediateTensors:
676
+ hidden_states = self.model(
677
+ input_ids, positions, intermediate_tensors, inputs_embeds
678
+ )
679
+ return hidden_states
680
+
681
+ def compute_logits(
682
+ self,
683
+ hidden_states: torch.Tensor,
684
+ ) -> torch.Tensor | None:
685
+ logits = self.logits_processor(self.lm_head, hidden_states)
686
+ return logits
687
+
688
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
689
+ loader = AutoWeightsLoader(self)
690
+ return loader.load_weights(weights)
691
+
692
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
693
+ return self.model.get_expert_mapping()