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
-    # sub-pool). Both default to the identity for the statically-partitioned
-    # pool, where req_to_token already holds physical ids.
-    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,
+    # 1: the two table pointers carry PAGE-granular, already kernel-facing
+    # read tables; emit verbatim -- no >>SHIFT, no v2p, no mapping gather.
+    SRC_IS_KERNEL_PAGE_TABLE: tl.constexpr = 0,
 ):
     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)
-    if page_size == 1:
+    # Compute column indices in the source tensor (token offset; page offset
+    # 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
-        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_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,
-    kernel_page_multiplier: int = 1,
+    src_is_read_table: bool = False,
+    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 via the legacy full->SWA map or the unified
-         SWA pool's independent page map
+      5. (optional) swa_page_table for sliding window attention
 
-    Step 4b is folded in rather than applied afterwards so the capture-stable
-    page_table is written already translated: no separate pass a caller could
-    forget, and no temporary to keep pointer-stable across cuda-graph replays.
-    Identity (None / 1) for the statically-partitioned pool.
+    Unified pool (``src_is_read_table=True``): ``req_to_token`` /
+    ``req_pool_indices`` carry the translator's PAGE-granular read table and its
+    row indices instead (entries already kernel-facing; ``swa_src_table`` is
+    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,
-            PAGE_MULT=kernel_page_multiplier,
-            SWA_MAPPING_IS_V2P=swa_uses_v2p,
+            SRC_IS_KERNEL_PAGE_TABLE=1 if src_is_read_table else 0,
             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
-        row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
-        if HAS_SWA:
-            swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_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 not SKIP_PAGE_TABLE:
+            req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
+            row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
+            row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
             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_row_out = (
+                    swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
                 )
-                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)
+            # 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)
+                    # 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
-    # views; FA3 / FlashInfer do not. EXCEPTION: the unified-memory MLA pool
-    # exposes each layer as a contiguous per-layer view
-    # (build_mla_views), which the paged MLA kernels consume directly,
-    # with their kv_indices / block tables remapped to kernel-facing ids. Names below
-    # are the RESOLVED ids from attention_backends_of: "flashinfer" is
-    # FlashInferMLAAttnBackend for an MLA model, "trtllm_mla" the trtllm
-    # decode kernel; "cutedsl_mla" and "tokenspeed_mla" subclass
-    # 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.
+    # Allow-list. Every backend below reads through the translator, so what
+    # gates one is only whether its kernels can address the per-layer views:
+    #   * MLA models: the full paged MLA family, incl. flashmla (ps=64
+    #     snap). cutlass_mla stays rejected (never exercised).
+    #   * MHA/SWA models: fa3 / fa4 / flashinfer / trtllm_mha alongside
+    #     Triton. fa4 is the fa3 class.
+    #   * Without the unified pool, plain page-major stays Triton-only.
+    # Names are the RESOLVED ids from attention_backends_of.
     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 "
-        "for the full-attention layers (unified-memory MLA also allows the "
-        f"paged MLA backends); got {sorted(backends)}, allowed "
-        f"{sorted(allowed_full)}. Pass a compatible --attention-backend."
+        "--enable-page-major-kv-layout: the resolved attention backends "
+        f"{sorted(backends)} are not in the allowed set "
+        f"{sorted(allowed_full)} for this configuration (unified memory "
+        "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
-        # are kernel-facing, so every page_table needs remapping. MLA-only -- the MHA/SWA
-        # 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.kv_index_translator = model_runner.kv_index_translator
+        self.kv_read_tables = None
         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_(
-                self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc)
-            )
+            if in_capture and self.kv_index_translator.is_translating:
+                # 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,
-                        req_pool_indices,
-                        self.decode_cuda_graph_metadata["strided_indices"],
+                        kv_view.ids,
+                        kv_view.row_ids,
                         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=(
-                            self._unified_hooks.v2p_page_table
-                            if self._unified_dense
-                            else None
-                        ),
-                        kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier,
+                        src_is_read_table=kv_view.is_translated,
+                        swa_src_table=kv_view.sliding_window_ids,
                     )
 
                 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,
-                        req_pool_indices,
-                        self.decode_cuda_graph_metadata["strided_indices"],
+                        kv_view.ids,
+                        kv_view.row_ids,
                         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=(
-                            self._unified_hooks.v2p_page_table
-                            if self._unified_dense
-                            else None
-                        ),
-                        kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier,
+                        src_is_read_table=kv_view.is_translated,
+                        swa_src_table=kv_view.sliding_window_ids,
                     )
 
                 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_(
-                self._swa_kv_pool.translate_loc_from_full_to_swa(
-                    forward_batch.out_cache_loc
+            if in_capture and self.kv_index_translator.is_translating:
+                # A runner-built capture batch never went through `init_new`,
+                # 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._swa_kv_pool.translate_loc_from_full_to_swa(
+            swa_out_cache_loc = self.kv_index_translator.sliding_window_write_loc_for(
                 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,
-                req_pool_indices,
+                src_table,
+                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,
-                    req_pool_indices,
+                    src_table,
+                    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,
-                req_pool_indices,
+                kv_view.ids,
+                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,
-                req_pool_indices,
+                kv_view.ids,
+                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[
-                (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),
-            )
+            if self.kv_index_translator.is_translating:
+                assert self.page_size == PAGE_SIZE
+                self.kv_index_translator.fill_read_table(
+                    out=block_kv_indices,
+                    req_pool_indices=forward_batch.req_pool_indices,
+                    seq_lens=forward_batch.seq_lens,
+                )
+            else:
+                create_flashmla_kv_indices_triton[
+                    (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[
-                (
-                    bs,
-                    get_num_kv_index_blocks_flashmla(
-                        self.cuda_graph_kv_indices.stride(0), PAGE_SIZE
-                    ),
+            if self.kv_index_translator.is_translating:
+                assert self.page_size == PAGE_SIZE
+                self.kv_index_translator.fill_read_table(
+                    out=self.cuda_graph_kv_indices,
+                    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).
-        if self._swa_kv_pool is not None:
+        # metadata kernel (gather + // page_size happen on device). The unified
+        # 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
-        metadata.page_table = torch.empty(
-            (batch_size, self.max_num_pages), dtype=torch.int32, device=device
-        )
-        metadata.swa_page_table = (
-            torch.empty(
+        kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
+        if kv_view.is_translated:
+            # No fill kernel: the kernels take the tensor's own width/stride
+            # and bound their reads by cache_seqlens.
+            metadata.page_table = kv_view.ids
+            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
-            else None
-        )
-        self._fill_page_table_device(
-            metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32
-        )
+            metadata.swa_page_table = (
+                torch.empty(
+                    (batch_size, self.max_num_pages), dtype=torch.int32, device=device
+                )
+                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
-        # holds VIRTUAL token ids; the block table needs kernel-facing page ids, so the
-        # kv-index kernels gather virtual->physical page through `_v2p_page_table`
-        # then scale by `_kernel_page_multiplier` (= num MLA layers). See
-        # build_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.
+        # [:n] view of a capture-stable buffer on the cuda-graph path; None on
+        # the eager path, which passes forward_batch.out_cache_loc 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[
-            (
-                batch_size,
-                get_num_kv_index_blocks_flashmla(max_blocks, self.page_size),
+        if self.kv_index_translator.is_translating:
+            self.kv_index_translator.fill_read_table(
+                out=block_kv_indices,
+                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[
-            (
-                bs,
-                get_num_kv_index_blocks_flashmla(
-                    metadata.block_kv_indices.shape[1], self.page_size
-                ),
+        if self.kv_index_translator.is_translating:
+            self.kv_index_translator.fill_read_table(
+                out=metadata.block_kv_indices,
+                req_pool_indices=req_pool_indices[:bs],
+                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
-        # buffer (both capture and each replay-prep run this out of the graph),
-        # so the in-graph set_mla_kv_buffer writes a dense loc without capturing
-        # a translate.
-        if self._unified_mla and (
+        # Out-of-graph on capture AND every replay-prep, so the in-graph
+        # set_mla_kv_buffer captures no translate.
+        if self.kv_index_translator.is_translating 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 = (
-                        self._decode_kernel_loc
-                        if self._decode_kernel_loc is not None
-                        else (
-                            None if self._unified_mla else forward_batch.out_cache_loc
-                        )
-                    )
+                    loc = self._resolve_fused_write_loc(forward_batch)
                     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
-      (``None`` otherwise).
-    - ``full_loc``: the full-attention-sub-pool location for the unified
-      memory pool (``None`` otherwise), carried in attention metadata
-      (``ForwardMetadata.out_cache_loc_full_physical``). Since the
-      construction-time rebind it is the SAME id space as ``loc``; the shared
-      full pool writes it directly and never translates.
+    - ``swa_loc``: the SWA-sub-pool location for hybrid SWA pools (``None``
+      otherwise); under the unified pool the translator derives it from the
+      same rebound loc (``sliding_window_write_loc_for``).
+    - ``full_loc``: OPTIONAL full-attention-sub-pool location. Since the
+      construction-time rebind it is the SAME id space as ``loc``, so pools
+      fall back to ``loc`` when it is ``None`` -- only triton's captured path
+      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:
-        token_to_kv_pool._full_translate = allocator.translate_kv_loc_for_kernel
+    # No full-KV translate hook is wired: both MLA doors now receive
+    # 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).
-        assert full_loc is not None, (
-            "UnifiedSWAKVPool.set_kv_buffer: full layer received no full_loc; "
-            "ForwardMetadata.out_cache_loc_full_physical must be precomputed for "
-            "the unified memory pool."
-        )
+        # Full layer: `loc` is already the full-side kernel-facing id, so an
+        # explicit full_loc is a same-space alias -- only triton's captured path
+        # passes one (its capture-stable buffer).
+        if full_loc is None:
+            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.