[Bugfix] Align vLLM YaRN with Transformers and stop re-scaling max_model_len (#56446)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Harry Mellor committed
Sep 11, 2026 at 14:01 UTC
c191787a6861868069bc4f6ed6f842af541de23a
11 files changed
+133
-134
tests/test_config.py
+48
@@ -1853,6 +1853,54 @@ def test_get_and_verify_max_len_with_nope_layers(
1853
assert hf_config.rope_parameters["full_attention"] is None
1854
1855
1856
+@pytest.mark.parametrize(
1857
+ ("rope_type", "factor", "expected_max_len"),
1858
+ [
1859
+ # TeleChat3-36B-Thinking: 32768 already scaled from 8192 by 4
1860
+ ("yarn", 4.0, 32768),
1861
+ # sarvam-105b: declares factor 40 but only serves 131072 of it
1862
+ ("deepseek_yarn", 40.0, 32768),
1863
+ ("deepseek_llama_scaling", 40.0, 32768),
1864
+ # Non-YaRN scaling still multiplies
1865
+ ("linear", 4.0, 131072),
1866
+ ],
1867
+)
1868
+def test_get_and_verify_max_len_yarn_is_already_scaled(
1869
+ rope_type, factor, expected_max_len
1870
+):
1871
+ """YaRN variants must not re-apply `factor` to max_position_embeddings.
1872
+
1873
+ Transformers treats max_position_embeddings as the final context length
1874
+ for every YaRN variant, so scaling it again overstates the limit and lets
1875
+ requests past the end of the cos/sin cache.
1876
+ """
1877
+ from transformers import PretrainedConfig
1878
+
1879
+ from vllm.config.model import _get_and_verify_max_len
1880
+ from vllm.transformers_utils.model_arch_config_convertor import (
1881
+ ModelArchConfigConvertorBase,
1882
+ )
1883
+
1884
+ hf_config = PretrainedConfig(max_position_embeddings=32768)
1885
+ hf_config.rope_parameters = {
1886
+ "rope_type": rope_type,
1887
+ "factor": factor,
1888
+ "original_max_position_embeddings": 8192,
1889
+ }
1890
+ model_arch_config = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
1891
+
1892
+ actual_max_len = _get_and_verify_max_len(
1893
+ hf_config=hf_config,
1894
+ model_arch_config=model_arch_config,
1895
+ tokenizer_config=None,
1896
+ max_model_len=None,
1897
+ disable_sliding_window=False,
1898
+ sliding_window=None,
1899
+ )
1900
+
1901
+ assert actual_max_len == expected_max_len
1902
+
1903
+
1904
class MockConfig:
1905
"""Simple mock object for testing maybe_pull_model_tokenizer_for_runai"""
1906
tests/transformers_utils/test_config.py
+26
@@ -6,6 +6,7 @@ only get the `eos_token_id` from the tokenizer as defined by
6
`BaseRenderer.get_eos_token_id`.
7
"""
8
9
+import math
10
from types import SimpleNamespace
11
from typing import cast
12
from unittest.mock import MagicMock, patch
@@ -52,6 +53,31 @@ def test_patch_legacy_rope_type_preserves_nope_layers():
53
}
54
55
56
+def test_patch_legacy_rope_type_normalizes_telechat3_yarn():
57
+ """TeleChat3's RoPE is YaRN with 0.07 in place of the usual 0.1.
58
+
59
+ Encoding that as a precomputed attention_factor keeps the config
60
+ plain YaRN, which Transformers and every "yarn" guard understand.
61
+ `mscale` cannot express it: Transformers only applies mscale when
62
+ mscale_all_dim is also truthy.
63
+ """
64
+ rope_parameters = {
65
+ "type": "telechat3-yarn",
66
+ "rope_type": "telechat3-yarn",
67
+ "factor": 4.0,
68
+ "original_max_position_embeddings": 8192,
69
+ }
70
+
71
+ patch_legacy_rope_type(rope_parameters)
72
+
73
+ assert rope_parameters == {
74
+ "rope_type": "yarn",
75
+ "factor": 4.0,
76
+ "original_max_position_embeddings": 8192,
77
+ "attention_factor": pytest.approx(0.07 * math.log(4.0) + 1.0),
78
+ }
79
+
80
+
81
def test_glm5_next_accepts_deepseek_sparse_attention_layers():
82
layer_types = ["linear_attention", "deepseek_sparse_attention"]
83
vllm/config/model.py
+11
-4
@@ -2465,13 +2465,20 @@ def _get_and_verify_max_len(
2465
# loading HF config
2466
rope_type = rp["rope_type"]
2467
2468
- if rope_type not in ("su", "longrope", "llama3"):
2468
+ # YaRN variants leave max_position_embeddings already scaled, as
2469
+ # Transformers' _compute_yarn_parameters assumes, so `factor` must
2470
+ # not be applied to it again.
2471
+ if rope_type not in (
2472
+ "su",
2473
+ "longrope",
2474
+ "llama3",
2475
+ "yarn",
2476
+ "deepseek_yarn",
2477
+ "deepseek_llama_scaling",
2478
+ ):
2479
# NOTE: rope_type == "default" does not define factor https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/modeling_rope_utils.py
2480
# NOTE: This assumes all layer types have the same scaling factor.
2481
scaling_factor = rp.get("factor", scaling_factor)
2472
-
2473
- if rope_type == "yarn":
2474
- derived_max_model_len = rp["original_max_position_embeddings"]
2482
if scaling_factor is None:
2483
# Fallback the factor to 1.0 if a user assigned `null`
2484
logger.warning_once(
vllm/model_executor/layers/rotary_embedding/__init__.py
+3
-37
@@ -24,7 +24,6 @@ from .mrope import MRotaryEmbedding
24
from .mrope_interleaved import MRotaryEmbeddingInterleaved
25
from .ntk_scaling_rope import NTKScalingRotaryEmbedding
26
from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding
27
-from .telechat3_scaling_rope import TeleChat3RoPEScaledRotaryEmbedding
27
from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
28
29
_ROPE_DICT: dict[tuple[Any, ...], RotaryEmbedding] = {}
@@ -246,16 +245,15 @@ def get_rope(
245
for k, v in rope_parameters.items()
246
if k
247
in (
249
- "extrapolation_factor",
250
- "attn_factor",
248
"beta_fast",
249
"beta_slow",
253
- "apply_yarn_scaling",
250
+ "mscale",
251
+ "mscale_all_dim",
252
+ "attention_factor",
253
"truncate",
254
)
255
}
256
if "mrope_section" in rope_parameters:
258
- extra_kwargs.pop("apply_yarn_scaling", None)
257
rotary_emb = MRotaryEmbedding(
258
head_size,
259
rotary_dim,
@@ -288,8 +286,6 @@ def get_rope(
286
for k, v in rope_parameters.items()
287
if k
288
in (
291
- "extrapolation_factor",
292
- "attn_factor",
289
"beta_fast",
290
"beta_slow",
291
"mscale",
@@ -346,36 +342,6 @@ def get_rope(
342
)
343
else:
344
raise ValueError("Pangu mrope lacks necessary parameters.")
349
- elif scaling_type == "telechat3-yarn":
350
- scaling_factor = rope_parameters["factor"]
351
- if "original_max_position_embeddings" in rope_parameters:
352
- original_max_position = rope_parameters["original_max_position_embeddings"]
353
- scaling_factor = max_position / original_max_position
354
- else:
355
- original_max_position = max_position
356
- extra_kwargs = {
357
- k: v
358
- for k, v in rope_parameters.items()
359
- if k
360
- in (
361
- "extrapolation_factor",
362
- "attn_factor",
363
- "beta_fast",
364
- "beta_slow",
365
- "mscale",
366
- "mscale_all_dim",
367
- )
368
- }
369
- rotary_emb = TeleChat3RoPEScaledRotaryEmbedding(
370
- head_size,
371
- rotary_dim,
372
- original_max_position,
373
- base,
374
- is_neox_style,
375
- scaling_factor,
376
- dtype,
377
- **extra_kwargs,
378
- )
345
else:
346
raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
347
_ROPE_DICT[key] = rotary_emb
vllm/model_executor/layers/rotary_embedding/common.py
+2
-2
@@ -70,10 +70,10 @@ def yarn_linear_ramp_mask(
70
return ramp_func
71
72
73
-def yarn_get_mscale(scale: float = 1) -> float:
73
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
74
if scale <= 1:
75
return 1.0
76
- return 0.1 * math.log(scale) + 1.0
76
+ return 0.1 * mscale * math.log(scale) + 1.0
77
78
79
def _flashinfer_rotary_embedding(
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+3
-9
@@ -39,8 +39,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
39
scaling_factor: float,
40
dtype: torch.dtype,
41
*,
42
- extrapolation_factor: float = 1,
43
- attn_factor: float = 1,
42
beta_fast: int = 32,
43
beta_slow: int = 1,
44
mscale: float = 1,
@@ -48,15 +46,12 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
46
init_cache: bool = True,
47
) -> None:
48
self.scaling_factor = scaling_factor
51
- self.extrapolation_factor = extrapolation_factor
52
- self.attn_factor = attn_factor
49
self.beta_fast = beta_fast
50
self.beta_slow = beta_slow
51
# Get n-d magnitude scaling corrected for interpolation.
52
self.mscale = float(
53
yarn_get_mscale(self.scaling_factor, float(mscale))
54
/ yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))
59
- * attn_factor
55
)
56
self.use_flashinfer = (
57
self.enabled()
@@ -96,10 +91,9 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
91
self.max_position_embeddings,
92
)
93
# Get n-d rotational scaling corrected for extrapolation
99
- inv_freq_mask = (
100
- 1
101
- - yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
102
- ) * self.extrapolation_factor
94
+ inv_freq_mask = 1 - yarn_linear_ramp_mask(
95
+ low, high, self.rotary_dim // 2, dtype=torch.float
96
+ )
97
inv_freq = (
98
inv_freq_interpolation * (1 - inv_freq_mask)
99
+ inv_freq_extrapolation * inv_freq_mask
vllm/model_executor/layers/rotary_embedding/mrope.py
+12
-5
@@ -252,21 +252,28 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
252
# YaRN parameters.
253
*,
254
scaling_factor: float | None = None,
255
- extrapolation_factor: float = 1,
256
- attn_factor: float = 1,
255
beta_fast: int = 32,
256
beta_slow: int = 1,
257
+ mscale: float | None = None,
258
+ mscale_all_dim: float | None = None,
259
+ attention_factor: float | None = None,
260
truncate: bool = True,
261
) -> None:
262
self.scaling_factor = scaling_factor
262
- self.extrapolation_factor = extrapolation_factor
263
- self.attn_factor = attn_factor
263
self.beta_fast = beta_fast
264
self.beta_slow = beta_slow
265
self.truncate = truncate
266
if self.scaling_factor is not None:
267
# Get n-d magnitude scaling corrected for interpolation
269
- self.mscale = float(yarn_get_mscale(self.scaling_factor) * attn_factor)
268
+ if attention_factor is not None:
269
+ self.mscale = float(attention_factor)
270
+ elif mscale and mscale_all_dim:
271
+ self.mscale = float(
272
+ yarn_get_mscale(self.scaling_factor, mscale)
273
+ / yarn_get_mscale(self.scaling_factor, mscale_all_dim)
274
+ )
275
+ else:
276
+ self.mscale = float(yarn_get_mscale(self.scaling_factor))
277
else:
278
self.mscale = 1.0
279
vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py
deleted
-57
@@ -1,57 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import math
5
-
6
-import torch
7
-
8
-from .base import RotaryEmbedding
9
-from .yarn_scaling_rope import YaRNScalingRotaryEmbedding
10
-
11
-
12
-class TeleChat3RoPEScaledRotaryEmbedding(YaRNScalingRotaryEmbedding):
13
- """TeleChat3 uses a variant of YaRN method.
14
-
15
- To achieve code reuse as much as possible, we have rewritten the
16
- `get_mscale` method in the initialization function
17
- """
18
-
19
- def __init__(
20
- self,
21
- head_size: int,
22
- rotary_dim: int,
23
- max_position_embeddings: int,
24
- base: int,
25
- is_neox_style: bool,
26
- scaling_factor: float,
27
- dtype: torch.dtype,
28
- *,
29
- extrapolation_factor: float = 1,
30
- attn_factor: float = 1,
31
- beta_fast: int = 32,
32
- beta_slow: int = 1,
33
- truncate: bool = True,
34
- ) -> None:
35
- self.scaling_factor = scaling_factor
36
- self.extrapolation_factor = extrapolation_factor
37
- self.attn_factor = attn_factor
38
- self.beta_fast = beta_fast
39
- self.beta_slow = beta_slow
40
- self.truncate = truncate
41
-
42
- def get_mscale(scale, mscale=1):
43
- if scale <= 1:
44
- return 1.0
45
- return 0.07 * mscale * math.log(scale) + 1.0
46
-
47
- self.mscale = float(get_mscale(self.scaling_factor) * attn_factor)
48
- # Initialization must be performed after mscale, otherwise mscale is useless
49
- RotaryEmbedding.__init__(
50
- self,
51
- head_size,
52
- rotary_dim,
53
- max_position_embeddings,
54
- base,
55
- is_neox_style,
56
- dtype,
57
- )
vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
+15
-14
@@ -23,25 +23,27 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
23
scaling_factor: float,
24
dtype: torch.dtype,
25
*,
26
- extrapolation_factor: float = 1,
27
- attn_factor: float = 1,
26
beta_fast: int = 32,
27
beta_slow: int = 1,
30
- apply_yarn_scaling: bool = True,
28
+ mscale: float | None = None,
29
+ mscale_all_dim: float | None = None,
30
+ attention_factor: float | None = None,
31
truncate: bool = True,
32
) -> None:
33
self.scaling_factor = scaling_factor
34
- self.extrapolation_factor = extrapolation_factor
35
- self.attn_factor = attn_factor
34
self.beta_fast = beta_fast
35
self.beta_slow = beta_slow
36
self.truncate = truncate
37
# Get n-d magnitude scaling corrected for interpolation
40
- self.mscale = (
41
- float(yarn_get_mscale(self.scaling_factor) * attn_factor)
42
- if apply_yarn_scaling
43
- else float(attn_factor)
44
- )
38
+ if attention_factor is not None:
39
+ self.mscale = float(attention_factor)
40
+ elif mscale and mscale_all_dim:
41
+ self.mscale = float(
42
+ yarn_get_mscale(self.scaling_factor, mscale)
43
+ / yarn_get_mscale(self.scaling_factor, mscale_all_dim)
44
+ )
45
+ else:
46
+ self.mscale = float(yarn_get_mscale(self.scaling_factor))
47
super().__init__(
48
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
49
)
@@ -62,10 +64,9 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
64
self.truncate,
65
)
66
# Get n-d rotational scaling corrected for extrapolation
65
- inv_freq_mask = (
66
- 1
67
- - yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
68
- ) * self.extrapolation_factor
67
+ inv_freq_mask = 1 - yarn_linear_ramp_mask(
68
+ low, high, self.rotary_dim // 2, dtype=torch.float
69
+ )
70
inv_freq = (
71
inv_freq_interpolation * (1 - inv_freq_mask)
72
+ inv_freq_extrapolation * inv_freq_mask
vllm/transformers_utils/config.py
+9
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
import os
6
from collections.abc import Callable, Iterator, Sequence
7
from contextlib import contextmanager
@@ -539,6 +540,14 @@ def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None:
540
)
541
rope_parameters["rope_type"] = "default"
542
logger.warning("Replacing legacy rope_type 'mrope' with 'default'")
543
+ elif rope_parameters["rope_type"] == "telechat3-yarn":
544
+ # TeleChat3 is YaRN with 0.07 in place of YaRN's 0.1 attention
545
+ # scaling coefficient. Precompute it so the config is plain YaRN.
546
+ factor = rope_parameters["factor"]
547
+ rope_parameters["rope_type"] = "yarn"
548
+ rope_parameters["attention_factor"] = 0.07 * math.log(factor) + 1.0
549
+ rope_parameters.pop("type", None)
550
+ logger.warning("Replacing rope_type 'telechat3-yarn' with 'yarn'")
551
552
# Handle nested rope_parameters in interleaved sliding attention models
553
if is_rope_parameters_nested(rope_parameters):
vllm/transformers_utils/configs/mistral.py
+4
-6
@@ -2,9 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import Any
4
5
-from packaging.version import Version
5
from transformers import PretrainedConfig, WhisperConfig
7
-from transformers import __version__ as TRANSFORMERS_VERSION
6
7
from vllm.logger import init_logger
8
@@ -119,7 +117,6 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
117
"original_max_position_embeddings": ("original_max_position_embeddings", int),
118
"beta": ("beta_fast", float),
119
"alpha": ("beta_slow", float),
122
- "apply_scale": ("apply_yarn_scaling", bool),
120
}
121
122
yarn_config = config.get("yarn") or {}
@@ -136,9 +133,10 @@ def _remap_mistral_yarn_args(config: dict) -> dict:
133
# Cast to remove Transformers > v5 type warnings
134
config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name))
135
139
- # Ignore apply_yarn_scaling in Transformers > v5 RoPE validation to remove warnings
140
- if Version(TRANSFORMERS_VERSION) >= Version("5.3.0.dev0"):
141
- config["ignore_keys_at_rope_validation"] = {"apply_yarn_scaling"}
136
+ # `apply_scale: false` means no magnitude correction, which Transformers
137
+ # spells as an explicit attention_factor of 1.
138
+ if not yarn_config.pop("apply_scale", True):
139
+ config["rope_parameters"]["attention_factor"] = 1.0
140
141
assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"
142