[unified-memory] Let Kimi-Linear use the paged MLA attention backends (#32972)

python/sglang/kernels/ops/kvcache/kv_indices.py

@@ -105,6 +105,15 @@ def create_flashmla_kv_indices_triton(
     req_to_token_ptr_stride: tl.constexpr,
     kv_indices_ptr_stride: tl.constexpr,
     PAGED_SIZE: tl.constexpr = 64,
+    # Unified-memory dense-view path (page-major envelope shared with the mamba
+    # sub-pool). req_to_token holds VIRTUAL token ids; the block table the MLA
+    # kernel consumes must hold DENSE page ids. When v2p_ptr is given, map each
+    # virtual page through it to the physical page, then scale by PAGE_MULT
+    # (= num MLA layers) so the entry addresses the layer's dense per-page block
+    # in the (num_pages*L, page_size, kv_cache_dim) reshaped view. Both default
+    # to the identity (v2p_ptr None, PAGE_MULT 1) for the static pool.
+    v2p_ptr=None,
+    PAGE_MULT: tl.constexpr = 1,
 ):
     NUM_PAGE_PER_BLOCK: tl.constexpr = (
         FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE
@@ -145,8 +154,13 @@ def create_flashmla_kv_indices_triton(
             + paged_offset,
             mask=mask,
         )
+        page = data // PAGED_SIZE
+        if v2p_ptr is not None:
+            # virtual page -> physical page (page-level v2p); masked so padded
+            # lanes never index the table out of bounds.
+            page = tl.load(v2p_ptr + page, mask=mask_out, other=0)
         tl.store(
             kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
-            data // PAGED_SIZE,
+            page * PAGE_MULT,
             mask=mask_out,
         )

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

@@ -66,6 +66,51 @@ if is_flashinfer_available():
     )
 
 
+@dataclass(frozen=True)
+class UnifiedMLAHooks:
+    """Allocator hooks the paged MLA backends need under the unified memory pool.
+
+    All-``None``/1/``False`` for the statically-partitioned pool, where
+    ``req_to_token`` already holds physical ids.
+    """
+
+    # Page-level virtual->physical table, gathered through by the block-table kernel.
+    v2p_page_table: Optional[torch.Tensor]
+    # Virtual token id -> DENSE kernel-facing id.
+    translate_kv_loc_dense: Optional[Callable[..., torch.Tensor]]
+    # Dense page stride scale (= number of full-attention MLA layers).
+    kernel_page_multiplier: int
+    enabled: bool
+
+
+def unified_mla_hooks(allocator) -> UnifiedMLAHooks:
+    """Probe ``allocator`` for the unified-pool dense-view hooks.
+
+    Detection keys on the page-level v2p table, NOT on
+    ``kernel_page_multiplier > 1``: a configuration with exactly ONE
+    full-attention layer (e.g. a pipeline-parallel rank that owns a single MLA
+    layer) has multiplier 1 while its ``req_to_token`` still holds VIRTUAL ids.
+    With multiplier 1 the dense id collapses onto the physical id, so the v2p
+    gather alone is the whole translation -- skipping it would leave the block
+    table and the KV write loc in virtual space and silently address the wrong
+    pages once virtual and physical diverge (e.g. after compaction).
+    """
+    v2p = getattr(allocator, "full_v2p_page_table", None)
+    if v2p is None:
+        return UnifiedMLAHooks(
+            v2p_page_table=None,
+            translate_kv_loc_dense=None,
+            kernel_page_multiplier=1,
+            enabled=False,
+        )
+    return UnifiedMLAHooks(
+        v2p_page_table=v2p,
+        translate_kv_loc_dense=getattr(allocator, "translate_kv_loc_dense", None),
+        kernel_page_multiplier=getattr(allocator, "kernel_page_multiplier", 1),
+        enabled=True,
+    )
+
+
 @dataclass
 class DecodeMetadata:
     decode_wrapper: BatchMLAPagedAttentionWrapper
@@ -675,6 +720,10 @@ class FlashInferMLAIndicesUpdaterDecode:
         self.kv_indptr = attn_backend.kv_indptr
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.q_indptr = attn_backend.q_indptr_decode
+        # Unified dense MLA pool: VIRTUAL -> DENSE kv_indices (see prefill updater).
+        self._translate_kv_loc_dense = unified_mla_hooks(
+            model_runner.token_to_kv_pool_allocator
+        ).translate_kv_loc_dense
 
     def update(
         self,
@@ -732,6 +781,21 @@ class FlashInferMLAIndicesUpdaterDecode:
                 kv_indices,
                 self.req_to_token.shape[1],
             )
+            # Unified pool: VIRTUAL -> DENSE, written back IN PLACE.
+            #
+            # On the cuda-graph replay path `kv_indices` IS the capture-stable
+            # buffer (fast_decode_kwargs["kv_indices"] == cuda_graph_kv_indices)
+            # that the captured wrapper reads, and `fast_mla_decode_plan` ignores
+            # the kv_indices argument entirely -- rebinding the local name to a
+            # fresh tensor would leave the graph reading VIRTUAL ids. Only the
+            # [:paged_kernel_lens_sum] prefix the index kernel just filled is
+            # translated; the stale tail is left alone so it can never index the
+            # v2p table out of bounds. The int64 translate result narrows back to
+            # the buffer's int32 on copy_ (flashinfer requires int32; dense ids
+            # fit comfortably).
+            if self._translate_kv_loc_dense is not None:
+                valid = kv_indices[:paged_kernel_lens_sum]
+                valid.copy_(self._translate_kv_loc_dense(valid))
 
             if get_parallel().dcp_enabled:
                 plan_dcp_decode_metadata(
@@ -797,6 +861,12 @@ class FlashInferMLAIndicesUpdaterPrefill:
         self.qo_indptr = attn_backend.qo_indptr
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
+        # Unified dense MLA pool: kv_indices built from req_to_token are VIRTUAL;
+        # the paged wrapper reads the dense per-layer view, so remap them to DENSE
+        # token ids. None (identity) unless the unified MLA pool is active.
+        self._translate_kv_loc_dense = unified_mla_hooks(
+            model_runner.token_to_kv_pool_allocator
+        ).translate_kv_loc_dense
 
     def update(
         self,
@@ -867,6 +937,12 @@ class FlashInferMLAIndicesUpdaterPrefill:
                 kv_indices,
                 self.req_to_token.shape[1],
             )
+            # Unified pool: VIRTUAL -> DENSE token ids for the paged wrapper.
+            # Prefill is not cuda-graph captured under unified memory, so an eager
+            # gather is safe. Dense ids fit int32 (max = full_slots*num_layers ~
+            # 1e7 << 2^31); the flashinfer wrapper requires int32.
+            if self._translate_kv_loc_dense is not None:
+                kv_indices = self._translate_kv_loc_dense(kv_indices).to(torch.int32)
             qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
             qo_indptr = qo_indptr[: bs + 1]
             custom_mask = None

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

@@ -34,6 +34,7 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.attention.flashinfer_mla_backend import (
     FlashInferMLAAttnBackend,
     FlashInferMLAMultiStepDraftBackend,
+    unified_mla_hooks,
 )
 from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -246,6 +247,23 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         # Tree-mask scratch is fetched from the target backend only.
         self.is_draft_runner = model_runner.is_draft_worker
 
+        # Unified-memory dense-view hooks (None on the static pool). req_to_token
+        # holds VIRTUAL token ids; the block table needs DENSE page ids, so the
+        # kv-index kernels gather virtual->physical page through `_v2p_page_table`
+        # then scale by `_kernel_page_multiplier` (= num MLA layers). See
+        # build_dense_mla_views / create_flashmla_kv_indices_triton.
+        _hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
+        self._v2p_page_table = _hooks.v2p_page_table
+        self._kernel_page_multiplier = _hooks.kernel_page_multiplier
+        self._unified_mla = _hooks.enabled
+        # virtual token id -> DENSE kernel-facing id, for the KV write loc.
+        self._translate_kv_loc_dense = _hooks.translate_kv_loc_dense
+        # Per-forward dense write loc ([:n] view of a capture-stable buffer),
+        # set by the cuda-graph out-graph hook; None on the eager path (where the
+        # write translates through the pool's _full_translate hook instead).
+        self._decode_dense_loc: Optional[torch.Tensor] = None
+        self.cuda_graph_out_cache_loc_dense: Optional[torch.Tensor] = None
+
     def _calc_padded_blocks(self, max_seq_len: int) -> int:
         """
         Calculate padded block count that satisfies both TRT-LLM and Triton constraints.
@@ -308,6 +326,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             self.req_to_token.stride(0),
             max_blocks,
             PAGED_SIZE=self.page_size,
+            v2p_ptr=self._v2p_page_table,
+            PAGE_MULT=self._kernel_page_multiplier,
         )
 
         return block_kv_indices
@@ -325,6 +345,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         self.decode_cuda_graph_kv_indices = torch.full(
             (max_bs, max_blocks_per_seq), -1, dtype=torch.int32, device=self.device
         )
+        # Unified pool: capture-stable buffer for the DENSE KV write loc, filled
+        # out-of-graph in init_forward_metadata_out_graph so the in-graph
+        # set_mla_kv_buffer captures no translate.
+        if self._unified_mla:
+            self.cuda_graph_out_cache_loc_dense = torch.zeros(
+                max_num_tokens, dtype=torch.int64, device=self.device
+            )
         num_tokens_per_req = max_num_tokens // max_bs
 
         if is_float4_e2m1fn_x2(self.data_type):
@@ -464,6 +491,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             self.req_to_token.stride(0),
             metadata.block_kv_indices.shape[1],
             PAGED_SIZE=self.page_size,
+            v2p_ptr=self._v2p_page_table,
+            PAGE_MULT=self._kernel_page_multiplier,
         )
 
     def get_cuda_graph_seq_len_fill_value(self) -> int:
@@ -522,8 +551,32 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 forward_mode=forward_mode,
             )
 
+        # Unified pool: precompute the DENSE KV write loc into the capture-stable
+        # buffer (both capture and each replay-prep run this out of the graph),
+        # so the in-graph set_mla_kv_buffer writes a dense loc without capturing a
+        # translate. Only decode writes KV under unified (spec is gated off).
+        if self._unified_mla and forward_mode.is_decode_or_idle():
+            out_cache_loc = forward_batch.out_cache_loc
+            n = out_cache_loc.shape[0]
+            dst = self.cuda_graph_out_cache_loc_dense[:n]
+            self._translate_kv_loc_dense(out_cache_loc, out=dst)
+            # Replay-prep receives the RAW (unpadded) out_cache_loc
+            # (build_replay_fb_view), but the captured write kernel consumes the
+            # full captured tier of this buffer. Zero the tail so pad rows write
+            # to the dense sink (row 0) instead of stale dense locs left by
+            # earlier larger replays — a stale tail scatters pad-row garbage into
+            # live KV pages. Mirrors the runner's PaddingPolicy.ZERO on its own
+            # out_cache_loc slot.
+            self.cuda_graph_out_cache_loc_dense[n:].zero_()
+            self._decode_dense_loc = dst
+        else:
+            self._decode_dense_loc = None
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Initialize the metadata for a forward pass."""
+        # Eager path: no capture-stable dense write loc; the pool's _full_translate
+        # hook translates the write loc (safe out of a cuda graph).
+        self._decode_dense_loc = None
         # Delegate to parent for non-decode modes.
         if (
             forward_batch.forward_mode.is_extend()
@@ -803,9 +856,17 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             assert (
                 k is not None and k_rope is not None
             ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
-            self.token_to_kv_pool.set_mla_kv_buffer(
-                layer, forward_batch.out_cache_loc, k, k_rope
-            )
+            if self._decode_dense_loc is not None:
+                # cuda-graph path: dense write loc precomputed out-of-graph, so
+                # the in-graph write captures no translate allocation.
+                self.token_to_kv_pool.set_mla_kv_buffer(
+                    layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
+                )
+            else:
+                # eager (or static pool): the pool's _full_translate handles it.
+                self.token_to_kv_pool.set_mla_kv_buffer(
+                    layer, forward_batch.out_cache_loc, k, k_rope
+                )
 
         # Prepare query tensor inline
         if merge_query:

python/sglang/srt/mem_cache/memory_pool.py

@@ -3836,12 +3836,17 @@ class HybridLinearKVPool(KVCache):
         loc: torch.Tensor,
         cache_k_nope: torch.Tensor,
         cache_k_rope: torch.Tensor,
+        loc_is_dense: bool = False,
     ):
         assert self.use_mla, "set_mla_kv_buffer called when use_mla is False"
         # Model-level MLA entry point: `loc` is a VIRTUAL loc under the unified
-        # pool (eager prefill only; the decode write goes through set_kv_buffer's
-        # pre-translated `full_loc`), so translate to the dense id space here.
-        loc = self._full_translate(loc)
+        # pool, so translate to the dense id space here.
+        #
+        # `loc_is_dense`: the caller already translated `loc` (the unified-pool
+        # cuda-graph decode precomputes it out-of-graph into a capture-stable
+        # buffer, so the in-graph write does not capture a translate allocation).
+        if not loc_is_dense:
+            loc = self._full_translate(loc)
         with self._transfer_id_context(layer):
             self.full_kv_pool.set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope)
 

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -1898,6 +1898,15 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def kernel_page_multiplier(self) -> int:
         return self.full_attn_allocator.kernel_page_multiplier
 
+    @property
+    def full_v2p_page_table(self) -> torch.Tensor:
+        """Page-level virtual->physical table of the full sub-pool. Kernels that
+        build the MLA block table directly from req_to_token (e.g. trtllm_mla,
+        flashmla) gather through this to turn a VIRTUAL page into a physical one,
+        then scale by `kernel_page_multiplier` to reach the dense per-page block.
+        """
+        return self.full_attn_allocator.virtual_to_physical
+
     def translate_kv_loc_dense(
         self,
         loc: torch.Tensor,

python/sglang/srt/server_args.py

@@ -7705,13 +7705,33 @@ class ServerArgs:
         if not self.enable_page_major_kv_layout:
             return
         # Only the Triton attention kernels read the strided 4-D envelope K/V
-        # views; FA3 / FlashInfer do not.
+        # views; FA3 / FlashInfer do not. EXCEPTION: the unified-memory MLA pool
+        # exposes each layer as a DENSE contiguous per-layer view
+        # (build_dense_mla_views), which the paged MLA kernels consume directly,
+        # with their kv_indices / block tables remapped to dense ids. Names below
+        # are the RESOLVED ids from _resolved_attention_backends: "flashinfer" is
+        # FlashInferMLAAttnBackend for an MLA model, "trtllm_mla" the trtllm
+        # decode kernel; "cutedsl_mla" and "tokenspeed_mla" subclass
+        # TRTLLMMLABackend and inherit its dense read/write path.
+        # flashmla / cutlass_mla share the create_flashmla block-table path and
+        # can be added the same way once exercised.
+        if self.enable_unified_memory and self.use_mla_backend():
+            allowed_full = {
+                "triton",
+                "trtllm_mla",
+                "flashinfer",
+                "cutedsl_mla",
+                "tokenspeed_mla",
+            }
+        else:
+            allowed_full = {"triton"}
         backends = set(self._resolved_attention_backends())
         backends.discard(None)
-        assert backends <= {"triton"}, (
+        assert backends <= allowed_full, (
             "--enable-page-major-kv-layout requires the Triton attention backend "
-            f"for the full-attention layers; got {sorted(backends)}. Pass "
-            "--attention-backend triton."
+            "for the full-attention layers (unified-memory MLA also allows the "
+            f"paged MLA backends); got {sorted(backends)}, allowed "
+            f"{sorted(allowed_full)}. Pass a compatible --attention-backend."
         )
         # The Mamba state is stored in envelope-strided views; only the
         # stride-aware Triton causal-conv / SSM kernels read them correctly.