[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(
+            if self._decode_dense_loc is not None:
-                layer, forward_batch.out_cache_loc, k, k_rope
+                # 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
+        # pool, so translate to the dense id space here.
-        # pre-translated `full_loc`), so translate to the dense id space here.
+        #
-        loc = self._full_translate(loc)
+        # `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 "
+            "for the full-attention layers (unified-memory MLA also allows the "
-            "--attention-backend triton."
+            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.

test/registered/models_e2e/test_kimi_linear_unified_memory.py

@@ -0,0 +1,77 @@
+"""Kimi-Linear (MLA full attention + KDA linear attention) served from the
+unified memory pool.
+
+`--enable-unified-memory` replaces the statically-partitioned hybrid pools with
+one byte buffer split dynamically between the full-attention KV sub-pool and the
+per-request KDA state sub-pool. For an MLA model the full side is exposed as
+DENSE per-layer views (`build_dense_mla_views`) and every loc the kernels see is
+a translated virtual id, so the whole read/write path differs from the static
+pool: `translate_kv_loc_dense` for kv_indices and the cuda-graph write loc,
+`HybridLinearKVPool._full_translate` for the model-level MLA entry points, and
+page-envelope relocation on allocator compaction.
+
+None of that is covered by the CPU/GPU unit tests, which pin the pool in
+isolation. This is the end-to-end guard: accuracy must match the static-pool
+baseline, and the prefix-cache branching case must still hit, since a radix hit
+replays virtual locs whose physical pages may have moved under compaction.
+
+Reference numbers on 2x H200 TP2, GSM8K 400 examples (2026-07-30):
+static pools 0.915, `--enable-unified-memory` 0.900 (1 sigma ~= 0.015) -- both with
+the attention backend pinned to triton, as this test runs it. For reference the
+paged MLA kernels land in the same band on a single B300 TP1, GSM8K 200, unified
+(2026-07-31): 0.915 with flashinfer prefill+decode, 0.900 with trtllm_mla. Those
+are not exercised here (see the comment on `other_args`).
+Nightly-only: it needs a second full 48B server launch, which is too much to add
+to per-PR CI on top of the existing Kimi-Linear e2e coverage.
+
+    python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory.py -v
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+
+register_cuda_ci(est_time=600, suite="nightly-4-gpu", nightly=True)
+
+KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
+
+
+class TestKimiLinearUnifiedMemory(
+    GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
+):
+    model = KIMI_LINEAR_MODEL
+    cache_chunk_size = 64
+    # Same bar as the static-pool Kimi-Linear e2e test: unified memory must not
+    # cost accuracy (measured 0.900 vs 0.915 static, see the module docstring).
+    gsm8k_score_threshold = 0.88
+    other_args = [
+        "--trust-remote-code",
+        "--tp-size",
+        "2",
+        "--chunked-prefill-size",
+        "2048",
+        "--enable-unified-memory",
+        # Pinned because the resolved default is not portable: on pre-Blackwell
+        # (this suite's runner is H100) an unspecified backend resolves to `fa3`,
+        # which cannot read the dense views at all, so the un-pinned form fails at
+        # startup with the page-major allowlist assertion. Unified memory on such a
+        # host currently REQUIRES an explicit compatible --attention-backend; that
+        # is a real usability gap, tracked separately, not something this test can
+        # paper over.
+        #
+        # Consequence to keep in mind: pinning triton means this test does NOT
+        # cover the paged MLA backends (trtllm_mla / flashinfer / cutedsl_mla /
+        # tokenspeed_mla), which is where dense-id translation bugs live -- a
+        # captured flashinfer decode reading untranslated virtual ids scored GSM8K
+        # 0.000 on a healthy server. Those paths are covered by the unit tests plus
+        # manual B300 runs; an sm100-gated case here would close the gap.
+        "--attention-backend",
+        "triton",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py

@@ -0,0 +1,375 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""DENSE block table / kv_indices for the paged MLA backends under the unified
+memory pool (Kimi-Linear).
+
+`req_to_token` holds VIRTUAL token ids, while the per-layer MLA views are dense
+(`build_dense_mla_views`). The paged MLA backends therefore need their page-level
+block table filled with DENSE page ids:
+
+    dense_page(virtual_page) = v2p[virtual_page] * layer_num
+
+`create_flashmla_kv_indices_triton` does that in-kernel via `v2p_ptr` / `PAGE_MULT`
+(trtllm_mla / cutedsl_mla / tokenspeed_mla), and the flashinfer_mla updaters do it
+by post-gathering `translate_kv_loc_dense` over the token-level kv_indices.
+
+Covered here:
+  - kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main;
+  - kernel dense mapping against the python reference, for several page sizes,
+    ragged sequence lengths and a non-identity v2p permutation;
+  - padded block-table lanes never index the v2p table out of bounds;
+  - the token-level dense translate the flashinfer updaters apply agrees with the
+    page-level block table the trtllm path builds.
+
+    python -m pytest test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py -v
+"""
+
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
+
+_HAS_CUDA = torch.cuda.is_available()
+_DEV = "cuda"
+_LAYERS = 24  # K3 MLA full-attention layer count
+
+
+def _fill_block_table(
+    req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult
+):
+    from sglang.kernels.ops.kvcache.kv_indices import (
+        create_flashmla_kv_indices_triton,
+        get_num_kv_index_blocks_flashmla,
+    )
+
+    bs = req_pool_indices.shape[0]
+    max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size
+    out = torch.full((bs, max_blocks), -1, dtype=torch.int32, device=_DEV)
+    grid = (bs, get_num_kv_index_blocks_flashmla(max_blocks, page_size))
+    create_flashmla_kv_indices_triton[grid](
+        req_to_token,
+        req_pool_indices,
+        seq_lens,
+        None,
+        out,
+        req_to_token.stride(0),
+        max_blocks,
+        PAGED_SIZE=page_size,
+        v2p_ptr=v2p,
+        PAGE_MULT=mult,
+    )
+    return out
+
+
+def _reference(req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult):
+    """Python reference: virtual token -> virtual page -> physical page -> dense."""
+    bs = req_pool_indices.shape[0]
+    max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size
+    ref = torch.full((bs, max_blocks), -1, dtype=torch.int64, device=_DEV)
+    for r in range(bs):
+        n_pages = (int(seq_lens[r].item()) + page_size - 1) // page_size
+        row = req_to_token[int(req_pool_indices[r].item())]
+        virt_pages = row[: n_pages * page_size : page_size] // page_size
+        pages = v2p[virt_pages.long()] if v2p is not None else virt_pages.long()
+        ref[r, :n_pages] = pages * mult
+    return ref
+
+
+@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
+class TestDenseBlockTable(unittest.TestCase):
+    def _make_batch(self, page_size, bs=5, max_ctx=2048, n_pages=512):
+        """Ragged batch with a non-identity virtual->physical page permutation."""
+        g = torch.Generator(device="cpu").manual_seed(97 + page_size)
+        seq_lens = torch.tensor(
+            [page_size, page_size + 1, 3 * page_size, 7 * page_size - 3, 1],
+            dtype=torch.int32,
+            device=_DEV,
+        )[:bs]
+        req_to_token = torch.zeros((bs, max_ctx), dtype=torch.int32, device=_DEV)
+        # Every request gets a distinct, non-monotonic run of virtual pages.
+        virt_page_perm = torch.randperm(n_pages - 1, generator=g)[: bs * 8] + 1
+        for r in range(bs):
+            n = int(seq_lens[r].item())
+            pages = virt_page_perm[r * 8 : r * 8 + 8].to(_DEV)
+            toks = (
+                pages[:, None] * page_size
+                + torch.arange(page_size, device=_DEV)[None, :]
+            ).reshape(-1)
+            req_to_token[r, :n] = toks[:n].to(torch.int32)
+        req_pool_indices = torch.arange(bs, dtype=torch.int32, device=_DEV)
+        # Non-identity page-level v2p, with a tombstone that no request references.
+        v2p = torch.randperm(n_pages, generator=g).to(_DEV).to(torch.int64)
+        v2p[0] = 0  # page 0 is the reserved sink
+        return req_to_token, req_pool_indices, seq_lens, v2p
+
+    def test_identity_when_hooks_absent(self):
+        """v2p_ptr=None / PAGE_MULT=1 must reproduce the pre-change behaviour."""
+        for page_size in (1, 32, 64):
+            rt, rpi, sl, _ = self._make_batch(page_size)
+            got = _fill_block_table(rt, rpi, sl, page_size, v2p=None, mult=1)
+            want = _reference(rt, rpi, sl, page_size, v2p=None, mult=1)
+            self.assertTrue(
+                torch.equal(got.long(), want), f"page_size={page_size}: {got} != {want}"
+            )
+
+    def test_dense_block_table_matches_reference(self):
+        for page_size in (1, 32, 64):
+            rt, rpi, sl, v2p = self._make_batch(page_size)
+            got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
+            want = _reference(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
+            self.assertTrue(
+                torch.equal(got.long(), want),
+                f"page_size={page_size}:\ngot ={got}\nwant={want}",
+            )
+
+    def test_single_full_attention_layer_still_maps_v2p(self):
+        """A config with exactly ONE full-attention layer (e.g. a PP rank owning a
+        single MLA layer) has `kernel_page_multiplier == 1`, but its req_to_token
+        still holds VIRTUAL ids. The dense id collapses onto the physical id, so
+        the v2p gather alone IS the whole translation -- it must not be skipped.
+
+        Regression guard for detecting the unified pool via `multiplier > 1`:
+        that predicate treats this config as a static pool and leaves the block
+        table in virtual id space.
+        """
+        for page_size in (1, 64):
+            rt, rpi, sl, v2p = self._make_batch(page_size)
+            got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=1).long()
+            want = _reference(rt, rpi, sl, page_size, v2p=v2p, mult=1)
+            self.assertTrue(torch.equal(got, want), f"page_size={page_size}")
+            # ... and the v2p permutation is non-trivial here, so a skipped
+            # translation would be visibly different rather than accidentally equal.
+            virtual = _reference(rt, rpi, sl, page_size, v2p=None, mult=1)
+            self.assertFalse(
+                torch.equal(want, virtual),
+                "test batch degenerated: v2p is the identity on the pages used",
+            )
+
+    def test_padded_lanes_stay_untouched(self):
+        """Lanes past a request's page count keep the -1 fill: the masked v2p load
+        must not write a translated value (nor read out of bounds)."""
+        page_size = 64
+        rt, rpi, sl, v2p = self._make_batch(page_size)
+        got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
+        for r in range(got.shape[0]):
+            n_pages = (int(sl[r].item()) + page_size - 1) // page_size
+            self.assertTrue(
+                torch.all(got[r, n_pages:] == -1),
+                f"row {r} padded lanes were written: {got[r]}",
+            )
+
+    def test_agrees_with_token_level_dense_translate(self):
+        """The flashinfer updaters translate TOKEN ids with
+        `translate_kv_loc_dense`; the trtllm path builds PAGE ids in-kernel. Both
+        must address the same dense page block."""
+        page_size = 64
+        rt, rpi, sl, v2p = self._make_batch(page_size)
+        block_table = _fill_block_table(
+            rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS
+        ).long()
+        for r in range(rt.shape[0]):
+            n = int(sl[r].item())
+            virt_tokens = rt[r, :n].long()
+            # translate_kv_loc_dense's formula, applied to token ids.
+            dense_tokens = (
+                v2p[virt_tokens // page_size] * (page_size * _LAYERS)
+                + virt_tokens % page_size
+            )
+            # The block-table entry scaled by page_size must be the dense id of
+            # each page's first token.
+            first_of_page = dense_tokens[::page_size]
+            n_pages = (n + page_size - 1) // page_size
+            self.assertTrue(
+                torch.equal(block_table[r, :n_pages] * page_size, first_of_page),
+                f"row {r}: block table and token translate disagree",
+            )
+
+
+class TestUnifiedMLAHookDetection(unittest.TestCase):
+    """`unified_mla_hooks` decides whether the paged MLA backends translate at
+    all. Getting the predicate wrong is silent: the block table and KV write loc
+    stay in virtual id space and address the wrong pages once virtual and
+    physical diverge (e.g. after compaction)."""
+
+    @staticmethod
+    def _probe(**attrs):
+        from sglang.srt.layers.attention.flashinfer_mla_backend import (
+            unified_mla_hooks,
+        )
+
+        class _Alloc:
+            pass
+
+        alloc = _Alloc()
+        for k, v in attrs.items():
+            setattr(alloc, k, v)
+        return unified_mla_hooks(alloc)
+
+    def test_static_pool_disables_every_hook(self):
+        """No v2p table -> statically-partitioned pool; req_to_token is already
+        physical, so all hooks must stay off (byte-identical to pre-change)."""
+        hooks = self._probe()
+        self.assertFalse(hooks.enabled)
+        self.assertIsNone(hooks.v2p_page_table)
+        self.assertIsNone(hooks.translate_kv_loc_dense)
+        self.assertEqual(hooks.kernel_page_multiplier, 1)
+
+    def test_multi_layer_unified_pool(self):
+        table = torch.arange(8)
+        hooks = self._probe(
+            full_v2p_page_table=table,
+            translate_kv_loc_dense=lambda x, **kw: x,
+            kernel_page_multiplier=_LAYERS,
+        )
+        self.assertTrue(hooks.enabled)
+        self.assertIs(hooks.v2p_page_table, table)
+        self.assertIsNotNone(hooks.translate_kv_loc_dense)
+        self.assertEqual(hooks.kernel_page_multiplier, _LAYERS)
+
+    def test_single_full_attention_layer_pool_is_still_unified(self):
+        """REGRESSION: `kernel_page_multiplier == 1` does NOT mean static.
+
+        A hybrid MLA config with exactly one full-attention layer (e.g. a
+        pipeline-parallel rank owning a single MLA layer) has multiplier 1, yet
+        its locs are still virtual. Detecting on `multiplier > 1` would disable
+        the v2p gather here and corrupt reads/writes after compaction.
+        """
+        table = torch.arange(8)
+        hooks = self._probe(
+            full_v2p_page_table=table,
+            translate_kv_loc_dense=lambda x, **kw: x,
+            kernel_page_multiplier=1,
+        )
+        self.assertTrue(hooks.enabled, "single-layer unified pool read as static")
+        self.assertIs(hooks.v2p_page_table, table)
+        self.assertIsNotNone(hooks.translate_kv_loc_dense)
+        # Multiplier stays 1: dense id == physical id, so the v2p gather alone is
+        # the whole translation and PAGE_MULT must not scale it.
+        self.assertEqual(hooks.kernel_page_multiplier, 1)
+
+
+@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
+class TestInPlaceKvIndicesTranslate(unittest.TestCase):
+    """The flashinfer decode updater must translate kv_indices IN PLACE.
+
+    Under cuda-graph replay the `kv_indices` it is handed IS the capture-stable
+    buffer the captured wrapper reads (`fast_decode_kwargs["kv_indices"]`), and
+    `fast_mla_decode_plan` ignores its `kv_indices` argument -- so rebinding the
+    local name to a fresh translated tensor leaves the graph reading VIRTUAL ids.
+    These pin the write-back contract that fix relies on.
+    """
+
+    def _allocator(self, page_size=1, n_full_tokens=4096):
+        from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
+        from sglang.srt.mem_cache.unified_memory_pool import (
+            MambaSubPoolSpec,
+            MLASubPoolSpec,
+            UnifiedKVPool,
+        )
+
+        full = MLASubPoolSpec(
+            name="full",
+            layer_num=_LAYERS,
+            kv_lora_rank=512,
+            qk_rope_head_dim=64,
+            store_dtype=torch.bfloat16,
+            grow_direction="down",
+        )
+        mamba = MambaSubPoolSpec(
+            name="mamba",
+            layer_num=2,
+            conv_state_shapes=((8, 16),),
+            conv_dtype=torch.bfloat16,
+            temporal_state_shape=(4, 8, 8),
+            temporal_dtype=torch.float32,
+            grow_direction="up",
+        )
+        pool = UnifiedKVPool(
+            total_bytes=full.entry_bytes() * n_full_tokens + mamba.entry_bytes() * 16,
+            sub_pool_specs=[full, mamba],
+            device=_DEV,
+            enable_memory_saver=False,
+            page_size=page_size,
+            view_tail_pad_bytes=page_size * full.entry_bytes(),
+        )
+
+        class _Stub:
+            def move_kv_cache(self, dst, src):
+                pass
+
+        full_alloc = MultiEndedAllocator(
+            kvcache=_Stub(),
+            unified_buffer=pool,
+            sub_pool_name="full",
+            device=_DEV,
+            is_id_owner=True,
+            page_size=page_size,
+            kernel_page_multiplier=_LAYERS,
+        )
+        mamba_alloc = MultiEndedAllocator(
+            kvcache=_Stub(),
+            unified_buffer=pool,
+            sub_pool_name="mamba",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        full_alloc.bind_peer(mamba_alloc)
+        mamba_alloc.bind_peer(full_alloc)
+        return full_alloc
+
+    def test_int32_buffer_prefix_translated_tail_untouched(self):
+        """Mirrors the updater: an int32 capture-stable buffer holding VIRTUAL
+        ids in [:n] gets the dense ids written back in place, narrowed to int32,
+        with the stale tail left alone (it must never index the v2p table)."""
+        alloc = self._allocator()
+        virt = alloc.alloc(64)
+        self.assertIsNotNone(virt)
+        n = virt.numel()
+
+        # Capture-stable int32 buffer: [:n] freshly filled with virtual ids by
+        # create_flashinfer_kv_indices_triton, tail = stale junk from a bigger replay.
+        buf = torch.full((n * 3,), 2**30, dtype=torch.int32, device=_DEV)
+        buf[:n] = virt.to(torch.int32)
+        tail_before = buf[n:].clone()
+
+        valid = buf[:n]
+        valid.copy_(alloc.translate_kv_loc_dense(valid))
+
+        expected = alloc.translate_kv_loc_dense(virt)
+        self.assertEqual(buf.dtype, torch.int32)
+        self.assertTrue(
+            torch.equal(buf[:n].long(), expected),
+            "in-place translate did not land dense ids in the stable buffer",
+        )
+        self.assertTrue(
+            torch.equal(buf[n:], tail_before),
+            "stale tail was modified -- it can hold ids outside the v2p table",
+        )
+
+    def test_dense_ids_differ_from_virtual(self):
+        """Guard the guard: if dense == virtual the in-place test proves nothing."""
+        alloc = self._allocator()
+        virt = alloc.alloc(64)
+        self.assertIsNotNone(virt)
+        self.assertFalse(
+            torch.equal(alloc.translate_kv_loc_dense(virt), virt),
+            "dense ids coincide with virtual ids; pick a different allocation",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/server_args/test_page_major_backend_allowlist.py

@@ -0,0 +1,111 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""`--enable-page-major-kv-layout` full-attention backend allowlist.
+
+The page-major envelope K/V views are strided, which only the Triton attention
+kernels read. The one exception is the unified-memory MLA pool: it exposes each
+layer as a DENSE contiguous view (`build_dense_mla_views`), so the paged MLA
+backends (`trtllm_mla` and its `cutedsl_mla` / `tokenspeed_mla` subclasses, plus
+`flashinfer`'s MLA backend) can read it directly once their kv_indices / block
+tables are remapped to dense ids.
+
+Pinned here so the exception cannot silently widen to a backend that has no
+dense-id remapping (`fa3`, `flashmla`, ...) or leak into the MHA path.
+
+    python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
+"""
+
+import unittest
+
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
+    """Run just `_handle_page_major_kv_layout` against a minimal stand-in.
+
+    ServerArgs' real constructor pulls in a model config; this exercises the
+    single handler under test with the fields it reads.
+    """
+    sa = ServerArgs.__new__(ServerArgs)
+    for name, value in {
+        "enable_unified_memory": unified,
+        # The unified pool sets this itself; without it the flag must be explicit
+        # or the handler returns before reaching the allowlist.
+        "enable_page_major_kv_layout": not unified,
+        "attention_backend": backend,
+        "prefill_attention_backend": None,
+        "decode_attention_backend": None,
+        "linear_attn_backend": "triton",
+        "linear_attn_decode_backend": None,
+        "linear_attn_prefill_backend": None,
+        "mamba_backend": "triton",
+    }.items():
+        object.__setattr__(sa, name, value)
+    sa.use_mla_backend = lambda: use_mla
+    sa._resolved_attention_backends = lambda: [backend]
+    try:
+        ServerArgs._handle_page_major_kv_layout(sa)
+        return True
+    except AssertionError:
+        return False
+
+
+class TestPageMajorBackendAllowlist(unittest.TestCase):
+    # Wired for the dense per-layer MLA views (see the module docstring).
+    DENSE_MLA_BACKENDS = ("trtllm_mla", "flashinfer", "cutedsl_mla", "tokenspeed_mla")
+    # No dense-id remapping: must stay rejected until they get one.
+    UNWIRED_BACKENDS = ("fa3", "flashmla", "cutlass_mla", "trtllm_mha", "aiter")
+
+    def test_triton_always_allowed(self):
+        for use_mla in (True, False):
+            self.assertTrue(_accepts("triton", use_mla=use_mla))
+
+    def test_dense_mla_backends_allowed_under_unified_mla(self):
+        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertTrue(
+                _accepts(backend, use_mla=True),
+                f"{backend} should be allowed with the unified-memory MLA pool",
+            )
+
+    def test_dense_mla_backends_rejected_for_mha(self):
+        """The dense-view exception is MLA-only -- MHA sub-pools stay strided."""
+        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertFalse(
+                _accepts(backend, use_mla=False),
+                f"{backend} must stay rejected for a non-MLA model",
+            )
+
+    def test_dense_mla_backends_rejected_without_unified_memory(self):
+        """Plain --enable-page-major-kv-layout (no unified pool) keeps the
+        strided views, so only Triton can read them."""
+        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertFalse(
+                _accepts(backend, use_mla=True, unified=False),
+                f"{backend} must stay rejected without --enable-unified-memory",
+            )
+
+    def test_unwired_backends_always_rejected(self):
+        for backend in self.UNWIRED_BACKENDS:
+            for use_mla in (True, False):
+                self.assertFalse(
+                    _accepts(backend, use_mla=use_mla),
+                    f"{backend} has no dense-id remapping and must be rejected",
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()