fix(unified-memory): forward the KV-index translator through every wrapper backend (#37307)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

python/sglang/srt/layers/attention/dots_hybrid_backend.py

@@ -140,6 +140,7 @@ class DotsSWAMLAAttnBackend(AttentionBackend):
         self._active_backend = backend
         self.token_to_kv_pool = backend.token_to_kv_pool
         self.req_to_token_pool = backend.req_to_token_pool
+        self.kv_index_translator = backend.kv_index_translator
         self.needs_cpu_seq_lens = True
         self._prefill_metadata: DotsSWAMLAPrefillMetadata | None = None
         self._dp_rebuilt_batch_id: int | None = None
@@ -404,6 +405,7 @@ class DotsHybridAttnBackend(AttentionBackend):
         self.swa_backend = swa_backend
         self.token_to_kv_pool = swa_backend.token_to_kv_pool
         self.req_to_token_pool = swa_backend.req_to_token_pool
+        self.kv_index_translator = swa_backend.kv_index_translator
         # SWA latent expansion uses host sequence-length mirrors.
         self.needs_cpu_seq_lens = True
         self._dp_rebuilt_batch_id: int | None = None

python/sglang/srt/layers/attention/hybrid_attn_backend.py

@@ -36,6 +36,7 @@ class HybridAttnBackend(AttentionBackend):
         self.data_type = model_runner.kv_cache_dtype
         self.token_to_kv_pool = model_runner.token_to_kv_pool
         self.req_to_token_pool = model_runner.req_to_token_pool
+        self.kv_index_translator = model_runner.kv_index_translator
         self.spec_attn_is_decode = get_spec().speculative_attention_mode == "decode"
         self.spec_attn_is_prefill = get_spec().speculative_attention_mode == "prefill"
         # Gates the FutureMap's per-step seq_lens D2H (decide_needs_cpu_seq_lens

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -997,6 +997,7 @@ class HybridLinearAttnBackend(AttentionBackend):
         self.attn_backend_list = [full_attn_backend, linear_attn_backend]
         self.token_to_kv_pool = full_attn_backend.token_to_kv_pool
         self.req_to_token_pool = full_attn_backend.req_to_token_pool
+        self.kv_index_translator = full_attn_backend.kv_index_translator
         self.max_context_len = getattr(full_attn_backend, "max_context_len", None)
         self.needs_cpu_seq_lens = (
             full_attn_backend.needs_cpu_seq_lens

python/sglang/srt/layers/attention/minimax_sparse_backend.py

@@ -1617,6 +1617,7 @@ class MiniMaxHybridAttnBackend(AttentionBackend):
     ):
         self.dense = dense_backend
         self.sparse = sparse_backend
+        self.kv_index_translator = dense_backend.kv_index_translator
         self.sparse_layer_ids = sparse_layer_ids
         # Let the sparse decode reuse the dense paged backend (page table + workspace).
         self.sparse.dense_backend = dense_backend

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -308,6 +308,22 @@ class KVIndexTranslator:
         self._index_table_memo = (weakref.ref(forward_batch), view)
         return view
 
+    def assert_backends_carry_translator(self, backends) -> None:
+        """Boot guard: under the unified pool every backend a forward can reach
+        must carry THIS translator."""
+        if not self.is_translating:
+            return
+        for backend in backends:
+            if backend is None:
+                continue
+            assert backend.kv_index_translator is self, (
+                f"{type(backend).__name__} does not carry the runner's "
+                "KVIndexTranslator. A backend (or wrapper) reachable under "
+                "--enable-unified-memory must forward `kv_index_translator`, or "
+                "read-index producers silently skip the virtual->kernel-facing "
+                "translation."
+            )
+
     # -- write loc (phase 1; phase 2 lives in build_index_table) ----------------
 
     def rebind_write_loc(self, forward_batch) -> None:

python/sglang/srt/model_executor/model_runner.py

@@ -1011,6 +1011,9 @@ class ModelRunner:
         self.attn_backend = backends.attn_backend
         self.decode_attn_backend = backends.decode_attn_backend
         self.decode_attn_backend_group = backends.decode_attn_backend_group
+        self.kv_index_translator.assert_backends_carry_translator(
+            [self.attn_backend, self.decode_attn_backend]
+        )
 
         if get_parallel().dcp_enabled and get_parallel().dcp_replicate_q_proj:
             self._prepare_replicated_q_proj()

test/registered/models_e2e/test_kimi_linear_models.py

@@ -15,7 +15,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=600, stage="base-b", runner_config="2-gpu-large")
+register_cuda_ci(est_time=900, stage="base-b", runner_config="2-gpu-large")
 
 KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
 
@@ -78,5 +78,26 @@ class TestKimiLinearExtraBuffer(
     ]
 
 
+class TestKimiLinearUnifiedMemory(
+    GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
+):
+    """BUG REGRESSION. The unified pool must keep GSM8K at the static-pool bar;
+    a wrapper backend that drops `kv_index_translator` silently skips the MLA
+    prefix translation and the model reads the wrong KV."""
+
+    model = KIMI_LINEAR_MODEL
+    cache_chunk_size = 64
+    gsm8k_score_threshold = 0.88
+    # No --attention-backend: the resolved default is the coverage.
+    other_args = [
+        "--trust-remote-code",
+        "--tp-size",
+        "2",
+        "--chunked-prefill-size",
+        "2048",
+        "--enable-unified-memory",
+    ]
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/layers/attention/test_kv_translate_ownership.py

@@ -17,11 +17,27 @@ consciously.
     python3 -m pytest test/registered/unit/layers/attention/test_kv_translate_ownership.py -v
 """
 
+import ast
 import os
 import re
 import unittest
 
 from sglang.srt.layers.attention import triton_backend as _anchor_module
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.dots_hybrid_backend import (
+    DotsHybridAttnBackend,
+    DotsSWAMLAAttnBackend,
+)
+from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    HybridLinearAttnBackend,
+    ShortConvHybridAttnBackend,
+)
+from sglang.srt.layers.attention.minimax_sparse_backend import (
+    MiniMaxHybridAttnBackend,
+)
+from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
+from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -79,5 +95,106 @@ class TestUnifiedTranslateBanned(CustomTestCase):
         self.assertEqual(hits, [])
 
 
+def _derive_wrapper_names():
+    """Wrapper classes, read off the source so a NEW one shows up the day it is
+    written: an AttentionBackend subclass whose own __init__ takes another
+    backend. Per class, not per file -- a file-wide scan passes as soon as any
+    one class in it forwards."""
+    backend = re.compile(r"Att(?:ention|n)Backend")
+    names = set()
+    for _rel, src in _iter_sources():
+        for node in ast.walk(ast.parse(src)):
+            if not isinstance(node, ast.ClassDef):
+                continue
+            if not any(backend.search(ast.unparse(b)) for b in node.bases):
+                continue
+            init = next(
+                (
+                    b
+                    for b in node.body
+                    if isinstance(b, ast.FunctionDef) and b.name == "__init__"
+                ),
+                None,
+            )
+            if init is None:
+                continue
+            if any(
+                a.annotation is not None and backend.search(ast.unparse(a.annotation))
+                for a in init.args.args[1:]
+            ):
+                names.add(node.name)
+    return names
+
+
+class _Inner(AttentionBackend):
+    """Stand-in for a wrapped backend, carrying only what the wrappers' __init__
+    bodies read off their inners."""
+
+    def __init__(self, translator=None):
+        self.kv_index_translator = translator
+        self.token_to_kv_pool = None
+        self.req_to_token_pool = None
+        self.needs_cpu_seq_lens = False
+        self.max_context_len = 8
+
+
+class _InnerModelConfig:
+    context_len = 8
+
+
+class _Runner:
+    """HybridAttnBackend takes its translator from the runner, not an inner."""
+
+    def __init__(self, translator):
+        self.kv_index_translator = translator
+        self.kv_cache_dtype = None
+        self.token_to_kv_pool = None
+        self.req_to_token_pool = None
+        self.model_config = _InnerModelConfig()
+
+
+def _build_wrappers(translator):
+    """One live instance per wrapper. Only the inner that MUST supply the
+    translator carries it; every other inner carries None, so a wrapper that
+    copies from the linear / sparse / DSA side ends up with None and fails."""
+    carrier = _Inner(translator)
+    # HybridAttnBackend reads the spec bag in __init__; the bag is unpublished
+    # outside a launched server.
+    with get_context().override_server_args(speculative_attention_mode="decode"):
+        hybrid = HybridAttnBackend(_Runner(translator), _Inner(), _Inner())
+    return {
+        "DotsSWAMLAAttnBackend": DotsSWAMLAAttnBackend(carrier),
+        "DotsHybridAttnBackend": DotsHybridAttnBackend(_Inner(), carrier),
+        "HybridAttnBackend": hybrid,
+        "HybridLinearAttnBackend": HybridLinearAttnBackend(carrier, _Inner(), [0]),
+        "ShortConvHybridAttnBackend": ShortConvHybridAttnBackend(
+            carrier, _Inner(), [0]
+        ),
+        "MiniMaxHybridAttnBackend": MiniMaxHybridAttnBackend(carrier, _Inner(), [0]),
+        "TboAttnBackend": TboAttnBackend(carrier, [_Inner()]),
+    }
+
+
+class TestWrapperBackendsForwardTranslator(CustomTestCase):
+    """BUG REGRESSION. `AttentionBackend.kv_index_translator` defaults to None,
+    so a wrapper that does not re-expose its inner's copy reads as "needs no
+    translation" and producers that fetch it off `get_attn_backend()` skip the
+    virtual->kernel-facing translation instead of failing."""
+
+    def test_every_wrapper_is_constructed_here(self):
+        self.assertEqual(
+            _derive_wrapper_names(),
+            set(_build_wrappers(object())),
+            "a wrapper backend has no instance in _build_wrappers; add one so "
+            "its translator forwarding is checked",
+        )
+
+    def test_wrappers_forward_the_translator(self):
+        translator = object()
+        for name, wrapper in _build_wrappers(translator).items():
+            with self.subTest(wrapper=name):
+                self.assertIs(wrapper.kv_index_translator, translator)
+
+
 if __name__ == "__main__":
     unittest.main()