[Fix] Preserve YaRN scaling when extending rotary caches (#38786)

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(

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
 

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(

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"]))