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>

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()