[rotary] Rebuild the shared RoPE cache entry when its buffers are dead (#33575)

Co-authored-by: mxz <mxz@fb.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

python/sglang/srt/layers/rotary_embedding/factory.py

@@ -60,6 +60,38 @@ if _use_aiter:
 _ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}
 
 
+def _get_live_rope_cache_entry(key: Tuple) -> Optional[RotaryEmbedding]:
+    """Return the cached module for ``key``, dropping it if its buffers are dead.
+
+    A cached module is shared process-wide and attached as a submodule of every
+    model that requests it, so a model teardown that frees CUDA storages and
+    re-points its own module tree at the meta device kills this entry for all
+    later models too. Nothing downstream can catch that: the in-place RoPE ops
+    take the cos/sin cache as an argument, so a meta tensor routes them to the
+    Meta backend, where they silently no-op and leave queries un-rotated.
+
+    A dead entry is indistinguishable from one a meta-device construction pass
+    built on purpose -- both are meta with no storage -- so the current device
+    is what separates them.
+    """
+    cached = _ROPE_DICT.get(key)
+    if cached is None:
+        return None
+    if torch.get_default_device().type == "meta":
+        return cached
+    for buf in cached.buffers():
+        if buf.device.type == "meta" or buf.untyped_storage().nbytes() == 0:
+            logger.warning(
+                "Discarding dead RoPE cache entry (key=%s): buffer on %s. "
+                "A shared RotaryEmbedding was freed by its owner.",
+                key,
+                buf.device,
+            )
+            del _ROPE_DICT[key]
+            return None
+    return cached
+
+
 def get_rope(
     head_size: int,
     rotary_dim: int,
@@ -103,8 +135,9 @@ def get_rope(
         dual_chunk_attention_args,
         dtype,
     )
-    if key in _ROPE_DICT:
-        return _ROPE_DICT[key]
+    cached = _get_live_rope_cache_entry(key)
+    if cached is not None:
+        return cached
 
     if dual_chunk_attention_config is not None:
         extra_kwargs = {
@@ -380,8 +413,9 @@ def get_rope_cpu(
         rope_scaling_args,
         dtype,
     )
-    if key in _ROPE_DICT:
-        return _ROPE_DICT[key]
+    cached = _get_live_rope_cache_entry(key)
+    if cached is not None:
+        return cached
 
     assert rope_scaling is not None
     scaling_type = rope_scaling["rope_type"]

test/registered/rotary/test_rope_cache_invalidation.py

@@ -0,0 +1,81 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers.rotary_embedding import get_rope
+from sglang.srt.layers.rotary_embedding.factory import _ROPE_DICT
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_ROPE_KWARGS = dict(
+    head_size=64,
+    rotary_dim=64,
+    max_position=256,
+    base=500000,
+    is_neox_style=False,
+)
+
+
+class TestRopeCacheInvalidation(CustomTestCase):
+    """`get_rope` hands out a process-wide shared module, so a model teardown
+    that meta-izes or frees its module tree also kills the cache entry for every
+    later model. Rebuilding is the only way that stays correct — a dead cos/sin
+    cache makes the in-place RoPE ops no-op instead of raising.
+    """
+
+    def setUp(self):
+        cpu_patch = patch("sglang.srt.layers.rotary_embedding.base._is_cpu", True)
+        cpu_patch.start()
+        self.addCleanup(cpu_patch.stop)
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+        _ROPE_DICT.clear()
+
+    def tearDown(self):
+        _ROPE_DICT.clear()
+
+    def test_rebuilds_after_meta_invalidation(self):
+        rope = get_rope(**_ROPE_KWARGS)
+        expected = rope.cos_sin_cache.clone()
+
+        rope.to(device="meta")
+        self.assertEqual(rope.cos_sin_cache.device.type, "meta")
+
+        rebuilt = get_rope(**_ROPE_KWARGS)
+        self.assertIsNot(rebuilt, rope)
+        self.assertNotEqual(rebuilt.cos_sin_cache.device.type, "meta")
+        self.assertTrue(torch.equal(rebuilt.cos_sin_cache, expected))
+
+    def test_rebuilds_after_storage_release(self):
+        rope = get_rope(**_ROPE_KWARGS)
+        expected = rope.cos_sin_cache.clone()
+
+        rope.cos_sin_cache.untyped_storage().resize_(0)
+
+        rebuilt = get_rope(**_ROPE_KWARGS)
+        self.assertIsNot(rebuilt, rope)
+        self.assertTrue(torch.equal(rebuilt.cos_sin_cache, expected))
+
+    def test_live_entry_is_still_shared(self):
+        rope = get_rope(**_ROPE_KWARGS)
+        self.assertIs(get_rope(**_ROPE_KWARGS), rope)
+
+    def test_meta_build_keeps_sharing(self):
+        # Meta-device construction passes build meta buffers on purpose; they
+        # must keep sharing one entry rather than rebuilding on every layer.
+        with torch.device("meta"):
+            rope = get_rope(**_ROPE_KWARGS)
+            self.assertEqual(rope.cos_sin_cache.device.type, "meta")
+            self.assertIs(get_rope(**_ROPE_KWARGS), rope)
+
+        # ...and the real build afterwards must not inherit the meta entry.
+        rebuilt = get_rope(**_ROPE_KWARGS)
+        self.assertIsNot(rebuilt, rope)
+        self.assertNotEqual(rebuilt.cos_sin_cache.device.type, "meta")
+
+
+if __name__ == "__main__":
+    unittest.main()