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