[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>
2026-08-04 17:24:42 -07:00
parent 87ed82ff7e
commit 211ee64249
2 changed files with 119 additions and 4 deletions
@@ -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()