A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 03:45 UTC
[Model] Add MiMo-V2-Flash support (#30836)
Signed-off-by: Abatom <abzhonghua@gmail.com> Signed-off-by: Jumiar <liuanqim10@126.com> Signed-off-by: Zyann7 <zyann7@outlook.com> Co-authored-by: Jumiar <liuanqim10@126.com> Co-authored-by: Zyann7 <zyann7@outlook.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Zhonghua Deng committed
Dec 20, 2025 at 01:17 UTC
969bbc7c6166ce10a48df2fcdc2fba508dbafcce
8 files changed
+789
-13
docs/models/supported_models.md
+1
@@ -415,6 +415,7 @@ th {
415
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
416
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
417
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
418
+| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | ︎| ✅︎ |
419
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
420
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
421
| `MiniMaxM2ForCausalLM` | MiniMax-M2 |`MiniMaxAI/MiniMax-M2`, etc. | | ✅︎ |
tests/models/registry.py
+3
@@ -459,6 +459,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
459
),
460
"Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"),
461
"MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL", trust_remote_code=True),
462
+ "MiMoV2FlashForCausalLM": _HfExamplesInfo(
463
+ "XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True
464
+ ),
465
"Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),
466
}
467
vllm/config/__init__.py
+2
@@ -18,6 +18,7 @@ from vllm.config.lora import LoRAConfig
18
from vllm.config.model import (
19
ModelConfig,
20
iter_architecture_defaults,
21
+ str_dtype_to_torch_dtype,
22
try_match_architecture_defaults,
23
)
24
from vllm.config.multimodal import MultiModalConfig
@@ -72,6 +73,7 @@ __all__ = [
73
# From vllm.config.model
74
"ModelConfig",
75
"iter_architecture_defaults",
76
+ "str_dtype_to_torch_dtype",
77
"try_match_architecture_defaults",
78
# From vllm.config.multimodal
79
"MultiModalConfig",
vllm/config/model.py
+5
@@ -1849,6 +1849,11 @@ _STR_DTYPE_TO_TORCH_DTYPE = {
1849
"bfloat16": torch.bfloat16,
1850
}
1851
1852
+
1853
+def str_dtype_to_torch_dtype(type: str):
1854
+ return _STR_DTYPE_TO_TORCH_DTYPE.get(type)
1855
+
1856
+
1857
# model_type -> reason
1858
_FLOAT16_NOT_SUPPORTED_MODELS = {
1859
"gemma2": "Numerical instability. Please use bfloat16 or float32 instead.",
vllm/model_executor/layers/linear.py
+49
-13
@@ -277,6 +277,7 @@ class LinearBase(CustomOp):
277
self.params_dtype = params_dtype
278
self.quant_config = quant_config
279
self.prefix = prefix
280
+ self.allow_fp8_block_shape_mismatch = False
281
if quant_config is None:
282
self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod()
283
else:
@@ -475,6 +476,7 @@ class ColumnParallelLinear(LinearBase):
476
disable_tp=disable_tp,
477
)
478
479
+ self._maybe_allow_fp8_block_shape_mismatch()
480
self.gather_output = gather_output
481
482
if output_sizes is None:
@@ -509,6 +511,33 @@ class ColumnParallelLinear(LinearBase):
511
self.register_parameter("bias", None)
512
self.update_param_tp_status()
513
514
+ def _maybe_allow_fp8_block_shape_mismatch(self) -> None:
515
+ quant_config = getattr(self, "quant_config", None)
516
+ weight_block = getattr(quant_config, "weight_block_size", None)
517
+ if (
518
+ weight_block is None
519
+ or len(weight_block) < 1
520
+ or len(self.output_partition_sizes) <= 1
521
+ ):
522
+ return
523
+
524
+ try:
525
+ block_n = int(weight_block[0])
526
+ except (ValueError, TypeError):
527
+ return
528
+
529
+ if block_n <= 0:
530
+ return
531
+
532
+ if any(size % block_n != 0 for size in self.output_partition_sizes):
533
+ self.allow_fp8_block_shape_mismatch = True
534
+ logger.debug(
535
+ "Allowing FP8 block shape mismatch for %s (block_n=%d, partitions=%s)",
536
+ getattr(self, "prefix", "<unknown>"),
537
+ block_n,
538
+ self.output_partition_sizes,
539
+ )
540
+
541
def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
542
output_dim = getattr(param, "output_dim", None)
543
@@ -906,9 +935,11 @@ class QKVParallelLinear(ColumnParallelLinear):
935
*,
936
return_bias: bool = True,
937
disable_tp: bool = False,
938
+ v_head_size: int | None = None,
939
):
940
self.hidden_size = hidden_size
941
self.head_size = head_size
942
+ self.v_head_size = v_head_size if v_head_size is not None else head_size
943
self.total_num_heads = total_num_heads
944
if total_num_kv_heads is None:
945
total_num_kv_heads = total_num_heads
@@ -924,12 +955,14 @@ class QKVParallelLinear(ColumnParallelLinear):
955
self.num_kv_head_replicas = 1
956
input_size = self.hidden_size
957
output_size = (
927
- (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size
928
- )
958
+ self.num_heads * self.head_size
959
+ + self.num_kv_heads * self.head_size
960
+ + self.num_kv_heads * self.v_head_size
961
+ ) * tp_size
962
self.output_sizes = [
963
self.num_heads * self.head_size * tp_size, # q_proj
964
self.num_kv_heads * self.head_size * tp_size, # k_proj
932
- self.num_kv_heads * self.head_size * tp_size, # v_proj
965
+ self.num_kv_heads * self.v_head_size * tp_size, # v_proj
966
]
967
968
super().__init__(
@@ -950,7 +983,8 @@ class QKVParallelLinear(ColumnParallelLinear):
983
"q": 0,
984
"k": self.num_heads * self.head_size,
985
"v": (self.num_heads + self.num_kv_heads) * self.head_size,
953
- "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size,
986
+ "total": (self.num_heads + self.num_kv_heads) * self.head_size
987
+ + self.num_kv_heads * self.v_head_size,
988
}
989
return shard_offset_mapping.get(loaded_shard_id)
990
@@ -958,7 +992,7 @@ class QKVParallelLinear(ColumnParallelLinear):
992
shard_size_mapping = {
993
"q": self.num_heads * self.head_size,
994
"k": self.num_kv_heads * self.head_size,
961
- "v": self.num_kv_heads * self.head_size,
995
+ "v": self.num_kv_heads * self.v_head_size,
996
}
997
return shard_size_mapping.get(loaded_shard_id)
998
@@ -985,7 +1019,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1019
(
1020
"v",
1021
(self.total_num_heads + self.total_num_kv_heads) * self.head_size,
988
- self.total_num_kv_heads * self.head_size,
1022
+ self.total_num_kv_heads * self.v_head_size,
1023
),
1024
]
1025
@@ -1110,7 +1144,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1144
(
1145
"v",
1146
(self.total_num_heads + self.total_num_kv_heads) * self.head_size,
1113
- self.total_num_kv_heads * self.head_size,
1147
+ self.total_num_kv_heads * self.v_head_size,
1148
),
1149
]
1150
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
@@ -1139,11 +1173,12 @@ class QKVParallelLinear(ColumnParallelLinear):
1173
"v": (
1174
(self.total_num_heads + self.total_num_kv_heads)
1175
* self.head_size,
1142
- self.total_num_kv_heads * self.head_size,
1176
+ self.total_num_kv_heads * self.v_head_size,
1177
),
1178
"total": (
1145
- (self.total_num_heads + 2 * self.total_num_kv_heads)
1146
- * self.head_size,
1179
+ (self.total_num_heads + self.total_num_kv_heads)
1180
+ * self.head_size
1181
+ + self.total_num_kv_heads * self.v_head_size,
1182
0,
1183
),
1184
}
@@ -1170,7 +1205,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1205
shard_size = self.num_kv_heads * self.head_size
1206
elif loaded_shard_id == "v":
1207
shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size
1173
- shard_size = self.num_kv_heads * self.head_size
1208
+ shard_size = self.num_kv_heads * self.v_head_size
1209
# Special case for Quantized Weights.
1210
# If quantized, we need to adjust the offset and size to account
1211
# for the packing.
@@ -1199,10 +1234,11 @@ class QKVParallelLinear(ColumnParallelLinear):
1234
),
1235
"v": (
1236
(self.num_heads + self.num_kv_heads) * self.head_size,
1202
- self.num_kv_heads * self.head_size,
1237
+ self.num_kv_heads * self.v_head_size,
1238
),
1239
"total": (
1205
- (self.num_heads + 2 * self.num_kv_heads) * self.head_size,
1240
+ (self.num_heads + self.num_kv_heads) * self.head_size
1241
+ + self.num_kv_heads * self.v_head_size,
1242
0,
1243
),
1244
}
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+8
@@ -1252,6 +1252,14 @@ def validate_fp8_block_shape(
1252
"""Validate block quantization shapes for tensor parallelism."""
1253
from vllm.distributed import get_tensor_model_parallel_world_size
1254
1255
+ if getattr(layer, "allow_fp8_block_shape_mismatch", False):
1256
+ logger.debug(
1257
+ "Skipping FP8 block shape validation for layer %s due to detected"
1258
+ " mismatch allowance.",
1259
+ getattr(layer, "prefix", "<unknown>"),
1260
+ )
1261
+ return
1262
+
1263
tp_size = getattr(layer, "tp_size", get_tensor_model_parallel_world_size())
1264
block_n, block_k = block_size[0], block_size[1]
1265
vllm/model_executor/models/mimo_v2_flash.py
new
+720
@@ -0,0 +1,720 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
+from itertools import islice
5
+
6
+import torch
7
+from torch import nn
8
+
9
+from vllm.attention.backends.abstract import AttentionType
10
+from vllm.attention.layer import Attention
11
+from vllm.config import (
12
+ CacheConfig,
13
+ VllmConfig,
14
+ get_current_vllm_config,
15
+ str_dtype_to_torch_dtype,
16
+)
17
+from vllm.distributed import (
18
+ get_ep_group,
19
+ get_pp_group,
20
+ get_tensor_model_parallel_rank,
21
+ get_tensor_model_parallel_world_size,
22
+ tensor_model_parallel_all_gather,
23
+)
24
+from vllm.logger import init_logger
25
+from vllm.model_executor.layers.activation import SiluAndMul
26
+from vllm.model_executor.layers.fused_moe import FusedMoE
27
+from vllm.model_executor.layers.layernorm import RMSNorm
28
+from vllm.model_executor.layers.linear import (
29
+ MergedColumnParallelLinear,
30
+ QKVParallelLinear,
31
+ RowParallelLinear,
32
+)
33
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
34
+from vllm.model_executor.layers.quantization import QuantizationConfig
35
+from vllm.model_executor.layers.rotary_embedding import get_rope
36
+from vllm.model_executor.layers.vocab_parallel_embedding import (
37
+ ParallelLMHead,
38
+ VocabParallelEmbedding,
39
+)
40
+from vllm.model_executor.model_loader.weight_utils import (
41
+ default_weight_loader,
42
+ maybe_remap_kv_scale_name,
43
+)
44
+from vllm.model_executor.models.utils import sequence_parallel_chunk
45
+from vllm.sequence import IntermediateTensors
46
+
47
+from .interfaces import MixtureOfExperts, SupportsPP
48
+from .utils import (
49
+ AutoWeightsLoader,
50
+ PPMissingLayer,
51
+ extract_layer_index,
52
+ is_pp_missing_parameter,
53
+ make_empty_intermediate_tensors_factory,
54
+ make_layers,
55
+ maybe_prefix,
56
+)
57
+
58
+logger = init_logger(__name__)
59
+
60
+
61
+class MiMoV2MLP(nn.Module):
62
+ def __init__(
63
+ self,
64
+ hidden_size: int,
65
+ intermediate_size: int,
66
+ hidden_act: str,
67
+ quant_config: QuantizationConfig | None = None,
68
+ reduce_results: bool = True,
69
+ prefix: str = "",
70
+ ) -> None:
71
+ super().__init__()
72
+ self.gate_up_proj = MergedColumnParallelLinear(
73
+ hidden_size,
74
+ [intermediate_size] * 2,
75
+ bias=False,
76
+ quant_config=quant_config,
77
+ prefix=f"{prefix}.gate_up_proj",
78
+ )
79
+ self.down_proj = RowParallelLinear(
80
+ intermediate_size,
81
+ hidden_size,
82
+ bias=False,
83
+ quant_config=quant_config,
84
+ reduce_results=reduce_results,
85
+ prefix=f"{prefix}.down_proj",
86
+ )
87
+ if hidden_act != "silu":
88
+ raise ValueError(
89
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
90
+ )
91
+ self.act_fn = SiluAndMul()
92
+
93
+ def forward(self, x):
94
+ gate_up, _ = self.gate_up_proj(x)
95
+ x = self.act_fn(gate_up)
96
+ x, _ = self.down_proj(x)
97
+ return x
98
+
99
+
100
+class MiMoV2MoE(nn.Module):
101
+ def __init__(
102
+ self,
103
+ vllm_config: VllmConfig,
104
+ prefix: str = "",
105
+ is_nextn: bool = False,
106
+ ):
107
+ super().__init__()
108
+
109
+ config = vllm_config.model_config.hf_text_config
110
+ parallel_config = vllm_config.parallel_config
111
+ quant_config = vllm_config.quant_config
112
+
113
+ self.tp_size = get_tensor_model_parallel_world_size()
114
+
115
+ self.ep_group = get_ep_group().device_group
116
+ self.ep_rank = get_ep_group().rank_in_group
117
+ self.ep_size = self.ep_group.size()
118
+ self.n_routed_experts = config.n_routed_experts
119
+
120
+ self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
121
+
122
+ if self.tp_size > config.n_routed_experts:
123
+ raise ValueError(
124
+ f"Tensor parallel size {self.tp_size} is greater than "
125
+ f"the number of experts {config.n_routed_experts}."
126
+ )
127
+
128
+ if config.hidden_act != "silu":
129
+ raise ValueError(
130
+ f"Unsupported activation: {config.hidden_act}. "
131
+ "Only silu is supported for now."
132
+ )
133
+
134
+ vllm_config = get_current_vllm_config()
135
+ eplb_config = vllm_config.parallel_config.eplb_config
136
+ self.enable_eplb = parallel_config.enable_eplb
137
+
138
+ self.n_logical_experts = self.n_routed_experts
139
+ self.n_redundant_experts = eplb_config.num_redundant_experts
140
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
141
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
142
+
143
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
144
+ self.physical_expert_end = (
145
+ self.physical_expert_start + self.n_local_physical_experts
146
+ )
147
+
148
+ dtype = getattr(config, "moe_router_dtype", "float32")
149
+ self.gate_dtype = str_dtype_to_torch_dtype(dtype)
150
+ self.gate = nn.Linear(
151
+ config.hidden_size,
152
+ config.n_routed_experts,
153
+ bias=False,
154
+ dtype=self.gate_dtype,
155
+ )
156
+ self.gate.e_score_correction_bias = nn.Parameter(
157
+ torch.empty(config.n_routed_experts, dtype=self.gate_dtype)
158
+ )
159
+
160
+ self.experts = FusedMoE(
161
+ num_experts=self.n_routed_experts,
162
+ top_k=config.num_experts_per_tok,
163
+ hidden_size=config.hidden_size,
164
+ intermediate_size=config.moe_intermediate_size,
165
+ reduce_results=True,
166
+ renormalize=config.norm_topk_prob,
167
+ quant_config=quant_config,
168
+ prefix=f"{prefix}.experts",
169
+ e_score_correction_bias=self.gate.e_score_correction_bias,
170
+ enable_eplb=self.enable_eplb,
171
+ num_redundant_experts=self.n_redundant_experts,
172
+ is_sequence_parallel=self.is_sequence_parallel,
173
+ use_grouped_topk=True,
174
+ num_expert_group=config.n_group,
175
+ topk_group=config.topk_group,
176
+ scoring_func="sigmoid",
177
+ )
178
+
179
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
180
+ assert hidden_states.dim() <= 2, "MiMoV2MoE only supports 1D or 2D inputs"
181
+ is_input_1d = hidden_states.dim() == 1
182
+ num_tokens, hidden_dim = hidden_states.shape
183
+ hidden_states = hidden_states.view(-1, hidden_dim)
184
+
185
+ if self.is_sequence_parallel:
186
+ hidden_states = sequence_parallel_chunk(hidden_states)
187
+
188
+ if self.gate_dtype is not None:
189
+ gate_input = hidden_states.to(self.gate_dtype)
190
+ else:
191
+ gate_input = hidden_states
192
+ router_logits = self.gate(gate_input)
193
+ final_hidden_states = self.experts(
194
+ hidden_states=hidden_states, router_logits=router_logits
195
+ )
196
+
197
+ if self.is_sequence_parallel:
198
+ final_hidden_states = tensor_model_parallel_all_gather(
199
+ final_hidden_states, 0
200
+ )
201
+ final_hidden_states = final_hidden_states[:num_tokens]
202
+
203
+ return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
204
+
205
+
206
+class MiMoV2Attention(nn.Module):
207
+ def __init__(
208
+ self,
209
+ hidden_size: int,
210
+ num_heads: int,
211
+ num_kv_heads: int,
212
+ head_dim: int,
213
+ v_head_dim: int | None = None,
214
+ sliding_window_size: int = -1,
215
+ attention_bias: bool = False,
216
+ add_swa_attention_sink_bias: bool = False,
217
+ layer_id: int = 0,
218
+ rope_theta: float = 1000000,
219
+ max_position_embeddings: int = 32768,
220
+ cache_config: CacheConfig | None = None,
221
+ quant_config: QuantizationConfig | None = None,
222
+ partial_rotary_factor: float = 1.0,
223
+ prefix: str = "",
224
+ ) -> None:
225
+ super().__init__()
226
+ self.hidden_size = hidden_size
227
+ self.layer_id = layer_id
228
+ tp_size = get_tensor_model_parallel_world_size()
229
+
230
+ self.total_num_heads = num_heads
231
+ self.num_heads = self.total_num_heads // tp_size
232
+
233
+ self.total_num_kv_heads = num_kv_heads
234
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
235
+
236
+ self.head_dim = head_dim
237
+
238
+ self.v_head_dim = v_head_dim if v_head_dim is not None else head_dim
239
+
240
+ self.q_size = self.num_heads * self.head_dim
241
+ self.k_size = self.num_kv_heads * self.head_dim
242
+ self.v_size = self.num_kv_heads * self.v_head_dim
243
+
244
+ self.scaling = self.head_dim**-0.5
245
+ self.rope_theta = rope_theta
246
+ self.max_position_embeddings = max_position_embeddings
247
+
248
+ self.qkv_proj = QKVParallelLinear(
249
+ hidden_size,
250
+ self.head_dim,
251
+ self.total_num_heads,
252
+ self.total_num_kv_heads,
253
+ bias=attention_bias,
254
+ quant_config=quant_config,
255
+ prefix=f"{prefix}.qkv_proj",
256
+ v_head_size=self.v_head_dim,
257
+ )
258
+
259
+ self.o_proj = RowParallelLinear(
260
+ self.total_num_heads * self.v_head_dim,
261
+ hidden_size,
262
+ bias=False,
263
+ quant_config=quant_config,
264
+ reduce_results=True,
265
+ prefix=f"{prefix}.o_proj",
266
+ )
267
+
268
+ self.rotary_emb = get_rope(
269
+ head_size=self.head_dim,
270
+ max_position=max_position_embeddings,
271
+ rope_parameters={
272
+ "rope_type": "default",
273
+ "rope_theta": rope_theta,
274
+ "partial_rotary_factor": partial_rotary_factor,
275
+ },
276
+ )
277
+
278
+ self.attention_sink_bias = (
279
+ torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False)
280
+ if add_swa_attention_sink_bias
281
+ else None
282
+ )
283
+
284
+ sliding_window = sliding_window_size if sliding_window_size > -1 else None
285
+ self.attn = Attention(
286
+ self.num_heads,
287
+ self.head_dim,
288
+ self.scaling,
289
+ num_kv_heads=self.num_kv_heads,
290
+ cache_config=cache_config,
291
+ quant_config=quant_config,
292
+ per_layer_sliding_window=sliding_window,
293
+ attn_type=AttentionType.DECODER,
294
+ prefix=f"{prefix}.attn",
295
+ sinks=self.attention_sink_bias,
296
+ )
297
+
298
+ def forward(
299
+ self,
300
+ positions: torch.Tensor,
301
+ hidden_states: torch.Tensor,
302
+ ) -> torch.Tensor:
303
+ qkv, _ = self.qkv_proj(hidden_states)
304
+ q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
305
+ q, k = self.rotary_emb(positions, q, k)
306
+
307
+ v = v.view(-1, self.num_kv_heads, self.v_head_dim)
308
+ v = torch.nn.functional.pad(v, [0, self.head_dim - self.v_head_dim], value=0)
309
+ v = v.view(-1, self.num_kv_heads * self.head_dim)
310
+
311
+ attn_output = self.attn(q, k, v)
312
+
313
+ attn_output = attn_output.view(-1, self.num_heads, self.head_dim)[
314
+ ..., : self.v_head_dim
315
+ ].reshape(-1, self.num_heads * self.v_head_dim)
316
+
317
+ output, _ = self.o_proj(attn_output)
318
+ return output
319
+
320
+
321
+class MiMoV2FlashDecoderLayer(nn.Module):
322
+ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
323
+ super().__init__()
324
+ config = vllm_config.model_config.hf_text_config
325
+ quant_config = vllm_config.quant_config
326
+ layer_id = extract_layer_index(prefix)
327
+
328
+ self.hidden_size = config.hidden_size
329
+ self.config = config
330
+ self.layer_id = layer_id
331
+
332
+ rope_theta = getattr(config, "rope_theta", 1000000)
333
+ max_position_embeddings = getattr(config, "max_position_embeddings", 32768)
334
+
335
+ if self.is_compressed_softmax_layer():
336
+ self.self_attn = MiMoV2Attention(
337
+ hidden_size=self.hidden_size,
338
+ num_heads=config.swa_num_attention_heads,
339
+ num_kv_heads=config.swa_num_key_value_heads,
340
+ head_dim=config.swa_head_dim,
341
+ v_head_dim=getattr(config, "swa_v_head_dim", None),
342
+ sliding_window_size=config.sliding_window_size,
343
+ attention_bias=config.attention_bias,
344
+ add_swa_attention_sink_bias=getattr(
345
+ config, "add_swa_attention_sink_bias", False
346
+ ),
347
+ layer_id=layer_id,
348
+ rope_theta=getattr(config, "swa_rope_theta", rope_theta),
349
+ max_position_embeddings=max_position_embeddings,
350
+ quant_config=quant_config,
351
+ partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0),
352
+ prefix=f"{prefix}.self_attn",
353
+ )
354
+ else:
355
+ self.self_attn = MiMoV2Attention(
356
+ hidden_size=self.hidden_size,
357
+ num_heads=config.num_attention_heads,
358
+ num_kv_heads=config.num_key_value_heads,
359
+ head_dim=config.head_dim,
360
+ v_head_dim=getattr(config, "v_head_dim", None),
361
+ sliding_window_size=-1, # normal attention
362
+ attention_bias=config.attention_bias,
363
+ layer_id=layer_id,
364
+ rope_theta=rope_theta,
365
+ max_position_embeddings=max_position_embeddings,
366
+ quant_config=quant_config,
367
+ partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0),
368
+ prefix=f"{prefix}.self_attn",
369
+ )
370
+
371
+ self.is_layer_sparse = self.is_moe_layer(layer_id)
372
+ if self.is_layer_sparse:
373
+ self.mlp = MiMoV2MoE(
374
+ vllm_config=vllm_config,
375
+ prefix=f"{prefix}.mlp",
376
+ )
377
+ else:
378
+ self.mlp = MiMoV2MLP(
379
+ hidden_size=self.hidden_size,
380
+ intermediate_size=config.intermediate_size,
381
+ hidden_act=config.hidden_act,
382
+ quant_config=quant_config,
383
+ prefix=f"{prefix}.mlp",
384
+ )
385
+
386
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
387
+ self.post_attention_layernorm = RMSNorm(
388
+ config.hidden_size, eps=config.layernorm_epsilon
389
+ )
390
+
391
+ def forward(
392
+ self,
393
+ positions: torch.Tensor,
394
+ hidden_states: torch.Tensor,
395
+ residual: torch.Tensor | None,
396
+ ) -> tuple[torch.Tensor, torch.Tensor]:
397
+ if residual is None:
398
+ residual = hidden_states
399
+ hidden_states = self.input_layernorm(hidden_states)
400
+ else:
401
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
402
+
403
+ hidden_states = self.self_attn(
404
+ positions=positions,
405
+ hidden_states=hidden_states,
406
+ )
407
+
408
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
409
+ hidden_states = self.mlp(hidden_states)
410
+ return hidden_states, residual
411
+
412
+ def is_moe_layer(self, layer_idx: int) -> bool:
413
+ return (
414
+ hasattr(self.config, "moe_layer_freq")
415
+ and layer_idx >= 0
416
+ and not isinstance(self.config.moe_layer_freq, int)
417
+ and self.config.moe_layer_freq[layer_idx]
418
+ )
419
+
420
+ def is_compressed_softmax_layer(self) -> bool:
421
+ return self.config.hybrid_layer_pattern[self.layer_id] == 1
422
+
423
+
424
+class MiMoV2Model(nn.Module):
425
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
426
+ super().__init__()
427
+
428
+ config = vllm_config.model_config.hf_config.get_text_config()
429
+ quant_config = vllm_config.quant_config
430
+ eplb_config = vllm_config.parallel_config.eplb_config
431
+
432
+ self.config = config
433
+ self.quant_config = quant_config
434
+ self.vocab_size = config.vocab_size
435
+ self.num_redundant_experts = eplb_config.num_redundant_experts
436
+ self.v_scale = getattr(config, "attention_value_scale", None)
437
+
438
+ if get_pp_group().is_first_rank or (
439
+ config.tie_word_embeddings and get_pp_group().is_last_rank
440
+ ):
441
+ self.embed_tokens = VocabParallelEmbedding(
442
+ config.vocab_size,
443
+ config.hidden_size,
444
+ quant_config=quant_config,
445
+ prefix=f"{prefix}.embed_tokens",
446
+ )
447
+ else:
448
+ self.embed_tokens = PPMissingLayer()
449
+
450
+ self.start_layer, self.end_layer, self.layers = make_layers(
451
+ config.num_hidden_layers,
452
+ lambda prefix: MiMoV2FlashDecoderLayer(
453
+ vllm_config=vllm_config,
454
+ prefix=prefix,
455
+ ),
456
+ prefix=f"{prefix}.layers",
457
+ )
458
+
459
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
460
+ ["hidden_states", "residual"], config.hidden_size
461
+ )
462
+ if get_pp_group().is_last_rank:
463
+ self.norm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
464
+ else:
465
+ self.norm = PPMissingLayer()
466
+
467
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
468
+ return self.embed_tokens(input_ids)
469
+
470
+ def forward(
471
+ self,
472
+ input_ids: torch.Tensor,
473
+ positions: torch.Tensor,
474
+ intermediate_tensors: IntermediateTensors | None = None,
475
+ inputs_embeds: torch.Tensor | None = None,
476
+ ) -> torch.Tensor | IntermediateTensors:
477
+ if get_pp_group().is_first_rank:
478
+ if inputs_embeds is not None:
479
+ hidden_states = inputs_embeds
480
+ else:
481
+ hidden_states = self.embed_input_ids(input_ids)
482
+ residual = None
483
+ else:
484
+ assert intermediate_tensors is not None
485
+ hidden_states = intermediate_tensors["hidden_states"]
486
+ residual = intermediate_tensors["residual"]
487
+
488
+ for idx, layer in enumerate(
489
+ islice(self.layers, self.start_layer, self.end_layer)
490
+ ):
491
+ hidden_states, residual = layer(positions, hidden_states, residual)
492
+
493
+ if not get_pp_group().is_last_rank:
494
+ return IntermediateTensors(
495
+ {"hidden_states": hidden_states, "residual": residual}
496
+ )
497
+
498
+ hidden_states, _ = self.norm(hidden_states, residual)
499
+
500
+ return hidden_states
501
+
502
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
503
+ # Params for weights, fp8 weight scales, fp8 activation scales
504
+ # (param_name, weight_name, expert_id, shard_id)
505
+ return FusedMoE.make_expert_params_mapping(
506
+ ckpt_gate_proj_name="gate_proj",
507
+ ckpt_down_proj_name="down_proj",
508
+ ckpt_up_proj_name="up_proj",
509
+ num_experts=self.config.n_routed_experts,
510
+ num_redundant_experts=self.num_redundant_experts,
511
+ )
512
+
513
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
514
+ stacked_params_mapping = [
515
+ # (param_name, shard_name, shard_id)
516
+ ("qkv_proj", "q_proj", "q"),
517
+ ("qkv_proj", "k_proj", "k"),
518
+ ("qkv_proj", "v_proj", "v"),
519
+ ("gate_up_proj", "gate_proj", 0),
520
+ ("gate_up_proj", "up_proj", 1),
521
+ ]
522
+
523
+ tp_rank = get_tensor_model_parallel_rank()
524
+ tp_size = get_tensor_model_parallel_world_size()
525
+
526
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
527
+ loaded_params: set[str] = set()
528
+ expert_params_mapping = self.get_expert_mapping()
529
+ for name, loaded_weight in weights:
530
+ if "rotary_emb.inv_freq" in name:
531
+ continue
532
+ if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
533
+ continue
534
+ if "mtp" in name:
535
+ continue
536
+
537
+ if self.quant_config is not None:
538
+ cache_scale_name = self.quant_config.get_cache_scale(name)
539
+ if cache_scale_name is not None and cache_scale_name in params_dict:
540
+ param = params_dict[cache_scale_name]
541
+ weight_loader = getattr(
542
+ param, "weight_loader", default_weight_loader
543
+ )
544
+
545
+ kv_scale = loaded_weight
546
+ if kv_scale.dim() > 0 and kv_scale.numel() > 1:
547
+ kv_scale = kv_scale.view(-1)[0]
548
+
549
+ weight_loader(param, kv_scale)
550
+ loaded_params.add(cache_scale_name)
551
+ continue
552
+
553
+ expert_matched = False
554
+ for param_name, weight_name, expert_id, shard_id in expert_params_mapping:
555
+ if weight_name not in name:
556
+ continue
557
+
558
+ name_rewritten = name.replace(weight_name, param_name)
559
+
560
+ if is_pp_missing_parameter(name_rewritten, self):
561
+ continue
562
+
563
+ if (
564
+ name_rewritten.endswith(".bias") or name_rewritten.endswith("_bias")
565
+ ) and name_rewritten not in params_dict:
566
+ continue
567
+
568
+ if name_rewritten not in params_dict:
569
+ continue
570
+
571
+ param = params_dict[name_rewritten]
572
+ weight_loader = param.weight_loader
573
+
574
+ weight_loader(
575
+ param,
576
+ loaded_weight,
577
+ name_rewritten,
578
+ shard_id=shard_id,
579
+ expert_id=expert_id,
580
+ )
581
+ loaded_params.add(name_rewritten)
582
+ expert_matched = True
583
+ break
584
+
585
+ if expert_matched:
586
+ continue
587
+
588
+ stacked_matched = False
589
+ for param_name, weight_name, shard_id in stacked_params_mapping:
590
+ if weight_name not in name:
591
+ continue
592
+ name_rewritten = name.replace(weight_name, param_name)
593
+
594
+ if (
595
+ name_rewritten.endswith(".bias")
596
+ and name_rewritten not in params_dict
597
+ ):
598
+ continue
599
+
600
+ if is_pp_missing_parameter(name_rewritten, self):
601
+ continue
602
+
603
+ if name_rewritten not in params_dict:
604
+ continue
605
+
606
+ param = params_dict[name_rewritten]
607
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
608
+
609
+ if param_name == "qkv_proj" and shard_id == "v":
610
+ v_scale = (
611
+ self.v_scale
612
+ if self.v_scale is not None
613
+ else getattr(self.config, "attention_value_scale", None)
614
+ )
615
+ if v_scale is not None and (
616
+ name.endswith("weight_scale_inv") or name.endswith(".bias")
617
+ ):
618
+ loaded_weight *= float(v_scale)
619
+
620
+ weight_loader(param, loaded_weight, shard_id)
621
+ loaded_params.add(name_rewritten)
622
+
623
+ stacked_matched = True
624
+ break
625
+
626
+ if stacked_matched:
627
+ continue
628
+
629
+ if name.endswith(".bias") and name not in params_dict:
630
+ continue
631
+
632
+ orig_name = name
633
+ mapped_name = maybe_remap_kv_scale_name(name, params_dict)
634
+ name = mapped_name if mapped_name is not None else orig_name
635
+
636
+ if name not in params_dict:
637
+ continue
638
+
639
+ param = params_dict[name]
640
+
641
+ if "attention_sink_bias" in name:
642
+ total_heads = loaded_weight.shape[0]
643
+ heads_per_rank = total_heads // tp_size
644
+ head_start = tp_rank * heads_per_rank
645
+ narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank)
646
+
647
+ param.data.copy_(narrow_weight)
648
+ loaded_params.add(name)
649
+ else:
650
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
651
+ weight_loader(param, loaded_weight)
652
+ loaded_params.add(name)
653
+
654
+ return loaded_params
655
+
656
+
657
+class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
658
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
659
+ super().__init__()
660
+ config = vllm_config.model_config.hf_config
661
+ quant_config = vllm_config.quant_config
662
+
663
+ self.config = config
664
+ self.quant_config = quant_config
665
+ self.model = MiMoV2Model(
666
+ vllm_config=vllm_config,
667
+ prefix=maybe_prefix(prefix, "model"),
668
+ )
669
+
670
+ if get_pp_group().is_last_rank:
671
+ self.lm_head = ParallelLMHead(
672
+ config.vocab_size,
673
+ config.hidden_size,
674
+ quant_config=quant_config,
675
+ prefix=maybe_prefix(prefix, "lm_head"),
676
+ )
677
+ else:
678
+ self.lm_head = PPMissingLayer()
679
+
680
+ self.logits_processor = LogitsProcessor(config.vocab_size)
681
+
682
+ self.make_empty_intermediate_tensors = (
683
+ self.model.make_empty_intermediate_tensors
684
+ )
685
+
686
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
687
+ self.model.aux_hidden_state_layers = layers
688
+
689
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
690
+ num_layers = len(self.model.layers)
691
+ return (2, num_layers // 2, num_layers - 3)
692
+
693
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
694
+ return self.model.embed_input_ids(input_ids)
695
+
696
+ def forward(
697
+ self,
698
+ input_ids: torch.Tensor,
699
+ positions: torch.Tensor,
700
+ intermediate_tensors: IntermediateTensors | None = None,
701
+ inputs_embeds: torch.Tensor | None = None,
702
+ ) -> torch.Tensor | IntermediateTensors:
703
+ hidden_states = self.model(
704
+ input_ids, positions, intermediate_tensors, inputs_embeds
705
+ )
706
+ return hidden_states
707
+
708
+ def compute_logits(
709
+ self,
710
+ hidden_states: torch.Tensor,
711
+ ) -> torch.Tensor | None:
712
+ logits = self.logits_processor(self.lm_head, hidden_states)
713
+ return logits
714
+
715
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
716
+ return self.model.get_expert_mapping()
717
+
718
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
719
+ loader = AutoWeightsLoader(self)
720
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -152,6 +152,7 @@ _TEXT_GENERATION_MODELS = {
152
"MptForCausalLM": ("mpt", "MPTForCausalLM"),
153
"MPTForCausalLM": ("mpt", "MPTForCausalLM"),
154
"MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
155
+ "MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),
156
"NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
157
"NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
158
"OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),