[Perf] Unified memory: close the DCP decode gap on Blackwell (#37926)

python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py

@@ -14,6 +14,8 @@ MTP inter-phase seam:
 This module fuses that chain into a single launch.
 """
 
+from typing import Optional
+
 import torch
 import triton
 import triton.language as tl
@@ -24,15 +26,20 @@ def _fused_replay_state_indices_kernel(
     req_pool_indices_ptr,  # (total_bs,) int64 — static replay buffer
     mamba_map_ptr,  # (req_pool_size,) int32 — req_index_to_mamba_index_mapping
     out_ptr,  # (total_bs,) int32 — state_indices_list[bs - 1]
+    v2p_ptr,  # (num_slots + 1,) int64 — mamba virtual->physical, or unused
     valid_bs,
     total_bs,
     BS_UPPER: tl.constexpr,
+    HAS_V2P: tl.constexpr,
 ):
     offs = tl.arange(0, BS_UPPER)
     in_range = offs < total_bs
     valid = offs < valid_bs
     req = tl.load(req_pool_indices_ptr + offs, mask=valid, other=0)
     idx = tl.load(mamba_map_ptr + req, mask=valid, other=0)
+    if HAS_V2P:
+        # Must gather before the padding sentinel, as the reference chain does.
+        idx = tl.load(v2p_ptr + idx, mask=valid, other=0)
     out_val = tl.where(valid, idx.to(tl.int32), -1)
     tl.store(out_ptr + offs, out_val, mask=in_range)
     # Preserve the reference chain's side effect: padded rows of the static
@@ -49,6 +56,7 @@ def fused_replay_state_indices(
     out_state_indices: torch.Tensor,
     valid_bs: int,
     total_bs: int,
+    v2p: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """Fill the captured replay state-indices buffer in one launch.
 
@@ -58,9 +66,11 @@ def fused_replay_state_indices(
     the ``-1`` sentinel (mamba kernels skip ``state_idx < 0``) and their
     ``req_pool_indices`` entries are zeroed.
 
-    Callers must supply an identity virtual->physical mapping (the static
-    hybrid pool); the unified pool's allocator translate is not a flat table
-    gather and has to take the reference chain.
+    ``v2p`` is the mamba virtual->physical slot table, for a pool whose slot
+    ids are virtual (the unified memory pool). Pass None when the mapping
+    already yields physical slots (the static hybrid pool). The unified
+    allocator runs the mamba sub-pool at page_size 1, so its translate is a
+    plain table gather and folds into this launch.
 
     Returns the filled ``out_state_indices[:total_bs]`` view.
     """
@@ -68,8 +78,10 @@ def fused_replay_state_indices(
         req_pool_indices,
         mamba_index_mapping,
         out_state_indices,
+        v2p,
         valid_bs,
         total_bs,
         BS_UPPER=triton.next_power_of_2(total_bs),
+        HAS_V2P=v2p is not None,
     )
     return out_state_indices[:total_bs]

python/sglang/kernels/ops/memory/__init__.py

@@ -18,6 +18,7 @@ _TRITON_KERNELS = [
     ("virtual_slot", "alloc_bind_inplace"),
     ("virtual_slot", "free_unbind_inplace"),
     ("virtual_slot", "bind_inplace"),
+    ("virtual_slot", "write_loc_to_kernel_ids"),
 ]
 for _mod, _fn in _TRITON_KERNELS:
     register_kernel(

python/sglang/kernels/ops/memory/virtual_slot.py

@@ -15,6 +15,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import triton
 import triton.language as tl
@@ -189,3 +191,134 @@ def bind_inplace(
         return
     grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),)
     bind_inplace_kernel[grid](v, p, v2p, p2v, N, BLOCK=ALLOC_BIND_BLOCK)
+
+
+WRITE_LOC_BLOCK = 512
+
+
+@triton.jit
+def write_loc_to_kernel_id_kernel(
+    loc_ptr,  # in:  [N] int64 — WIDENED virtual token ids
+    v2p_ptr,  # in:  [num_pages + 1] int64 — virtual->physical page table
+    out_ptr,  # out: [N] int64 — kernel-facing ids
+    N,  # runtime: live element count
+    W,  # runtime: lanes to write; [N, W) get 0
+    stride,  # runtime: pool_page_size * kernel_page_multiplier
+    PAGE_SIZE: tl.constexpr,
+    DCP_SIZE: tl.constexpr,
+    DCP_RANK: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """``kernel_id(t) = v2p[t // ps] * ps * mult + t % ps``, clamped at 0.
+
+    Under DCP the incoming id is WIDENED: ``loc % dcp_size`` names its owner
+    and ``loc // dcp_size`` is the row. Ids this rank does not own resolve to
+    kernel id 0, the padding sink every write kernel skips.
+
+    Triton truncates ``//`` toward zero where torch floors it, so a negative
+    loc is tested explicitly rather than left to the division; it resolves to
+    0, as the torch path does.
+
+    Writing ``W > N`` lanes fills ``[N, W)`` with 0, the padding sink, so a
+    caller may hand in a capture-stable buffer wider than this batch and have
+    the stale tail cleared in the same launch.
+    """
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    in_range = offs < W
+    mask = offs < N
+    loc = tl.load(loc_ptr + offs, mask=mask, other=0).to(tl.int64)
+
+    keep = mask & (loc >= 0)
+    if DCP_SIZE > 1:
+        keep = keep & ((loc % DCP_SIZE) == DCP_RANK)
+        loc = loc // DCP_SIZE
+
+    page = loc // PAGE_SIZE if PAGE_SIZE > 1 else loc
+    offset = loc % PAGE_SIZE if PAGE_SIZE > 1 else 0
+    # `keep` already excludes negatives, so the gather index is in range.
+    phys = tl.load(v2p_ptr + tl.where(keep, page, 0), mask=mask, other=0).to(tl.int64)
+    ids = tl.maximum(phys * stride + offset, 0)
+    tl.store(out_ptr + offs, tl.where(keep, ids, 0), mask=in_range)
+
+
+def write_loc_to_kernel_ids(
+    *,
+    loc: torch.Tensor,
+    v2p: torch.Tensor,
+    page_size: int,
+    stride: int,
+    dcp_size: int = 1,
+    dcp_rank: int = 0,
+    out: Optional[torch.Tensor] = None,
+    out_width: Optional[int] = None,
+) -> torch.Tensor:
+    """One launch for the whole write-loc conversion; see the kernel.
+
+    ``out`` is written in place when given (a captured graph records the
+    gather against a fixed ``data_ptr``), else a fresh int64 tensor is
+    returned. Cuda-graph safe: no ``.item()``, no host sync, no allocation on
+    the ``out=`` path.
+
+    ``out_width`` writes that many lanes rather than ``loc.numel()``, zeroing
+    the ones past the batch; pass the captured tier's width to clear a stale
+    tail here.
+    """
+    N = int(loc.numel())
+    # Flat-indexed as `ptr + offs`, so a strided view is mis-addressed.
+    assert loc.is_contiguous(), (
+        f"write_loc_to_kernel_ids: loc must be contiguous, got shape "
+        f"{tuple(loc.shape)} stride {tuple(loc.stride())}"
+    )
+    if out is None:
+        out = torch.empty_like(loc, dtype=torch.int64)
+    width = N if out_width is None else int(out_width)
+    assert out.dtype == torch.int64, (
+        f"write_loc_to_kernel_ids: out dtype must be int64 (matches v2p), "
+        f"got {out.dtype}"
+    )
+    if out_width is None:
+        # `out` mirrors `loc` whatever its shape; a 2-D page table is legal.
+        assert out.shape == loc.shape and out.is_contiguous(), (
+            f"write_loc_to_kernel_ids: out shape {tuple(out.shape)} must match "
+            f"loc shape {tuple(loc.shape)}"
+        )
+    else:
+        assert out.dim() == 1 and out.is_contiguous() and out.numel() >= width, (
+            f"write_loc_to_kernel_ids: out_width needs a packed 1-D out of at "
+            f"least {width}, got {tuple(out.shape)}"
+        )
+        assert width >= N, (
+            f"write_loc_to_kernel_ids: out_width {width} is under the batch's "
+            f"{N} locs, which would drop live rows"
+        )
+    if width == 0:
+        return out
+    if not loc.is_cuda:
+        # Pure-torch reference; the allocator's unit tests run on CPU.
+        big = loc.to(torch.int64)
+        keep = big >= 0
+        if dcp_size > 1:
+            keep = keep & (big % dcp_size == dcp_rank)
+            big = torch.div(big, dcp_size, rounding_mode="floor")
+        page = torch.where(keep, torch.div(big, page_size, rounding_mode="floor"), 0)
+        offset = big % page_size if page_size > 1 else 0
+        ids = (v2p[page] * stride + offset).clamp_(min=0)
+        out[:N].copy_(torch.where(keep, ids, torch.zeros_like(ids)))
+        if width > N:
+            out[N:width].zero_()
+        return out
+    grid = (triton.cdiv(width, WRITE_LOC_BLOCK),)
+    write_loc_to_kernel_id_kernel[grid](
+        loc,
+        v2p,
+        out,
+        N,
+        width,
+        stride,
+        PAGE_SIZE=page_size,
+        DCP_SIZE=dcp_size,
+        DCP_RANK=dcp_rank,
+        BLOCK=WRITE_LOC_BLOCK,
+    )
+    return out

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -249,6 +249,12 @@ def handle_unified_memory_pool(server_args: Any) -> None:
             "not translate speculative verify indices to the unified "
             "pool's kernel-facing space yet."
         )
+    assert not cfg.enable_two_batch_overlap, (
+        "--enable-unified-memory does not support --enable-two-batch-overlap: "
+        "TBO's replay split hands each child a view without the pre-translate "
+        "write loc, so a captured decode replay raises. "
+        "TODO(ch-wan): carry out_cache_loc_virtual into the child view."
+    )
     assert not (cfg.enable_hierarchical_cache or cfg.enable_lmcache), (
         "--enable-unified-memory is not yet compatible with hierarchical / "
         "host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): "

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -693,6 +693,11 @@ class TboForwardBatchPreparer:
             )
             output_dict[key] = old_value[start_token_index:end_token_index]
 
+        if batch.out_cache_loc_virtual is not None:
+            output_dict["out_cache_loc_virtual"] = batch.out_cache_loc_virtual[
+                start_token_index:end_token_index
+            ]
+
         attention_tp_size = get_parallel().attn_tp_size
         _tbo_padded_len = (
             (end_token_index - start_token_index - 1) // attention_tp_size + 1

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

@@ -70,15 +70,11 @@ class MambaAttnBackendBase(AttentionBackend):
         # backends on runners without a real model_config.
         self._model_runner = model_runner
         self._mamba_chunk_size: Optional[int] = None
-        # Fused replay-prep state-indices fast path (fused_replay_state_indices):
-        # requires the static hybrid pool whose v2p translate is the identity —
-        # the unified pool overrides translate_mamba_indices with an allocator
-        # lookup that is not a flat table gather.
+        pool = self.req_to_token_pool
         self._fused_state_indices_ok = (
-            str(self.device).startswith("cuda")
-            and isinstance(self.req_to_token_pool, HybridReqToTokenPool)
-            and type(self.req_to_token_pool).translate_mamba_indices
-            is HybridReqToTokenPool.translate_mamba_indices
+            torch.device(self.device).type == "cuda"
+            and isinstance(pool, HybridReqToTokenPool)
+            and pool.mamba_translate_is_fusable
         )
         self.forward_metadata: ForwardMetadata = None
         self.state_indices_list = []
@@ -643,6 +639,7 @@ class MambaAttnBackendBase(AttentionBackend):
                 out_state_indices=self.state_indices_list[bs - 1],
                 valid_bs=bs - int(num_padding),
                 total_bs=bs,
+                v2p=self.req_to_token_pool.mamba_v2p_table,
             )
         else:
             # Make sure forward metadata is correctly handled for padding reqs

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

@@ -748,22 +748,16 @@ class TritonAttnBackend(AttentionBackend):
     def _fill_cuda_graph_write_locs(
         self, forward_batch: ForwardBatch, bs: int
     ) -> Optional[torch.Tensor]:
-        """Copy the cuda-graph WRITE loc into the capture-stable buffer and
-        return the ``[:n]`` view; no-op for non-unified pools.
-
-        Runs BEFORE graph.replay() so it reads the live post-compaction v2p.
-        The capture batch is runner-built with zeros, which is safe because
-        slot 0 is the reserved sink in every id space.
-        """
+        """Runs BEFORE graph.replay(), so it reads the live post-compaction
+        v2p; no-op for non-unified pools."""
+        # The buffer exists only for a translating pool; return before naming it.
         if not self.kv_index_translator.is_translating:
             return None
-        out_cache_loc = forward_batch.out_cache_loc
-        n = out_cache_loc.shape[0]
-        # Zero the padded tail first: a smaller replay batch leaves [n:] holding
-        # stale ids that the captured store would write; send them to slot 0 (sink).
-        self.cuda_graph_out_cache_loc_full_physical[n:].zero_()
-        self.cuda_graph_out_cache_loc_full_physical[:n].copy_(out_cache_loc)
-        return self.cuda_graph_out_cache_loc_full_physical[:n]
+        return self.kv_index_translator.fill_capture_write_loc(
+            out=self.cuda_graph_out_cache_loc_full_physical,
+            forward_batch=forward_batch,
+            width=self.cuda_graph_out_cache_loc_full_physical.numel(),
+        )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Init auxiliary variables for triton attention backend."""

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

@@ -766,19 +766,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         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
-            n = out_cache_loc.shape[0]
-            dst = self.cuda_graph_out_cache_loc_kernel[:n]
-            dst.copy_(out_cache_loc)
-            # Replay-prep receives the RAW (unpadded) out_cache_loc
-            # (build_replay_fb_view), but the captured write kernel consumes the
-            # full captured tier of this buffer. Zero the tail so pad rows write
-            # to the sink (row 0) instead of stale kernel-facing locs left by
-            # earlier larger replays — a stale tail scatters pad-row garbage into
-            # live KV pages. Mirrors the runner's PaddingPolicy.ZERO on its own
-            # out_cache_loc slot.
-            self.cuda_graph_out_cache_loc_kernel[n:].zero_()
-            self._decode_kernel_loc = dst
+            # The captured kernel consumes the whole buffer, so the tail a
+            # shorter replay leaves must go to slot 0 rather than live pages.
+            self._decode_kernel_loc = self.kv_index_translator.fill_capture_write_loc(
+                out=self.cuda_graph_out_cache_loc_kernel,
+                forward_batch=forward_batch,
+                width=self.cuda_graph_out_cache_loc_kernel.numel(),
+            )
         else:
             self._decode_kernel_loc = None
 

python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py

@@ -331,10 +331,13 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         loc: torch.Tensor,
         *,
         out: Optional[torch.Tensor] = None,
+        out_width: Optional[int] = None,
     ) -> torch.Tensor:
         """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this
         composite at argument validation, so it coincides with the read translate."""
-        return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
+        return self.full_attn_allocator.translate_write_loc_for_kernel(
+            loc, out=out, out_width=out_width
+        )
 
     @property
     def swa_kernel_page_multiplier(self) -> int:

python/sglang/srt/mem_cache/allocator/unified_mamba.py

@@ -260,9 +260,12 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         loc: torch.Tensor,
         *,
         out: Optional[torch.Tensor] = None,
+        out_width: Optional[int] = None,
     ) -> torch.Tensor:
         """Widened virtual WRITE loc -> DENSE id; see the sub-allocator's copy."""
-        return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
+        return self.full_attn_allocator.translate_write_loc_for_kernel(
+            loc, out=out, out_width=out_width
+        )
 
     def translate_kv_indices_for_transfer(
         self, kv_indices: torch.Tensor

python/sglang/srt/mem_cache/allocator/unified_sub_pool.py

@@ -43,6 +43,7 @@ from sglang.kernels.ops.memory.virtual_slot import (
     alloc_bind_inplace,
     bind_inplace,
     free_unbind_inplace,
+    write_loc_to_kernel_ids,
 )
 from sglang.srt.environ import envs
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -990,39 +991,47 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         clamp to kernel-facing id 0, the page-0 sink. int64 out; a consumer whose
         kernel ABI wants int32 narrows where it fills that buffer.
         """
-        ps = self.pool_page_size
-        stride = ps * self.kernel_page_multiplier
         with record_function("MultiEndedAlloc.translate_kv_loc_for_kernel"):
-            pages = virt_tokens if ps == 1 else virt_tokens // ps
-            offsets = None if ps == 1 else virt_tokens % ps
-            if out is None:
-                phys = self.virtual_to_physical[pages]
-                ids = phys * stride if offsets is None else phys * stride + offsets
-                return ids.clamp_(min=0)
+            return self._translate_loc_fused(virt_tokens, dcp_size=1, out=out)
+
+    def _translate_loc_fused(
+        self,
+        loc: torch.Tensor,
+        *,
+        dcp_size: int,
+        dcp_rank: int = 0,
+        out: Optional[torch.Tensor] = None,
+        out_width: Optional[int] = None,
+    ) -> torch.Tensor:
+        """One launch for the read and write conversions alike; see
+        `write_loc_to_kernel_ids`."""
+        if out is not None:
             assert out.dtype == torch.int64, (
                 f"translate_kv_loc_for_kernel: out= dtype must be int64 (matches v2p), "
                 f"got {out.dtype}"
             )
-            assert out.shape == virt_tokens.shape, (
-                f"translate_kv_loc_for_kernel: out= shape {tuple(out.shape)} must "
-                f"match virt_tokens shape {tuple(virt_tokens.shape)}"
-            )
-            if pages.dtype != torch.int64:
-                pages = pages.to(torch.int64)
-            if pages is virt_tokens:
-                out.copy_(torch.take(self.virtual_to_physical, pages))
-            else:
-                torch.take(self.virtual_to_physical, pages, out=out)
-            out.mul_(stride)
-            if offsets is not None:
-                out.add_(offsets)
-            return out.clamp_(min=0)
+            if out_width is None:
+                assert out.shape == loc.shape, (
+                    f"translate_kv_loc_for_kernel: out= shape {tuple(out.shape)} must "
+                    f"match virt_tokens shape {tuple(loc.shape)}"
+                )
+        return write_loc_to_kernel_ids(
+            loc=loc,
+            v2p=self.virtual_to_physical,
+            page_size=self.pool_page_size,
+            stride=self.pool_page_size * self.kernel_page_multiplier,
+            dcp_size=dcp_size,
+            dcp_rank=dcp_rank,
+            out=out,
+            out_width=out_width,
+        )
 
     def translate_write_loc_for_kernel(
         self,
         widened_loc: torch.Tensor,
         *,
         out: Optional[torch.Tensor] = None,
+        out_width: Optional[int] = None,
     ) -> torch.Tensor:
         """Widened virtual WRITE loc (`out_cache_loc`) -> kernel-facing id.
 
@@ -1032,16 +1041,14 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
         """
         parallel = get_parallel()
         dcp_size = parallel.attn_dcp_size if self.shards_under_dcp else 1
-        if dcp_size == 1:
-            return self.translate_kv_loc_for_kernel(widened_loc, out=out)
         with record_function("MultiEndedAlloc.translate_write_loc_for_kernel"):
-            owned = (widened_loc % dcp_size) == parallel.attn_dcp_rank
-            dense = self.translate_kv_loc_for_kernel(widened_loc // dcp_size)
-            dense = torch.where(owned, dense, torch.zeros_like(dense))
-            if out is not None:
-                out.copy_(dense)
-                return out
-            return dense
+            return self._translate_loc_fused(
+                widened_loc,
+                dcp_size=dcp_size,
+                dcp_rank=parallel.attn_dcp_rank,
+                out=out,
+                out_width=out_width,
+            )
 
     # -- alloc --
 

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -432,20 +432,55 @@ class KVIndexTranslator:
 
     def rebind_write_loc(self, forward_batch) -> None:
         """Phase 1 of the WRITE contract: translate the batch's write loc to
-        FULL-side kernel-facing ids exactly once, at ForwardBatch
-        construction. No-op on non-unified pools.
+        FULL-side kernel-facing ids, once, at ForwardBatch construction.
 
         REBIND, never mutate: the translate returns a FRESH tensor, so the
         ScheduleBatch's aliased tensor stays VIRTUAL for the radix / accept /
-        in-flight machinery that reads it.
+        in-flight machinery that reads it. The pre-translate tensor stays on
+        the batch for `fill_capture_write_loc`.
         """
         self._index_table_memo = None
         if not self.is_translating or forward_batch.out_cache_loc is None:
             return
+        forward_batch.out_cache_loc_virtual = forward_batch.out_cache_loc
         forward_batch.out_cache_loc = self._translate_write_full(
             forward_batch.out_cache_loc
         )
 
+    def fill_capture_write_loc(
+        self,
+        *,
+        out: torch.Tensor,
+        forward_batch,
+        width: Optional[int] = None,
+    ) -> Optional[torch.Tensor]:
+        """Translate this batch's WRITE loc straight into ``out``, a backend's
+        capture-stable buffer, and return the live ``[:n]`` view. One launch
+        fills the live prefix and clears the tail a shorter replay leaves;
+        None when this pool needs no translation.
+
+        Must run at metadata-init time: `out` is reused every step, so filling
+        it sooner would race a still-pending previous step under overlap
+        scheduling.
+        """
+        if not self.is_translating:
+            return None
+        virtual = forward_batch.out_cache_loc_virtual
+        if virtual is None:
+            loc = forward_batch.out_cache_loc
+            if loc is None:
+                return None
+            # The runner builds the capture batch outside `init_new`, so no
+            # rebind marked its virtual source; bake it holding sink ids.
+            width = int(loc.numel()) if width is None else int(width)
+            out[:width].zero_()
+            return out[: int(loc.numel())]
+        n = int(virtual.numel())
+        width = n if width is None else int(width)
+        buf = out[:width]
+        self._translate_write_full(virtual, out=buf, out_width=width)
+        return buf[:n]
+
     def sliding_window_write_loc_for(
         self, out_cache_loc: Optional[torch.Tensor]
     ) -> Optional[torch.Tensor]:

python/sglang/srt/mem_cache/memory_pool.py

@@ -1385,6 +1385,28 @@ class HybridReqToTokenPool(ReqToTokenPool):
     def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor:
         return self.req_index_to_mamba_index_mapping[req_indices]
 
+    @property
+    def mamba_v2p_table(self) -> Optional[torch.Tensor]:
+        """The mamba virtual->physical slot table, or None when the ids this
+        pool hands out are already physical."""
+        return None
+
+    @property
+    def mamba_translate_is_fusable(self) -> bool:
+        """Whether `fused_replay_state_indices` can reproduce this pool's
+        `translate_mamba_indices` in its own launch.
+
+        The kernel expresses exactly two shapes: the identity, and one gather
+        through `mamba_v2p_table`. A subclass that replaces the translate with
+        anything else is excluded here rather than silently mis-served.
+        """
+        if self.mamba_v2p_table is not None:
+            return True
+        return (
+            type(self).translate_mamba_indices
+            is HybridReqToTokenPool.translate_mamba_indices
+        )
+
     def translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor:
         """Virtual->physical mamba-slot translate. Identity for a static pool
         (slots are physical); UnifiedHybridReqToTokenPool overrides it for the

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -1107,6 +1107,15 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
                 )
             )
 
+    @property
+    def mamba_v2p_table(self) -> Optional[torch.Tensor]:
+        """This pool's ids ARE virtual; page_size is 1, so the translate is the
+        plain gather this table serves, which keeps `mamba_translate_is_fusable`
+        true despite the override."""
+        if self.mamba_allocator is None:
+            return None
+        return self.mamba_allocator.virtual_to_physical
+
     def translate_mamba_indices(self, virtual_ids: torch.Tensor) -> torch.Tensor:
         """Virtual mamba ids -> physical slot ids."""
         return self.mamba_allocator.translate(virtual_ids).to(torch.int32)

python/sglang/srt/model_executor/forward_batch_info.py

@@ -414,6 +414,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     # The original sequence length without being chunked. Qwen-1M related.
     orig_seq_lens: Optional[torch.Tensor] = None
 
+    # The write loc before `rebind_write_loc` replaced it with kernel-facing
+    # ids; a backend re-derives from it into its capture-stable buffer.
+    out_cache_loc_virtual: Optional[torch.Tensor] = None
     # DSV4-NPU only: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator,
     # consumed by the Ascend backend for PA_ND block tables. None elsewhere.
     out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None
@@ -1836,6 +1839,8 @@ def build_inner_fb_view(
         seq_lens_cpu=forward_batch.seq_lens_cpu,
         encoder_lens=encoder_lens,
         out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
+        # A caller may hand in another view that does not carry this field.
+        out_cache_loc_virtual=getattr(forward_batch, "out_cache_loc_virtual", None),
         out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
         spec_info=forward_batch.spec_info,
     )

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -192,6 +192,7 @@ def build_replay_fb_view(
         num_padding=bs - raw_bs,
         encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None,
         out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
+        out_cache_loc_virtual=forward_batch.out_cache_loc_virtual,
         out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
         # The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE
         # for the backend, which copies the result into its own static buffer and