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