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 13, 17:00 UTC
Add llama 4 scaling support (#28145)
Signed-off-by: Julien Denize <julien.denize@mistral.ai>
Julien Denize committed
Nov 6, 2025 at 19:55 UTC
7a8375f8a0cd37ee7e3d2be8498b5ae543694179
4 files changed
+59
-8
vllm/model_executor/layers/rotary_embedding/__init__.py
+8
-1
@@ -191,9 +191,16 @@ def get_rope(
191
k: v
192
for k, v in rope_scaling.items()
193
if k
194
- in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow")
194
+ in (
195
+ "extrapolation_factor",
196
+ "attn_factor",
197
+ "beta_fast",
198
+ "beta_slow",
199
+ "apply_yarn_scaling",
200
+ )
201
}
202
if "mrope_section" in rope_scaling:
203
+ extra_kwargs.pop("apply_yarn_scaling", None)
204
rotary_emb = MRotaryEmbedding(
205
head_size,
206
rotary_dim,
vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
+6
-1
@@ -27,6 +27,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
27
attn_factor: float = 1,
28
beta_fast: int = 32,
29
beta_slow: int = 1,
30
+ apply_yarn_scaling: bool = True,
31
) -> None:
32
self.scaling_factor = scaling_factor
33
self.extrapolation_factor = extrapolation_factor
@@ -34,7 +35,11 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
35
self.beta_fast = beta_fast
36
self.beta_slow = beta_slow
37
# Get n-d magnitude scaling corrected for interpolation
37
- self.mscale = float(yarn_get_mscale(self.scaling_factor) * attn_factor)
38
+ self.mscale = (
39
+ float(yarn_get_mscale(self.scaling_factor) * attn_factor)
40
+ if apply_yarn_scaling
41
+ else float(attn_factor)
42
+ )
43
super().__init__(
44
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
45
)
vllm/model_executor/models/llama.py
+22
@@ -160,6 +160,14 @@ class LlamaAttention(nn.Module):
160
self.rope_theta = rope_theta
161
self.max_position_embeddings = max_position_embeddings
162
163
+ llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
164
+ self.do_llama_4_scaling = llama_4_scaling_config is not None
165
+ if self.do_llama_4_scaling:
166
+ self.llama_4_scaling_original_max_position_embeddings = (
167
+ llama_4_scaling_config["original_max_position_embeddings"]
168
+ )
169
+ self.llama_4_scaling_beta = llama_4_scaling_config["beta"]
170
+
171
self.qkv_proj = QKVParallelLinear(
172
hidden_size=hidden_size,
173
head_size=self.head_dim,
@@ -221,6 +229,17 @@ class LlamaAttention(nn.Module):
229
prefix=f"{prefix}.attn",
230
)
231
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
+
243
def forward(
244
self,
245
positions: torch.Tensor,
@@ -229,6 +248,9 @@ class LlamaAttention(nn.Module):
248
qkv, _ = self.qkv_proj(hidden_states)
249
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
250
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)
254
attn_output = self.attn(q, k, v)
255
output, _ = self.o_proj(attn_output)
256
return output
vllm/transformers_utils/configs/mistral.py
+23
-6
@@ -24,6 +24,18 @@ def adapt_config_dict(config_dict: dict[str, Any], **kwargs) -> PretrainedConfig
24
if bool(config_dict.get("yarn")):
25
config_dict = _remap_mistral_yarn_args(config_dict)
26
27
+ if bool(config_dict.get("llama_4_scaling")):
28
+ llama_4_scaling_config_keys = ["original_max_position_embeddings", "beta"]
29
+ assert all(
30
+ [
31
+ key in config_dict["llama_4_scaling"]
32
+ for key in llama_4_scaling_config_keys
33
+ ]
34
+ ), (
35
+ "llama_4_scaling config should define the keys: "
36
+ f"{','.join(llama_4_scaling_config_keys)}"
37
+ )
38
+
39
is_vision = (config_dict.get("multimodal") or {}).get(
40
"vision_encoder_args"
41
) or config_dict.get("vision_encoder")
@@ -66,19 +78,24 @@ def _remap_mistral_vision_args(config: dict) -> dict:
78
79
80
def _remap_mistral_yarn_args(config: dict) -> dict:
69
- # Direct remaps: yarn.X -> rope_scaling.Y
70
- # Source keys are from mistral.model.args.YarnArgs
71
- _map = {
81
+ yarn_config_map = {
82
+ "factor": "factor",
83
+ "original_max_position_embeddings": "original_max_position_embeddings",
84
"beta": "beta_fast",
85
"alpha": "beta_slow",
86
+ "apply_scale": "apply_yarn_scaling",
87
}
88
yarn_config = config.get("yarn") or {}
76
- renamed_yarn_config = {_map.get(k, k): v for k, v in yarn_config.items()}
89
config["rope_scaling"] = {
90
"rope_type": "yarn",
79
- "mscale_all_dim": 1, # We hardcoded this to 1
80
- **renamed_yarn_config,
91
+ "mscale_all_dim": 1,
92
}
93
+ for old_name, new_name in yarn_config_map.items():
94
+ if old_name in yarn_config:
95
+ config["rope_scaling"][new_name] = yarn_config.pop(old_name)
96
+
97
+ assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"
98
+
99
return config
100
101