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 12, 21:00 UTC
Add truncate arg to yarn to match openai implementation of gpt-oss (#28244)
Signed-off-by: ashors1 <ashors@nvidia.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>
Anna Shors committed
Nov 20, 2025 at 02:53 UTC
6eb745d9bdf5b69bb63f897b32465c62ecb9e14a
4 files changed
+12
-7
vllm/model_executor/layers/rotary_embedding/__init__.py
+1
@@ -197,6 +197,7 @@ def get_rope(
197
"beta_fast",
198
"beta_slow",
199
"apply_yarn_scaling",
200
+ "truncate",
201
)
202
}
203
if "mrope_section" in rope_parameters:
vllm/model_executor/layers/rotary_embedding/common.py
+7
-7
@@ -117,13 +117,13 @@ def yarn_find_correction_range(
117
dim: int,
118
base: float = 10000,
119
max_position_embeddings: int = 2048,
120
-) -> tuple[int, int]:
121
- low = math.floor(
122
- yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
123
- )
124
- high = math.ceil(
125
- yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
126
- )
120
+ truncate: bool = True,
121
+) -> tuple[float | int, float | int]:
122
+ low = yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
123
+ high = yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
124
+ if truncate:
125
+ low = math.floor(low)
126
+ high = math.ceil(high)
127
return max(low, 0), min(high, dim - 1) # Clamp values just in case
128
129
vllm/model_executor/layers/rotary_embedding/yarn_scaling_rope.py
+3
@@ -28,12 +28,14 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
28
beta_fast: int = 32,
29
beta_slow: int = 1,
30
apply_yarn_scaling: bool = True,
31
+ truncate: bool = True,
32
) -> None:
33
self.scaling_factor = scaling_factor
34
self.extrapolation_factor = extrapolation_factor
35
self.attn_factor = attn_factor
36
self.beta_fast = beta_fast
37
self.beta_slow = beta_slow
38
+ self.truncate = truncate
39
# Get n-d magnitude scaling corrected for interpolation
40
self.mscale = (
41
float(yarn_get_mscale(self.scaling_factor) * attn_factor)
@@ -57,6 +59,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):
59
self.rotary_dim,
60
self.base,
61
self.max_position_embeddings,
62
+ self.truncate,
63
)
64
# Get n-d rotational scaling corrected for extrapolation
65
inv_freq_mask = (
vllm/model_executor/models/gpt_oss.py
+1
@@ -78,6 +78,7 @@ class OAIAttention(nn.Module):
78
],
79
"beta_fast": config.rope_parameters["beta_fast"],
80
"beta_slow": config.rope_parameters["beta_slow"],
81
+ "truncate": config.rope_parameters.get("truncate", True),
82
},
83
is_neox_style=True,
84
)