feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla (#34613)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

python/sglang/kernels/ops/attention/metadata.py

@@ -193,13 +193,9 @@ def _fused_metadata_kernel_general(
     use_swa: tl.constexpr,
     SHIFT: tl.constexpr,
     BLOCK_COLS: tl.constexpr,
-    # Unified-memory per-layer-view path (page-major envelope shared with the mamba
+    # 1: the two table pointers carry PAGE-granular, already kernel-facing
-    # sub-pool). Both default to the identity for the statically-partitioned
+    # read tables; emit verbatim -- no >>SHIFT, no v2p, no mapping gather.
-    # pool, where req_to_token already holds physical ids.
+    SRC_IS_KERNEL_PAGE_TABLE: tl.constexpr = 0,
-    v2p_ptr=None,
-    PAGE_MULT: tl.constexpr = 1,
-    # Unified SWA puts its independent v2p table in the legacy mapping slot.
-    SWA_MAPPING_IS_V2P: tl.constexpr = False,
 ):
     pid_b = tl.program_id(0)  # batch index
     pid_c = tl.program_id(1)  # column chunk index
@@ -240,8 +236,11 @@ def _fused_metadata_kernel_general(
     col_offsets = col_start + tl.arange(0, BLOCK_COLS)
     mask = col_offsets < num_live_pages
 
-    # Compute column indices in the source tensor (token offset)
+    # Compute column indices in the source tensor (token offset; page offset
-    if page_size == 1:
+    # when the source is already the page-granular canonical)
+    if SRC_IS_KERNEL_PAGE_TABLE:
+        col_idx = col_offsets
+    elif page_size == 1:
         col_idx = col_offsets
     else:
         col_idx = col_offsets << SHIFT  # faster than multiplication for power-of-two
@@ -253,40 +252,38 @@ def _fused_metadata_kernel_general(
     )
 
     # Compute page_table
-    if page_size == 1:
+    if SRC_IS_KERNEL_PAGE_TABLE:
+        page_table_val = page_index  # read-table entries are the page ids
+    elif page_size == 1:
         page_table_val = page_index
     else:
         page_table_val = page_index >> SHIFT
 
-    # Unified memory: virtual page -> physical page -> that layer's dense block.
-    # Derived from page_table_val, NOT page_index, which the SWA branch below
-    # still needs in virtual space. Masked so padded lanes never index the table.
-    if v2p_ptr is not None:
-        page_table_val = tl.load(v2p_ptr + page_table_val, mask=mask, other=0)
-    page_table_val = page_table_val * PAGE_MULT
-
     # Store to page_table
     pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
     tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg")
 
     if use_swa:
-        if SWA_MAPPING_IS_V2P:
+        if SRC_IS_KERNEL_PAGE_TABLE:
+            # The swa canonical shares the full canonical's shape and strides,
+            # so the SAME rt_offsets address the matching swa entry.
+            swa_val = tl.load(
+                full_to_swa_mapping + rt_offsets,
+                mask=mask,
+                other=0,
+                cache_modifier=".cg",
+            )
+        else:
+            swa_slot = tl.load(
+                full_to_swa_mapping + page_index * full_to_swa_mapping_stride_0,
+                mask=mask,
+                other=0,
+                cache_modifier=".cg",
+            )
             if page_size == 1:
-                swa_mapping_index = page_index
+                swa_val = swa_slot
             else:
-                swa_mapping_index = page_index >> SHIFT
+                swa_val = swa_slot >> SHIFT
-        else:
-            swa_mapping_index = page_index * full_to_swa_mapping_stride_0
-        swa_slot = tl.load(
-            full_to_swa_mapping + swa_mapping_index,
-            mask=mask,
-            other=0,
-            cache_modifier=".cg",
-        )
-        if page_size == 1 or SWA_MAPPING_IS_V2P:
-            swa_val = swa_slot
-        else:
-            swa_val = swa_slot >> SHIFT
         swa_offsets = (
             i * swa_page_table_stride_0 + col_offsets * swa_page_table_stride_1
         )
@@ -316,9 +313,6 @@ def _fused_metadata_kernel_ps1_no_swa(
     max_seq_pages,
     seq_len_delta: tl.constexpr,
     BLOCK_COLS: tl.constexpr,
-    # Unified-memory per-layer-view path; identity defaults for the static pool.
-    v2p_ptr=None,
-    PAGE_MULT: tl.constexpr = 1,
 ):
     pid_b = tl.program_id(0)  # batch index
     pid_c = tl.program_id(1)  # column chunk index
@@ -361,11 +355,6 @@ def _fused_metadata_kernel_ps1_no_swa(
         req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
     )
 
-    # page_table = page_index // 1 = page_index
-    # Unified memory: at page_size 1 the virtual token id IS the virtual page id.
-    if v2p_ptr is not None:
-        page_index = tl.load(v2p_ptr + page_index, mask=mask, other=0)
-    page_index = page_index * PAGE_MULT
     pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
     tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg")
 
@@ -586,15 +575,14 @@ def normal_decode_set_metadata(
     page_table: torch.Tensor,
     req_to_token: torch.Tensor,
     req_pool_indices: torch.Tensor,
-    strided_indices: torch.Tensor,
     max_seq_pages: torch.Tensor,
     seq_lens: torch.Tensor,
     seq_len_delta: int,
     page_size: int,
     swa_page_table: Optional[torch.Tensor] = None,
     token_to_kv_pool: Optional["SWAKVPool"] = None,
-    v2p_page_table: Optional[torch.Tensor] = None,
+    src_is_read_table: bool = False,
-    kernel_page_multiplier: int = 1,
+    swa_src_table: Optional[torch.Tensor] = None,
 ):
     """
     Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels:
@@ -602,14 +590,15 @@ def normal_decode_set_metadata(
       2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum)
       3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather)
       4. page_table = page_indices // page_size (floor-divide)
-      4b. (unified memory) page_table = v2p_page_table[page] * kernel_page_multiplier
+      5. (optional) swa_page_table for sliding window attention
-      5. (optional) swa_page_table via the legacy full->SWA map or the unified
-         SWA pool's independent page map
 
-    Step 4b is folded in rather than applied afterwards so the capture-stable
+    Unified pool (``src_is_read_table=True``): ``req_to_token`` /
-    page_table is written already translated: no separate pass a caller could
+    ``req_pool_indices`` carry the translator's PAGE-granular read table and its
-    forget, and no temporary to keep pointer-stable across cuda-graph replays.
+    row indices instead (entries already kernel-facing; ``swa_src_table`` is
-    Identity (None / 1) for the statically-partitioned pool.
+    the swa canonical, same shape and strides); steps 3-5 become verbatim
+    copies of the read table's rows' live prefixes, folded into the same launch
+    so the capture-stable page_table is written translated with no separate
+    pass a caller could forget.
 
     Achieves ~5.2x speedup on H200 hardware for typical decode workloads.
 
@@ -639,7 +628,9 @@ def normal_decode_set_metadata(
     page_table_stride_0 = page_table.stride(0)
     page_table_stride_1 = page_table.stride(1)
 
-    use_swa = swa_page_table is not None and token_to_kv_pool is not None
+    use_swa = swa_page_table is not None and (
+        token_to_kv_pool is not None or swa_src_table is not None
+    )
 
     # Unified SWA uses an independent SWA v2p table.
     swa_v2p_page_table = None
@@ -681,15 +672,26 @@ def normal_decode_set_metadata(
             max_seq_pages,
             seq_len_delta,
             BLOCK_COLS=BLOCK_COLS,
-            v2p_ptr=v2p_page_table,
-            PAGE_MULT=kernel_page_multiplier,
             num_warps=8,
             num_stages=3,
         )
     else:
         # General kernel for page_size > 1 or SWA cases
         # SWA parameters
-        if use_swa:
+        if use_swa and src_is_read_table:
+            # Unified pool: the swa canonical rides in the mapping slot; the
+            # kernel addresses it with the SAME row/col offsets as the full
+            # canonical, so their layouts must match exactly.
+            assert swa_src_table is not None
+            assert (
+                swa_src_table.stride() == req_to_token.stride()
+            ), "swa canonical must share the full canonical's strides"
+            swa_page_table = swa_page_table.contiguous()
+            swa_page_table_stride_0 = swa_page_table.stride(0)
+            swa_page_table_stride_1 = swa_page_table.stride(1)
+            full_to_swa_mapping = swa_src_table
+            full_to_swa_mapping_stride_0 = 0  # unused under the canonical source
+        elif use_swa:
             from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 
             swa_page_table = swa_page_table.contiguous()
@@ -749,9 +751,7 @@ def normal_decode_set_metadata(
             use_swa,
             shift,
             BLOCK_COLS=BLOCK_COLS,
-            v2p_ptr=v2p_page_table,
+            SRC_IS_KERNEL_PAGE_TABLE=1 if src_is_read_table else 0,
-            PAGE_MULT=kernel_page_multiplier,
-            SWA_MAPPING_IS_V2P=swa_uses_v2p,
             num_warps=4,
             num_stages=3,
         )

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

@@ -129,16 +129,9 @@ 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 per-layer-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 kernel-facing 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 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,
 ):
+    # Static-pool builder only: token ids here are physical, entry = token//ps.
+    # The unified pool's block table is filled by KVIndexTranslator instead.
     NUM_PAGE_PER_BLOCK: tl.constexpr = (
         FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE
     )
@@ -178,13 +171,8 @@ 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,
-            page * PAGE_MULT,
+            data // PAGED_SIZE,
             mask=mask_out,
         )

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

@@ -62,38 +62,44 @@ def update_trtllm_mha_graph_metadata_kernel(
     Q_MODE: tl.constexpr,
     PAGE_BLOCK: tl.constexpr,
     BS_BLOCK: tl.constexpr,
+    # 1: the page tables are refreshed out-of-graph, so rebuild only the
+    # seqlen metadata. 0 = static pool, full rebuild.
+    SKIP_PAGE_TABLE: tl.constexpr = 0,
 ):
     pid = tl.program_id(axis=0)
 
     if pid < bs:
         # One program per batch row: cache_seqlens + page table row(s).
-        req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
         seqlen = (tl.load(seq_lens_ptr + pid) + seqlen_offset).to(tl.int32)
         tl.store(cache_seqlens_ptr + pid, seqlen)
 
-        row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
+        if not SKIP_PAGE_TABLE:
-        row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
+            req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
-        if HAS_SWA:
+            row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
-            swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
+            row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
-        # Self-guard on the device-side seqlen: pages past cdiv(cache_seqlen,
-        # PAGE_SIZE) keep stale values the attention kernels never read.
-        num_live_pages = tl.minimum(tl.cdiv(seqlen, PAGE_SIZE), max_seq_pages)
-        for i in range(tl.cdiv(num_live_pages, PAGE_BLOCK)):
-            page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
-            mask = page_idx < num_live_pages
-            token = tl.load(
-                row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0
-            )
-            tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask)
             if HAS_SWA:
-                token64 = token.to(tl.int64)
+                swa_row_out = (
-                # Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror
+                    swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
-                # the swa_out_cache_loc -1 sentinel (uniform handling, no wrap).
-                swa_token = tl.load(
-                    swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1
                 )
-                swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE)
+            # Self-guard on the device-side seqlen: pages past cdiv(cache_seqlen,
-                tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask)
+            # PAGE_SIZE) keep stale values the attention kernels never read.
+            num_live_pages = tl.minimum(tl.cdiv(seqlen, PAGE_SIZE), max_seq_pages)
+            for i in range(tl.cdiv(num_live_pages, PAGE_BLOCK)):
+                page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
+                mask = page_idx < num_live_pages
+                token = tl.load(
+                    row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0
+                )
+                tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask)
+                if HAS_SWA:
+                    token64 = token.to(tl.int64)
+                    # Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror
+                    # the swa_out_cache_loc -1 sentinel (uniform handling, no wrap).
+                    swa_token = tl.load(
+                        swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1
+                    )
+                    swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE)
+                    tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask)
     elif pid == bs:
         # Single program: cu_seqlens_k (+ optional cu_seqlens_q) cumsum.
         offs = tl.arange(0, BS_BLOCK)
@@ -145,12 +151,18 @@ def update_trtllm_mha_graph_metadata(
     qlens=None,
     q_stride: int = 0,
     q_mode: int = Q_MODE_NONE,
+    skip_page_table: bool = False,
 ):
     """Launch the fused metadata update (one kernel for the whole replay init).
 
     Contract: only the live prefix (cdiv(cache_seqlens, page_size) pages) of each
     page_table / swa_page_table row is (re)written; the tail keeps stale values
     across replays, so consumers must bound reads by cache_seqlens.
+
+    ``skip_page_table=True`` (the unified memory pool): the seqlen metadata is
+    still rebuilt in one launch, but the page-table writes are compiled out --
+    the bound tables are capture-stable read tables the translator
+    refreshes out-of-graph. ``page_table`` may then be None.
     """
     if bs == 0:
         return
@@ -160,6 +172,10 @@ def update_trtllm_mha_graph_metadata(
     # set small enough to stay off the register-pressure / occupancy cliff while
     # being wide enough to cover the static page-table width in few iterations.
     PAGE_BLOCK = 512
+    if skip_page_table:
+        # Dead pointers under SKIP_PAGE_TABLE=1; pass a valid dummy for codegen.
+        page_table = cache_seqlens
+        swa_page_table = None
     has_swa = swa_page_table is not None
     has_swa_out = swa_out_cache_loc is not None
 
@@ -203,4 +219,5 @@ def update_trtllm_mha_graph_metadata(
         Q_MODE=q_mode,
         PAGE_BLOCK=PAGE_BLOCK,
         BS_BLOCK=triton.next_power_of_2(bs),
+        SKIP_PAGE_TABLE=1 if skip_page_table else 0,
     )

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -318,18 +318,14 @@ def handle_page_major_kv_layout(server_args: Any):
         "pool's per-layer views require a uniform row width; run "
         "this model without --enable-unified-memory."
     )
-    # Only the Triton attention kernels read the strided 4-D envelope K/V
+    # Allow-list. Every backend below reads through the translator, so what
-    # views; FA3 / FlashInfer do not. EXCEPTION: the unified-memory MLA pool
+    # gates one is only whether its kernels can address the per-layer views:
-    # exposes each layer as a contiguous per-layer view
+    #   * MLA models: the full paged MLA family, incl. flashmla (ps=64
-    # (build_mla_views), which the paged MLA kernels consume directly,
+    #     snap). cutlass_mla stays rejected (never exercised).
-    # with their kv_indices / block tables remapped to kernel-facing ids. Names below
+    #   * MHA/SWA models: fa3 / fa4 / flashinfer / trtllm_mha alongside
-    # are the RESOLVED ids from attention_backends_of: "flashinfer" is
+    #     Triton. fa4 is the fa3 class.
-    # FlashInferMLAAttnBackend for an MLA model, "trtllm_mla" the trtllm
+    #   * Without the unified pool, plain page-major stays Triton-only.
-    # decode kernel; "cutedsl_mla" and "tokenspeed_mla" subclass
+    # Names are the RESOLVED ids from attention_backends_of.
-    # TRTLLMMLABackend and inherit its read/write path; "fa3" remaps its
-    # page_table (in-kernel for captured decode, one funnel for eager).
-    # flashmla / cutlass_mla share the create_flashmla block-table path and
-    # can be added the same way once exercised.
     if cfg.enable_unified_memory and use_mla_backend(server_args):
         allowed_full = {
             "triton",
@@ -338,16 +334,27 @@ def handle_page_major_kv_layout(server_args: Any):
             "flashinfer",
             "cutedsl_mla",
             "tokenspeed_mla",
+            "flashmla",
+        }
+    elif cfg.enable_unified_memory:
+        allowed_full = {
+            "triton",
+            "fa3",
+            "fa4",
+            "flashinfer",
+            "trtllm_mha",
         }
     else:
         allowed_full = {"triton"}
     backends = set(attention_backends_of(resolved_view(server_args)))
     backends.discard(None)
     assert backends <= allowed_full, (
-        "--enable-page-major-kv-layout requires the Triton attention backend "
+        "--enable-page-major-kv-layout: the resolved attention backends "
-        "for the full-attention layers (unified-memory MLA also allows the "
+        f"{sorted(backends)} are not in the allowed set "
-        f"paged MLA backends); got {sorted(backends)}, allowed "
+        f"{sorted(allowed_full)} for this configuration (unified memory "
-        f"{sorted(allowed_full)}. Pass a compatible --attention-backend."
+        "allows the per-layer-view families; plain page-major keeps the "
+        "envelope-strided views only Triton reads). Pass a compatible "
+        "--attention-backend."
     )
     # The Mamba/KDA state is stored in envelope-strided views; only
     # stride-audited kernels may read it (Stage 4 audit, per slot):

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

@@ -79,6 +79,12 @@ class AttentionBackend(ABC):
     # (metadata glue graph) can read it off any backend without hasattr.
     forward_metadata: Optional[object] = None
 
+    # The runner's KVIndexTranslator; backends that read through it set the
+    # instance attribute in __init__. None means "no translate" -- a backend
+    # that never set it cannot serve the unified pool, which the server-args
+    # allow-list enforces.
+    kv_index_translator = None
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Eager entry point. Default = ``_out_graph(fb) + _in_graph(fb)``.
 

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

@@ -364,7 +364,7 @@ class CuteDslMLABackend(TRTLLMMLABackend):
                 if (
                     save_kv_cache
                     and self._fused_set_kv_concat_q_fp8
-                    and not self._unified_mla
+                    and not self.kv_index_translator.is_translating
                 ):
                     # Static pool: out_cache_loc is already the physical loc.
                     # Fused: bf16->fp8 quantize + KV scatter + q concat in one

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

@@ -18,7 +18,6 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
 )
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks
 from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_v2_active
@@ -192,11 +191,8 @@ class FlashAttentionBackend(AttentionBackend):
         # seq_lens_cpu / seq_lens_sum D2H sync is ever needed.
         self.needs_cpu_seq_lens = False
         self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
-        # Unified pool: req_to_token holds VIRTUAL ids but the MLA per-layer views
+        self.kv_index_translator = model_runner.kv_index_translator
-        # are kernel-facing, so every page_table needs remapping. MLA-only -- the MHA/SWA
+        self.kv_read_tables = None
-        # sub-pools keep the strided envelope layout FA3 cannot read at all.
-        self._unified_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
-        self._unified_dense = self._unified_hooks.enabled and self.use_mla
         self.skip_prefill = skip_prefill
         self.attn_cp_size = model_runner.ps.attn_cp_size
         self._verify_mask = None
@@ -229,6 +225,16 @@ class FlashAttentionBackend(AttentionBackend):
 
         # Local attention settings
         self.has_local_attention = model_runner.model_config.is_local_attention_model
+        # Local (chunked) attention derives its page table by re-translating
+        # metadata.page_table through the static full->swa map -- meaningless
+        # on the unified pool's kernel-facing tables, and no unified-eligible
+        # model uses it. Fail loud rather than silently double-translate.
+        assert not (
+            self.kv_index_translator.is_translating and self.has_local_attention
+        ), (
+            "--enable-unified-memory does not support local-attention models "
+            "on the fa3/fa4 backend."
+        )
         if self.has_local_attention:
             assert (
                 model_runner.attention_chunk_size is not None
@@ -248,6 +254,12 @@ class FlashAttentionBackend(AttentionBackend):
                 "Prefill-aware SWA requires page_size=1, "
                 f"got page_size={self.page_size}"
             )
+            # Its page-table builder indexes prefill_lens by POOL SLOT,
+            # incompatible with the batch-row canonical source.
+            assert not self.kv_index_translator.is_translating, (
+                "--enable-unified-memory does not support the prefill-aware "
+                "SWA decode mode; disable it for this model."
+            )
             # Indexed by raw req_pool_idx values (see the write below and
             # _build_pa_page_table), which range over [0, size] (row 0 is the
             # reserved padding slot) -- so this needs size+1, not size.
@@ -527,6 +539,7 @@ class FlashAttentionBackend(AttentionBackend):
                 spec_info=spec_info,
                 seq_lens_cpu=seq_lens_cpu,
                 out_cache_loc=out_cache_loc,
+                in_capture=True,
             )
 
             if forward_mode.is_decode_or_idle() and spec_info is None:
@@ -1081,7 +1094,25 @@ class FlashAttentionBackend(AttentionBackend):
                 text_row, text_col
             ]
 
-        if self.use_sliding_window_kv_pool:
+        # Safe to rebind: every eager branch above produced a fresh tensor.
+        _unified_read = (
+            self.kv_index_translator.is_translating and metadata.page_table is not None
+        )
+        if _unified_read:
+            kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
+            metadata.page_table = kv_view.ids
+            if self.use_sliding_window_kv_pool:
+                metadata.swa_page_table = kv_view.sliding_window_ids
+                if forward_batch.out_cache_loc is not None:
+                    # The swa write loc was computed from the still-VIRTUAL
+                    # loc at ForwardBatch construction; re-running the
+                    # full->swa map on the kernel-facing loc would be garbage.
+                    metadata.swa_out_cache_loc = (
+                        self.kv_index_translator.sliding_window_write_loc_for(
+                            forward_batch.out_cache_loc
+                        )
+                    )
+        elif self.use_sliding_window_kv_pool:
             # FA3 requires an int32 page_table.
             metadata.swa_page_table = (
                 self.token_to_kv_pool.translate_loc_from_full_to_swa(
@@ -1095,28 +1126,8 @@ class FlashAttentionBackend(AttentionBackend):
                     )
                 )
 
-        # Unified pool: one remap for every eager branch above, which all filled
-        # page_table with VIRTUAL token ids. Rebinding is safe here because those
-        # branches each produced a fresh tensor; the captured path instead folds
-        # the remap into normal_decode_set_metadata, which must write in place.
-        #
-        # Placed BEFORE the `// page_size` reduction, in token space: since
-        # kernel_id(t) = phys_page * (ps * L) + t % ps, dense(page_start) // ps is
-        # phys_page * L, the dense page id the kernel wants. One site then serves
-        # both page sizes, and it inherits translate_kv_loc_for_kernel's tombstone
-        # clamp so an unwritten req_to_token slot lands in the page-0 sink.
-        if self._unified_dense and metadata.page_table is not None:
-            # Flattened: the page_size == 1 translate path uses index_select,
-            # which rejects a 2-D index.
-            pt = metadata.page_table
-            metadata.page_table = (
-                self._unified_hooks.translate_kv_loc_for_kernel(pt.reshape(-1))
-                .to(torch.int32)
-                .view(pt.shape)
-            )
-
         # Convert the page table to a strided format which is needed by FA3 API
-        if self.page_size > 1:
+        if self.page_size > 1 and not _unified_read:
             self.strided_indices = torch.arange(
                 0, metadata.page_table.shape[1], self.page_size, device=self.device
             )
@@ -2157,6 +2168,12 @@ class FlashAttentionBackend(AttentionBackend):
         """
         max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size
 
+        if self.kv_index_translator.is_translating:
+            # Zero-filled: slot 0 is the reserved sink in every id space.
+            self.kv_read_tables = self.kv_index_translator.make_capture_tables(
+                max_bs=max_bs, max_context_len=self.max_context_len
+            )
+
         # This is being used by normal decode and draft decode when topk == 1
         self.decode_cuda_graph_metadata = {
             "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
@@ -2709,6 +2726,7 @@ class FlashAttentionBackend(AttentionBackend):
         spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
         out_cache_loc: Optional[torch.Tensor] = None,
+        in_capture: bool = False,
     ):
         """Shared capture+replay body for the cuda-graph init path.
 
@@ -2731,9 +2749,14 @@ class FlashAttentionBackend(AttentionBackend):
         if self.use_sliding_window_kv_pool and out_cache_loc is not None:
             n = out_cache_loc.shape[0]
             self.cuda_graph_swa_out_cache_loc[n:].zero_()
-            self.cuda_graph_swa_out_cache_loc[:n].copy_(
+            if in_capture and self.kv_index_translator.is_translating:
-                self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc)
+                # A capture batch never went through `init_new`, so there is no
-            )
+                # rebound write loc; zeros are the page-0 sink.
+                self.cuda_graph_swa_out_cache_loc[:n].zero_()
+            else:
+                self.cuda_graph_swa_out_cache_loc[:n].copy_(
+                    self.kv_index_translator.sliding_window_write_loc_for(out_cache_loc)
+                )
 
         if forward_mode.is_decode_or_idle():
             if spec_info is not None:
@@ -2744,13 +2767,19 @@ class FlashAttentionBackend(AttentionBackend):
                     # Page table built on-device (self-guards on cache_seqlens);
                     # max_seq_len_k left unset -- unread here (scheduler_metadata
                     # is normal-decode-only).
+                    # Spec is asserted off under the unified pool, so this
+                    # captured view is always the passthrough (req_to_token).
+                    kv_view = self.kv_index_translator.build_index_table(
+                        req_pool_indices=req_pool_indices,
+                        seq_lens=seq_lens,
+                        into=self.kv_read_tables,
+                    )
                     normal_decode_set_metadata(
                         metadata.cache_seqlens_int32,
                         metadata.cu_seqlens_k,
                         metadata.page_table,
-                        self.req_to_token,
+                        kv_view.ids,
-                        req_pool_indices,
+                        kv_view.row_ids,
-                        self.decode_cuda_graph_metadata["strided_indices"],
                         self.max_num_pages,
                         seq_lens,
                         self.speculative_step_id + 1,
@@ -2761,12 +2790,8 @@ class FlashAttentionBackend(AttentionBackend):
                             if self.use_sliding_window_kv_pool
                             else None
                         ),
-                        v2p_page_table=(
+                        src_is_read_table=kv_view.is_translated,
-                            self._unified_hooks.v2p_page_table
+                        swa_src_table=kv_view.sliding_window_ids,
-                            if self._unified_dense
-                            else None
-                        ),
-                        kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier,
                     )
 
                 else:
@@ -2866,13 +2891,17 @@ class FlashAttentionBackend(AttentionBackend):
                         if seq_lens_cpu is not None
                         else self.max_context_len
                     )
+                    kv_view = self.kv_index_translator.build_index_table(
+                        req_pool_indices=req_pool_indices,
+                        seq_lens=seq_lens,
+                        into=self.kv_read_tables,
+                    )
                     normal_decode_set_metadata(
                         metadata.cache_seqlens_int32,
                         metadata.cu_seqlens_k,
                         metadata.page_table,
-                        self.req_to_token,
+                        kv_view.ids,
-                        req_pool_indices,
+                        kv_view.row_ids,
-                        self.decode_cuda_graph_metadata["strided_indices"],
                         self.max_num_pages,
                         seq_lens,
                         0,
@@ -2883,12 +2912,8 @@ class FlashAttentionBackend(AttentionBackend):
                             if self.use_sliding_window_kv_pool
                             else None
                         ),
-                        v2p_page_table=(
+                        src_is_read_table=kv_view.is_translated,
-                            self._unified_hooks.v2p_page_table
+                        swa_src_table=kv_view.sliding_window_ids,
-                            if self._unified_dense
-                            else None
-                        ),
-                        kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier,
                     )
 
                 self._maybe_update_local_attn_metadata_for_replay(

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

@@ -36,6 +36,7 @@ from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
 )
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTable
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
@@ -310,6 +311,8 @@ class FlashInferAttnBackend(AttentionBackend):
 
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.token_to_kv_pool = model_runner.token_to_kv_pool
+        self.kv_index_translator = model_runner.kv_index_translator
+        self.kv_read_tables = None
         self._swa_kv_pool: Optional[BaseSWAKVPool] = self._resolve_swa_kv_pool(
             model_runner
         )
@@ -721,9 +724,16 @@ class FlashInferAttnBackend(AttentionBackend):
             num_tokens = forward_batch.positions.numel()
             self._prepare_cuda_graph_metadata(bs, num_tokens, forward_mode, spec_info)
 
+        # All flashinfer gathers run OUT-of-graph (plan time), so the
+        # capture-stable read table is buffer reuse, not pointer stability.
+        kv_view = self.kv_index_translator.build_index_table(
+            req_pool_indices=req_pool_indices[:bs],
+            seq_lens=seq_lens[:bs],
+            into=self.kv_read_tables,
+        )
+
         if forward_mode.is_decode_or_idle():
             self.indices_updater_decode.update(
-                req_pool_indices[:bs],
                 seq_lens[:bs],
                 seq_lens_cpu[:bs] if seq_lens_cpu is not None else None,
                 seq_lens_sum,
@@ -732,6 +742,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 spec_info=spec_info,
                 fixed_split_size=None,
                 disable_split_kv=self.disable_cuda_graph_kv_split,
+                kv_view=kv_view,
             )
         elif forward_mode.is_target_verify():
             self.indices_updater_prefill.update(
@@ -744,6 +755,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 use_ragged=False,
                 encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
                 spec_info=spec_info,
+                kv_view=kv_view,
             )
         elif forward_mode.is_dllm_extend():
             self.indices_updater_prefill.update(
@@ -756,6 +768,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 use_ragged=not self.use_paged,
                 encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
                 spec_info=None,
+                kv_view=kv_view,
             )
         elif forward_mode.is_draft_extend_v2():
             self.indices_updater_prefill.update(
@@ -768,6 +781,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 use_ragged=False,
                 encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
                 spec_info=spec_info,
+                kv_view=kv_view,
             )
         elif forward_mode.is_extend():
             # Plain EXTEND under full prefill CUDA graph. plan() runs
@@ -786,6 +800,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 use_ragged=False,
                 encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
                 spec_info=None,
+                kv_view=kv_view,
             )
         else:
             raise ValueError("Invalid forward mode")
@@ -840,14 +855,19 @@ class FlashInferAttnBackend(AttentionBackend):
         # Refill the SWA write-target buffer from the live out_cache_loc before
         # replay (bound onto the metadata at capture below).
         if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
-            assert self._swa_kv_pool is not None
             n = forward_batch.out_cache_loc.shape[0]
             self.cuda_graph_swa_out_cache_loc[n:].zero_()
-            self.cuda_graph_swa_out_cache_loc[:n].copy_(
+            if in_capture and self.kv_index_translator.is_translating:
-                self._swa_kv_pool.translate_loc_from_full_to_swa(
+                # A runner-built capture batch never went through `init_new`,
-                    forward_batch.out_cache_loc
+                # so there is no prepared write loc to resolve -- and zeros are the
+                # page-0 sink in every id space. Replay refills below.
+                self.cuda_graph_swa_out_cache_loc[:n].zero_()
+            else:
+                self.cuda_graph_swa_out_cache_loc[:n].copy_(
+                    self.kv_index_translator.sliding_window_write_loc_for(
+                        forward_batch.out_cache_loc
+                    )
                 )
-            )
             if in_capture:
                 self.forward_metadata.swa_out_cache_loc = (
                     self.cuda_graph_swa_out_cache_loc[:n]
@@ -934,16 +954,15 @@ class FlashInferAttnBackend(AttentionBackend):
         return layer.k_scale, layer.v_scale
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
+        kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
         swa_out_cache_loc = None
         if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
-            assert self._swa_kv_pool is not None
+            swa_out_cache_loc = self.kv_index_translator.sliding_window_write_loc_for(
-            swa_out_cache_loc = self._swa_kv_pool.translate_loc_from_full_to_swa(
                 forward_batch.out_cache_loc
             )
 
         if forward_batch.forward_mode.is_decode_or_idle():
             self.indices_updater_decode.update(
-                forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
                 forward_batch.seq_lens_cpu,
                 forward_batch.seq_lens_sum,
@@ -952,6 +971,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 spec_info=forward_batch.spec_info,
                 fixed_split_size=self.decode_split_tile_size,
                 disable_split_kv=False,
+                kv_view=kv_view,
             )
             self.forward_metadata = DecodeMetadata(
                 self.decode_wrappers, swa_out_cache_loc=swa_out_cache_loc
@@ -967,6 +987,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 use_ragged=False,
                 encoder_lens=forward_batch.encoder_lens,
                 spec_info=forward_batch.spec_info,
+                kv_view=kv_view,
             )
             self.forward_metadata = PrefillMetadata(
                 self.prefill_wrappers_verify,
@@ -1020,6 +1041,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 cross_attention_custom_mask=forward_batch.cross_attention_custom_mask,
                 extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
                 custom_kv_indices=self.dq_page_table,
+                kv_view=kv_view,
             )
             self.forward_metadata = PrefillMetadata(
                 self.prefill_wrappers_paged,
@@ -1035,6 +1057,9 @@ class FlashInferAttnBackend(AttentionBackend):
         max_num_tokens: int,
         kv_indices_buf: Optional[torch.Tensor] = None,
     ):
+        self.kv_read_tables = self.kv_index_translator.make_capture_tables(
+            max_bs=max_bs, max_context_len=self.max_context_len
+        )
         if kv_indices_buf is None:
             cuda_graph_kv_indices = torch.zeros(
                 (max_num_tokens * self.max_context_len,),
@@ -1539,7 +1564,6 @@ class FlashInferIndicesUpdaterDecode:
         # Buffers and wrappers
         self.kv_indptr = attn_backend.kv_indptr
         self.kv_last_page_len = attn_backend.kv_last_page_len
-        self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self._swa_kv_pool = attn_backend._swa_kv_pool
 
         # Dispatch the update function
@@ -1553,7 +1577,6 @@ class FlashInferIndicesUpdaterDecode:
 
     def update(
         self,
-        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
@@ -1562,13 +1585,14 @@ class FlashInferIndicesUpdaterDecode:
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
         disable_split_kv: Optional[bool] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
 
     def update_single_wrapper(
         self,
-        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
@@ -1577,11 +1601,12 @@ class FlashInferIndicesUpdaterDecode:
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
         disable_split_kv: Optional[bool] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         decode_wrappers = decode_wrappers or self.decode_wrappers
         self.call_begin_forward(
             decode_wrappers[0],
-            req_pool_indices,
             seq_lens,
             seq_lens_sum,
             self.kv_indptr[0],
@@ -1590,11 +1615,11 @@ class FlashInferIndicesUpdaterDecode:
             seq_lens_cpu,
             fixed_split_size=fixed_split_size,
             disable_split_kv=disable_split_kv,
+            kv_view=kv_view,
         )
 
     def update_sliding_window(
         self,
-        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
@@ -1603,6 +1628,8 @@ class FlashInferIndicesUpdaterDecode:
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
         disable_split_kv: Optional[bool] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         assert self.sliding_window_size is not None
         for wrapper_id in range(2):
@@ -1632,7 +1659,6 @@ class FlashInferIndicesUpdaterDecode:
 
             self.call_begin_forward(
                 decode_wrappers[wrapper_id],
-                req_pool_indices,
                 paged_kernel_lens_tmp,
                 paged_kernel_lens_sum_tmp,
                 self.kv_indptr[wrapper_id],
@@ -1642,11 +1668,11 @@ class FlashInferIndicesUpdaterDecode:
                 use_sliding_window_kv_pool=use_sliding_window_kv_pool,
                 fixed_split_size=fixed_split_size,
                 disable_split_kv=disable_split_kv,
+                kv_view=kv_view,
             )
 
     def update_cross_attention(
         self,
-        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_cpu: Optional[torch.Tensor],
         seq_lens_sum: int,
@@ -1655,6 +1681,8 @@ class FlashInferIndicesUpdaterDecode:
         spec_info: Optional[SpecInput],
         fixed_split_size: Optional[int] = None,
         disable_split_kv: Optional[bool] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         # Cache encoder_lens on CPU to avoid GPU→CPU transfer per call
         encoder_lens_cpu = encoder_lens.cpu() if encoder_lens is not None else None
@@ -1672,7 +1700,6 @@ class FlashInferIndicesUpdaterDecode:
 
             self.call_begin_forward(
                 decode_wrappers[wrapper_id],
-                req_pool_indices,
                 paged_kernel_lens,
                 seq_lens_sum,
                 self.kv_indptr[wrapper_id],
@@ -1681,12 +1708,12 @@ class FlashInferIndicesUpdaterDecode:
                 seq_lens_cpu=kv_lens_cpu,
                 fixed_split_size=fixed_split_size,
                 disable_split_kv=disable_split_kv,
+                kv_view=kv_view,
             )
 
     def call_begin_forward(
         self,
         wrapper: BatchDecodeWithPagedKVCacheWrapper,
-        req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
         kv_indptr: torch.Tensor,
@@ -1696,9 +1723,15 @@ class FlashInferIndicesUpdaterDecode:
         use_sliding_window_kv_pool: bool = False,
         fixed_split_size: Optional[int] = None,
         disable_split_kv: Optional[bool] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
+        # Unified SWA wrapper-0: gather from the swa canonical directly -- its
+        # entries are already swa-side kernel-facing ids, so the in-place
+        # full->swa translate below must not run on top of them.
+        use_swa_source = use_sliding_window_kv_pool and kv_view.is_translated
         if spec_info is None or getattr(spec_info, "kv_indptr", None) is None:
-            bs = len(req_pool_indices)
+            bs = len(paged_kernel_lens)
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
 
@@ -1710,20 +1743,26 @@ class FlashInferIndicesUpdaterDecode:
                     paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
                 )
 
+            if use_swa_source:
+                assert kv_view.sliding_window_ids is not None
+                src_table = kv_view.sliding_window_ids
+            else:
+                src_table = kv_view.ids
             create_flashinfer_kv_indices_triton[(bs,)](
-                self.req_to_token,
+                src_table,
-                req_pool_indices,
+                kv_view.row_ids,
                 paged_kernel_lens,
                 kv_indptr,
                 kv_start_idx,
                 kv_indices,
-                self.req_to_token.shape[1],
+                kv_view.row_stride,
+                ENTRY_PAGE_SIZE=kv_view.entry_page_size,
             )
         else:
             kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
             bs = kv_indptr.shape[0] - 1
 
-        if use_sliding_window_kv_pool:
+        if use_sliding_window_kv_pool and not use_swa_source:
             assert self._swa_kv_pool is not None
             kv_last_index = kv_indptr[-1]
             kv_indices[:kv_last_index] = (
@@ -1811,6 +1850,9 @@ class FlashInferIndicesUpdaterPrefill:
         self.kv_indptr = attn_backend.kv_indptr
         self.kv_last_page_len = attn_backend.kv_last_page_len
         self.qo_indptr = attn_backend.qo_indptr
+        # Kept ONLY for the spec-info branches (generate_attn_arg_prefill),
+        # which are static-pool-only: unified memory asserts spec off. The
+        # normal builders source from the per-batch KVIndexTable.
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self._swa_kv_pool = attn_backend._swa_kv_pool
         self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
@@ -1840,6 +1882,8 @@ class FlashInferIndicesUpdaterPrefill:
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
@@ -1860,6 +1904,8 @@ class FlashInferIndicesUpdaterPrefill:
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         if use_ragged:
             assert prefix_lens is not None
@@ -1890,6 +1936,7 @@ class FlashInferIndicesUpdaterPrefill:
             multi_item_params=multi_item_params,
             seq_lens_cpu=seq_lens_cpu,
             custom_kv_indices=custom_kv_indices,
+            kv_view=kv_view,
         )
 
     def update_sliding_window(
@@ -1908,6 +1955,8 @@ class FlashInferIndicesUpdaterPrefill:
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         if custom_kv_indices is not None:
             raise RuntimeError(
@@ -1983,6 +2032,7 @@ class FlashInferIndicesUpdaterPrefill:
                     if (wrapper_id == 0 and not use_ragged and spec_info is None)
                     else -1
                 ),
+                kv_view=kv_view,
             )
 
     def _build_swa_prefix_custom_mask(
@@ -2042,6 +2092,8 @@ class FlashInferIndicesUpdaterPrefill:
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
+        *,
+        kv_view: KVIndexTable,
     ):
         if custom_kv_indices is not None:
             raise RuntimeError(
@@ -2077,6 +2129,7 @@ class FlashInferIndicesUpdaterPrefill:
                 cross_attention_custom_mask=(
                     cross_attention_custom_mask if wrapper_id == 1 else None
                 ),
+                kv_view=kv_view,
             )
 
     def call_begin_forward(
@@ -2100,8 +2153,14 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens_cpu: Optional[torch.Tensor] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
         window_left: int = -1,
+        *,
+        kv_view: KVIndexTable,
     ):
         bs = len(seq_lens)
+        # Unified SWA wrapper-0: gather from the swa canonical directly -- its
+        # entries are already swa-side kernel-facing ids, so the in-place
+        # full->swa translate below must not run on top of them.
+        use_swa_source = use_sliding_window_kv_pool and kv_view.is_translated
         if spec_info is None:
             assert prefix_lens is not None
             assert len(seq_lens) == len(req_pool_indices)
@@ -2127,14 +2186,20 @@ class FlashInferIndicesUpdaterPrefill:
                     dtype=torch.int32,
                     device=req_pool_indices.device,
                 )
+                if use_swa_source:
+                    assert kv_view.sliding_window_ids is not None
+                    src_table = kv_view.sliding_window_ids
+                else:
+                    src_table = kv_view.ids
                 create_flashinfer_kv_indices_triton[(bs,)](
-                    self.req_to_token,
+                    src_table,
-                    req_pool_indices,
+                    kv_view.row_ids,
                     paged_kernel_lens,
                     kv_indptr,
                     kv_start_idx,
                     kv_indices,
-                    self.req_to_token.shape[1],
+                    kv_view.row_stride,
+                    ENTRY_PAGE_SIZE=kv_view.entry_page_size,
                 )
             qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
             qo_indptr = qo_indptr[: bs + 1]
@@ -2173,7 +2238,7 @@ class FlashInferIndicesUpdaterPrefill:
                 q_data_type=self.q_data_type,
             )
 
-        if use_sliding_window_kv_pool:
+        if use_sliding_window_kv_pool and not use_swa_source:
             assert self._swa_kv_pool is not None
             kv_last_index = kv_indptr[-1]
             kv_indices[:kv_last_index] = (

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

@@ -30,12 +30,12 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.flashinfer_backend import (
     create_flashinfer_kv_indices_triton,
 )
-from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks
 from sglang.srt.layers.dcp import (
     DecodeContextParallelMetadata,
     update_local_kv_lens_for_dcp,
 )
 from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTable
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
     is_in_breakable_cuda_graph,
@@ -238,6 +238,8 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         # corresponding ForwardBatch fields.
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.token_to_kv_pool = model_runner.token_to_kv_pool
+        self.kv_index_translator = model_runner.kv_index_translator
+        self.kv_read_tables = None
         self.enable_chunk_kv = (
             not skip_prefill
             and get_disagg().disaggregation_mode != "decode"
@@ -342,6 +344,14 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         forward_mode = forward_batch.forward_mode
         spec_info = forward_batch.spec_info
 
+        # All flashinfer gathers run OUT-of-graph (plan time), so the
+        # capture-stable table is buffer reuse, not pointer stability.
+        kv_view = self.kv_index_translator.build_index_table(
+            req_pool_indices=req_pool_indices[:bs],
+            seq_lens=seq_lens[:bs],
+            into=self.kv_read_tables,
+        )
+
         if in_capture:
             num_tokens = forward_batch.positions.numel()
             seq_lens_sum = seq_lens.sum().item()
@@ -358,12 +368,12 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                     backend="auto",
                 )
                 self.indices_updater_decode.update(
-                    req_pool_indices,
                     seq_lens,
                     seq_lens_sum,
                     decode_wrapper=decode_wrapper,
                     init_metadata_replay=False,
                     spec_info=spec_info,
+                    kv_view=kv_view,
                 )
                 self.decode_cuda_graph_metadata[bs] = decode_wrapper
                 self.forward_metadata = DecodeMetadata(decode_wrapper)
@@ -396,6 +406,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 spec_info=spec_info,
                 seq_lens_cpu=seq_lens_cpu,
                 in_capture=True,
+                kv_view=kv_view,
             )
             if forward_mode.is_target_verify() and (
                 spec_info is None
@@ -412,16 +423,18 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 forward_mode=forward_mode,
                 spec_info=spec_info,
                 seq_lens_cpu=forward_batch.seq_lens_cpu,
+                kv_view=kv_view,
             )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
+        kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
         if forward_batch.forward_mode.is_decode_or_idle():
             self.indices_updater_decode.update(
-                forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
                 forward_batch.seq_lens_sum,
                 decode_wrapper=self.decode_wrapper,
                 init_metadata_replay=False,
+                kv_view=kv_view,
             )
             self.forward_metadata = DecodeMetadata(self.decode_wrapper)
         elif forward_batch.forward_mode.is_target_verify():
@@ -433,6 +446,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 prefill_wrapper_paged=self.prefill_wrapper_verify,
                 use_ragged=False,
                 spec_info=forward_batch.spec_info,
+                kv_view=kv_view,
             )
             self.forward_metadata = PrefillMetadata(self.prefill_wrapper_verify, False)
         else:
@@ -481,6 +495,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 qo_indptr_cpu=qo_indptr_cpu,
                 kv_indptr_cpu=kv_indptr_cpu,
                 kv_len_arr_cpu=kv_len_arr_cpu,
+                kv_view=kv_view,
             )
             self.forward_metadata = PrefillMetadata(
                 self.prefill_wrapper_paged, use_ragged
@@ -492,6 +507,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         max_num_tokens: int,
         kv_indices_buf: Optional[torch.Tensor] = None,
     ):
+        self.kv_read_tables = self.kv_index_translator.make_capture_tables(
+            max_bs=max_bs, max_context_len=self.max_context_len
+        )
         if kv_indices_buf is None:
             cuda_graph_kv_indices = torch.zeros(
                 (max_bs * self.max_context_len,),
@@ -535,6 +553,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         forward_mode: ForwardMode,
         spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
+        kv_view: KVIndexTable,
         in_capture: bool = False,
     ):
         """Shared capture+replay body for the cuda-graph init path.
@@ -557,12 +576,12 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             )
 
             self.indices_updater_decode.update(
-                req_pool_indices[:bs],
                 seq_lens[:bs],
                 seq_lens_sum,
                 decode_wrapper=self.decode_cuda_graph_metadata[bs],
                 init_metadata_replay=True,
                 spec_info=spec_info,
+                kv_view=kv_view,
                 **self.fast_decode_kwargs,
             )
         elif forward_mode.is_target_verify():
@@ -613,6 +632,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                     if use_generic_fast_plan
                     else None
                 ),
+                kv_view=kv_view,
             )
         else:
             raise ValueError(f"Invalid forward mode: {forward_mode=}")
@@ -845,84 +865,70 @@ class FlashInferMLAIndicesUpdaterDecode:
 
         # Buffers and wrappers
         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_for_kernel
 
     def update(
         self,
-        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         decode_wrapper: BatchMLAPagedAttentionWrapper,
         init_metadata_replay: bool = False,
         spec_info: Optional[SpecInput] = None,
+        *,
+        kv_view: KVIndexTable,
         **fast_decode_kwargs,
     ):
         decode_wrapper = decode_wrapper or self.decode_wrapper
         self.call_begin_forward(
             decode_wrapper,
-            req_pool_indices,
             seq_lens,
             seq_lens_sum,
             self.q_indptr,
             self.kv_indptr,
             init_metadata_replay,
             spec_info,
+            kv_view=kv_view,
             **fast_decode_kwargs,
         )
 
     def call_begin_forward(
         self,
         wrapper: BatchMLAPagedAttentionWrapper,
-        req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
         q_indptr: torch.Tensor,
         kv_indptr: torch.Tensor,
         init_metadata_replay: bool = False,
         spec_info: Optional[SpecInput] = None,
+        *,
+        kv_view: KVIndexTable,
         **fast_decode_kwargs,
     ):
-        bs = len(req_pool_indices)
+        bs = len(paged_kernel_lens)
         q_indptr = q_indptr[: bs + 1]
         kv_lens = paged_kernel_lens.to(torch.int32)
         sm_scale = self.scaling
         if spec_info is None:
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
+            # On replay `kv_indices` IS the capture-stable buffer the captured
+            # wrapper reads -- never rebind it. The builder fills only the
+            # [:paged_kernel_lens_sum] prefix; the stale tail is unread.
             kv_indices = (
                 torch.empty(paged_kernel_lens_sum, dtype=torch.int32, device="cuda")
                 if not init_metadata_replay
                 else fast_decode_kwargs["kv_indices"]
             )
             create_flashinfer_kv_indices_triton[(bs,)](
-                self.req_to_token,
+                kv_view.ids,
-                req_pool_indices,
+                kv_view.row_ids,
                 paged_kernel_lens,
                 kv_indptr,
                 None,
                 kv_indices,
-                self.req_to_token.shape[1],
+                kv_view.row_stride,
+                ENTRY_PAGE_SIZE=kv_view.entry_page_size,
             )
-            # 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; kernel-facing 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(
@@ -986,14 +992,11 @@ class FlashInferMLAIndicesUpdaterPrefill:
         # Buffers and wrappers
         self.kv_indptr = attn_backend.kv_indptr
         self.qo_indptr = attn_backend.qo_indptr
+        # Kept ONLY for the spec-info branch (generate_attn_arg_prefill), which
+        # is static-pool-only: unified memory asserts spec off. The normal
+        # builder sources from the per-batch KVIndexTable.
         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 per-layer view, so remap them to kernel-facing
-        # 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_for_kernel
 
     def update(
         self,
@@ -1006,6 +1009,8 @@ class FlashInferMLAIndicesUpdaterPrefill:
         spec_info: Optional[SpecInput] = None,
         attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
         fast_verify_plan_kwargs: Optional[dict] = None,
+        *,
+        kv_view: KVIndexTable,
         qo_indptr_cpu: Optional[torch.Tensor] = None,
         kv_indptr_cpu: Optional[torch.Tensor] = None,
         kv_len_arr_cpu: Optional[torch.Tensor] = None,
@@ -1034,6 +1039,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
             qo_indptr_cpu=qo_indptr_cpu,
             kv_indptr_cpu=kv_indptr_cpu,
             kv_len_arr_cpu=kv_len_arr_cpu,
+            kv_view=kv_view,
         )
 
     def call_begin_forward(
@@ -1051,6 +1057,8 @@ class FlashInferMLAIndicesUpdaterPrefill:
         spec_info: Optional[SpecInput] = None,
         attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
         fast_verify_plan_kwargs: Optional[dict] = None,
+        *,
+        kv_view: KVIndexTable,
         qo_indptr_cpu: Optional[torch.Tensor] = None,
         kv_indptr_cpu: Optional[torch.Tensor] = None,
         kv_len_arr_cpu: Optional[torch.Tensor] = None,
@@ -1068,20 +1076,15 @@ class FlashInferMLAIndicesUpdaterPrefill:
                 device=req_pool_indices.device,
             )
             create_flashinfer_kv_indices_triton[(bs,)](
-                self.req_to_token,
+                kv_view.ids,
-                req_pool_indices,
+                kv_view.row_ids,
                 paged_kernel_lens,
                 kv_indptr,
                 None,
                 kv_indices,
-                self.req_to_token.shape[1],
+                kv_view.row_stride,
+                ENTRY_PAGE_SIZE=kv_view.entry_page_size,
             )
-            # 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/flashmla_backend.py

@@ -162,17 +162,25 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         if forward_batch.forward_mode.is_decode_or_idle():
             max_seqlen_pad = triton.cdiv(eager_max_k, PAGE_SIZE)
             block_kv_indices = self._eager_block_kv_indices(bs, max_seqlen_pad)
-            create_flashmla_kv_indices_triton[
+            if self.kv_index_translator.is_translating:
-                (bs, get_num_kv_index_blocks_flashmla(max_seqlen_pad, PAGE_SIZE))
+                assert self.page_size == PAGE_SIZE
-            ](
+                self.kv_index_translator.fill_read_table(
-                self.req_to_token,
+                    out=block_kv_indices,
-                forward_batch.req_pool_indices,
+                    req_pool_indices=forward_batch.req_pool_indices,
-                forward_batch.seq_lens,
+                    seq_lens=forward_batch.seq_lens,
-                None,
+                )
-                block_kv_indices,
+            else:
-                self.req_to_token.stride(0),
+                create_flashmla_kv_indices_triton[
-                block_kv_indices.stride(0),
+                    (bs, get_num_kv_index_blocks_flashmla(max_seqlen_pad, PAGE_SIZE))
-            )
+                ](
+                    self.req_to_token,
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    None,
+                    block_kv_indices,
+                    self.req_to_token.stride(0),
+                    block_kv_indices.stride(0),
+                )
             mla_metadata, num_splits = get_mla_metadata(
                 forward_batch.seq_lens.to(torch.int32),
                 self.num_q_heads,
@@ -328,22 +336,30 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
             else:
                 max_seqlen_pad = self.cuda_graph_kv_indices.shape[1]
 
-            create_flashmla_kv_indices_triton[
+            if self.kv_index_translator.is_translating:
-                (
+                assert self.page_size == PAGE_SIZE
-                    bs,
+                self.kv_index_translator.fill_read_table(
-                    get_num_kv_index_blocks_flashmla(
+                    out=self.cuda_graph_kv_indices,
-                        self.cuda_graph_kv_indices.stride(0), PAGE_SIZE
+                    req_pool_indices=req_pool_indices[:bs],
-                    ),
+                    seq_lens=seq_lens,
+                )
+            else:
+                create_flashmla_kv_indices_triton[
+                    (
+                        bs,
+                        get_num_kv_index_blocks_flashmla(
+                            self.cuda_graph_kv_indices.stride(0), PAGE_SIZE
+                        ),
+                    )
+                ](
+                    self.req_to_token,
+                    req_pool_indices[:bs],
+                    seq_lens,
+                    None,
+                    self.cuda_graph_kv_indices,
+                    self.req_to_token.stride(0),
+                    self.cuda_graph_kv_indices.stride(0),
                 )
-            ](
-                self.req_to_token,
-                req_pool_indices[:bs],
-                seq_lens,
-                None,
-                self.cuda_graph_kv_indices,
-                self.req_to_token.stride(0),
-                self.cuda_graph_kv_indices.stride(0),
-            )
 
             q_head_mult = (
                 self.num_draft_tokens

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

@@ -23,6 +23,7 @@ class TboAttnBackend(AttentionBackend):
         # reads through TboAttnBackend resolve to the underlying pool.
         self.token_to_kv_pool = primary.token_to_kv_pool
         self.req_to_token_pool = primary.req_to_token_pool
+        self.kv_index_translator = primary.kv_index_translator
         self.extend_dummy_seqs_capped_by_req_pool = getattr(
             primary, "extend_dummy_seqs_capped_by_req_pool", False
         )

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

@@ -198,8 +198,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         # separate index spaces; SWA layers need a translated page_table.
         self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner)
         # Raw full->swa index mapping tensor for the fused cuda-graph
-        # metadata kernel (gather + // page_size happen on device).
+        # metadata kernel (gather + // page_size happen on device). The unified
-        if self._swa_kv_pool is not None:
+        # pool has no token-level mapping, so this is a static-pool mechanism.
+        if (
+            self._swa_kv_pool is not None
+            and not self.kv_index_translator.is_translating
+        ):
             self._swa_full_to_swa_mapping = self._swa_kv_pool.full_to_swa_index_mapping
             assert self._swa_full_to_swa_mapping is not None, (
                 "SWA pool must register full_to_swa_index_mapping before "
@@ -465,6 +469,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         kv_indices_buf: Optional[torch.Tensor] = None,
     ):
         """Initialize CUDA graph state for TRTLLM MHA."""
+        self.kv_read_tables = self.kv_index_translator.make_capture_tables(
+            max_bs=max_bs, max_context_len=self.max_context_len
+        )
         max_num_pages = self.max_num_pages
         self.decode_cuda_graph_metadata = {
             "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
@@ -766,25 +773,27 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         # bounds real KV reads by cache_seqlens, so this is a fixed loop
         # bound only — never a host max / seq_lens_cpu D2H sync.
         max_seq_pages = self.max_num_pages
+        unified = self.kv_index_translator.is_translating
         update_trtllm_mha_graph_metadata(
             req_pool_indices=req_pool_indices,
             seq_lens=seq_lens,
             req_to_token=self.req_to_token,
             cache_seqlens=metadata.cache_seqlens_int32,
             cu_seqlens_k=metadata.cu_seqlens_k,
-            page_table=metadata.page_table,
+            page_table=None if unified else metadata.page_table,
             bs=bs,
             seqlen_offset=seqlen_offset,
             max_seq_pages=max_seq_pages,
             page_size=self.page_size,
             swa_mapping=self._swa_full_to_swa_mapping,
-            swa_page_table=metadata.swa_page_table,
+            swa_page_table=None if unified else metadata.swa_page_table,
             out_cache_loc=out_cache_loc,
-            swa_out_cache_loc=metadata.swa_out_cache_loc,
+            swa_out_cache_loc=None if unified else metadata.swa_out_cache_loc,
             cu_seqlens_q=cu_seqlens_q,
             qlens=qlens,
             q_stride=q_stride,
             q_mode=q_mode,
+            skip_page_table=unified,
         )
 
         if self._needs_encoder_only_expand(forward_mode, metadata):
@@ -888,6 +897,39 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 f"Invalid forward mode: {forward_mode=} for CUDA Graph replay."
             )
 
+        if self.kv_index_translator.is_translating:
+            # Unified pool: refresh the capture-stable read table (this runs
+            # out-of-graph on BOTH capture and every replay-prep; the recorded
+            # fused kernel skips its page-table writes so the graph reads the
+            # refreshed content through pointers baked at capture).
+            kv_view = self.kv_index_translator.build_index_table(
+                req_pool_indices=forward_batch.req_pool_indices[:bs],
+                seq_lens=forward_batch.seq_lens[:bs],
+                into=self.kv_read_tables,
+            )
+            metadata = self.forward_metadata
+            if in_capture:
+                # Bind ONCE: the attention kernels bake these pointers at capture.
+                metadata.page_table = kv_view.ids[:bs]
+                if kv_view.sliding_window_ids is not None:
+                    metadata.swa_page_table = kv_view.sliding_window_ids[:bs]
+            # A capture batch carries no prepared write loc; zeros are the
+            # page-0 sink.
+            if (
+                self.use_sliding_window_kv_pool
+                and forward_batch.out_cache_loc is not None
+            ):
+                n = forward_batch.out_cache_loc.shape[0]
+                self.cuda_graph_swa_out_cache_loc[n:].zero_()
+                if in_capture and self.kv_index_translator.is_translating:
+                    self.cuda_graph_swa_out_cache_loc[:n].zero_()
+                else:
+                    self.cuda_graph_swa_out_cache_loc[:n].copy_(
+                        self.kv_index_translator.sliding_window_write_loc_for(
+                            forward_batch.out_cache_loc
+                        )
+                    )
+
     def _assert_ragged_verify_supported(self) -> None:
         if self.is_xqa_impl:
             raise NotImplementedError(
@@ -1035,20 +1077,27 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             else:
                 metadata.cu_seqlens_q = metadata.cu_seqlens_k
 
-        has_swa = self._swa_kv_pool is not None
+        kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
-        metadata.page_table = torch.empty(
+        if kv_view.is_translated:
-            (batch_size, self.max_num_pages), dtype=torch.int32, device=device
+            # No fill kernel: the kernels take the tensor's own width/stride
-        )
+            # and bound their reads by cache_seqlens.
-        metadata.swa_page_table = (
+            metadata.page_table = kv_view.ids
-            torch.empty(
+            metadata.swa_page_table = kv_view.sliding_window_ids
+        else:
+            has_swa = self._swa_kv_pool is not None
+            metadata.page_table = torch.empty(
                 (batch_size, self.max_num_pages), dtype=torch.int32, device=device
             )
-            if has_swa
+            metadata.swa_page_table = (
-            else None
+                torch.empty(
-        )
+                    (batch_size, self.max_num_pages), dtype=torch.int32, device=device
-        self._fill_page_table_device(
+                )
-            metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32
+                if has_swa
-        )
+                else None
+            )
+            self._fill_page_table_device(
+                metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32
+            )
         self._maybe_build_cp_zigzag_page_tables(metadata, forward_batch)
 
         if self._needs_encoder_only_expand(forward_batch.forward_mode, metadata):
@@ -1063,7 +1112,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         # int64 scatter index (unlike the int32 read page table above).
         if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
             metadata.swa_out_cache_loc = (
-                self.token_to_kv_pool.translate_loc_from_full_to_swa(
+                self.kv_index_translator.sliding_window_write_loc_for(
                     forward_batch.out_cache_loc
                 )
             )

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

@@ -49,7 +49,6 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import (
     FlashInferMLAAttnBackend,
     FlashInferMLAMultiStepDraftBackend,
 )
-from sglang.srt.layers.attention.unified_mem_hooks import 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
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
@@ -285,17 +284,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         # Tree-mask scratch is fetched from the target backend only.
         self.is_draft_runner = model_runner.is_draft_worker
 
-        # Unified-memory per-layer-view hooks (None on the static pool). req_to_token
+        # [:n] view of a capture-stable buffer on the cuda-graph path; None on
-        # holds VIRTUAL token ids; the block table needs kernel-facing page ids, so the
+        # the eager path, which passes forward_batch.out_cache_loc through.
-        # kv-index kernels gather virtual->physical page through `_v2p_page_table`
-        # then scale by `_kernel_page_multiplier` (= num MLA layers). See
-        # build_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
-        # Per-forward kernel-facing write loc ([:n] view of a capture-stable buffer);
-        # None on the eager path, which passes out_cache_loc straight through.
         self._decode_kernel_loc: Optional[torch.Tensor] = None
         self.cuda_graph_out_cache_loc_kernel: Optional[torch.Tensor] = None
         # Fused KV-scatter + q-concat on the decode dense-loc path (one launch
@@ -368,23 +358,28 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             (batch_size, max_blocks), -1, dtype=torch.int32, device=device
         )
 
-        create_flashmla_kv_indices_triton[
+        if self.kv_index_translator.is_translating:
-            (
+            self.kv_index_translator.fill_read_table(
-                batch_size,
+                out=block_kv_indices,
-                get_num_kv_index_blocks_flashmla(max_blocks, self.page_size),
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens,
+            )
+        else:
+            create_flashmla_kv_indices_triton[
+                (
+                    batch_size,
+                    get_num_kv_index_blocks_flashmla(max_blocks, self.page_size),
+                )
+            ](
+                self.req_to_token,
+                req_pool_indices,
+                seq_lens,
+                None,
+                block_kv_indices,
+                self.req_to_token.stride(0),
+                max_blocks,
+                PAGED_SIZE=self.page_size,
             )
-        ](
-            self.req_to_token,
-            req_pool_indices,
-            seq_lens,
-            None,
-            block_kv_indices,
-            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
 
@@ -404,7 +399,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         # 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:
+        if self.kv_index_translator.is_translating:
             self.cuda_graph_out_cache_loc_kernel = torch.zeros(
                 max_num_tokens, dtype=torch.int64, device=self.device
             )
@@ -545,25 +540,30 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             metadata.seq_lens_k.copy_(seq_lens[:bs])
 
         # Update block indices for new sequences.
-        create_flashmla_kv_indices_triton[
+        if self.kv_index_translator.is_translating:
-            (
+            self.kv_index_translator.fill_read_table(
-                bs,
+                out=metadata.block_kv_indices,
-                get_num_kv_index_blocks_flashmla(
+                req_pool_indices=req_pool_indices[:bs],
-                    metadata.block_kv_indices.shape[1], self.page_size
+                seq_lens=seq_lens,
-                ),
+            )
+        else:
+            create_flashmla_kv_indices_triton[
+                (
+                    bs,
+                    get_num_kv_index_blocks_flashmla(
+                        metadata.block_kv_indices.shape[1], self.page_size
+                    ),
+                )
+            ](
+                self.req_to_token,
+                req_pool_indices[:bs],
+                seq_lens,
+                None,
+                metadata.block_kv_indices,
+                self.req_to_token.stride(0),
+                metadata.block_kv_indices.shape[1],
+                PAGED_SIZE=self.page_size,
             )
-        ](
-            self.req_to_token,
-            req_pool_indices[:bs],
-            seq_lens,
-            None,
-            metadata.block_kv_indices,
-            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:
         """Get the fill value for sequence lengths in CUDA graph."""
@@ -623,11 +623,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 forward_mode=forward_mode,
             )
 
-        # Unified pool: precompute the DENSE KV write loc into the capture-stable
+        # Out-of-graph on capture AND every replay-prep, so the in-graph
-        # buffer (both capture and each replay-prep run this out of the graph),
+        # set_mla_kv_buffer captures no translate.
-        # so the in-graph set_mla_kv_buffer writes a dense loc without capturing
+        if self.kv_index_translator.is_translating and (
-        # a translate.
-        if self._unified_mla and (
             forward_mode.is_decode_or_idle() or forward_mode.is_target_verify()
         ):
             out_cache_loc = forward_batch.out_cache_loc
@@ -646,6 +644,24 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         else:
             self._decode_kernel_loc = None
 
+    def _resolve_fused_write_loc(
+        self, forward_batch: ForwardBatch
+    ) -> Optional[torch.Tensor]:
+        """Write loc for the fused fp8 KV scatter, or None when this batch is
+        not covered by it.
+
+        Captured decode refills `_decode_kernel_loc` out of the graph, and the
+        captured kernel must read that buffer. Eager decode on a unified pool
+        has no such buffer, and the caller falls back to the unfused path.
+        """
+        if self._decode_kernel_loc is not None:
+            return self._decode_kernel_loc
+        return (
+            None
+            if self.kv_index_translator.is_translating
+            else forward_batch.out_cache_loc
+        )
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Initialize the metadata for a forward pass."""
         self._decode_kernel_loc = None
@@ -1050,13 +1066,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             assert q_rope is not None and k_rope is not None
             if cos_sin_cache is None:
                 if save_kv_cache and self._fused_set_kv_concat_q_fp8:
-                    loc = (
+                    loc = self._resolve_fused_write_loc(forward_batch)
-                        self._decode_kernel_loc
-                        if self._decode_kernel_loc is not None
-                        else (
-                            None if self._unified_mla else forward_batch.out_cache_loc
-                        )
-                    )
                     if loc is not None:
                         # Fused: bf16->fp8 quantize + KV scatter + q concat
                         # in one launch; None when not covered.
@@ -1113,7 +1123,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 if (
                     merge_query
                     and self._fused_set_kv_concat_q
-                    and not self._unified_mla
+                    and not self.kv_index_translator.is_translating
                 ):
                     # Static pool only, conservatively.
                     query = self._set_kv_and_concat_q_fused(

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

@@ -1,72 +0,0 @@
-# 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.
-# ==============================================================================
-"""Allocator hooks the paged MLA attention backends need under the unified
-memory pool.
-
-Lives in its own module because three unrelated backend families consume it
-(fa3, flashinfer_mla, and trtllm_mla with its cutedsl_mla / tokenspeed_mla
-subclasses) and none of them should have to import another's module to get it.
-"""
-
-from __future__ import annotations
-
-from typing import Callable, Optional
-
-import msgspec
-import torch
-
-
-class UnifiedMLAHooks(msgspec.Struct, frozen=True):
-    """Dense-view hooks for one KV allocator.
-
-    All-``None``/1/``False`` for the statically-partitioned pool, where
-    ``req_to_token`` already holds physical ids and no translation is needed.
-    """
-
-    # Page-level virtual->physical table, gathered through by block-table kernels.
-    v2p_page_table: Optional[torch.Tensor]
-    # Virtual token id -> DENSE kernel-facing id (tombstones clamped to the sink).
-    translate_kv_loc_for_kernel: Optional[Callable[..., torch.Tensor]]
-    # Dense page stride scale (= number of full-attention MLA layers).
-    kernel_page_multiplier: int
-    enabled: bool
-
-
-_STATIC_POOL = UnifiedMLAHooks(
-    v2p_page_table=None,
-    translate_kv_loc_for_kernel=None,
-    kernel_page_multiplier=1,
-    enabled=False,
-)
-
-
-def unified_mla_hooks(allocator) -> UnifiedMLAHooks:
-    """Probe ``allocator`` for the unified-pool per-layer-view hooks.
-
-    Detection keys on the v2p table, NOT on ``kernel_page_multiplier > 1``: a
-    rank owning exactly ONE full-attention layer has multiplier 1 while its
-    ``req_to_token`` is still virtual. There the kernel-facing id collapses onto the
-    physical id, so the v2p gather alone is the whole translation.
-    """
-    v2p = getattr(allocator, "full_v2p_page_table", None)
-    if v2p is None:
-        return _STATIC_POOL
-    return UnifiedMLAHooks(
-        v2p_page_table=v2p,
-        translate_kv_loc_for_kernel=getattr(
-            allocator, "translate_kv_loc_for_kernel", None
-        ),
-        kernel_page_multiplier=getattr(allocator, "kernel_page_multiplier", 1),
-        enabled=True,
-    )

python/sglang/srt/mem_cache/memory_pool.py

@@ -1602,13 +1602,15 @@ class KVWriteLoc:
       KERNEL-FACING on every pool: physical by allocation on non-unified
       pools, rebound at ForwardBatch construction (``rebind_write_loc``) on
       the unified pool.
-    - ``swa_loc``: the pre-resolved SWA-sub-pool location for hybrid SWA pools
+    - ``swa_loc``: the SWA-sub-pool location for hybrid SWA pools (``None``
-      (``None`` otherwise).
+      otherwise); under the unified pool the translator derives it from the
-    - ``full_loc``: the full-attention-sub-pool location for the unified
+      same rebound loc (``sliding_window_write_loc_for``).
-      memory pool (``None`` otherwise), carried in attention metadata
+    - ``full_loc``: OPTIONAL full-attention-sub-pool location. Since the
-      (``ForwardMetadata.out_cache_loc_full_physical``). Since the
+      construction-time rebind it is the SAME id space as ``loc``, so pools
-      construction-time rebind it is the SAME id space as ``loc``; the shared
+      fall back to ``loc`` when it is ``None`` -- only triton's captured path
-      full pool writes it directly and never translates.
+      still passes its capture-stable
+      ``ForwardMetadata.out_cache_loc_full_physical`` buffer here (a
+      same-space alias slated for collapse).
 
     ``swa_loc`` and ``full_loc`` are the parallel pair (each a pre-resolved
     loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``);
@@ -3665,11 +3667,6 @@ class HybridLinearKVPool(KVCache):
         # virtual->physical mamba-slot translate for the HiCache offload path;
         # identity for a static pool, the allocator's `translate` for the unified pool.
         self._mamba_translate = lambda ids: ids
-        # The MLA doors take DIFFERENT id spaces: `get_mla_kv_buffer` gets
-        # ForwardBatch-built read indices (prefix_chunk_kv_indices /
-        # fetch_mha_one_shot_kv_indices), still VIRTUAL, so it translates;
-        # `set_mla_kv_buffer` gets out_cache_loc, already kernel-facing.
-        self._full_translate = lambda ids: ids
         self.use_mla = use_mla
         if full_kv_pool is not None:
             # Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool
@@ -3967,7 +3964,10 @@ class HybridLinearKVPool(KVCache):
         dst_dtype: Optional[torch.dtype] = None,
     ):
         assert self.use_mla, "get_mla_kv_buffer called when use_mla is False"
-        loc = self._full_translate(loc)
+        # Read door -- same kernel-facing contract as the write door: `loc` is
+        # a read-index tensor already translated at its production site
+        # (fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
+        # pool never translates.
         with self._transfer_id_context(layer):
             return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype)
 

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -1244,8 +1244,9 @@ def init_unified_mamba_pools(
     # `_mamba_translate` feeds the HiCache offload path, GATED OFF here — wired but inert.
     req_to_token_pool.mamba_allocator = mamba_slot_allocator
     token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate
-    if use_mla_backend:
+    # No full-KV translate hook is wired: both MLA doors now receive
-        token_to_kv_pool._full_translate = allocator.translate_kv_loc_for_kernel
+    # KERNEL-FACING ids -- writes from the ForwardBatch rebind, reads
+    # translated at their production sites.
 
     logger.info(
         "[unified-memory-pool] ============================================================"
@@ -1466,7 +1467,7 @@ class UnifiedSWAKVPool(SWAKVPool):
         """Route to the right sub-pool. Both `swa_loc` and `full_loc` are PHYSICAL
         (pre-translated once per forward by the attention backend); never translates here.
         """
-        _, swa_loc, full_loc = unwrap_write_loc(loc_info)
+        loc, swa_loc, full_loc = unwrap_write_loc(loc_info)
         layer_id = layer.layer_id
         pool_layer_id, is_swa = self.layers_mapping[layer_id]
         if is_swa:
@@ -1486,12 +1487,11 @@ class UnifiedSWAKVPool(SWAKVPool):
                 layer_id_override=pool_layer_id,
             )
             return
-        # Full layer: full_loc is full-physical, always precomputed (eager + cuda-graph).
+        # Full layer: `loc` is already the full-side kernel-facing id, so an
-        assert full_loc is not None, (
+        # explicit full_loc is a same-space alias -- only triton's captured path
-            "UnifiedSWAKVPool.set_kv_buffer: full layer received no full_loc; "
+        # passes one (its capture-stable buffer).
-            "ForwardMetadata.out_cache_loc_full_physical must be precomputed for "
+        if full_loc is None:
-            "the unified memory pool."
+            full_loc = loc
-        )
         self.full_kv_pool.set_kv_buffer(
             None,
             full_loc,

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -12,6 +12,7 @@ from sglang.kernels.ops.kvcache.kv_indices import (
 from sglang.srt.environ import envs
 from sglang.srt.layers.dcp.layout import filter_dcp_local_chunk_kv_indices
 from sglang.srt.model_executor.forward_context import (
+    get_attn_backend,
     get_req_to_token_pool,
     get_token_to_kv_pool,
 )
@@ -90,6 +91,10 @@ class ForwardBatchDeepSeekMHAMixin:
                 self.prefix_chunk_starts_cpu[idx],
                 self.prefix_chunk_seq_lens_cpu[idx],
             )
+            # None on a backend that never bound a translator.
+            src = get_attn_backend().kv_index_translator
+            if src is not None:
+                chunk_kv_indices = src.translate_full_attn_ids(chunk_kv_indices)
             self.prefix_chunk_kv_indices.append(chunk_kv_indices)
 
     # Here we suppose the length of each chunk is equal
@@ -235,5 +240,9 @@ class ForwardBatchDeepSeekMHAMixin:
             kv_indices,
             req_to_token.shape[1],
         )
+        # None on a backend that never bound a translator.
+        src = get_attn_backend().kv_index_translator
+        if src is not None:
+            kv_indices = src.translate_full_attn_ids(kv_indices)
         self.mha_one_shot_kv_indices = kv_indices
         return kv_indices

python/sglang/srt/server_args.py

@@ -3762,7 +3762,7 @@ class ServerArgs:
 
     LANGUAGE_MODEL_ONLY_ARCHITECTURES = ("MuseGlimmerForConditionalGeneration",)
 
-    # The strided-layout Triton requirement is enforced via
+    # The attention-backend allow-list is enforced via
     # --enable-page-major-kv-layout (implied by the unified pool in
     # _handle_page_major_kv_layout); the model-family gate is enforced at pool
     # construction in model_runner_kv_cache_mixin._init_pools.

test/registered/attention/test_normal_decode_set_metadata.py

@@ -31,7 +31,6 @@ def reference_normal_decode_set_metadata(
     page_table: torch.Tensor,
     req_to_token: torch.Tensor,
     req_pool_indices: torch.Tensor,
-    strided_indices: torch.Tensor,
     max_seq_pages: int,
     seq_lens: torch.Tensor,
     seq_len_delta: int,
@@ -45,6 +44,11 @@ def reference_normal_decode_set_metadata(
     """
     cache_seqlens_int32.copy_(seq_lens + seq_len_delta)
     cu_seqlens_k[1:].copy_(torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32))
+    # Page-start columns, derived internally (the wrapper's dead
+    # strided_indices parameter was removed alongside its v2p args).
+    strided_indices = torch.arange(
+        0, req_to_token.shape[1], page_size, device=req_to_token.device
+    )
     page_indices = req_to_token[
         req_pool_indices[:, None],
         strided_indices[:max_seq_pages][None, :],
@@ -213,7 +217,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
             ref_data["page_table"],
             test_data["req_to_token"],
             test_data["req_pool_indices"],
-            test_data["strided_indices"],
             test_data["max_seq_pages"],
             test_data["seq_lens"],
             test_data["seq_len_delta"],
@@ -229,7 +232,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
             test_data["page_table"],
             test_data["req_to_token"],
             test_data["req_pool_indices"],
-            test_data["strided_indices"],
             test_data["max_seq_pages"],
             test_data["seq_lens"],
             test_data["seq_len_delta"],
@@ -351,7 +353,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
             test_data["page_table"],
             test_data["req_to_token"],
             test_data["req_pool_indices"],
-            test_data["strided_indices"],
             test_data["max_seq_pages"],
             test_data["seq_lens"],
             test_data["seq_len_delta"],
@@ -408,7 +409,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
             ref_data["page_table"],
             test_data["req_to_token"],
             test_data["req_pool_indices"],
-            test_data["strided_indices"],
             test_data["max_seq_pages"],
             test_data["seq_lens"],
             0,
@@ -423,7 +423,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
             test_data["page_table"],
             test_data["req_to_token"],
             test_data["req_pool_indices"],
-            test_data["strided_indices"],
             test_data["max_seq_pages"],
             test_data["seq_lens"],
             0,

test/registered/attention/test_trtllm_mha_graph_metadata.py

@@ -30,6 +30,8 @@ PAGE_SIZE = 128
 
 
 def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
+    from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
+
     backend = TRTLLMHAAttnBackend.__new__(TRTLLMHAAttnBackend)
     backend.device = torch.device("cpu")
     backend.max_context_len = 1024
@@ -45,6 +47,15 @@ def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
     backend.decode_cuda_graph_metadata = {}
     backend.target_verify_metadata = {}
     backend.draft_extend_metadata = {}
+    # Passthrough source (static pool): every unified-arm branch stays off,
+    # matching the real __init__'s parent binding.
+    backend.kv_index_translator = KVIndexTranslator(
+        req_to_token=backend.req_to_token,
+        token_to_kv_pool_allocator=SimpleNamespace(),
+        token_to_kv_pool=SimpleNamespace(),
+        page_size=PAGE_SIZE,
+        device="cpu",
+    )
     backend.init_cuda_graph_state(max_bs=4, max_num_tokens=16)
     return backend
 
@@ -543,6 +554,67 @@ def test_metadata_correctness(bs, seqlen_offset, q_mode, with_swa, static_width)
         torch.testing.assert_close(swa_out_cache_loc, out_ref, rtol=0, atol=0)
 
 
+@pytest.mark.parametrize("pass_tables", [False, True])
+def test_skip_page_table_updates_seqlens_only(pass_tables):
+    """The unified-memory arm: skip_page_table=True must still rebuild the
+    seqlen metadata in-graph but leave every page-table byte alone -- the bound
+    tables are capture-stable read tables the translator refreshes
+    out-of-graph, and an in-graph write would clobber them with virtual-derived
+    pages. Covers both call shapes: page_table=None (what the backend passes)
+    and a real sentinel-filled table (pins that the writes are compiled out,
+    not just unpassed)."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+
+    bs, seqlen_offset, seed = 5, 1, 4242
+    pool_size, max_num_pages = 64, 16
+    seq_max = (max_num_pages - 2) * PAGE_SIZE
+    (
+        req_to_token,
+        req_pool_indices,
+        seq_lens,
+        _stride,
+        _cap,
+    ) = _build_inputs(bs, pool_size, max_num_pages, None, seq_max, seed)
+
+    cache_seqlens = torch.zeros(bs, dtype=torch.int32, device=DEVICE)
+    cu_seqlens_k = torch.zeros(bs + 1, dtype=torch.int32, device=DEVICE)
+    sentinel_pt = None
+    sentinel_swa = None
+    if pass_tables:
+        sentinel_pt = torch.full(
+            (bs, max_num_pages), 777, dtype=torch.int32, device=DEVICE
+        )
+        sentinel_swa = torch.full(
+            (bs, max_num_pages), 888, dtype=torch.int32, device=DEVICE
+        )
+
+    update_trtllm_mha_graph_metadata(
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        req_to_token=req_to_token,
+        cache_seqlens=cache_seqlens,
+        cu_seqlens_k=cu_seqlens_k,
+        page_table=sentinel_pt,
+        bs=bs,
+        seqlen_offset=seqlen_offset,
+        max_seq_pages=max_num_pages,
+        page_size=PAGE_SIZE,
+        swa_page_table=sentinel_swa,
+        skip_page_table=True,
+    )
+    torch.cuda.synchronize()
+
+    cache_seqlens_ref = _ref_cache_seqlens(seq_lens, seqlen_offset)
+    torch.testing.assert_close(cache_seqlens, cache_seqlens_ref, rtol=0, atol=0)
+    cu_k_ref = torch.zeros(bs + 1, dtype=torch.int32, device=DEVICE)
+    cu_k_ref[1:] = torch.cumsum(cache_seqlens_ref, dim=0, dtype=torch.int32)
+    torch.testing.assert_close(cu_seqlens_k, cu_k_ref, rtol=0, atol=0)
+    if pass_tables:
+        assert bool((sentinel_pt == 777).all()), "page_table written despite skip"
+        assert bool((sentinel_swa == 888).all()), "swa_page_table written despite skip"
+
+
 def test_bs_zero_noop():
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")

test/registered/models_e2e/test_kimi_linear_unified_memory.py

@@ -32,7 +32,7 @@ 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=570, stage="nightly", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=1200, stage="nightly", runner_config="4-gpu-h100")
 
 KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
 
@@ -55,5 +55,19 @@ class TestKimiLinearUnifiedMemory(
     ]
 
 
+class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
+    """flashmla at its ps=64 snap: the canonical block-table route
+    (KVIndexTranslator.build_into into flashmla's padded tables) plus the ps=64
+    sub-pool sizing (64-token sink floor, dense-view tail pad) end to end.
+    Hopper-only, like the rest of this nightly suite."""
+
+    other_args = TestKimiLinearUnifiedMemory.other_args + [
+        "--attention-backend",
+        "flashmla",
+        "--page-size",
+        "64",
+    ]
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/page_major/test_page_major_gpt_oss.py

@@ -1,15 +1,12 @@
-"""
+"""Unified memory pool on a hybrid-SWA MoE model, across the backend matrix.
-End-to-end accuracy test for the unified memory pool on a hybrid-SWA MoE model.
 
-Launches gpt-oss-20b with ``--enable-unified-memory`` on the Triton attention
+gpt-oss-20b is uniform-row hybrid-SWA, so its MHA and SWA sub-pools are
-backend and checks that GSM8K accuracy holds. This exercises the SWA +
+per-layer views and the fa3 cell reads them through the translator's read
-full-attention KV sub-pools stored as per-layer views in the unified
+tables. The resolved-default cell pins the no-pin path, since a pinned
-page-major envelope.
+backend hides default-resolution breakage by construction. flashinfer is
+absent on purpose: gpt-oss uses attention sinks, which it does not support.
 
 Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
-
-Usage:
-    python3 -m unittest test_page_major_gpt_oss
 """
 
 import unittest
@@ -20,7 +17,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
 
-register_cuda_ci(est_time=420, stage="extra-a", runner_config="1-gpu-large")
+register_cuda_ci(est_time=1500, stage="extra-a", runner_config="1-gpu-large")
 
 _UNIFIED_COMMON_ARGS = [
     "--enable-unified-memory",
@@ -64,5 +61,19 @@ class TestUnifiedGptOssTriton(DefaultServerBase):
         self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
 
 
+class TestUnifiedGptOssFa3(TestUnifiedGptOssTriton):
+    """fa3 pinned: the per-layer views read through the translator's read
+    tables (eager direct-bind + captured fused copy)."""
+
+    other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
+
+
+class TestUnifiedGptOssResolvedDefault(TestUnifiedGptOssTriton):
+    """No backend pin: whatever the host resolves must be in the allow-list,
+    or the server fails to boot under its own defaults."""
+
+    other_args = _UNIFIED_COMMON_ARGS
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/page_major/test_page_major_qwen_hybrid.py

@@ -1,17 +1,14 @@
-"""
+"""Unified memory pool on a GDN-hybrid model, across the backend matrix.
-End-to-end accuracy test for the unified memory pool on a GDN-hybrid model.
 
-Launches Qwen3.5-4B (a gated-delta-net / linear-attention hybrid) with
+Qwen3.5-4B is a gated-delta-net / linear-attention hybrid, which exercises the
-``--enable-unified-memory`` on the Triton attention + linear-attn + Mamba
+path most prone to subtle bugs: the Mamba conv/SSM state stays a strided
-backends and checks that GSM8K accuracy holds. This exercises the unified
+envelope view (its kernels are stride-aware by design) while the
-envelope's most bug-prone path: the Mamba conv/SSM state stored as a strided
+full-attention KV is per-layer views, which the fa3 / flashinfer cells read
-envelope view, plus the full-attention KV stored as per-layer views,
+through the translator's read tables. The resolved-default cell pins the
-both read/written by the GDN prefill and decode kernels.
+no-pin path, since a pinned backend hides default-resolution breakage by
+construction.
 
 Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
-
-Usage:
-    python3 -m unittest test_page_major_qwen_hybrid
 """
 
 import unittest
@@ -22,7 +19,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
 
-register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-large")
+register_cuda_ci(est_time=1600, stage="extra-a", runner_config="1-gpu-large")
 
 _UNIFIED_COMMON_ARGS = [
     "--trust-remote-code",
@@ -74,5 +71,25 @@ class TestUnifiedQwenHybridTriton(DefaultServerBase):
         self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
 
 
+class TestUnifiedQwenHybridFa3(TestUnifiedQwenHybridTriton):
+    """fa3 pinned: read tables, eager direct-bind + captured fused copy."""
+
+    other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
+
+
+class TestUnifiedQwenHybridFlashinfer(TestUnifiedQwenHybridTriton):
+    """flashinfer pinned: token ids reconstructed from the read table by the
+    ENTRY_PAGE_SIZE CSR builder."""
+
+    other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "flashinfer"]
+
+
+class TestUnifiedQwenHybridResolvedDefault(TestUnifiedQwenHybridTriton):
+    """No backend pin: whatever the host resolves must be in the allow-list,
+    or the server fails to boot under its own defaults."""
+
+    other_args = _UNIFIED_COMMON_ARGS
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/layers/attention/test_flashattention_graph_metadata.py

@@ -4,6 +4,7 @@ from types import SimpleNamespace
 import torch
 
 from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -46,6 +47,15 @@ class TestFlashAttentionGraphMetadata(CustomTestCase):
         backend.req_to_token_pool = SimpleNamespace(
             req_to_token=torch.zeros((1, 16), dtype=torch.int32)
         )
+        # A real source over the stub pool: the probe disables it, giving the
+        # backend the strict passthrough view it now reads its tables from.
+        backend.kv_index_translator = KVIndexTranslator(
+            req_to_token=backend.req_to_token_pool.req_to_token,
+            token_to_kv_pool_allocator=SimpleNamespace(),
+            token_to_kv_pool=SimpleNamespace(),
+            page_size=1,
+            device="cpu",
+        )
         backend.is_prefill_aware_swa = False
         backend.has_swa = False
         backend.use_sliding_window_kv_pool = False

test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py

@@ -24,6 +24,7 @@ import torch
 
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
 from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -62,15 +63,24 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
         enable_prefill_cp=False,
         enable_dp_attention=False,
     )
+    token_to_kv_pool = object()
+    token_to_kv_pool_allocator = object()
     return SimpleNamespace(
         sliding_window_size=None,
         model_config=model_config,
         device=device,
         req_to_token_pool=req_to_token_pool,
-        token_to_kv_pool=object(),  # not a SWAKVPool instance -> use_sliding_window_kv_pool=False
+        token_to_kv_pool=token_to_kv_pool,  # not a SWAKVPool -> use_sliding_window_kv_pool=False
-        # getattr(..., "full_v2p_page_table", None) is None -> unified_mla_hooks
+        token_to_kv_pool_allocator=token_to_kv_pool_allocator,
-        # falls back to the static (disabled) hook set.
+        # A real KVIndexTranslator over this non-unified pool: the probe finds no
-        token_to_kv_pool_allocator=object(),
+        # unified composite, so it is the strict passthrough the backend reads.
+        kv_index_translator=KVIndexTranslator(
+            req_to_token=req_to_token_pool.req_to_token,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+            token_to_kv_pool=token_to_kv_pool,
+            page_size=1,
+            device=device,
+        ),
         kv_cache_dtype=torch.float16,
         kv_cache_dtype_str="auto",
         page_size=1,

test/registered/unit/layers/attention/test_kv_translate_ownership.py

@@ -0,0 +1,83 @@
+"""Nothing under layers/attention may translate KV ids for itself.
+
+Ownership is exactly two places: `KVIndexTranslator` for READS (indices are
+born kernel-facing, backends consume its tables) and the ForwardBatch rebind
+(`rebind_write_loc`) for WRITES. Virtual and physical ids share a value range,
+so a backend that forgets a translate -- or does one twice -- reads the wrong
+rows and nothing crashes. This scan makes both unrepresentable.
+
+Out of scope, deliberately: the allocator-internal implementations
+(`multi_ended_allocator` / `unified_memory_pool`), which ARE the mechanism the
+translator calls; the PD transfer plane's `translate_kv_indices_for_transfer`,
+which stages for RDMA outside the forward path; and the STATIC SWA pool's
+legacy full->swa slot map, a different mapping kind with no virtual/physical
+ambiguity -- its call sites are count-pinned below so new ones are added
+consciously.
+
+    python3 -m pytest test/registered/unit/layers/attention/test_kv_translate_ownership.py -v
+"""
+
+import os
+import re
+import unittest
+
+from sglang.srt.layers.attention import triton_backend as _anchor_module
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+# The attention package is a namespace package (no __init__), so anchor the
+# scan on a concrete module inside it.
+_ATTN_DIR = os.path.dirname(os.path.abspath(_anchor_module.__file__))
+
+
+def _iter_sources():
+    for root, _dirs, files in os.walk(_ATTN_DIR):
+        for name in sorted(files):
+            if not name.endswith(".py"):
+                continue
+            path = os.path.join(root, name)
+            with open(path, encoding="utf-8") as fh:
+                yield os.path.relpath(path, _ATTN_DIR), fh.read()
+
+
+class TestUnifiedTranslateBanned(CustomTestCase):
+    def test_no_unified_translate_calls(self):
+        """No backend calls the unified translate surfaces. A hit here means
+        a backend re-grew its own id-space transition -- the design whose two
+        failure modes (forgotten translate, duplicated translate) this scan
+        exists to prevent. Route reads through KVIndexTranslator views and
+        writes through the ForwardBatch rebind instead."""
+        banned = re.compile(r"\.translate_kv_loc(_kernel_id)?\(")
+        hits = [
+            f"{rel}: {m.group(0)}"
+            for rel, src in _iter_sources()
+            for m in banned.finditer(src)
+        ]
+        self.assertEqual(hits, [])
+
+    def test_no_translate_capability_probing(self):
+        """No backend probes an allocator for translate capability -- the
+        getattr-hook pattern is how per-backend translation grew the first
+        time."""
+        probing = re.compile(r"""getattr\([^)]*['"]translate_kv_loc""")
+        hits = [rel for rel, src in _iter_sources() if probing.search(src)]
+        self.assertEqual(hits, [])
+
+    def test_hooks_module_deleted_and_unimported(self):
+        """The per-backend hooks module (the previous owner of backend-side
+        v2p knowledge) stays deleted, and nothing imports it."""
+        self.assertFalse(
+            os.path.exists(os.path.join(_ATTN_DIR, "unified_mem_hooks.py"))
+        )
+        hits = [
+            rel
+            for rel, src in _iter_sources()
+            if "unified_mem_hooks" in src or "unified_mla_hooks" in src
+        ]
+        self.assertEqual(hits, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_full_loc_fast_path.py

@@ -59,9 +59,11 @@ class _RecordingPool:
 
 
 class TestUnifiedSWARouting(unittest.TestCase):
-    """`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write the full-physical
+    """`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write `full_loc`
-    `full_loc`; SWA layers write the swa-physical `swa_loc`. Both come from the
+    when present (triton's capture-stable buffer), else the rebound generic
-    write metadata; the pool never translates."""
+    `loc` -- the same id space once the loc is rebound; SWA layers write the swa-physical
+    `swa_loc`, which has no fallback (a different id space). The pool never
+    translates."""
 
     def _make_bare_pool(self):
         from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -96,21 +98,29 @@ class TestUnifiedSWARouting(unittest.TestCase):
         self.assertIsNot(forwarded, virtual_loc)
         self.assertNotIn("already_physical", kwargs)
 
-    def test_full_layer_requires_full_loc(self):
+    def test_full_layer_falls_back_to_generic_loc(self):
+        """Bug regression: fa3 x unified-SWA crashed at gpt-oss
+        cuda-graph capture because every backend except triton bundles the
+        2-arg KVWriteLoc(loc, swa) and the full-layer door demanded an explicit
+        full_loc. Once the loc is rebound the generic `loc` IS the full-side kernel-facing id
+        (rebind_write_loc runs at ForwardBatch construction),
+        so the door must fall back to it -- the pool still never translates."""
         pool = self._make_bare_pool()
-        virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
+        rebound_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
         swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
 
         layer = types.SimpleNamespace(layer_id=0)
-        # No full_loc precomputed -> fail loud (the unified memory pool must precompute
+        pool.set_kv_buffer(
-        # out_cache_loc_full_physical) rather than write a virtual loc as physical.
+            layer,
-        with self.assertRaises(AssertionError):
+            _loc_info(rebound_loc, swa_phys),
-            pool.set_kv_buffer(
+            torch.zeros(3, 4, 8),
-                layer,
+            torch.zeros(3, 4, 8),
-                _loc_info(virtual_loc, swa_phys),
+        )
-                torch.zeros(3, 4, 8),
+
-                torch.zeros(3, 4, 8),
+        self.assertEqual(len(pool.full_kv_pool.calls), 1)
-            )
+        forwarded, kwargs = pool.full_kv_pool.calls[0]
+        self.assertIs(forwarded, rebound_loc)
+        self.assertNotIn("already_physical", kwargs)
 
     def test_swa_layer_writes_swa_loc(self):
         pool = self._make_bare_pool()
@@ -264,19 +274,17 @@ class TestHybridLinearMLARouting(unittest.TestCase):
     - `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
       pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
       carries the DENSE loc), else the raw `loc` (static pool, already physical).
-    - `set_mla_kv_buffer` forwards `loc` untouched (kernel-facing since the
+    - `set_mla_kv_buffer` / `get_mla_kv_buffer` forward `loc` untouched:
-      ForwardBatch rebind); `get_mla_kv_buffer` applies `_full_translate`
+      writes are kernel-facing since the ForwardBatch rebind, and read
-      exactly once (its indices are req_to_token-produced, virtual under the
+      indices are translated at their production sites."""
-      unified pool)."""
 
-    def _make_bare_pool(self, translate=None):
+    def _make_bare_pool(self):
         from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
 
         pool = object.__new__(HybridLinearKVPool)
         pool.full_kv_pool = _RecordingMLAPool()
         pool.use_mla = True
         pool.full_attention_layer_id_mapping = {0: 0}
-        pool._full_translate = translate if translate is not None else (lambda ids: ids)
         return pool
 
     def test_mla_writes_full_loc_from_write_loc(self):
@@ -328,28 +336,20 @@ class TestHybridLinearMLARouting(unittest.TestCase):
         self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
         self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
 
-    def test_get_mla_kv_buffer_translates_exactly_once(self):
+    def test_get_mla_kv_buffer_door_never_translates(self):
-        """READ door: `loc` is produced from req_to_token (VIRTUAL under the
+        """Kernel-facing contract, read side: `loc` is a read-index tensor
-        unified pool), so the get side still translates here — exactly once.
+        already translated at its production site
-        The WRITE door (case above) never translates: the split is the write
+        (fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
-        flip's contract."""
+        door forwards it UNTOUCHED -- a re-added door translate would
-        calls = []
+        double-translate every unified MLA prefix read."""
-
+        pool = self._make_bare_pool()
-        def translate(ids):
+        loc = torch.tensor([104, 105], dtype=torch.int64)
-            calls.append(ids)
-            return ids + 100
-
-        pool = self._make_bare_pool(translate=translate)
-        virtual_loc = torch.tensor([4, 5], dtype=torch.int64)
         layer = types.SimpleNamespace(layer_id=0)
 
-        pool.get_mla_kv_buffer(layer, virtual_loc)
+        pool.get_mla_kv_buffer(layer, loc)
 
-        self.assertEqual(len(calls), 1)
         self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
-        self.assertTrue(
+        self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
-            torch.all(pool.full_kv_pool.mla_get_calls[0] == virtual_loc + 100)
-        )
 
 
 if __name__ == "__main__":

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -31,11 +31,13 @@ import torch
 
 from sglang.srt.mem_cache.multi_ended_allocator import (
     MultiEndedAllocator,
+    UnifiedMambaTokenToKVPoolAllocator,
     UnifiedSWATokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.unified_memory_pool import (
     MambaSubPoolSpec,
     MHASubPoolSpec,
+    MLASubPoolSpec,
     UnifiedKVPool,
 )
 
@@ -2746,5 +2748,79 @@ class TestSWACompositeDenseSurface(unittest.TestCase):
         self.assertTrue(bool((got[in_tomb] == 0).all().item()))
 
 
+class TestPs64MLACompositeFeasibility(unittest.TestCase):
+    """The Kimi/flashmla shape: MLA + mamba composite at page_size=64 (the
+    flashmla arg snap). Large pages stress every sizing derivation at once —
+    the 64-token sink-page floor, the ps*entry_bytes per-layer-view tail pad, and
+    the page-granular alloc — so this pins that the factory-shaped
+    construction stays FEASIBLE and the dense surface stays on-formula when
+    the page size jumps from the usual 1..4 to 64."""
+
+    PS = 64
+    LAYERS = 3
+
+    def _build(self):
+        full = MLASubPoolSpec(
+            name="full",
+            layer_num=self.LAYERS,
+            kv_lora_rank=64,
+            qk_rope_head_dim=16,
+            store_dtype=torch.float16,
+            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",
+        )
+        n_full = 8 * self.PS  # 8 pages incl. the sink page
+        total = n_full * full.entry_bytes() + 16 * mamba.entry_bytes()
+        pool = UnifiedKVPool(
+            total_bytes=total,
+            sub_pool_specs=[full, mamba],
+            device=_DEV,
+            enable_memory_saver=False,
+            page_size=self.PS,
+        )
+        full_kv = _FakeKVCache(pool.max_slots("full"))
+        full_kv.attach_allocator = lambda allocator: None
+        mamba_kv = _FakeKVCache(pool.max_slots("mamba"))
+        mamba_kv.attach_allocator = lambda allocator: None
+        mamba_kv._copy_from_physical = lambda src, dst: None
+
+        class _FakeHybridLinearKVPool:
+            full_kv_pool = full_kv
+            mamba_pool = mamba_kv
+
+        return UnifiedMambaTokenToKVPoolAllocator(
+            unified_buffer=pool,
+            kvcache=_FakeHybridLinearKVPool(),
+            device=_DEV,
+            page_size=self.PS,
+            need_sort=False,
+            forward_stream=None,
+        )
+
+    def test_construction_alloc_and_dense_formula(self):
+        a = self._build()
+        # MLA: one latent row per layer, so the spec reports LAYERS blocks.
+        self.assertEqual(a.kernel_page_multiplier, self.LAYERS)
+        v = a.alloc(2 * self.PS)
+        self.assertIsNotNone(v, "2-page alloc infeasible at ps=64")
+        # Page-aligned virtual run (page-granular allocator invariant).
+        self.assertEqual(int(v[0].item()) % self.PS, 0)
+        # Dense translate follows the affine formula at ps=64, and every id
+        # fits int32 (the canonical narrows on store).
+        v2p = a.full_v2p_page_table
+        want = v2p[v // self.PS] * (self.PS * self.LAYERS) + v % self.PS
+        got = a.translate_kv_loc_for_kernel(v)
+        self.assertTrue(torch.equal(got, want), "kernel-facing formula broke at ps=64")
+        self.assertTrue(bool((got < 2**31).all().item()))
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_mla_block_table.py

@@ -20,20 +20,26 @@ block table filled with kernel-facing page ids:
 
     dense_page(virtual_page) = v2p[virtual_page] * layer_num
 
-Three backend families reach that same formula by different routes:
+Since the read-path translator, ONE builder computes that formula for every
-  - `create_flashmla_kv_indices_triton` in-kernel via `v2p_ptr` / `PAGE_MULT`
+family — `build_kv_read_table` (the canonical) — and the backends only
-    (trtllm_mla / cutedsl_mla / tokenspeed_mla);
+differ in how they consume it:
-  - the flashinfer_mla updaters, post-gathering `translate_kv_loc_for_kernel` over the
+  - trtllm_mla / cutedsl_mla / tokenspeed_mla / flashmla: rows filled straight
-    token-level kv_indices;
+    into their padded block tables (`KVIndexTranslator.build_into`, prefix-only so
-  - `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table.
+    the backends' own -1 / stale tail sentinels survive);
+  - the flashinfer updaters: token ids reconstructed from the canonical by
+    `create_flashinfer_kv_indices_triton[ENTRY_PAGE_SIZE=ps]`;
+  - fa3's captured decode: `normal_decode_set_metadata` copies the canonical
+    rows' live prefixes (src_is_read_table=True).
 
 Covered here:
-  - kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main;
+  - the static `create_flashmla_kv_indices_triton` (no id-space knowledge left)
-  - kernel dense mapping against the python reference, for several page sizes,
+    still matches the plain token//ps reference;
-    ragged sequence lengths and a non-identity v2p permutation;
+  - the canonical route against the python dense reference, for several page
-  - padded block-table lanes never index the v2p table out of bounds;
+    sizes, ragged sequence lengths and a non-identity v2p permutation;
-  - the token-level kernel-facing translate the flashinfer updaters apply agrees with the
+  - lanes past a row's live prefix keep the backend's -1 sentinel (prefix-only
-    page-level block table the trtllm path builds;
+    discipline — the trtllm/flashmla tail contract);
+  - the token-level kernel-facing translate the flashinfer updaters used to apply
+    agrees with the canonical page table (page-affinity of the id space);
   - fa3's fused metadata kernels agree with the same reference, on both the
     page_size == 1 fast path (which is what Kimi-Linear takes: fa3 imposes no
     page-size constraint) and the general path.
@@ -57,6 +63,28 @@ _LAYERS = 24  # K3 MLA full-attention layer count
 def _fill_block_table(
     req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult
 ):
+    """The unified route: canonical builder into a -1-filled block table
+    (exactly what KVIndexTranslator.build_into does for trtllm_mla/flashmla)."""
+    from sglang.kernels.ops.kvcache.kv_read_table import build_kv_read_table
+
+    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)
+    build_kv_read_table(
+        req_to_token=req_to_token,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens.to(torch.int64),
+        v2p=v2p,
+        multiplier=mult,
+        page_size=page_size,
+        max_pages=max_blocks,
+        out=out,
+    )
+    return out
+
+
+def _fill_block_table_static(req_to_token, req_pool_indices, seq_lens, page_size):
+    """The static-pool route: the stripped flashmla kernel, token//ps verbatim."""
     from sglang.kernels.ops.kvcache.kv_indices import (
         create_flashmla_kv_indices_triton,
         get_num_kv_index_blocks_flashmla,
@@ -75,8 +103,6 @@ def _fill_block_table(
         req_to_token.stride(0),
         max_blocks,
         PAGED_SIZE=page_size,
-        v2p_ptr=v2p,
-        PAGE_MULT=mult,
     )
     return out
 
@@ -122,11 +148,12 @@ class TestDenseBlockTable(unittest.TestCase):
         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):
+    def test_static_kernel_matches_reference(self):
-        """v2p_ptr=None / PAGE_MULT=1 must reproduce the pre-change behaviour."""
+        """The stripped (id-space-free) flashmla kernel is byte-identical to the
+        plain token//ps reference -- guards the v2p-arg removal itself."""
         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)
+            got = _fill_block_table_static(rt, rpi, sl, page_size)
             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}"
@@ -166,8 +193,9 @@ class TestDenseBlockTable(unittest.TestCase):
             )
 
     def test_padded_lanes_stay_untouched(self):
-        """Lanes past a request's page count keep the -1 fill: the masked v2p load
+        """Lanes past a request's page count keep the -1 fill: the prefix-only
-        must not write a translated value (nor read out of bounds)."""
+        canonical build must never write a backend's tail sentinel (the
+        trtllm/flashmla block-table contract)."""
         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)
@@ -207,43 +235,77 @@ class TestDenseBlockTable(unittest.TestCase):
 
 @unittest.skipUnless(_HAS_CUDA, "requires CUDA")
 class TestFa3MetadataDenseBlockTable(unittest.TestCase):
-    """fa3 folds the unified remap into `normal_decode_set_metadata`, the fused
+    """fa3's captured-decode page table is written by `normal_decode_set_metadata`
-    gather that writes its captured-decode page table, so the kernel itself has
+    fed with the translator's read table kernel page table
-    to get the mapping right. Two kernels back it: a page_size == 1 / no-SWA fast
+    (src_is_read_table=True): the fused kernel copies the canonical
-    path (what Kimi-Linear takes, since fa3 imposes no page-size constraint) and
+    rows' live prefixes into the capture-stable buffer. Pinned END-TO-END:
-    a general one.
+    build_kv_read_table -> wrapper -> page_table must equal the python
+    reference of the kernel-facing formula, on both the page_size == 1 / no-SWA fast
+    path (what Kimi-Linear takes) and the general kernel. The static call
+    (no source flag) stays byte-identical to the pre-translator kernel.
     """
 
     def _run(self, page_size, *, v2p, mult, bs=5, max_ctx=2048):
         from sglang.kernels.ops.attention.metadata import normal_decode_set_metadata
+        from sglang.kernels.ops.kvcache.kv_read_table import (
+            build_kv_read_table,
+        )
 
         maker = TestDenseBlockTable._make_batch
         rt, rpi, sl, v2p_full = maker(self, page_size, bs=bs, max_ctx=max_ctx)
-        v2p_arg = v2p_full if v2p else None
 
         max_pages = (max_ctx + page_size - 1) // page_size
         page_table = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
         cache_seqlens = torch.zeros((bs,), dtype=torch.int32, device=_DEV)
         cu_seqlens_k = torch.zeros((bs + 1,), dtype=torch.int32, device=_DEV)
-        strided = torch.arange(0, max_ctx, page_size, device=_DEV)
         max_seq_pages = (int(sl.max().item()) + page_size - 1) // page_size
 
-        normal_decode_set_metadata(
+        if v2p:
-            cache_seqlens,
+            # The translator's canonical, then the wrapper copies its rows.
-            cu_seqlens_k,
+            canonical = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
-            page_table,
+            build_kv_read_table(
-            rt,
+                req_to_token=rt,
-            rpi,
+                req_pool_indices=rpi,
-            strided,
+                seq_lens=sl.to(torch.int64),
-            max_seq_pages,
+                v2p=v2p_full,
-            sl.to(torch.int64),
+                multiplier=mult,
-            0,
+                page_size=page_size,
-            page_size,
+                max_pages=max_pages,
-            v2p_page_table=v2p_arg,
+                out=canonical,
-            kernel_page_multiplier=mult,
+            )
-        )
+            rows = torch.arange(bs, dtype=torch.int64, device=_DEV)
+            normal_decode_set_metadata(
+                cache_seqlens,
+                cu_seqlens_k,
+                page_table,
+                canonical,
+                rows,
+                max_seq_pages,
+                sl.to(torch.int64),
+                0,
+                page_size,
+                None,
+                None,
+                src_is_read_table=True,
+            )
+        else:
+            normal_decode_set_metadata(
+                cache_seqlens,
+                cu_seqlens_k,
+                page_table,
+                rt,
+                rpi,
+                max_seq_pages,
+                sl.to(torch.int64),
+                0,
+                page_size,
+                None,
+                None,
+            )
         torch.cuda.synchronize()
-        want = _reference(rt, rpi, sl, page_size, v2p=v2p_arg, mult=mult)
+        want = _reference(
+            rt, rpi, sl, page_size, v2p=(v2p_full if v2p else None), mult=mult
+        )
         return page_table, want, sl
 
     def _assert_live_prefix(self, got, want, sl, page_size):
@@ -287,192 +349,9 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase):
                 "test batch degenerated: v2p is the identity on the pages used",
             )
 
-    def test_agrees_with_flashmla_block_table(self):
+    # (The old fa3<->flashmla agreement case is gone: both families now
-        """fa3 and trtllm_mla build the same table two different ways; a
+    # consume the SAME canonical builder, so cross-family agreement holds by
-        disagreement means one family is addressing the wrong pages."""
+    # construction and the per-family cases above cover the two consumers.)
-        for page_size in (1, 64):
-            got, _, sl = self._run(page_size, v2p=True, mult=_LAYERS)
-            rt, rpi, sl2, v2p = TestDenseBlockTable._make_batch(self, page_size)
-            other = _fill_block_table(
-                rt, rpi, sl2, page_size, v2p=v2p, mult=_LAYERS
-            ).long()
-            for r in range(got.shape[0]):
-                n_pages = (int(sl[r].item()) + page_size - 1) // page_size
-                self.assertTrue(
-                    torch.equal(got[r, :n_pages].long(), other[r, :n_pages]),
-                    f"fa3 and flashmla block tables disagree (row {r}, ps={page_size})",
-                )
-
-
-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.unified_mem_hooks 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_for_kernel)
-        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_for_kernel=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_for_kernel)
-        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_for_kernel=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_for_kernel)
-        # Multiplier stays 1: kernel-facing 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,
-        )
-
-        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 kernel-facing 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_for_kernel(valid))
-
-        expected = alloc.translate_kv_loc_for_kernel(virt)
-        self.assertEqual(buf.dtype, torch.int32)
-        self.assertTrue(
-            torch.equal(buf[:n].long(), expected),
-            "in-place translate did not land kernel-facing 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_for_kernel(virt), virt),
-            "kernel-facing ids coincide with virtual ids; pick a different allocation",
-        )
 
 
 if __name__ == "__main__":

test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py

@@ -27,6 +27,7 @@ import inspect
 import textwrap
 import unittest
 from types import SimpleNamespace
+from unittest.mock import create_autospec
 
 import torch
 
@@ -177,5 +178,104 @@ class TestPadComposesWithDerivation(CustomTestCase):
         self.assertEqual(src._swa_write_loc_unified(fb.out_cache_loc).numel(), 0)
 
 
+class TestReadRailTranslatesAtProduction(CustomTestCase):
+    """The model-door READ indices (req_to_token-derived, VIRTUAL under the
+    unified pool) are translated at their PRODUCTION site -- the cache then
+    holds the kernel-facing result and the pool door never translates."""
+
+    def _fb_for_one_shot(self):
+        fb = _make_fb(torch.tensor([1, 2], dtype=torch.int64))
+        fb.batch_size = 2
+        fb.seq_lens = torch.tensor([2, 3], dtype=torch.int64)
+        fb.seq_lens_cpu = torch.tensor([2, 3], dtype=torch.int32)
+        fb.req_pool_indices = torch.tensor([0, 1], dtype=torch.int64)
+        return fb
+
+    def test_one_shot_indices_translated_once_and_cached(self):
+        from unittest.mock import patch
+
+        from sglang.srt.model_executor import forward_batch_deepseek_mha_mixin as mix
+
+        calls = []
+        sentinel = torch.arange(5, dtype=torch.int64) + 5000
+
+        def translate(t):
+            calls.append(t)
+            return sentinel
+
+        fb = self._fb_for_one_shot()
+        fake_pool = SimpleNamespace(
+            req_to_token=torch.zeros((4, 16), dtype=torch.int32)
+        )
+        # autospec, not a bare namespace: setting a name the translator does
+        # not have raises, so renaming the method breaks this test loudly.
+        fake_translator = create_autospec(KVIndexTranslator, instance=True)
+        fake_translator.translate_full_attn_ids = translate
+        fake_backend = SimpleNamespace(kv_index_translator=fake_translator)
+        with (
+            patch.object(mix, "get_req_to_token_pool", return_value=fake_pool),
+            patch.object(mix, "get_attn_backend", return_value=fake_backend),
+            patch.object(mix, "create_flashinfer_kv_indices_triton"),
+        ):
+            r1 = fb.fetch_mha_one_shot_kv_indices()
+            r2 = fb.fetch_mha_one_shot_kv_indices()
+
+        self.assertIs(r1, sentinel)  # production site translated
+        self.assertIs(r2, sentinel)  # cache holds the TRANSLATED result
+        self.assertEqual(len(calls), 1)  # translated exactly once
+        self.assertEqual(calls[0].dtype, torch.int32)  # raw producer output
+
+    def test_one_shot_indices_noop_on_unmigrated_backend(self):
+        from unittest.mock import patch
+
+        from sglang.srt.model_executor import forward_batch_deepseek_mha_mixin as mix
+
+        fb = self._fb_for_one_shot()
+        fake_pool = SimpleNamespace(
+            req_to_token=torch.zeros((4, 16), dtype=torch.int32)
+        )
+        # A backend that never set the attribute inherits the base-class None.
+        fake_backend = SimpleNamespace(kv_index_translator=None)
+        with (
+            patch.object(mix, "get_req_to_token_pool", return_value=fake_pool),
+            patch.object(mix, "get_attn_backend", return_value=fake_backend),
+            patch.object(mix, "create_flashinfer_kv_indices_triton"),
+        ):
+            r = fb.fetch_mha_one_shot_kv_indices()
+        # The raw int32 producer output passes through untouched.
+        self.assertEqual(r.dtype, torch.int32)
+
+    def test_get_mla_kv_buffer_door_passes_loc_untranslated(self):
+        from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
+
+        recorded = {}
+
+        class _RecordingLeafPool:
+            def get_mla_kv_buffer(self, layer, loc, dst_dtype):
+                recorded["loc"] = loc
+                return None, None
+
+            def get_kv_size_bytes(self):
+                return 0
+
+        pool = HybridLinearKVPool(
+            size=16,
+            dtype=torch.float16,
+            page_size=1,
+            head_num=1,
+            head_dim=8,
+            full_attention_layer_ids=[0],
+            device=_DEV,
+            mamba_pool=SimpleNamespace(get_size_per_token=lambda: 0),
+            enable_memory_saver=False,
+            use_mla=True,
+            start_layer=0,
+            full_kv_pool=_RecordingLeafPool(),
+        )
+        loc = torch.tensor([9, 10], dtype=torch.int64)
+        pool.get_mla_kv_buffer(SimpleNamespace(layer_id=0), loc, torch.float16)
+        self.assertIs(recorded["loc"], loc)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/server_args/test_page_major_backend_allowlist.py

@@ -13,18 +13,28 @@
 # ==============================================================================
 """`--enable-page-major-kv-layout` full-attention backend allowlist.
 
-The page-major envelope K/V views are strided, which only the Triton attention
+Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool
-kernels read. The one exception is the unified-memory MLA pool: it exposes each
+exposes per-layer views and nothing else:
-layer as a contiguous view (`build_mla_views`), so the paged MLA
+  * unified-memory MLA models (`build_mla_views`) allow the whole wired
-backends can read it directly once their kv_indices / block tables are remapped
+    paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with
-to kernel-facing ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its
+    its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64
-`cutedsl_mla` / `tokenspeed_mla` subclasses.
+    snap);
+  * unified-memory MHA/SWA models (`build_mha_views`) allow `fa3` /
+    `fa4` / `flashinfer` / `trtllm_mha` alongside Triton;
+  * plain `--enable-page-major-kv-layout` without the unified pool keeps the
+    envelope-strided 4-D views only the stride-aware Triton kernels read.
 
-Pinned here so the exception cannot silently widen to a backend that has no
+The same handler also screens the pool itself: the dense MHA/SWA views need
-dense-id remapping (`flashmla`, `cutlass_mla`, ...) or leak into the MHA path.
+uniform K/V rows, so an asymmetric-K/V model (MiMoV2: head_dim 192 !=
-`fa3` matters most: it is the resolved default on pre-Blackwell hosts, so it is
+v_head_dim 128) cannot run `--enable-unified-memory` at all and is rejected on
-the one entry whose absence used to make `--enable-unified-memory` fail to boot
+EVERY backend, Triton included. MLA models are exempt -- their sub-pool keeps
-under its own default configuration.
+one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72,
+v_head_dim 128) report asymmetric dims while running the unified pool today.
+
+Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`,
+`aiter`) and no arm silently narrows: `fa3` is the resolved default on
+pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory`
+fail to boot under its own default configuration.
 
     python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
 """
@@ -96,13 +106,23 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
         "flashinfer",
         "cutedsl_mla",
         "tokenspeed_mla",
+        "flashmla",
     )
-    # No dense-id remapping: must stay rejected until they get one.
+    # Wired for the dense per-layer MHA/SWA views (uniform-row models).
-    UNWIRED_BACKENDS = ("flashmla", "cutlass_mla", "trtllm_mha", "aiter")
+    DENSE_MHA_BACKENDS = ("fa3", "fa4", "flashinfer", "trtllm_mha")
+    # MLA-family kernels that must never leak into the MHA arm.
+    MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla")
+    # No dense-id wiring anywhere: must stay rejected until they get one.
+    UNWIRED_BACKENDS = ("cutlass_mla", "aiter")
 
-    def test_triton_always_allowed(self):
+    def test_triton_allowed_on_every_arm(self):
+        """Triton reads both view families, so it is the one backend neither
+        the MLA nor the MHA arm can narrow away."""
         for use_mla in (True, False):
             self.assertTrue(_accepts("triton", use_mla=use_mla))
+        # The uniform-row screen is a property of the model, not of the
+        # backend, so it rejects even Triton.
+        self.assertFalse(_accepts("triton", use_mla=False, has_asymmetric_kv=True))
 
     def test_dense_mla_backends_allowed_under_unified_mla(self):
         for backend in self.DENSE_MLA_BACKENDS:
@@ -111,21 +131,18 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
                 f"{backend} should be allowed with the unified-memory MLA pool",
             )
 
-    def test_dense_mla_backends_rejected_for_mha(self):
+    def test_dense_mha_backends_allowed_for_uniform_row_models(self):
-        """The per-layer-view exception is MLA-only -- MHA sub-pools stay strided."""
+        for backend in self.DENSE_MHA_BACKENDS:
-        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertTrue(
-            self.assertFalse(
                 _accepts(backend, use_mla=False),
-                f"{backend} must stay rejected for a non-MLA model",
+                f"{backend} should be allowed for a uniform-row MHA model",
             )
 
-    def test_dense_mla_backends_rejected_without_unified_memory(self):
+    def test_mla_only_backends_rejected_for_mha(self):
-        """Plain --enable-page-major-kv-layout (no unified pool) keeps the
+        for backend in self.MLA_ONLY_BACKENDS:
-        strided views, so only Triton can read them."""
-        for backend in self.DENSE_MLA_BACKENDS:
             self.assertFalse(
-                _accepts(backend, use_mla=True, unified=False),
+                _accepts(backend, use_mla=False),
-                f"{backend} must stay rejected without --enable-unified-memory",
+                f"{backend} is an MLA kernel and must stay out of the MHA arm",
             )
 
     def test_plain_page_major_arm_is_gated_at_boot(self):
@@ -143,7 +160,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
         """head_dim != v_head_dim (MiMoV2): no uniform rows, so no per-layer views
         and no unified pool. The rejection is the POOL's, not a backend's, so
         it must fire on every backend -- Triton included."""
-        for backend in ("triton",) + self.DENSE_MLA_BACKENDS:
+        for backend in ("triton",) + self.DENSE_MHA_BACKENDS:
             self.assertFalse(
                 _accepts(backend, use_mla=False, has_asymmetric_kv=True),
                 f"--enable-unified-memory + {backend} must be rejected for an "
@@ -162,12 +179,25 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
                 "K/V head dims",
             )
 
+    def test_page_major_rejected_without_unified_memory(self):
+        """--enable-page-major-kv-layout without the unified pool is rejected
+        outright, Triton included: the static page-major arm went away with
+        the strided views and awaits its per-layer-view reimplementation."""
+        for backend in ("triton",) + tuple(
+            set(self.DENSE_MLA_BACKENDS + self.DENSE_MHA_BACKENDS)
+        ):
+            for use_mla in (True, False):
+                self.assertFalse(
+                    _accepts(backend, use_mla=use_mla, 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",
+                    f"{backend} has no dense-id wiring and must be rejected",
                 )
 
     def test_helion_linear_attention_is_kda_only(self):