From f5f3c38aad65d70e5ee07bb1b5dce7bb4c6d2d49 Mon Sep 17 00:00:00 2001 From: skyler-apdx Date: Mon, 21 Sep 2026 14:14:05 +0800 Subject: [PATCH] [Fix] Preserve YaRN scaling when extending rotary caches (#38786) --- .../srt/layers/rotary_embedding/mrope.py | 9 ++++ .../layers/rotary_embedding/rope_variant.py | 20 ++++++++ .../srt/layers/rotary_embedding/yarn.py | 36 ++++++++++++- .../unit/layers/test_yarn_cache_extension.py | 50 +++++++++++++++++++ 4 files changed, 114 insertions(+), 1 deletion(-) create mode 100644 test/registered/unit/layers/test_yarn_cache_extension.py diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index b5710d340..5f42f89fc 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -14,6 +14,7 @@ from sglang.kernels.ops.attention.rotary_triton import ( from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.rotary_embedding.yarn import ( + _extend_yarn_cache, yarn_find_correction_range, yarn_get_mscale_simple, yarn_linear_ramp_mask, @@ -488,6 +489,14 @@ class YaRNScalingMRotaryEmbedding(MRotaryEmbedding): ) return inv_freq + def _ensure_cos_sin_cache_length(self, needed_max_pos: int): + self.cos_sin_cache, _ = _extend_yarn_cache( + cache=self.cos_sin_cache, + compute_inv_freq=lambda: self._compute_inv_freq(self.scaling_factor), + mscale=self.mscale, + needed_max_pos=needed_max_pos, + ) + def _compute_cos_sin_cache(self) -> torch.Tensor: inv_freq = self._compute_inv_freq(self.scaling_factor) t = torch.arange( diff --git a/python/sglang/srt/layers/rotary_embedding/rope_variant.py b/python/sglang/srt/layers/rotary_embedding/rope_variant.py index 795e1105b..cc0b61a26 100644 --- a/python/sglang/srt/layers/rotary_embedding/rope_variant.py +++ b/python/sglang/srt/layers/rotary_embedding/rope_variant.py @@ -18,6 +18,7 @@ from sglang.srt.layers.rotary_embedding.utils import ( rotate_neox, ) from sglang.srt.layers.rotary_embedding.yarn import ( + _extend_yarn_cache, yarn_find_correction_range, yarn_get_mscale, yarn_linear_ramp_mask, @@ -419,6 +420,25 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): self.sin_cached_total = torch.sin(emb) * self.mscale return cache + def _ensure_cos_sin_cache_length(self, needed_max_pos: int): + self.cos_sin_cache, rows = _extend_yarn_cache( + cache=self.cos_sin_cache, + compute_inv_freq=lambda: self._compute_inv_freq(self.scaling_factor), + mscale=self.mscale, + needed_max_pos=needed_max_pos, + ) + # NPU also consumes full-width cos/sin tables, built before dtype casting. + if rows is not None and self.cos_cached_total is not None: + cos, sin = rows.chunk(2, dim=-1) + self.cos_cached_total = torch.cat( + (self.cos_cached_total, cos.repeat(1, 2).to(self.cos_cached_total)), + dim=0, + ) + self.sin_cached_total = torch.cat( + (self.sin_cached_total, sin.repeat(1, 2).to(self.sin_cached_total)), + dim=0, + ) + def get_cos_cached_total(self): return self.cos_cached_total diff --git a/python/sglang/srt/layers/rotary_embedding/yarn.py b/python/sglang/srt/layers/rotary_embedding/yarn.py index e2ccb82f5..440a786c0 100644 --- a/python/sglang/srt/layers/rotary_embedding/yarn.py +++ b/python/sglang/srt/layers/rotary_embedding/yarn.py @@ -3,10 +3,11 @@ from __future__ import annotations import math -from typing import Tuple +from typing import Callable, Optional, Tuple import torch +from sglang.srt.environ import envs from sglang.srt.layers.rotary_embedding.base import RotaryEmbedding @@ -62,6 +63,31 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: return 0.1 * mscale * math.log(scale) + 1.0 +def _extend_yarn_cache( + cache: torch.Tensor, + compute_inv_freq: Callable[[], torch.Tensor], + mscale: float, + needed_max_pos: int, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Return the extended cache and uncast rows for auxiliary tables. + + A no-op returns the original cache and None without computing frequencies. + The caller supplies YaRN frequencies using its scaling factor and unchanged + correction-range bound, rather than the base extension's theta argument. + """ + if needed_max_pos < cache.shape[0]: + return cache, None + align = envs.SGLANG_ROPE_CACHE_ALIGN.get() + new_len = ((needed_max_pos + align) // align) * align + inv_freq = compute_inv_freq().to(cache.device) + positions = torch.arange( + cache.shape[0], new_len, dtype=inv_freq.dtype, device=cache.device + ) + freqs = torch.einsum("i,j->ij", positions, inv_freq) + rows = torch.cat((freqs.cos() * mscale, freqs.sin() * mscale), dim=-1) + return torch.cat((cache, rows.to(cache.dtype)), dim=0), rows + + class YaRNScalingRotaryEmbedding(RotaryEmbedding): """RotaryEmbedding extended with YaRN method. @@ -134,6 +160,14 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding): ) return inv_freq + def _ensure_cos_sin_cache_length(self, needed_max_pos: int): + self.cos_sin_cache, _ = _extend_yarn_cache( + cache=self.cos_sin_cache, + compute_inv_freq=lambda: self._compute_inv_freq(self.scaling_factor), + mscale=self.mscale, + needed_max_pos=needed_max_pos, + ) + def _compute_cos_sin_cache(self) -> torch.Tensor: inv_freq = self._compute_inv_freq(self.scaling_factor) t = torch.arange( diff --git a/test/registered/unit/layers/test_yarn_cache_extension.py b/test/registered/unit/layers/test_yarn_cache_extension.py new file mode 100644 index 000000000..361c22f8d --- /dev/null +++ b/test/registered/unit/layers/test_yarn_cache_extension.py @@ -0,0 +1,50 @@ +import sys + +import pytest +import torch +from torch.testing import assert_close + +from sglang.srt.layers.rotary_embedding import base, factory, rope_variant +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +@pytest.mark.parametrize("kind", ("yarn", "deepseek_yarn", "mrope")) +@pytest.mark.parametrize("factor", (1.0, 2.0)) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_cache_extension(monkeypatch, kind, factor, dtype): + monkeypatch.setattr(base, "_is_cpu", True) + monkeypatch.setattr(base, "publish_role", lambda: None) + monkeypatch.setattr(rope_variant, "get_device", lambda: "cpu") + monkeypatch.setattr(rope_variant, "_is_npu", True) + monkeypatch.setattr(factory, "_ROPE_DICT", {}) + scaling = {"rope_type": kind, "factor": factor, "attn_factor": 1.3} + if kind == "mrope": + scaling.update(rope_type="yarn", mrope_section=[16, 8, 8]) + rope = factory.get_rope( + head_size=64, + rotary_dim=64, + max_position=4096, + base=10000, + dtype=torch.float32, + rope_scaling=scaling, + ) + rope.cos_sin_cache = rope.cos_sin_cache.to(dtype) + prefix = rope.cos_sin_cache.clone() + for needed in (len(prefix) - 1, len(prefix), len(prefix) + 256): + rope._ensure_cos_sin_cache_length(needed) + assert len(rope.cos_sin_cache) > needed + assert rope.max_position_embeddings == 4096 + positions = torch.arange(len(rope.cos_sin_cache), dtype=torch.float32) + phase = torch.outer(positions, rope._compute_inv_freq(factor)) + expected = torch.cat((phase.cos(), phase.sin()), dim=-1) * rope.mscale + assert_close(rope.cos_sin_cache, expected.to(dtype), rtol=0, atol=0) + assert torch.equal(rope.cos_sin_cache[: len(prefix)], prefix) + if kind == "deepseek_yarn": + assert_close(rope.cos_cached_total, expected[:, :32].repeat(1, 2)) + assert_close(rope.sin_cached_total, expected[:, 32:].repeat(1, 2)) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-q"]))