diff --git a/python/sglang/srt/layers/attention/dots_hybrid_backend.py b/python/sglang/srt/layers/attention/dots_hybrid_backend.py index 887f70b4e..cfcff5b4f 100644 --- a/python/sglang/srt/layers/attention/dots_hybrid_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 359bdc851..210323127 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index bfcda1965..126e48bdd 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 814a70d79..a49cd02c7 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 0e2399cec..2962392e4 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/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: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index a02221cbb..9f66e95ab 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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() diff --git a/test/registered/models_e2e/test_kimi_linear_models.py b/test/registered/models_e2e/test_kimi_linear_models.py index 97dff7a5c..9494085bf 100644 --- a/test/registered/models_e2e/test_kimi_linear_models.py +++ b/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() diff --git a/test/registered/unit/layers/attention/test_kv_translate_ownership.py b/test/registered/unit/layers/attention/test_kv_translate_ownership.py index d81b55af0..51c603576 100644 --- a/test/registered/unit/layers/attention/test_kv_translate_ownership.py +++ b/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()