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 14, 05:30 UTC
[feat] Support MRoPE + YaRN (#25384)
Signed-off-by: liuye.hj <liuye.hj@alibaba-inc.com> Co-authored-by: liuye.hj <liuye.hj@alibaba-inc.com>
JJJYmmm committed
Sep 23, 2025 at 11:04 UTC
fc97733da801f1d6464e706f41ebd5180f2abb51
2 files changed
+48
-5
vllm/model_executor/layers/rotary_embedding/__init__.py
+17
-5
@@ -153,11 +153,23 @@ def get_rope(
153
if k in ("extrapolation_factor", "attn_factor", "beta_fast",
154
"beta_slow")
155
}
156
- rotary_emb = YaRNScalingRotaryEmbedding(head_size, rotary_dim,
157
- original_max_position,
158
- base, is_neox_style,
159
- scaling_factor, dtype,
160
- **extra_kwargs)
156
+ if "mrope_section" in rope_scaling:
157
+ rotary_emb = MRotaryEmbedding(
158
+ head_size,
159
+ rotary_dim,
160
+ original_max_position,
161
+ base,
162
+ is_neox_style,
163
+ dtype,
164
+ mrope_section=rope_scaling["mrope_section"],
165
+ mrope_interleaved=rope_scaling.get("mrope_interleaved",
166
+ False),
167
+ scaling_factor=scaling_factor,
168
+ **extra_kwargs)
169
+ else:
170
+ rotary_emb = YaRNScalingRotaryEmbedding(
171
+ head_size, rotary_dim, original_max_position, base,
172
+ is_neox_style, scaling_factor, dtype, **extra_kwargs)
173
elif scaling_type == "deepseek_yarn":
174
scaling_factor = rope_scaling["factor"]
175
original_max_position = rope_scaling[
vllm/model_executor/layers/rotary_embedding/mrope.py
+31
@@ -12,6 +12,7 @@ from vllm.triton_utils import tl, triton
12
13
from .base import RotaryEmbedding
14
from .common import apply_rotary_emb_dispatch
15
+from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale
16
17
18
@triton.jit
@@ -213,7 +214,27 @@ class MRotaryEmbedding(RotaryEmbedding):
214
dtype: torch.dtype,
215
mrope_section: Optional[list[int]] = None,
216
mrope_interleaved: bool = False,
217
+ # YaRN parameters.
218
+ *,
219
+ scaling_factor: Optional[float] = None,
220
+ extrapolation_factor: float = 1,
221
+ attn_factor: float = 1,
222
+ beta_fast: int = 32,
223
+ beta_slow: int = 1,
224
) -> None:
225
+
226
+ self.scaling_factor = scaling_factor
227
+ self.extrapolation_factor = extrapolation_factor
228
+ self.attn_factor = attn_factor
229
+ self.beta_fast = beta_fast
230
+ self.beta_slow = beta_slow
231
+ if self.scaling_factor is not None:
232
+ # Get n-d magnitude scaling corrected for interpolation
233
+ self.mscale = float(
234
+ yarn_get_mscale(self.scaling_factor) * attn_factor)
235
+ else:
236
+ self.mscale = 1.0
237
+
238
# In Qwen2.5-VL, the maximum index value is related to the duration of
239
# the input video. We enlarge max_position_embeddings to 4 times to get
240
# a larger the cos and sin cache.
@@ -226,6 +247,16 @@ class MRotaryEmbedding(RotaryEmbedding):
247
if self.mrope_section:
248
assert sum(self.mrope_section) == rotary_dim // 2
249
250
+ def _compute_inv_freq(self, base: float) -> torch.Tensor:
251
+ if self.scaling_factor is None:
252
+ return super()._compute_inv_freq(base)
253
+ return YaRNScalingRotaryEmbedding._compute_inv_freq(self, base)
254
+
255
+ def _compute_cos_sin_cache(self) -> torch.Tensor:
256
+ if self.scaling_factor is None:
257
+ return super()._compute_cos_sin_cache()
258
+ return YaRNScalingRotaryEmbedding._compute_cos_sin_cache(self)
259
+
260
def forward_native(
261
self,
262
positions: torch.Tensor,