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 9, 20:45 UTC
[Llama.py -> mistral.py] Extract mistral-only relevant code into separate file (#32780)
Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
Patrick von Platen committed
Jan 22, 2026 at 06:14 UTC
1579c9b5fd0f6f213be3bbc55c88b1d215245ed6
3 files changed
+248
-115
vllm/model_executor/models/llama.py
+5
-114
@@ -152,24 +152,14 @@ class LlamaAttention(nn.Module):
152
# the KV heads across multiple tensor parallel GPUs.
153
assert tp_size % self.total_num_kv_heads == 0
154
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
155
- # MistralConfig has an optional head_dim introduced by Mistral-Nemo
155
+
156
head_dim = getattr(config, "head_dim", None)
157
- if head_dim is None:
158
- head_dim = self.hidden_size // self.total_num_heads
159
- self.head_dim = head_dim
157
+ self.head_dim = head_dim or self.hidden_size // self.total_num_heads
158
self.q_size = self.num_heads * self.head_dim
159
self.kv_size = self.num_kv_heads * self.head_dim
160
self.scaling = self.head_dim**-0.5
161
self.max_position_embeddings = max_position_embeddings
162
165
- llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
166
- self.do_llama_4_scaling = llama_4_scaling_config is not None
167
- if self.do_llama_4_scaling:
168
- self.llama_4_scaling_original_max_position_embeddings = (
169
- llama_4_scaling_config["original_max_position_embeddings"]
170
- )
171
- self.llama_4_scaling_beta = llama_4_scaling_config["beta"]
172
-
163
self.qkv_proj = QKVParallelLinear(
164
hidden_size=hidden_size,
165
head_size=self.head_dim,
@@ -229,17 +219,6 @@ class LlamaAttention(nn.Module):
219
prefix=f"{prefix}.attn",
220
)
221
232
- def _get_llama_4_attn_scale(self, positions: torch.Tensor) -> torch.Tensor:
233
- # Llama4 scaling
234
- scaling = 1 + self.llama_4_scaling_beta * torch.log(
235
- 1
236
- + torch.floor(
237
- positions / self.llama_4_scaling_original_max_position_embeddings
238
- )
239
- )
240
- # Broadcast over head_dim
241
- return scaling.unsqueeze(-1)
242
-
222
def forward(
223
self,
224
positions: torch.Tensor,
@@ -248,9 +227,6 @@ class LlamaAttention(nn.Module):
227
qkv, _ = self.qkv_proj(hidden_states)
228
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
229
q, k = self.rotary_emb(positions, q, k)
251
- if self.do_llama_4_scaling:
252
- attn_scale = self._get_llama_4_attn_scale(positions)
253
- q = (q * attn_scale).to(q.dtype)
230
attn_output = self.attn(q, k, v)
231
output, _ = self.o_proj(attn_output)
232
return output
@@ -279,6 +255,7 @@ class LlamaDecoderLayer(nn.Module):
255
vllm_config: VllmConfig,
256
prefix: str = "",
257
config: LlamaConfig | None = None,
258
+ attn_layer_type: type[nn.Module] = LlamaAttention,
259
) -> None:
260
super().__init__()
261
@@ -307,7 +284,7 @@ class LlamaDecoderLayer(nn.Module):
284
else:
285
attn_type = AttentionType.ENCODER_ONLY
286
310
- self.self_attn = LlamaAttention(
287
+ self.self_attn = attn_layer_type(
288
config=config,
289
hidden_size=self.hidden_size,
290
num_heads=config.num_attention_heads,
@@ -537,32 +514,6 @@ class LlamaForCausalLM(
514
"lm_head": "output_embeddings",
515
}
516
540
- # Mistral/Llama models can also be loaded with --load-format mistral
541
- # from consolidated.safetensors checkpoints
542
- mistral_mapping = {
543
- "layers": "model.layers",
544
- "attention": "self_attn",
545
- "qscale_act": "input_scale",
546
- "qscale_weight": "weight_scale",
547
- "kv_fake_quantizer.qscale_act": "kv_scale",
548
- "q_fake_quantizer.qscale_act": "attn.q_scale",
549
- "k_fake_quantizer.qscale_act": "k_scale",
550
- "v_fake_quantizer.qscale_act": "v_scale",
551
- "wq": "q_proj",
552
- "wk": "k_proj",
553
- "wv": "v_proj",
554
- "wo": "o_proj",
555
- "attention_norm": "input_layernorm",
556
- "feed_forward": "mlp",
557
- "w1": "gate_proj",
558
- "w2": "down_proj",
559
- "w3": "up_proj",
560
- "ffn_norm": "post_attention_layernorm",
561
- "tok_embeddings": "model.embed_tokens",
562
- "output": "lm_head",
563
- "norm": "model.norm",
564
- }
565
-
517
def __init__(
518
self,
519
*,
@@ -649,67 +600,7 @@ class LlamaForCausalLM(
600
self,
601
skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
602
)
652
- return loader.load_weights(
653
- self.maybe_remap_mistral(name, loaded_weight)
654
- for name, loaded_weight in weights
655
- )
656
-
657
- # This function is used to remap the mistral format as
658
- # used by Mistral and Llama <=2
659
- def maybe_remap_mistral(
660
- self,
661
- name: str,
662
- loaded_weight: torch.Tensor,
663
- ) -> tuple[str, torch.Tensor]:
664
- def permute(w: torch.Tensor, n_heads: int, attn_out: int):
665
- attn_in = self.config.head_dim * n_heads
666
-
667
- return (
668
- w.view(n_heads, attn_in // n_heads // 2, 2, attn_out)
669
- .transpose(1, 2)
670
- .reshape(attn_in, attn_out)
671
- )
672
-
673
- mapping = self.mistral_mapping
674
- modules = name.split(".")
675
-
676
- # rotary embeds should be sliced
677
- # If using quantized model in mistral format,
678
- # quantization scales (qscale_weight) also need to be sliced
679
- if "wk" in modules and modules[-1] == "weight":
680
- loaded_weight = permute(
681
- loaded_weight, self.config.num_key_value_heads, self.config.hidden_size
682
- )
683
- elif (
684
- "wk" in modules
685
- and modules[-1] == "qscale_weight"
686
- and loaded_weight.numel() > 1
687
- ):
688
- loaded_weight = permute(loaded_weight, self.config.num_key_value_heads, 1)
689
- elif "wq" in modules and modules[-1] == "weight":
690
- loaded_weight = permute(
691
- loaded_weight, self.config.num_attention_heads, self.config.hidden_size
692
- )
693
- elif (
694
- "wq" in modules
695
- and modules[-1] == "qscale_weight"
696
- and loaded_weight.numel() > 1
697
- ):
698
- loaded_weight = permute(loaded_weight, self.config.num_attention_heads, 1)
699
-
700
- num_modules = len(modules)
701
- for i in range(num_modules):
702
- item = modules[i]
703
- next_item = modules[i + 1] if i < num_modules - 1 else None
704
-
705
- combined_item = f"{item}.{next_item}" if next_item is not None else None
706
-
707
- if combined_item in mapping:
708
- name = name.replace(combined_item, mapping[combined_item])
709
- elif item in mapping and mapping[item] not in name:
710
- name = name.replace(item, mapping[item])
711
-
712
- return name, loaded_weight
603
+ return loader.load_weights(weights)
604
605
606
class LlamaBidirectionalForSequenceClassification(as_seq_cls_model(LlamaForCausalLM)):
vllm/model_executor/models/mistral.py
new
+242
@@ -0,0 +1,242 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Mistral adaptation of the LLaMA architecture."""
4
+
5
+from collections.abc import Iterable
6
+
7
+import torch
8
+from torch import nn
9
+from transformers import LlamaConfig
10
+
11
+from vllm.compilation.decorators import support_torch_compile
12
+from vllm.config import CacheConfig, VllmConfig
13
+from vllm.model_executor.layers.quantization import QuantizationConfig
14
+from vllm.model_executor.models.llama import (
15
+ LlamaAttention,
16
+ LlamaDecoderLayer,
17
+ LlamaForCausalLM,
18
+ LlamaModel,
19
+)
20
+from vllm.v1.attention.backend import AttentionType
21
+
22
+from .utils import AutoWeightsLoader
23
+
24
+
25
+class MistralAttention(LlamaAttention):
26
+ def __init__(
27
+ self,
28
+ config: LlamaConfig,
29
+ hidden_size: int,
30
+ num_heads: int,
31
+ num_kv_heads: int,
32
+ max_position_embeddings: int = 8192,
33
+ quant_config: QuantizationConfig | None = None,
34
+ bias: bool = False,
35
+ bias_o_proj: bool = False,
36
+ cache_config: CacheConfig | None = None,
37
+ prefix: str = "",
38
+ attn_type: str = AttentionType.DECODER,
39
+ ) -> None:
40
+ super().__init__(
41
+ config=config,
42
+ hidden_size=hidden_size,
43
+ num_heads=num_heads,
44
+ num_kv_heads=num_kv_heads,
45
+ max_position_embeddings=max_position_embeddings,
46
+ quant_config=quant_config,
47
+ bias=bias,
48
+ bias_o_proj=bias_o_proj,
49
+ cache_config=cache_config,
50
+ prefix=prefix,
51
+ attn_type=attn_type,
52
+ )
53
+
54
+ llama_4_scaling_config: dict[str, int | float | str] | None = getattr(
55
+ config, "llama_4_scaling", None
56
+ )
57
+ self.do_llama_4_scaling = llama_4_scaling_config is not None
58
+ if self.do_llama_4_scaling:
59
+ assert llama_4_scaling_config is not None
60
+ self.llama_4_scaling_original_max_position_embeddings = (
61
+ llama_4_scaling_config["original_max_position_embeddings"]
62
+ )
63
+ self.llama_4_scaling_beta = llama_4_scaling_config["beta"]
64
+
65
+ def _get_llama_4_attn_scale(self, positions: torch.Tensor) -> torch.Tensor:
66
+ # Llama4 scaling
67
+ scaling = 1 + self.llama_4_scaling_beta * torch.log(
68
+ 1
69
+ + torch.floor(
70
+ positions / self.llama_4_scaling_original_max_position_embeddings
71
+ )
72
+ )
73
+ # Broadcast over head_dim
74
+ return scaling.unsqueeze(-1)
75
+
76
+ def forward(
77
+ self,
78
+ positions: torch.Tensor,
79
+ hidden_states: torch.Tensor,
80
+ ) -> torch.Tensor:
81
+ qkv, _ = self.qkv_proj(hidden_states)
82
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
83
+ q, k = self.rotary_emb(positions, q, k)
84
+ if self.do_llama_4_scaling:
85
+ attn_scale = self._get_llama_4_attn_scale(positions)
86
+ q = (q * attn_scale).to(q.dtype)
87
+ attn_output = self.attn(q, k, v)
88
+ output, _ = self.o_proj(attn_output)
89
+ return output
90
+
91
+
92
+class MistralDecoderLayer(LlamaDecoderLayer):
93
+ def __init__(
94
+ self,
95
+ vllm_config: VllmConfig,
96
+ prefix: str = "",
97
+ config: LlamaConfig | None = None,
98
+ ) -> None:
99
+ super().__init__(
100
+ vllm_config=vllm_config,
101
+ prefix=prefix,
102
+ config=config,
103
+ attn_layer_type=MistralAttention,
104
+ )
105
+
106
+ self.layer_idx = int(prefix.split(sep=".")[-1])
107
+ quant_config = self.get_quant_config(vllm_config)
108
+ config = config or vllm_config.model_config.hf_config
109
+
110
+ do_fusion = getattr(
111
+ quant_config, "enable_quantization_scaling_fusion", False
112
+ ) and vllm_config.cache_config.cache_dtype.startswith("fp8")
113
+ if do_fusion:
114
+ self.input_layernorm.quant_scaling_from = self.self_attn.qkv_proj
115
+ self.post_attention_layernorm.quant_scaling_from = self.mlp.gate_up_proj
116
+
117
+
118
+@support_torch_compile
119
+class MistralModel(LlamaModel):
120
+ def __init__(
121
+ self,
122
+ *,
123
+ vllm_config: VllmConfig,
124
+ prefix: str = "",
125
+ layer_type: type[nn.Module] = MistralDecoderLayer,
126
+ ):
127
+ super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
128
+
129
+
130
+class MistralForCausalLM(LlamaForCausalLM):
131
+ # Mistral: We don't support LoRA on the embedding layers.
132
+ embedding_modules: dict[str, str] = {}
133
+
134
+ # Mistral/Llama models can also be loaded with --load-format mistral
135
+ # from consolidated.safetensors checkpoints
136
+ mistral_mapping = {
137
+ "layers": "model.layers",
138
+ "attention": "self_attn",
139
+ "qscale_act": "input_scale",
140
+ "qscale_weight": "weight_scale",
141
+ "kv_fake_quantizer.qscale_act": "kv_scale",
142
+ "q_fake_quantizer.qscale_act": "attn.q_scale",
143
+ "k_fake_quantizer.qscale_act": "k_scale",
144
+ "v_fake_quantizer.qscale_act": "v_scale",
145
+ "wq": "q_proj",
146
+ "wk": "k_proj",
147
+ "wv": "v_proj",
148
+ "wo": "o_proj",
149
+ "attention_norm": "input_layernorm",
150
+ "feed_forward": "mlp",
151
+ "w1": "gate_proj",
152
+ "w2": "down_proj",
153
+ "w3": "up_proj",
154
+ "ffn_norm": "post_attention_layernorm",
155
+ "tok_embeddings": "model.embed_tokens",
156
+ "output": "lm_head",
157
+ "norm": "model.norm",
158
+ }
159
+
160
+ def __init__(
161
+ self,
162
+ *,
163
+ vllm_config: VllmConfig,
164
+ prefix: str = "",
165
+ layer_type: type[nn.Module] = MistralDecoderLayer,
166
+ ):
167
+ super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
168
+
169
+ def _init_model(
170
+ self,
171
+ vllm_config: VllmConfig,
172
+ prefix: str = "",
173
+ layer_type: type[nn.Module] = MistralDecoderLayer,
174
+ ):
175
+ return MistralModel(
176
+ vllm_config=vllm_config, prefix=prefix, layer_type=layer_type
177
+ )
178
+
179
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
180
+ loader = AutoWeightsLoader(
181
+ self,
182
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
183
+ )
184
+ return loader.load_weights(
185
+ self.maybe_remap_mistral(name, loaded_weight)
186
+ for name, loaded_weight in weights
187
+ )
188
+
189
+ def maybe_remap_mistral(
190
+ self,
191
+ name: str,
192
+ loaded_weight: torch.Tensor,
193
+ ) -> tuple[str, torch.Tensor]:
194
+ def permute(w: torch.Tensor, n_heads: int, attn_out: int):
195
+ attn_in = self.config.head_dim * n_heads
196
+
197
+ return (
198
+ w.view(n_heads, attn_in // n_heads // 2, 2, attn_out)
199
+ .transpose(1, 2)
200
+ .reshape(attn_in, attn_out)
201
+ )
202
+
203
+ mapping = self.mistral_mapping
204
+ modules = name.split(".")
205
+
206
+ # rotary embeds should be sliced
207
+ # If using quantized model in mistral format,
208
+ # quantization scales (qscale_weight) also need to be sliced
209
+ if "wk" in modules and modules[-1] == "weight":
210
+ loaded_weight = permute(
211
+ loaded_weight, self.config.num_key_value_heads, self.config.hidden_size
212
+ )
213
+ elif (
214
+ "wk" in modules
215
+ and modules[-1] == "qscale_weight"
216
+ and loaded_weight.numel() > 1
217
+ ):
218
+ loaded_weight = permute(loaded_weight, self.config.num_key_value_heads, 1)
219
+ elif "wq" in modules and modules[-1] == "weight":
220
+ loaded_weight = permute(
221
+ loaded_weight, self.config.num_attention_heads, self.config.hidden_size
222
+ )
223
+ elif (
224
+ "wq" in modules
225
+ and modules[-1] == "qscale_weight"
226
+ and loaded_weight.numel() > 1
227
+ ):
228
+ loaded_weight = permute(loaded_weight, self.config.num_attention_heads, 1)
229
+
230
+ num_modules = len(modules)
231
+ for i in range(num_modules):
232
+ item = modules[i]
233
+ next_item = modules[i + 1] if i < num_modules - 1 else None
234
+
235
+ combined_item = f"{item}.{next_item}" if next_item is not None else None
236
+
237
+ if combined_item in mapping:
238
+ name = name.replace(combined_item, mapping[combined_item])
239
+ elif item in mapping and mapping[item] not in name:
240
+ name = name.replace(item, mapping[item])
241
+
242
+ return name, loaded_weight
vllm/model_executor/models/registry.py
+1
-1
@@ -153,7 +153,7 @@ _TEXT_GENERATION_MODELS = {
153
"MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
154
"MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
155
"MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
156
- "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
156
+ "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
157
"MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
158
"MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
159
# transformers's mpt class has lower case