minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

refactor(unified-memory): translate the KV write location once, at ForwardBatch construction (#35245)

Browse Source 
Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
caihuali95
2026-08-30 23:52:14 -07:00
committed by GitHub
co-authored by Caihua Li Cheng Wan
parent 4f997a432a
commit 29578d5578
28 changed files with 1837 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/kvcache/__init__.py
+1
View File
@@ -62,6 +62,7 @@ _TRITON_KERNELS = [
("cache_ops", "concat_and_cast_mha_k_triton"),
("cache_ops", "launch_reshape_and_cache_flash"),
("pd_dcp_gather", "copy_mla_rows_into_pack"),
("kv_read_table", "build_kv_read_table"),
("kv_indices", "create_flashinfer_kv_indices_triton"),
("kv_indices", "create_flashmla_kv_indices_triton"),
("kv_indices", "create_chunked_prefix_cache_kv_indices"),
python/sglang/kernels/ops/kvcache/kv_indices.py
+34 -10
View File
@@ -7,19 +7,33 @@ FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SI
@triton.jit
def create_flashinfer_kv_indices_triton(
req_to_token_ptr, # [max_batch, max_context_len]
req_to_token_ptr, # [max_batch, max_context_len] token table; at
# ENTRY_PAGE_SIZE > 1 a PAGE-granular table (the unified pool's read table)
req_pool_indices_ptr,
page_kernel_lens_ptr,
kv_indptr,
kv_start_idx,
kv_indices_ptr,
req_to_token_ptr_stride: tl.constexpr,
# Runtime, not constexpr: the translator's eager table is allocated at the
# batch's live width, so a constexpr stride would JIT-specialize per width
# (a recompile every few decode steps at small page sizes).
req_to_token_ptr_stride,
ENTRY_PAGE_SIZE: tl.constexpr = 1,
):
"""Gather per-request token ids into a flat CSR kv_indices stream.
``ENTRY_PAGE_SIZE == 1`` (default): the source table is token-granular and
entries are emitted verbatim -- byte-identical to the historical kernel.
``ENTRY_PAGE_SIZE == ps``: the source is the translator's PAGE-granular
read table (entries already kernel-facing page ids); token ids are rebuilt
as ``token = entry * ps + pos % ps``, exact because converting an id keeps
its offset inside the page.
"""
BLOCK_SIZE: tl.constexpr = 512
pid = tl.program_id(axis=0)
# find the req pool idx, this is for batch to token
req_pool_index = tl.load(req_pool_indices_ptr + pid)
req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
kv_indices_offset = tl.load(kv_indptr + pid)
kv_start = 0
@@ -34,13 +48,23 @@ def create_flashinfer_kv_indices_triton(
# index into req_to_token_ptr needs to be int64
offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE
mask = offset < kv_end - kv_start
data = tl.load(
req_to_token_ptr
+ req_pool_index * req_to_token_ptr_stride
+ kv_start
+ offset,
mask=mask,
)
if ENTRY_PAGE_SIZE == 1:
data = tl.load(
req_to_token_ptr
+ req_pool_index * req_to_token_ptr_stride
+ kv_start
+ offset,
mask=mask,
)
else:
pos = kv_start + offset
entry = tl.load(
req_to_token_ptr
+ req_pool_index * req_to_token_ptr_stride
+ pos // ENTRY_PAGE_SIZE,
mask=mask,
)
data = entry.to(tl.int64) * ENTRY_PAGE_SIZE + pos % ENTRY_PAGE_SIZE
tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)
python/sglang/kernels/ops/kvcache/kv_read_table.py
+140
View File
@@ -0,0 +1,140 @@
# 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.
# ==============================================================================
"""Builds the per-batch read table for the unified memory pool.
One fused gather-and-translate. For each request row it reads the virtual ids
out of `req_to_token`, converts each to the id the kernels can use, and writes
the result into `out`:
out[b, c] = clamp(v2p[req_to_token[req[b], c * ps] // ps] * multiplier, 0)
for c < ceil(seq_lens[b] / ps) -- the row's LIVE prefix
`v2p` is the pool's virtual->physical page table and `multiplier` scales a
physical page into the id space the per-layer views use (1 when they are not
dense). Since only the page number is rewritten, a token-level consumer can
rebuild flat ids as `entry * ps + offset`.
PREFIX-ONLY per row: columns past the live prefix are never written, so a
caller-owned buffer keeps what it had there -- which is what lets a captured
cuda-graph buffer be refreshed in place. Readers bound themselves by
`cache_seqlens` and never look past the prefix.
A `-1` in `req_to_token` and a freed (`-1`) v2p row both clamp to entry 0, the
reserved padding slot, so a kernel dereferences padding, not a wild address.
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
_BLOCK_COLS = 256
@triton.jit
def build_kv_read_table_kernel(
req_to_token_ptr, # in: [max_reqs, max_context] -- VIRTUAL token ids
req_pool_indices_ptr, # in: [bs] -- row per batch lane
seq_lens_ptr, # in: [bs]
v2p_ptr, # in: [num_pages + 1] int64 -- virtual->physical page table
out_ptr, # out: [>=bs, >=max_pages] int32 -- the read table
req_stride, # runtime: req_to_token row stride (elements)
out_stride, # runtime: out row stride (elements)
mult, # runtime: kernel_page_multiplier of the target sub-pool
PAGE_SIZE: tl.constexpr,
BLOCK: tl.constexpr,
):
bid = tl.program_id(0)
blk = tl.program_id(1)
req = tl.load(req_pool_indices_ptr + bid).to(tl.int64)
seqlen = tl.load(seq_lens_ptr + bid)
n_pages = (seqlen + PAGE_SIZE - 1) // PAGE_SIZE
cols = blk * BLOCK + tl.arange(0, BLOCK)
mask = cols < n_pages
tok = tl.load(
req_to_token_ptr + req * req_stride + cols.to(tl.int64) * PAGE_SIZE,
mask=mask,
other=0,
).to(tl.int64)
# Triton's `//` truncates toward zero, so `-1 // ps` is 0 for ps > 1 but
# -1 at ps == 1, which would read one element BEFORE `v2p`.
page = tl.where(tok < 0, 0, tok // PAGE_SIZE)
phys = tl.load(v2p_ptr + page, mask=mask, other=0)
entry = tl.maximum(phys * mult, 0).to(tl.int32)
tl.store(out_ptr + bid.to(tl.int64) * out_stride + cols, entry, mask=mask)
def build_kv_read_table(
*,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
v2p: torch.Tensor,
multiplier: int,
page_size: int,
max_pages: int,
out: torch.Tensor,
) -> torch.Tensor:
"""Fill ``out``'s live prefix with read-table entries.
``out`` is caller-owned (fresh zeros for the eager path, the module's
capture-stable buffer for replay) and only its ``[:bs, :max_pages]``
region's live prefix is written -- never rebound, never tail-cleared.
"""
bs = int(req_pool_indices.numel())
assert (
out.dtype == torch.int32
), f"build_kv_read_table: out must be int32, got {out.dtype}"
assert out.dim() == 2 and out.shape[0] >= bs and out.shape[1] >= max_pages, (
f"build_kv_read_table: out {tuple(out.shape)} cannot hold "
f"(bs={bs}, max_pages={max_pages})"
)
assert out.stride(1) == 1, "build_kv_read_table: out rows must be packed"
assert (max_pages - 1) * page_size < req_to_token.shape[1], (
f"build_kv_read_table: max_pages={max_pages} x ps={page_size} "
f"exceeds req_to_token width {req_to_token.shape[1]}"
)
if bs == 0 or max_pages == 0:
return out
if not req_to_token.is_cuda:
cols = torch.arange(max_pages, device=req_to_token.device)
live = cols[None, :] < (
(seq_lens[:bs, None].to(torch.int64) + page_size - 1) // page_size
)
tok = req_to_token[
req_pool_indices[:bs, None].to(torch.int64), (cols * page_size)[None, :]
].to(torch.int64)
pages = torch.where(tok < 0, 0, tok // page_size)
entry = (v2p[pages] * multiplier).clamp(min=0).to(torch.int32)
dst = out[:bs, :max_pages]
dst.copy_(torch.where(live, entry, dst))
return out
grid = (bs, triton.cdiv(max_pages, _BLOCK_COLS))
build_kv_read_table_kernel[grid](
req_to_token,
req_pool_indices,
seq_lens,
v2p,
out,
req_to_token.stride(0),
out.stride(0),
multiplier,
PAGE_SIZE=page_size,
BLOCK=_BLOCK_COLS,
)
return out
python/sglang/srt/arg_groups/kv_cache_hook.py
+15 -5
View File
@@ -262,11 +262,21 @@ def handle_unified_memory_pool(server_args: Any) -> None:
# Only monolithic decode cuda-graph capture is wired; piecewise prefill
# capture is not. Guard when the user opts into it.
_cg_cfg = cfg.cuda_graph_config
if _cg_cfg is not None and _cg_cfg.prefill.backend == Backend.TC_PIECEWISE:
raise ValueError(
"--enable-unified-memory supports monolithic (decode) "
"cuda-graph capture only; disable piecewise prefill capture "
"(e.g. --cuda-graph-backend-prefill=disabled)."
if _cg_cfg is not None and _cg_cfg.prefill.backend != Backend.DISABLED:
if cfg.cuda_graph_backend_prefill is not None:
raise ValueError(
"--enable-unified-memory supports decode cuda-graph "
"capture only; prefill capture is not wired (the prefill "
"graph runner bypasses the unified virtual->physical loc "
"rebind). Got --cuda-graph-backend-prefill="
f"{cfg.cuda_graph_backend_prefill!r}; pass "
"--cuda-graph-backend-prefill=disabled."
)
_cg_cfg.prefill.backend = Backend.DISABLED
logger.warning(
"--enable-unified-memory: disabling prefill cuda-graph "
"capture (not wired for the unified pool's loc rebind); "
"decode capture is unaffected."
)
python/sglang/srt/layers/attention/triton_backend.py
+90 -141
View File
@@ -29,6 +29,7 @@ from sglang.srt.layers.dcp import (
get_dcp_lens,
)
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.cuda_graph_config import (
@@ -208,14 +209,8 @@ class TritonAttnBackend(AttentionBackend):
# Lets the Triton wrappers specialize on PAGE_SIZE; page_size=1 is
# byte-identical to the slot-based envelope.
self.page_size = getattr(model_runner, "page_size", 1) or 1
# Unified pool v2p hook (None = no-op): req_to_token holds VIRTUAL ids but
# kernels need the kernel-facing id space — PHYSICAL for MHA, DENSE for the
# per-layer-view MLA pool (translate_kv_loc_for_kernel falls back to the physical
# translate when kernel_page_multiplier == 1, so preferring it is exact for
# both). Applied eagerly so the captured graph has no translate.
self._translate_kv_loc = getattr(
self.token_to_kv_pool_allocator, "translate_kv_loc_for_kernel", None
) or getattr(self.token_to_kv_pool_allocator, "translate_kv_loc", None)
self.kv_index_translator = model_runner.kv_index_translator
self.kv_read_tables = None
self.num_draft_tokens = get_spec().speculative_num_draft_tokens
self.speculative_num_steps = get_spec().speculative_num_steps
self.topk = get_spec().speculative_eagle_topk or 0
@@ -466,19 +461,20 @@ class TritonAttnBackend(AttentionBackend):
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
index_table,
kv_indices: torch.Tensor,
) -> torch.Tensor:
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1:] = torch.cumsum(seq_lens, dim=0)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
index_table.ids,
index_table.row_ids,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
index_table.row_stride,
ENTRY_PAGE_SIZE=index_table.entry_page_size,
)
return kv_indptr
@@ -487,6 +483,7 @@ class TritonAttnBackend(AttentionBackend):
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
index_table,
):
"""Fill KV (and SWA) cuda-graph buffers for decode/idle mode.
@@ -495,6 +492,9 @@ class TritonAttnBackend(AttentionBackend):
``num_kv_splits_lens`` is the per-request length used to size kv splits
(per-DCP-rank length clamped to >=1 when DCP is enabled, full seq_lens
otherwise).
``index_table`` is the captured read-index view: under the unified pool the
gathers below read the converted tables.
"""
seq_lens = seq_lens[:bs]
req_pool_indices = req_pool_indices[:bs]
@@ -512,27 +512,20 @@ class TritonAttnBackend(AttentionBackend):
num_kv_splits_lens = dcp_seq_lens.clamp_min(1)
else:
kv_indptr = self._fill_kv_indptr_and_indices(
bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
bs, seq_lens, index_table, self.cuda_graph_kv_indices
)
# Unified pool: VIRTUAL ids written here are translated to PHYSICAL in
# init_forward_metadata_out_graph (replay-prep) so the captured graph
# carries zero translate nodes.
num_kv_splits_lens = seq_lens
window_kv_indptr = self.window_kv_indptr
window_kv_lens = None
if self.sliding_window_size is not None and self.sliding_window_size > 0:
# Unified pool: leave the window VIRTUAL too (translated alongside the
# full kv_indices later); baseline SWA keeps the eager window translate.
window_kv_indptr, _, window_kv_lens, _ = update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
index_table,
self.sliding_window_size,
seq_lens,
req_pool_indices,
bs,
token_to_kv_pool=self.token_to_kv_pool,
window_kv_indices=self.cuda_graph_window_kv_indices,
skip_full_to_swa_translation=(self._translate_kv_loc is not None),
)
return kv_indptr, window_kv_indptr, window_kv_lens, num_kv_splits_lens
@@ -540,8 +533,8 @@ class TritonAttnBackend(AttentionBackend):
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
spec_info,
index_table,
):
"""Fill all cuda-graph buffers for target_verify mode."""
# Prefer the spec_info's per-request query length (DSpark draft propose
@@ -561,7 +554,7 @@ class TritonAttnBackend(AttentionBackend):
device=self.device,
)
kv_indptr = self._fill_kv_indptr_and_indices(
bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
bs, seq_lens, index_table, self.cuda_graph_kv_indices
)
window_kv_indptr = self.window_kv_indptr
window_kv_indices = None
@@ -574,10 +567,9 @@ class TritonAttnBackend(AttentionBackend):
window_kv_indptr, window_kv_indices, _, window_kv_offsets[:bs] = (
update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
index_table,
self.sliding_window_size,
seq_lens[:bs],
req_pool_indices,
bs,
token_to_kv_pool=self.token_to_kv_pool,
window_kv_indices=window_kv_indices,
@@ -611,9 +603,9 @@ class TritonAttnBackend(AttentionBackend):
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
index_table,
):
"""Fill QO + KV cuda-graph buffers for draft_extend mode."""
seq_lens = seq_lens[:bs]
@@ -644,7 +636,7 @@ class TritonAttnBackend(AttentionBackend):
extend_seq_lens = torch.zeros(bs, dtype=torch.int32, device=seq_lens.device)
kv_lens = torch.clamp(seq_lens - extend_seq_lens, min=0).to(torch.int32)
kv_indptr = self._fill_kv_indptr_and_indices(
bs, kv_lens, req_pool_indices, self.cuda_graph_kv_indices
bs, kv_lens, index_table, self.cuda_graph_kv_indices
)
return qo_indptr, kv_indptr, num_tokens_per_req
@@ -689,7 +681,7 @@ class TritonAttnBackend(AttentionBackend):
forward_mode=forward_mode,
spec_info=spec_info,
)
out_cache_loc_full_physical = self._translate_cuda_graph_shared_pool_locs(
out_cache_loc_full_physical = self._fill_cuda_graph_write_locs(
forward_batch, bs
)
swa_out_cache_loc = self._fill_cuda_graph_swa_out_cache_loc(forward_batch)
@@ -709,15 +701,16 @@ class TritonAttnBackend(AttentionBackend):
spec_info=spec_info,
)
# Metadata view is reused from capture; just refill the buffers.
self._translate_cuda_graph_shared_pool_locs(forward_batch, bs)
self._fill_cuda_graph_write_locs(forward_batch, bs)
self._fill_cuda_graph_swa_out_cache_loc(forward_batch)
def _fill_cuda_graph_swa_out_cache_loc(
self, forward_batch: ForwardBatch
) -> Optional[torch.Tensor]:
"""Refill the SWA write-target buffer from live out_cache_loc, returning the
[:n] view (None for non-SWA / multi-step draft) so the captured store reads
fresh slots on replay."""
"""Refill the SWA write-target buffer from the batch's derived
sliding-window write loc, returning the [:n] view (None for non-SWA /
multi-step draft) so the captured store reads fresh slots on replay.
"""
if not self.use_sliding_window_kv_pool:
return None
out_cache_loc = forward_batch.out_cache_loc
@@ -726,68 +719,32 @@ class TritonAttnBackend(AttentionBackend):
or out_cache_loc.shape[0] > self.cuda_graph_swa_out_cache_loc.shape[0]
):
return None
swa_write_loc = self.kv_index_translator.sliding_window_write_loc_for(
out_cache_loc
)
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)
)
self.cuda_graph_swa_out_cache_loc[:n].copy_(swa_write_loc)
return self.cuda_graph_swa_out_cache_loc[:n]
def _translate_cuda_graph_shared_pool_locs(
def _fill_cuda_graph_write_locs(
self, forward_batch: ForwardBatch, bs: int
) -> Optional[torch.Tensor]:
"""Unified pool: eager v2p translate of the cuda-graph read+write LOC buffers,
run BEFORE graph.replay() reading the live post-compaction v2p, so the
captured graph carries zero translate nodes. No-op for non-unified pools.
"""Copy the cuda-graph WRITE loc into the capture-stable buffer and
return the ``[:n]`` view; no-op for non-unified pools.
Read buffers (full kv_indices, SWA window) are translated IN PLACE; the
full-attn WRITE loc is RETURNED as the [:n] view of the backend-owned
out_cache_loc_full_physical buffer. Eager .item() bounds are fine here
(out-of-graph), so no in-graph translate variant is needed.
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.
"""
if self._translate_kv_loc is None:
if not self.kv_index_translator.is_translating:
return None
# seq_lens_sum is the reliable "mirror present" signal: it is
# None-preserving into the replay view, unlike seq_lens_cpu (always a
# non-None but stale slice for gpu_only batches). None -> fall back to a
# per-step D2H `.item()` on the indptr.
have_cpu_mirror = forward_batch.seq_lens_sum is not None
# Full-attention read path. kv_indptr[bs] == seq_lens_sum.
n_kv = (
forward_batch.seq_lens_sum
if have_cpu_mirror
else int(self.kv_indptr[bs].item())
)
if n_kv > 0:
self.cuda_graph_kv_indices[:n_kv] = self._translate_kv_loc(
self.cuda_graph_kv_indices[:n_kv]
)
# SWA window read path. window_kv_indptr[bs] == sum(min(seq_len, window)).
if self.sliding_window_size is not None and self.sliding_window_size > 0:
if have_cpu_mirror:
n_win = int(
forward_batch.seq_lens_cpu[:bs]
.clamp(max=self.sliding_window_size)
.sum()
)
else:
n_win = int(self.window_kv_indptr[bs].item())
if n_win > 0:
self.cuda_graph_window_kv_indices[:n_win] = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.cuda_graph_window_kv_indices[:n_win]
)
)
# Full-attention write path: translate out_cache_loc -> physical into the
# capture-stable buffer and RETURN the [:n] view.
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_(
self._translate_kv_loc(out_cache_loc)
)
self.cuda_graph_out_cache_loc_full_physical[:n].copy_(out_cache_loc)
return self.cuda_graph_out_cache_loc_full_physical[:n]
def init_forward_metadata(self, forward_batch: ForwardBatch):
@@ -806,6 +763,9 @@ class TritonAttnBackend(AttentionBackend):
if forward_batch.forward_mode.is_decode_or_idle():
if spec_info is None or spec_info.kv_indptr is None:
index_table = self.kv_index_translator.index_table_for_batch(
forward_batch
)
# kv_indptr is None for draft-extend's idle batch; build from seq_lens.
if self.dcp_size > 1:
# DCP: per-rank sharded KV indices, else each rank reads the
@@ -826,11 +786,9 @@ class TritonAttnBackend(AttentionBackend):
kv_indptr = self._fill_kv_indptr_and_indices(
bs,
forward_batch.seq_lens,
forward_batch.req_pool_indices,
index_table,
kv_indices,
)
if self._translate_kv_loc is not None:
kv_indices = self._translate_kv_loc(kv_indices)
if (
self.sliding_window_size is not None
and self.sliding_window_size > 0
@@ -838,10 +796,9 @@ class TritonAttnBackend(AttentionBackend):
window_kv_indptr, window_kv_indices, window_kv_lens, _ = (
update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
index_table,
self.sliding_window_size,
forward_batch.seq_lens,
forward_batch.req_pool_indices,
bs,
self.device,
self.token_to_kv_pool,
@@ -931,10 +888,11 @@ class TritonAttnBackend(AttentionBackend):
kv_indices = torch.empty(
seq_lens_sum, dtype=torch.int64, device=self.device
)
index_table = self.kv_index_translator.index_table_for_batch(forward_batch)
kv_indptr = self._fill_kv_indptr_and_indices(
bs,
forward_batch.seq_lens,
forward_batch.req_pool_indices,
index_table,
kv_indices,
)
@@ -947,10 +905,9 @@ class TritonAttnBackend(AttentionBackend):
window_kv_offsets,
) = update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
index_table,
self.sliding_window_size,
forward_batch.seq_lens,
forward_batch.req_pool_indices,
bs,
self.device,
self.token_to_kv_pool,
@@ -969,6 +926,7 @@ class TritonAttnBackend(AttentionBackend):
attn_lse = None
else:
index_table = self.kv_index_translator.index_table_for_batch(forward_batch)
if self.dcp_size > 1:
kv_indptr, kv_indices, _ = self._dcp_kv_indices(
forward_batch.req_pool_indices,
@@ -989,11 +947,9 @@ class TritonAttnBackend(AttentionBackend):
kv_indptr = self._fill_kv_indptr_and_indices(
bs,
forward_batch.extend_prefix_lens,
forward_batch.req_pool_indices,
index_table,
kv_indices,
)
if self._translate_kv_loc is not None:
kv_indices = self._translate_kv_loc(kv_indices)
if self.sliding_window_size is not None and self.sliding_window_size > 0:
(
window_kv_indptr,
@@ -1002,10 +958,9 @@ class TritonAttnBackend(AttentionBackend):
window_kv_offsets,
) = update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
index_table,
self.sliding_window_size,
forward_batch.extend_prefix_lens,
forward_batch.req_pool_indices,
bs,
self.device,
self.token_to_kv_pool,
@@ -1027,18 +982,7 @@ class TritonAttnBackend(AttentionBackend):
swa_out_cache_loc = None
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
swa_out_cache_loc = self.token_to_kv_pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc
)
# Unified pool full-attention WRITE loc (virtual out_cache_loc -> physical),
# carried in the metadata (-> KVWriteLoc.full_loc). None for non-unified pools.
out_cache_loc_full_physical = None
if (
self._translate_kv_loc is not None
and forward_batch.out_cache_loc is not None
):
out_cache_loc_full_physical = self._translate_kv_loc(
swa_out_cache_loc = self.kv_index_translator.sliding_window_write_loc_for(
forward_batch.out_cache_loc
)
@@ -1058,7 +1002,11 @@ class TritonAttnBackend(AttentionBackend):
window_kv_offsets,
swa_attn_logits=swa_attn_logits,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
out_cache_loc_full_physical=(
forward_batch.out_cache_loc
if self.kv_index_translator.is_translating
else None
),
lean_Mp=lean_Mp,
lean_Lp=lean_Lp,
lean_Op=lean_Op,
@@ -1179,7 +1127,7 @@ class TritonAttnBackend(AttentionBackend):
device=self.device,
)
if self._translate_kv_loc is not None:
if self.kv_index_translator.is_translating:
# Unified pool full-attention write-target buffer, refilled at replay
# (-> KVWriteLoc.full_loc). Capture-stable, mirrors cuda_graph_swa_out_cache_loc.
self.cuda_graph_out_cache_loc_full_physical = torch.zeros(
@@ -1187,6 +1135,9 @@ class TritonAttnBackend(AttentionBackend):
dtype=torch.int64,
device=self.device,
)
self.kv_read_tables = self.kv_index_translator.make_capture_tables(
max_bs=max_bs, max_context_len=self.max_context_len
)
def _build_cuda_graph_forward_metadata(
self,
@@ -1300,16 +1251,21 @@ class TritonAttnBackend(AttentionBackend):
seq_lens: torch.Tensor,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
):
) -> None:
"""Shared capture+replay body for the cuda-graph init path.
Public entry: :py:meth:`init_forward_metadata_out_graph`.
"""
# NOTE: encoder_lens expected to be zeros or None
index_table = self.kv_index_translator.build_index_table(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
into=self.kv_read_tables,
)
if forward_mode.is_decode_or_idle():
assert spec_info is None, "Multi-step cuda graph init is not done here."
_, _, window_kv_lens, num_kv_splits_lens = self._update_decode_kv_buffers(
bs, seq_lens, req_pool_indices
bs, seq_lens, req_pool_indices, index_table
)
self.get_num_kv_splits(
self.cuda_graph_num_kv_splits[:bs], num_kv_splits_lens[:bs]
@@ -1320,12 +1276,10 @@ class TritonAttnBackend(AttentionBackend):
)
elif forward_mode.is_target_verify():
bs = len(req_pool_indices)
self._update_target_verify_buffers(
bs, seq_lens, req_pool_indices, spec_info
)
self._update_target_verify_buffers(bs, seq_lens, spec_info, index_table)
elif forward_mode.is_draft_extend_v2():
self._update_draft_extend_buffers(
bs, seq_lens, req_pool_indices, forward_mode, spec_info
bs, seq_lens, forward_mode, spec_info, index_table
)
else:
raise ValueError(
@@ -1402,10 +1356,11 @@ class TritonAttnBackend(AttentionBackend):
pool = self.token_to_kv_pool
cache_loc = forward_batch.out_cache_loc
if isinstance(pool, SWAKVPool) and pool.layers_mapping[layer.layer_id][1]:
cache_loc = pool.translate_loc_from_full_to_swa(cache_loc)
elif self._translate_kv_loc is not None:
# Unified pool: buffers are indexed in the kernel-facing id space.
cache_loc = self._translate_kv_loc(cache_loc)
assert self.forward_metadata.swa_out_cache_loc is not None, (
"window-layer read-back before the metadata carried a "
"sliding-window write loc"
)
cache_loc = self.forward_metadata.swa_out_cache_loc
k_buffer, v_buffer = pool.get_kv_buffer(layer.layer_id)
k = k_buffer[cache_loc]
v = v_buffer[cache_loc]
@@ -1776,15 +1731,12 @@ class TritonAttnBackend(AttentionBackend):
and isinstance(pool, SWAKVPool)
and pool.layers_mapping[layer.layer_id][1]
):
# Consumes VIRTUAL ids, so it must see out_cache_loc untranslated.
extend_kv_indices = pool.translate_loc_from_full_to_swa(extend_kv_indices)
extend_kv_indices = self.forward_metadata.swa_out_cache_loc
assert extend_kv_indices is not None, (
"window-layer extend before the metadata carried a "
"sliding-window write loc"
)
elif self.forward_metadata.out_cache_loc_full_physical is not None:
# Unified pool: this kernel reads the extend half OUT OF THE POOL (the
# 2-stage path takes it from the k/v arguments), so it needs the same
# translated loc the KV write uses -- otherwise the prefix is read at
# physical ids and the extend tokens at virtual ones. Reuse the
# per-forward translation rather than re-translating: this runs once
# per layer.
extend_kv_indices = self.forward_metadata.out_cache_loc_full_physical
# Handle cases where extend_seq_lens or extend_start_loc might not be set
@@ -2246,15 +2198,13 @@ class TritonMultiStepDraftBackend:
def update_sliding_window_buffer(
window_kv_indptr,
req_to_token,
index_table,
sliding_window_size,
seq_lens,
req_pool_indices,
bs,
device=None,
token_to_kv_pool=None,
window_kv_indices=None,
skip_full_to_swa_translation=False,
):
"""Fill window KV buffers for sliding-window attention.
@@ -2262,13 +2212,12 @@ def update_sliding_window_buffer(
path); omit it (or pass ``None``) to allocate a fresh tensor (eager path,
requires ``device``).
``skip_full_to_swa_translation=True`` leaves ``window_kv_indices`` as VIRTUAL
full-token ids (no eager full->swa translate). The unified-memory-pool cuda-graph
builder passes this so the window translate is deferred to
``TritonAttnBackend._translate_cuda_graph_shared_pool_locs`` (run in
``init_forward_metadata_out_graph``, BEFORE ``graph.replay()``), which reads
the live v2p and rewrites the static window buffer to swa-physical in place;
baseline SWA leaves it False (eager).
``index_table`` is the batch's read-index source view. Unified pool: the
gather reads the parallel SWA array (built directly from virtual ids
through the swa side's own v2p), so the window indices come out
already swa-side ids -- no translate here, eager or captured. Static SWA
pools gather full-token ids from req_to_token and keep the legacy
full->swa translate below.
"""
window_kv_lens = torch.minimum(
seq_lens,
@@ -2281,18 +2230,18 @@ def update_sliding_window_buffer(
window_kv_indptr[-1], dtype=torch.int64, device=device
)
window_kv_start_idx = seq_lens - window_kv_lens
source_ids = index_table.sliding_window_read_ids()
create_flashinfer_kv_indices_triton[(bs,)](
req_to_token,
req_pool_indices,
source_ids,
index_table.row_ids,
window_kv_lens,
window_kv_indptr,
window_kv_start_idx,
window_kv_indices,
req_to_token.stride(0),
source_ids.stride(0),
ENTRY_PAGE_SIZE=index_table.entry_page_size,
)
if not skip_full_to_swa_translation and hasattr(
token_to_kv_pool, "translate_loc_from_full_to_swa"
):
if not index_table.is_translated and isinstance(token_to_kv_pool, BaseSWAKVPool):
kv_last_index = window_kv_indptr[-1]
window_kv_indices[:kv_last_index] = (
token_to_kv_pool.translate_loc_from_full_to_swa(
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+10 -24
View File
@@ -294,11 +294,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self._v2p_page_table = _hooks.v2p_page_table
self._kernel_page_multiplier = _hooks.kernel_page_multiplier
self._unified_mla = _hooks.enabled
# virtual token id -> DENSE kernel-facing id, for the KV write loc.
self._translate_kv_loc_dense = _hooks.translate_kv_loc_for_kernel
# Per-forward kernel-facing write loc ([:n] view of a capture-stable buffer),
# set by the cuda-graph out-graph hook; None on the eager path (where the
# write translates through the pool's _full_translate hook instead).
# Per-forward kernel-facing write loc ([:n] view of a capture-stable buffer);
# None on the eager path, which passes out_cache_loc straight through.
self._decode_kernel_loc: Optional[torch.Tensor] = None
self.cuda_graph_out_cache_loc_kernel: Optional[torch.Tensor] = None
# Fused KV-scatter + q-concat on the decode dense-loc path (one launch
@@ -636,7 +633,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
out_cache_loc = forward_batch.out_cache_loc
n = out_cache_loc.shape[0]
dst = self.cuda_graph_out_cache_loc_kernel[:n]
self._translate_kv_loc_dense(out_cache_loc, out=dst)
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
@@ -651,8 +648,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass."""
# Eager path: no capture-stable kernel-facing write loc; the pool's _full_translate
# hook translates the write loc (safe out of a cuda graph).
self._decode_kernel_loc = None
# Delegate to parent for non-decode modes.
if (
@@ -950,11 +945,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
q: torch.Tensor,
q_rope: torch.Tensor,
) -> Optional[torch.Tensor]:
"""Decode: scatter the KV row at ``loc`` (already physical — the
dense-loc buffer on the unified pool, or out_cache_loc on the static
pool where ``_full_translate`` is identity) and build the
[q_nope | q_rope] fmha query in one kernel launch (saves one launch
per MLA layer and keeps the PDL chain intact).
"""Decode: scatter the KV row at ``loc`` (already kernel-facing) and
build the [q_nope | q_rope] fmha query in one kernel launch (saves one
launch per MLA layer and keeps the PDL chain intact).
Returns the concatenated query, or None when the fused kernel does
not cover the inputs (caller falls back to the two-kernel path).
@@ -1100,8 +1093,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
k is not None and k_rope is not None
), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
if self._decode_kernel_loc is not None:
# cuda-graph path: kernel-facing write loc precomputed out-of-graph, so
# the in-graph write captures no translate allocation.
if merge_query and self._fused_set_kv_concat_q:
# Fused: KV scatter + [q_nope | q_rope] concat in one
# launch; None when the inputs are not covered.
@@ -1115,21 +1106,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
)
if query is None:
self.token_to_kv_pool.set_mla_kv_buffer(
layer,
self._decode_kernel_loc,
k,
k_rope,
loc_is_kernel_facing=True,
layer, self._decode_kernel_loc, k, k_rope
)
else:
# eager (or static pool): the pool's _full_translate handles it.
# eager (or static pool): out_cache_loc is kernel-facing.
if (
merge_query
and self._fused_set_kv_concat_q
and not self._unified_mla
):
# Static pool: _full_translate is identity, so
# out_cache_loc is already the physical write loc.
# Static pool only, conservatively.
query = self._set_kv_and_concat_q_fused(
layer=layer,
loc=forward_batch.out_cache_loc,
@@ -1265,7 +1251,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
if self._decode_kernel_loc is not None:
self.token_to_kv_pool.set_mla_kv_buffer(
layer, self._decode_kernel_loc, k, k_rope, loc_is_kernel_facing=True
layer, self._decode_kernel_loc, k, k_rope
)
else:
self.token_to_kv_pool.set_mla_kv_buffer(
python/sglang/srt/mem_cache/kv_index_translator.py
+358
View File
@@ -0,0 +1,358 @@
# 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.
# ==============================================================================
"""Turns the KV ids stored in `req_to_token` into ids attention kernels can use.
A KV slot can be named in three id spaces:
* **virtual** - what `req_to_token` stores. Keeps naming the same logical
slot even after the pool moves data around.
* **physical** - where that slot sits in the pool right now.
* **kernel-facing** - what a kernel can index the per-layer K/V tensors
with. Same as physical on a plain pool; under the unified pool it is the
physical page scaled by the per-page block count.
All three coincide on a plain pool, so nothing here does any work there.
Backends get a `KVIndexTable`, which answers "what do I gather from, and
which row is mine?":
ids[row_ids[b], pos]
plain pool : ids = req_to_token, row_ids = req_pool_indices (those very
objects - no copy, no kernel)
unified : ids = a built array of kernel-facing ids,
row_ids = arange(batch_size)
Backends call their own copy a *page table* (fa3) or a *block table*
(trtllm); here it is the **index table**.
Converting only ever rewrites the page number and keeps the in-page offset, so
one page-granular table serves both kinds of consumer: a block-table backend
uses its rows as-is, and one that wants flat per-token ids rebuilds them as
token_id = entry * entry_page_size + pos % entry_page_size
WRITES, IN TWO PHASES. The full-side write loc is rebound to kernel-facing
ids at ForwardBatch construction - the earliest consumer can snapshot it
right after. The sliding-window write loc is derived at the same moment as
read table, into the same index table.
"""
from __future__ import annotations
import weakref
from typing import Optional, Tuple
import msgspec
import torch
from sglang.kernels.ops.kvcache.kv_read_table import build_kv_read_table
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
class KVReadTables(msgspec.Struct, frozen=True):
"""One capture-stable destination, however many id spaces the pool has.
A backend holds this and hands it back to `build_index_table(into=...)`;
it never has to know whether there is a sliding-window space behind it.
"""
full: torch.Tensor
sliding_window: Optional[torch.Tensor]
class KVIndexTable(msgspec.Struct, frozen=True):
"""Collection of what one batch gathers from."""
ids: torch.Tensor # 2-D array of KV ids to gather from
row_ids: torch.Tensor # which row belongs to batch lane b
row_stride: int # stride between rows of `ids`, in elements
entry_page_size: int # what one entry covers: 1 = a token, N = a page of N
is_translated: bool # entries are already kernel-facing ids
sliding_window_ids: Optional[torch.Tensor] # SWA models: the parallel swa array
def sliding_window_read_ids(self) -> torch.Tensor:
"""Which array a sliding-window gather reads: the parallel swa array
when translated, else the full-attention array, which the caller maps
through the pool's own full->swa map."""
return self.sliding_window_ids if self.is_translated else self.ids
class KVIndexTranslator:
"""Built once per ModelRunner."""
def __init__(
self,
*,
req_to_token: torch.Tensor,
token_to_kv_pool_allocator,
token_to_kv_pool,
page_size: int,
device: str,
):
self.req_to_token = req_to_token
self.page_size = page_size
self.device = device
self.is_translating = (
isinstance(
token_to_kv_pool_allocator,
(UnifiedMambaTokenToKVPoolAllocator, UnifiedSWATokenToKVPoolAllocator),
)
and token_to_kv_pool_allocator.get_kvcache() is token_to_kv_pool
)
if self.is_translating:
alloc = token_to_kv_pool_allocator
self._full_v2p_table = alloc.full_v2p_page_table
self._full_p2v_table = alloc.full_p2v_page_table
self._full_page_multiplier = alloc.kernel_page_multiplier
self._translate_full = alloc.translate_kv_loc_for_kernel
if isinstance(alloc, UnifiedSWATokenToKVPoolAllocator):
self._swa_v2p_table = alloc.swa_v2p_page_table
self._swa_page_multiplier = alloc.swa_kernel_page_multiplier
self._swa_write_loc_from_full = self._swa_write_loc_unified
else:
self._swa_v2p_table = None
self._swa_page_multiplier = 1
self._swa_write_loc_from_full = None
else:
self._full_v2p_table = None
self._full_p2v_table = None
self._full_page_multiplier = 1
self._translate_full = None
self._swa_v2p_table = None
self._swa_page_multiplier = 1
self._swa_write_loc_from_full = (
token_to_kv_pool.translate_loc_from_full_to_swa
if isinstance(token_to_kv_pool, SWAKVPool)
else None
)
self._rows: Optional[torch.Tensor] = (
torch.arange(req_to_token.shape[0], dtype=torch.int64, device=device)
if self.is_translating
else None
)
self._index_table_memo: Optional[Tuple[weakref.ref, KVIndexTable]] = None
def make_capture_tables(
self, *, max_bs: int, max_context_len: int
) -> Optional[KVReadTables]:
"""Capture-stable destinations for a backend to own, or None when this
pool needs no translation and the backend will never fill any.
Zero-filled: entry 0 is the reserved padding slot in every id space, so
a captured graph replaying before its first refresh reads padding, not
garbage.
"""
if not self.is_translating:
return None
max_pages = -(-max_context_len // self.page_size)
def _zeros():
return torch.zeros(
(max_bs, max_pages), dtype=torch.int32, device=self.device
)
return KVReadTables(
full=_zeros(),
sliding_window=_zeros() if self._swa_v2p_table is not None else None,
)
# -- per-batch view --------------------------------------------------------
def build_index_table(
self,
*,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
max_pages: Optional[int] = None,
into: Optional[KVReadTables] = None,
) -> KVIndexTable:
"""The one per-batch entry point.
Non-unified: the raw ``(req_to_token, req_pool_indices)`` passthrough,
no tensor ops and no copies. Unified: fills each row's live prefix and
returns the table WHOLE, so a caller needing a stable pointer (a
captured graph bakes it) passes its own tables in ``into``;
``into=None`` allocates of width ``max_pages`` instead.
"""
if not self.is_translating:
return KVIndexTable(
ids=self.req_to_token,
row_ids=req_pool_indices,
row_stride=self.req_to_token.stride(0),
entry_page_size=1,
is_translated=False,
sliding_window_ids=None,
)
bs = int(req_pool_indices.numel())
if into is not None:
out_full = into.full
out_swa = into.sliding_window
# A caller-owned table may be padded wider than req_to_token's span
# (trtllm_mla / flashmla pad to a page-count bound); the columns
# past it have no source to read, so stop there.
width = min(
out_full.shape[1] if max_pages is None else max_pages,
-(-self.req_to_token.shape[1] // self.page_size),
)
else:
assert max_pages is not None, (
"KVIndexTranslator.build_index_table: allocating needs max_pages "
"(from the batch's seq_lens_cpu max)"
)
width = max_pages
out_full = torch.zeros((bs, width), dtype=torch.int32, device=self.device)
out_swa = (
torch.zeros((bs, width), dtype=torch.int32, device=self.device)
if self._swa_v2p_table is not None
else None
)
build_kv_read_table(
req_to_token=self.req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
v2p=self._full_v2p_table,
multiplier=self._full_page_multiplier,
page_size=self.page_size,
max_pages=width,
out=out_full,
)
if out_swa is not None:
build_kv_read_table(
req_to_token=self.req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
v2p=self._swa_v2p_table,
multiplier=self._swa_page_multiplier,
page_size=self.page_size,
max_pages=width,
out=out_swa,
)
return KVIndexTable(
ids=out_full,
row_ids=self._rows[:bs],
row_stride=out_full.stride(0),
entry_page_size=self.page_size,
is_translated=True,
sliding_window_ids=out_swa,
)
def fill_read_table(
self,
*,
out: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
) -> None:
"""`build_index_table(into=...)` for a caller that owns a bare block
table rather than a KVReadTables: trtllm_mla / flashmla consume that
table directly, its rows already being the index table's rows.
"""
assert (
self.is_translating
), "KVIndexTranslator.fill_read_table on a pool that needs no translation"
self.build_index_table(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
into=KVReadTables(full=out, sliding_window=None),
)
def index_table_for_batch(self, forward_batch) -> KVIndexTable:
"""Eager per-batch view, memoized in one slot keyed by batch identity
so multi-consumer metadata builds share a build. The next batch
replaces the slot; consumers only read during their own build. The
captured path does not memoize -- it refreshes its buffers per
replay."""
memo = self._index_table_memo
if memo is not None and memo[0]() is forward_batch:
return memo[1]
max_pages = None
if self.is_translating:
# `seq_lens_cpu` is a non-None but STALE slice on a gpu_only
# batch; `seq_lens_sum` is the signal that it is live. A stale max
# under-sizes the table and the tail then reads as the sink.
slc = forward_batch.seq_lens_cpu
if (
forward_batch.seq_lens_sum is not None
and slc is not None
and slc.numel() > 0
):
max_seq = int(slc.max())
else:
max_seq = self.req_to_token.shape[1]
max_pages = max(-(-max_seq // self.page_size), 1)
view = self.build_index_table(
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
max_pages=max_pages,
)
self._index_table_memo = (weakref.ref(forward_batch), view)
return view
# -- write loc (phase 1; phase 2 lives in build_index_table) ----------------
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.
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.
"""
self._index_table_memo = None
if not self.is_translating or forward_batch.out_cache_loc is None:
return
forward_batch.out_cache_loc = self._translate_full(forward_batch.out_cache_loc)
def sliding_window_write_loc_for(
self, out_cache_loc: Optional[torch.Tensor]
) -> Optional[torch.Tensor]:
"""This batch's sliding-window write loc, or None when there is no loc
this forward or the pool has no sliding-window id space."""
if out_cache_loc is None or self._swa_write_loc_from_full is None:
return None
return self._swa_write_loc_from_full(out_cache_loc)
def _swa_write_loc_unified(self, kernel_loc: torch.Tensor) -> torch.Tensor:
"""Sliding-window write loc, derived pointwise from FULL-side
kernel-facing values (phase 2 of the write contract).
"""
full_stride = self.page_size * self._full_page_multiplier
offset = kernel_loc % full_stride # == virtual_token % page_size
# An unmapped physical page reads back as -1; clamp it rather than let
# the gather wrap onto the v2p table's last element.
virt_page = self._full_p2v_table[kernel_loc // full_stride].clamp_(min=0)
swa_stride = self.page_size * self._swa_page_multiplier
return (self._swa_v2p_table[virt_page] * swa_stride + offset).clamp_(min=0)
# -- token-level translate surface (the mixin / local-attn consumers) ------
def translate_full_attn_ids(
self, kv_indices: torch.Tensor, *, out: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""Virtual token ids -> kernel-facing full-attention ids (the identity
when no translation is needed, so callers never branch)."""
if not self.is_translating:
assert out is None, "passthrough translate takes no out="
return kv_indices
return self._translate_full(kv_indices, out=out)
python/sglang/srt/mem_cache/memory_pool.py
+35 -26
View File
@@ -81,7 +81,10 @@ from sglang.srt.utils import (
is_npu,
next_power_of_2,
)
from sglang.srt.utils.async_probe import maybe_detect_oob
from sglang.srt.utils.async_probe import (
maybe_detect_kernel_facing_loc,
maybe_detect_oob,
)
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
if TYPE_CHECKING:
@@ -1595,21 +1598,22 @@ class KVWriteLoc:
"""Write target(s) for ``KVCache.set_kv_buffer``.
All location info lives here (in the attention metadata), NOT in the pool:
- ``loc``: the generic per-token write location (the allocated
``out_cache_loc``). VIRTUAL under the unified memory pool (it indexes the
virtual slot space); already physical for a non-unified memory pool.
- ``swa_loc``: the pre-translated SWA-sub-pool PHYSICAL location for hybrid
SWA pools (``None`` otherwise).
- ``full_loc``: the pre-translated full-attention-sub-pool PHYSICAL location
for the unified memory pool (``None`` otherwise), computed once per forward in
attention metadata (``ForwardMetadata.out_cache_loc_full_physical``). The
shared full pool writes it directly; the pool never translates (replacing
the former per-layer v2p gather / ``set_full_loc`` pin).
- ``loc``: the generic per-token write location (``out_cache_loc``).
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`` and ``full_loc`` are the parallel pair (each a pre-resolved
PHYSICAL loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``);
``loc`` is the generic, possibly-virtual fallback. Bundling them lets a
backend issue one ``set_kv_buffer`` call regardless of pool type.
loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``);
``loc`` is the generic fallback. Bundling them lets a backend issue one
``set_kv_buffer`` call regardless of pool type.
"""
loc: torch.Tensor
@@ -1690,6 +1694,10 @@ class KVCache(abc.ABC):
):
self.size = size
self.page_size = page_size
# Row-blocks one page holds in this pool's kernel-facing id space; >1
# only where the per-layer views are dense (the unified pool), and then
# a write loc must have been translated into that space first.
self.kernel_page_blocks = 1
self.dtype = dtype
self.device = device
if dtype in (torch.float8_e5m2, torch.float8_e4m3fn, torch.float8_e4m3fnuz):
@@ -2389,6 +2397,9 @@ class MHATokenToKVPool(KVCache):
# Catch stale slot ids here instead of as illegal-addr / silent KV
# corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT).
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)")
maybe_detect_kernel_facing_loc(
loc, self.page_size, self.kernel_page_blocks, "set_kv_buffer (MHA)"
)
layer_id = (
layer_id_override if layer_id_override is not None else layer.layer_id
)
@@ -3654,9 +3665,10 @@ 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
# virtual->kernel-facing full-KV translate for the model-level MLA entry points
# (`set_mla_kv_buffer` / `get_mla_kv_buffer` receive VIRTUAL locs);
# identity for a static pool, `translate_kv_loc_for_kernel` for the unified pool.
# 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:
@@ -3943,17 +3955,8 @@ class HybridLinearKVPool(KVCache):
loc: torch.Tensor,
cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor,
loc_is_kernel_facing: bool = False,
):
assert self.use_mla, "set_mla_kv_buffer called when use_mla is False"
# Model-level MLA entry point: `loc` is a VIRTUAL loc under the unified
# pool, so translate to the kernel-facing id space here.
#
# `loc_is_kernel_facing`: the caller already translated `loc` (the unified-pool
# cuda-graph decode precomputes it out-of-graph into a capture-stable
# buffer, so the in-graph write does not capture a translate allocation).
if not loc_is_kernel_facing:
loc = self._full_translate(loc)
with self._transfer_id_context(layer):
self.full_kv_pool.set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope)
@@ -4092,6 +4095,9 @@ class MLATokenToKVPool(KVCache):
):
loc, _, _ = unwrap_write_loc(loc_info)
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)")
maybe_detect_kernel_facing_loc(
loc, self.page_size, self.kernel_page_blocks, "set_kv_buffer (MLA)"
)
layer_id = (
layer_id_override if layer_id_override is not None else layer.layer_id
)
@@ -4176,6 +4182,9 @@ class MLATokenToKVPool(KVCache):
(self.size + self.page_size) * get_parallel().attn_dcp_size,
"set_mla_kv_buffer (MLA)",
)
maybe_detect_kernel_facing_loc(
loc, self.page_size, self.kernel_page_blocks, "set_mla_kv_buffer (MLA)"
)
layer_id = (
layer_id_override if layer_id_override is not None else layer.layer_id
)
python/sglang/srt/mem_cache/multi_ended_allocator.py
+10
View File
@@ -1906,6 +1906,11 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
"""
return self.full_attn_allocator.virtual_to_physical
@property
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,
@@ -2257,6 +2262,11 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
"""Page-level virtual->physical table of the full sub-pool."""
return self.full_attn_allocator.virtual_to_physical
@property
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,
python/sglang/srt/mem_cache/unified_memory_pool.py
+3 -4
View File
@@ -548,6 +548,7 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
enable_kv_cache_copy=False,
kv_cache_layout="page_major",
)
self.kernel_page_blocks = spec.blocks_per_page()
def _create_buffers(self):
self.k_buffer = self._k_views
@@ -641,7 +642,7 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
max_slots = unified_buffer.max_slots(sub_pool_name)
self._num_pages = max_slots // page_size
self._page_bytes = page_size * spec.entry_bytes()
self._view_rows = self._num_pages * spec.layer_num * page_size
self._view_rows = self._num_pages * spec.blocks_per_page() * page_size
super().__init__(
# OOB checks bound locs by `size + page_size`; kernel-facing ids run to
@@ -655,6 +656,7 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
device=unified_buffer.device,
enable_memory_saver=False, # buffer owned by UnifiedKVPool
)
self.kernel_page_blocks = spec.blocks_per_page()
def _create_buffers(self):
self.kv_buffer = self._kv_views
@@ -1243,9 +1245,6 @@ def init_unified_mamba_pools(
req_to_token_pool.mamba_allocator = mamba_slot_allocator
token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate
if use_mla_backend:
# Model-level MLA entry points (`set_mla_kv_buffer` / `get_mla_kv_buffer`)
# receive VIRTUAL locs and translate to the kernel-facing space internally
# (eager-prefill-only paths; never captured in a cuda graph).
token_to_kv_pool._full_translate = allocator.translate_kv_loc_for_kernel
logger.info(
python/sglang/srt/model_executor/forward_batch_info.py
+2
View File
@@ -827,6 +827,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
device = model_runner.device
model_runner.kv_index_translator.rebind_write_loc(ret)
if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
hashed = _hash_rids_to_tensor(
rids=[req.rid for req in batch.reqs],
python/sglang/srt/model_executor/model_runner.py
+15
View File
@@ -89,6 +89,7 @@ from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.kv_cache_configurator import (
KVCacheConfigurator,
)
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
from sglang.srt.model_executor.cuda_graph_config import (
cuda_graph_fully_disabled,
@@ -856,6 +857,18 @@ class ModelRunner:
return
self.pre_model_load_memory += preloaded_weights_bytes / (1 << 30)
def init_kv_index_translator(self):
"""The one object that converts KV ids for this runner: attention
backends build their read indices from the table it hands them instead
of probing the pool's id spaces themselves."""
self.kv_index_translator = KVIndexTranslator(
req_to_token=self.req_to_token_pool.req_to_token,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
token_to_kv_pool=self.token_to_kv_pool,
page_size=self.page_size or 1,
device=self.device,
)
def alloc_memory_pool(self, memory_pool_config: Optional[MemoryPoolConfig] = None):
"""Allocate KV cache memory pools only (no backends or cuda graphs)."""
if memory_pool_config is not None:
@@ -882,6 +895,8 @@ class ModelRunner:
def _init_post_memory_pool_components(self):
"""Post-pool component wiring, split out of alloc_memory_pool so forks
that build bespoke memory pools can reuse it after allocating them."""
self.init_kv_index_translator()
# Must be called AFTER init_memory_pool so the pool object exists for
# canary to monkey-patch, and BEFORE init_decode_cuda_graph so warmup
# forwards captured into the graph see the patched pool methods.
python/sglang/srt/utils/async_probe.py
+21
View File
@@ -141,6 +141,27 @@ def maybe_detect_oob(indices: Optional[torch.Tensor], low: int, high: int, msg:
)
def maybe_detect_kernel_facing_loc(
indices: Optional[torch.Tensor], page_size: int, blocks_per_page: int, msg: str
):
"""Async check that a write loc is in the pool's KERNEL-FACING id space.
A kernel-facing id is `phys_page * (page_size * blocks_per_page) + offset`
with `offset < page_size`, so its remainder modulo the page stride is
below page_size; a VIRTUAL id satisfies that only in the first block.
Vacuous at blocks_per_page 1. Virtual ids are in range for the OOB probe,
so this is the only check that separates them.
"""
if blocks_per_page <= 1 or not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
return
if indices is None or indices.numel() == 0:
return
torch._assert_async(
(indices % (page_size * blocks_per_page) < page_size).all(),
f"write loc outside the kernel-facing id space (virtual ids?): {msg}",
)
def maybe_detect_page_aligned(
indices: Optional[torch.Tensor], page_size: int, msg: str
):
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+1
View File
@@ -411,6 +411,7 @@ class MockModelRunner(ModelRunner):
page_size=case.page_size,
get_kvcache=lambda: self.token_to_kv_pool,
)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py
+1
View File
@@ -404,6 +404,7 @@ class DSAMockModelRunner(ModelRunner):
kv_cache_dim=pool_kv_cache_dim,
)
self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py
+1
View File
@@ -384,6 +384,7 @@ class DualChunkMockModelRunner(ModelRunner):
enable_alt_stream=False,
)
self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py
+1
View File
@@ -318,6 +318,7 @@ class MockGDNModelRunner(ModelRunner):
page_size=case.page_size,
get_kvcache=lambda: self.token_to_kv_pool,
)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py
+1
View File
@@ -317,6 +317,7 @@ class MockKDAModelRunner(ModelRunner):
enable_alt_stream=False,
)
self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py
+1
View File
@@ -326,6 +326,7 @@ class MockLightningModelRunner(ModelRunner):
enable_alt_stream=False,
)
self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py
+1
View File
@@ -453,6 +453,7 @@ class MockMamba2ModelRunner(ModelRunner):
enable_alt_stream=False,
)
self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None
python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py
+1
View File
@@ -310,6 +310,7 @@ class MockMLAModelRunner(ModelRunner):
enable_memory_saver=False,
)
self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size)
self.init_kv_index_translator()
self.attn_cp_size = 1
self.attention_chunk_size = None
self.hisparse_coordinator = None