[unified-memory] Stride the KV translate kernel and route every translate through it (6/7)

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

@@ -15,7 +15,7 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import triton
@@ -203,9 +203,16 @@ def write_loc_to_kernel_id_kernel(
     out_ptr,  # out: [N] int64 — physical token ids
     N,  # runtime: live element count
     W,  # runtime: lanes to write; [N, W) get 0
+    num_v_pages,  # runtime: rows of v2p; a page at or past it is unmapped
+    cols,  # runtime: row width of the (rows, cols) view; STRIDED only
+    loc_row_stride,  # runtime: element strides of `loc`; STRIDED only
+    loc_col_stride,
+    out_row_stride,  # runtime: element strides of `out`; STRIDED only
+    out_col_stride,
     PAGE_SIZE: tl.constexpr,
     DCP_SIZE: tl.constexpr,
     DCP_RANK: tl.constexpr,
+    STRIDED: tl.constexpr,
     BLOCK: tl.constexpr,
 ):
     """``physical(t) = v2p[t // ps] * ps + t % ps``, clamped at 0.
@@ -221,12 +228,25 @@ def write_loc_to_kernel_id_kernel(
     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.
+
+    ``STRIDED`` walks `loc` and `out` as a ``(rows, cols)`` view, so a column
+    slice of a wider page table addresses correctly; a contiguous pair takes
+    the flat path and the stride arguments go unread.
     """
     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)
+    if STRIDED:
+        wide = offs.to(tl.int64)
+        row = wide // cols
+        col = wide - row * cols
+        loc_off = row * loc_row_stride + col * loc_col_stride
+        out_off = row * out_row_stride + col * out_col_stride
+    else:
+        loc_off = offs
+        out_off = offs
+    loc = tl.load(loc_ptr + loc_off, mask=mask, other=0).to(tl.int64)
 
     keep = mask & (loc >= 0)
     if DCP_SIZE > 1:
@@ -235,10 +255,19 @@ def write_loc_to_kernel_id_kernel(
 
     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.
+    # Torch's gather bounds-checks and `tl.load` does not, so an id past the
+    # table is excluded here rather than read out of bounds.
+    keep = keep & (page < num_v_pages)
     phys = tl.load(v2p_ptr + tl.where(keep, page, 0), mask=mask, other=0).to(tl.int64)
     ids = tl.maximum(phys * PAGE_SIZE + offset, 0)
-    tl.store(out_ptr + offs, tl.where(keep, ids, 0), mask=in_range)
+    tl.store(out_ptr + out_off, tl.where(keep, ids, 0), mask=in_range)
+
+
+def _row_col_strides(t: torch.Tensor) -> Tuple[int, int]:
+    """Element strides of `t` seen as ``(rows, cols)``; 0-d and 1-D are one row."""
+    if t.dim() == 2:
+        return t.stride(0), t.stride(1)
+    return 0, (t.stride(0) if t.dim() == 1 else 1)
 
 
 def write_loc_to_kernel_ids(
@@ -260,28 +289,28 @@ def write_loc_to_kernel_ids(
 
     ``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.
+    tail here. That mode addresses a packed lane range, so it takes a
+    contiguous ``loc``; everywhere else a 1-D or 2-D strided view is legal.
     """
     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)
+        out = torch.empty(loc.shape, dtype=torch.int64, device=loc.device)
     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(), (
+        assert out.shape == loc.shape, (
             f"write_loc_to_kernel_ids: out shape {tuple(out.shape)} must match "
             f"loc shape {tuple(loc.shape)}"
         )
     else:
+        assert loc.is_contiguous(), (
+            f"write_loc_to_kernel_ids: out_width writes a packed lane range, "
+            f"so loc must be packed; got stride {tuple(loc.stride())}"
+        )
         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)}"
@@ -293,29 +322,55 @@ def write_loc_to_kernel_ids(
     if width == 0:
         return out
     if not loc.is_cuda:
-        # Pure-torch reference; the allocator's unit tests run on CPU.
+        # Pure-torch reference; taken by the CPU unit tests and by every
+        # non-CUDA accelerator.
         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)
+        page = torch.div(big, page_size, rounding_mode="floor")
+        keep = keep & (page < v2p.numel())
+        page = torch.where(keep, page, 0)
         offset = big % page_size if page_size > 1 else 0
         ids = (v2p[page] * page_size + offset).clamp_(min=0)
-        out[:N].copy_(torch.where(keep, ids, torch.zeros_like(ids)))
-        if width > N:
-            out[N:width].zero_()
+        ids = torch.where(keep, ids, torch.zeros_like(ids))
+        if out_width is None:
+            out.copy_(ids)
+        else:
+            out[:N].copy_(ids)
+            if width > N:
+                out[N:width].zero_()
         return out
+    flat = loc.is_contiguous() and out.is_contiguous()
+    if flat:
+        cols = loc_row_stride = loc_col_stride = 1
+        out_row_stride = out_col_stride = 1
+    else:
+        assert loc.dim() <= 2 and out.dim() <= 2, (
+            f"write_loc_to_kernel_ids: a strided loc/out is walked as 2-D, got "
+            f"{loc.dim()}-D loc and {out.dim()}-D out; pass a contiguous view"
+        )
+        cols = loc.shape[-1] if loc.dim() == 2 else N
+        loc_row_stride, loc_col_stride = _row_col_strides(loc)
+        out_row_stride, out_col_stride = _row_col_strides(out)
     grid = (triton.cdiv(width, WRITE_LOC_BLOCK),)
     write_loc_to_kernel_id_kernel[grid](
         loc,
         v2p,
         out,
         N,
         width,
+        int(v2p.numel()),
+        cols,
+        loc_row_stride,
+        loc_col_stride,
+        out_row_stride,
+        out_col_stride,
         PAGE_SIZE=page_size,
         DCP_SIZE=dcp_size,
         DCP_RANK=dcp_rank,
+        STRIDED=not flat,
         BLOCK=WRITE_LOC_BLOCK,
     )
     return out

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

@@ -788,10 +788,11 @@ def _init_full_cg_prefill_metadata(
                 "full-CG prefill SWA write-location buffer",
             )
             # Under the unified pool `out_cache_loc` was rebound to FULL-side
-            # KERNEL-FACING ids at ForwardBatch construction, so the full->swa
-            # map cannot be re-run on it -- those values index far past the swa
-            # v2p table (a device-side "index out of bounds" assert). Phase 2 of
-            # the write contract derives the swa loc from them instead.
+            # physical ids at ForwardBatch construction, so the full->swa map
+            # cannot be re-run on it: a full physical page number is also a
+            # valid virtual page of the swa v2p, so the map would silently
+            # resolve it to an unrelated token's swa slot. Phase 2 of the write
+            # contract derives the swa loc from them instead.
             swa_write_loc = (
                 self.kv_index_translator.sliding_window_write_loc_for(
                     forward_batch.out_cache_loc

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

@@ -165,7 +165,7 @@ def __init__(
         # Full-attention transfers use virtual IDs. SWA transfers already use
         # kernel-facing IDs from translate_loc_from_full_to_swa.
         kvcache.full_kv_pool.host_transfer_translate = (
-            self.full_attn_allocator.translate_kv_loc_for_kernel
+            self.full_attn_allocator.translate_kv_loc
         )
 
         self.free_group = None
@@ -287,8 +287,8 @@ def translate_loc_from_full_to_swa(
         *,
         out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids."""
-        return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out)
+        """SWA-layer read path: virtual TOKEN ids -> swa-physical TOKEN ids."""
+        return self.swa_attn_allocator.translate_kv_loc(kv_indices, out=out)
 
     @property
     def full_v2p_page_table(self) -> torch.Tensor:
@@ -316,14 +316,9 @@ def translate_kv_indices_for_transfer(
     def translate_swa_indices_for_transfer(
         self, kv_indices: torch.Tensor
     ) -> torch.Tensor:
-        """Virtual TOKEN ids -> swa-sub-pool PHYSICAL token ids.
-
-        The SWA counterpart of the above. `translate_loc_from_full_to_swa`
-        cannot serve here: it returns KERNEL-FACING ids (the physical page
-        scaled by the sub-pool's per-page block count), which index the
-        per-layer views, whereas the SWA state component is registered as whole
-        page envelopes and addressed by physical page.
-        """
+        """Virtual TOKEN ids -> swa-sub-pool PHYSICAL token ids, for the PD
+        transfer engine: the same translate as `translate_loc_from_full_to_swa`,
+        on int64 ids."""
         return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
 
     def _move_gate_targets(self):
@@ -377,15 +372,6 @@ def set_host_transfer_move_gate(self, gate: Callable[[], bool]) -> None:
             lazy_compaction=self.lazy_compaction,
         )
 
-    def translate_kv_loc_for_kernel(
-        self,
-        loc: torch.Tensor,
-        *,
-        out: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Full-pool virtual TOKEN ids -> kernel-facing ids."""
-        return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
-
     def translate_write_loc_for_kernel(
         self,
         loc: torch.Tensor,

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

@@ -129,7 +129,7 @@ def __init__(
 
         # HiCache indexes the full sub-pool's per-layer views with kernel-facing IDs.
         kvcache.full_kv_pool.host_transfer_translate = (
-            self.full_attn_allocator.translate_kv_loc_for_kernel
+            self.full_attn_allocator.translate_kv_loc
         )
 
     # -- size: dynamic --
@@ -255,8 +255,9 @@ def translate_kv_loc(
         *,
         out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """Full-pool virtual TOKEN ids -> physical TOKEN ids; `-1` passes through as
-        `-1` (padding downstream). ``out=`` supports cuda-graph buffer stability."""
+        """Full-pool virtual TOKEN ids -> physical TOKEN ids; an unmapped or
+        negative id lands on 0, the sink. ``out=`` supports cuda-graph buffer
+        stability."""
         result = self.full_attn_allocator.translate_kv_loc(loc, out=out)
         return result
 
@@ -272,15 +273,6 @@ def full_p2v_page_table(self) -> torch.Tensor:
         """Page-level physical->virtual table of the full sub-pool."""
         return self.full_attn_allocator.physical_to_virtual
 
-    def translate_kv_loc_for_kernel(
-        self,
-        loc: torch.Tensor,
-        *,
-        out: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Full-pool virtual TOKEN ids -> kernel-facing ids."""
-        return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
-
     def translate_write_loc_for_kernel(
         self,
         loc: torch.Tensor,

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

@@ -993,76 +993,21 @@ def translate_kv_loc(
         *,
         out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """Translate token-granular virtual ids to physical ids.
+        """Virtual token ids -> physical token ids, which under the token-major
+        views ARE the kernel-facing ids:
+
+            kernel_id(t) = v2p[t // ps] * ps + t % ps
 
         Under DCP the input is the DCP-collapsed id (`widened // dcp_size`, what
         `KVIndexTranslator.translate_dcp_read_ids` hands down), so this works on
-        `pool_page_size`. ``out=`` writes in-place into a caller-owned buffer,
-        required under cuda-graph capture: the captured graph records the gather
-        against a fixed ``data_ptr``.
+        `pool_page_size`. An unmapped, out-of-range or negative id resolves to
+        0, the page-0 sink every kernel skips. ``out=`` writes in-place into a
+        caller-owned buffer, required under cuda-graph capture: the captured
+        graph records the gather against a fixed ``data_ptr``. int64 out; a
+        consumer whose kernel ABI wants int32 narrows where it fills that
+        buffer.
         """
-        if out is not None:
-            assert out.dtype == torch.int64, (
-                f"translate_kv_loc: out= dtype must be int64 (matches v2p), "
-                f"got {out.dtype}"
-            )
-            assert out.shape == virt_tokens.shape, (
-                f"translate_kv_loc: out= shape {tuple(out.shape)} must match "
-                f"virt_tokens shape {tuple(virt_tokens.shape)}"
-            )
         with record_function("MultiEndedAlloc.translate_kv_loc"):
-            return self._translate_kv_loc_impl(virt_tokens, out)
-
-    def _translate_kv_loc_impl(
-        self,
-        virt_tokens: torch.Tensor,
-        out: Optional[torch.Tensor],
-    ) -> torch.Tensor:
-        # Tombstone-safety clamp: a tombstoned v2p entry (-1) must not reach
-        # `k_buffer[-1]` (illegal access under captured graph replay). Clamping to
-        # 0 routes it to physical slot 0, reserved sink space holding no real data.
-        ps = self.pool_page_size
-        if ps == 1:
-            if out is not None:
-                # `index_select(out=out)` forbids index/out aliasing, but the
-                # canonical caller passes `out=kv_indices` in place.
-                tmp = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
-                tmp = torch.clamp_min(tmp, 0)
-                out.copy_(tmp)
-                return out
-            result = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
-            return torch.clamp_min(result, 0)
-        # ps > 1: page math. `virt_pages`/`offsets` are fresh, so they
-        # cannot alias `out` -- `index_select(out=out)` is safe.
-        virt_pages = virt_tokens // ps
-        offsets = virt_tokens % ps
-        if out is not None:
-            torch.index_select(self.virtual_to_physical, 0, virt_pages, out=out)
-            out.mul_(ps)
-            out.add_(offsets)
-            out.clamp_(min=0)  # tombstoned page: -1*ps + offset in [-ps, -1]
-            return out
-        phys_pages = self.virtual_to_physical[virt_pages]
-        result = phys_pages * ps + offsets
-        return torch.clamp_min(result, 0)
-
-    def translate_kv_loc_for_kernel(
-        self,
-        virt_tokens: torch.Tensor,
-        *,
-        out: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Virtual token ids -> kernel-facing ids, which under the token-major
-        views are the physical token ids:
-
-            kernel_id(t) = v2p[t // ps] * ps + t % ps
-
-        Same id as `translate_kv_loc` for any mapped virtual token; this one
-        is one Triton launch, and it is the path that sends an unmapped or
-        negative loc to id 0, the page-0 sink. int64 out; a consumer whose
-        kernel ABI wants int32 narrows where it fills that buffer.
-        """
-        with record_function("MultiEndedAlloc.translate_kv_loc_for_kernel"):
             return self._translate_loc_fused(virt_tokens, dcp_size=1, out=out)
 
     def _translate_loc_fused(
@@ -1075,17 +1020,8 @@ def _translate_loc_fused(
         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}"
-            )
-            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)}"
-                )
+        `write_loc_to_kernel_ids`, which owns the `out=` dtype and shape
+        contract for both."""
         return write_loc_to_kernel_ids(
             loc=loc,
             v2p=self.virtual_to_physical,

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -124,7 +124,7 @@ def __init__(
             self._capture_page_size = alloc.page_size
             self._full_v2p_table = alloc.full_v2p_page_table
             self._full_p2v_table = alloc.full_p2v_page_table
-            self._translate_full = alloc.translate_kv_loc_for_kernel
+            self._translate_full = alloc.translate_kv_loc
             # The WRITE loc is the one id that arrives DCP-WIDENED: read indices
             # are collapsed by the DCP index kernels, `out_cache_loc` still
             # carries the owner rule in `loc % dcp_size`. Identity with the read