diff --git a/python/sglang/srt/layers/rotary_embedding/factory.py b/python/sglang/srt/layers/rotary_embedding/factory.py index a5f4d9bb3..22f28f571 100644 --- a/python/sglang/srt/layers/rotary_embedding/factory.py +++ b/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"] diff --git a/test/registered/rotary/test_rope_cache_invalidation.py b/test/registered/rotary/test_rope_cache_invalidation.py new file mode 100644 index 000000000..3fd170e63 --- /dev/null +++ b/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()