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

feat(unified-memory): dense KV views for uniform-row MHA/SWA models (#34602)

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 15:10:12 -07:00
committed by GitHub
co-authored by Caihua Li Cheng Wan
parent 007ef5e23a
commit 4bea51d885
30 changed files with 1310 additions and 1819 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/decode_attention.py
+13 -54
View File
@@ -195,53 +195,18 @@ def _mla_launch_plan(
def _extract_kv_strides(buf, page_size: int):
"""Extract (slot_stride, head_stride, page_stride, tok_stride) for a
KV buffer that may be:
- 3-D ``[max_slots, head_num, head_dim]`` (legacy / non-shared) — the
contiguous layout most callers use. page/tok strides are synthesized
so the kernel's PAGE_SIZE>1 math collapses to ``kv_loc * stride(0)``.
- 4-D ``[num_pages, page_size, head_num, head_dim]`` (shared
pool). page/tok strides come from stride(0)/stride(1) directly;
legacy ``stride_bs`` is set to 0 (unused at PAGE_SIZE>1).
"""Extract (slot_stride, head_stride, page_stride, tok_stride) for a 3-D
``[max_slots, head_num, head_dim]`` KV buffer.
Returns a 4-tuple of ints suitable for passing as ``stride_buf_*bs``,
``stride_buf_*h``, ``stride_buf_*page``, ``stride_buf_*tok``.
"""
if buf.ndim == 4:
# 4-D view ``[num_pages, page_size, head_num, head_dim]``.
# stride(0) = per-PAGE stride (page_bytes/itemsize)
# stride(1) = within-page per-TOKEN stride (k_row/v_row bytes/itemsize)
# The PAGE_SIZE>1 kernel branch uses page_stride/tok_stride and does
# NOT read slot_stride. slot_stride is consumed ONLY by the
# PAGE_SIZE==1 branch (``offs = kv_loc * stride_buf_*bs``), where one
# page holds exactly one slot, so the per-slot stride is the per-page
# stride — NOT the within-page token stride. Concretely the per-slot
# stride is ``page_stride // page_size`` (= entry_bytes/itemsize),
# which at ps=1 equals page_stride. Using ``tok_stride`` here (one
# layer's k_row) would make the ps=1 read address ``kv_loc * k_row``
# instead of ``kv_loc * entry_bytes`` and read the wrong slot.
page_stride = buf.stride(0)
tok_stride = buf.stride(1)
head_stride = buf.stride(2)
slot_stride = (
page_stride // page_size
) # per-slot stride; == page_stride at ps=1
assert buf.shape[1] == page_size, (
f"4-D KV buffer's dim-1 must equal page_size; got "
f"shape[1]={buf.shape[1]}, page_size={page_size}"
)
elif buf.ndim == 3:
# Legacy 3-D ``[N, head, dim]``. Synthesize page/tok strides such
# that ``(kv_loc // ps) * page_stride + (kv_loc % ps) * tok_stride
# == kv_loc * slot_stride`` for the page-aware branch — this lets
# the same kernel handle non-shared paged-allocator buffers without
# any caller adjustment.
slot_stride = buf.stride(0)
head_stride = buf.stride(1)
page_stride = slot_stride * page_size
tok_stride = slot_stride
else: # pragma: no cover
if buf.ndim != 3:
raise ValueError(f"unexpected KV buffer ndim={buf.ndim}, shape={buf.shape}")
slot_stride = buf.stride(0)
head_stride = buf.stride(1)
page_stride = slot_stride * page_size
tok_stride = slot_stride
return slot_stride, head_stride, page_stride, tok_stride
@@ -464,9 +429,6 @@ def _decode_att_m_fwd(
Lk = k_buffer.shape[-1]
Lv = v_buffer.shape[-1]
# head_num lives in the dim immediately before the head_dim. For 3-D
# ``[N, head_num, head_dim]`` that's dim 1; for 4-D
# ``[num_pages, page_size, head_num, head_dim]`` that's dim 2.
kv_head_num = k_buffer.shape[-2]
batch, head_num = q.shape[0], q.shape[1]
@@ -1204,8 +1166,7 @@ def decode_attention_fwd(
# schedule from kv_indptr on-device, so this path involves no host sync and is safe to
# capture in a CUDA graph. Whether Lean pays off for a given shape is decided cheaply by
# the backend's host-side seqlen gate (lean_decode_seqlen_gate) before we get here.
# Lean supports both the contiguous 3-D [N, head, dim] and paged 4-D
# [num_pages, page_size, head, dim] KV layouts (page-aware address math in the kernel).
# Lean handles both page sizes: the kernel does page-aware address math.
# ROCm/AMD only: Lean is validated on MI300X/MI355X; CUDA/NVIDIA uses the standard kernel.
if (
_is_hip
@@ -1587,8 +1548,8 @@ def _lean_attention_decode_kernel(
# Load K transposed: [BLOCK_DMODEL, BLOCK_N] so qk = q @ k directly.
# Page-aware KV address math (mirrors the standard grouped kernel): at
# PAGE_SIZE==1 the slot index addresses directly; otherwise it splits into
# (page_id, tok_in_p) for a [num_pages, page_size, head, dim] paged buffer.
# PAGE_SIZE==1 the slot index addresses directly; otherwise it splits
# into (page_id, tok_in_p).
if PAGE_SIZE == 1:
offs_buf_k = (
kv_loc[None, :] * stride_buf_kbs
@@ -1968,13 +1929,11 @@ def _decode_lean_attention_fwd(
``total_programs`` is the fixed persistent-grid size (2× device CU count). The kernel
derives its own tile schedule from ``kv_indptr`` on-device, so no host sync is needed and
the launch is CUDA-graph capturable. ``Mp``, ``Lp``, ``Op``, ``locks`` are pre-allocated
persistent-grid partial-result buffers reused across decode steps. ``page_size`` selects
the KV address math: 1 for a contiguous ``[N, head, dim]`` buffer, >1 for a paged
``[num_pages, page_size, head, dim]`` buffer (strides via ``_extract_kv_strides``).
persistent-grid partial-result buffers reused across decode steps. ``page_size``
selects the KV address math over the ``[N, head, dim]`` buffer (strides via
``_extract_kv_strides``).
"""
batch, head_num = q.shape[0], q.shape[1]
# head_num lives at dim -2 for both the 3-D [N, head, dim] and 4-D paged
# [num_pages, page_size, head, dim] layouts.
num_kv_heads = k_buffer.shape[-2]
Lk = k_buffer.shape[-1]
Lv = v_buffer.shape[-1]
python/sglang/kernels/ops/attention/metadata.py
+2 -2
View File
@@ -193,7 +193,7 @@ def _fused_metadata_kernel_general(
use_swa: tl.constexpr,
SHIFT: tl.constexpr,
BLOCK_COLS: tl.constexpr,
# Unified-memory dense-view path (page-major envelope shared with the mamba
# Unified-memory per-layer-view path (page-major envelope shared with the mamba
# sub-pool). Both default to the identity for the statically-partitioned
# pool, where req_to_token already holds physical ids.
v2p_ptr=None,
@@ -316,7 +316,7 @@ def _fused_metadata_kernel_ps1_no_swa(
max_seq_pages,
seq_len_delta: tl.constexpr,
BLOCK_COLS: tl.constexpr,
# Unified-memory dense-view path; identity defaults for the static pool.
# Unified-memory per-layer-view path; identity defaults for the static pool.
v2p_ptr=None,
PAGE_MULT: tl.constexpr = 1,
):
python/sglang/kernels/ops/kvcache/__init__.py
-1
View File
@@ -72,7 +72,6 @@ _TRITON_KERNELS = [
("trtllm_mha_graph_metadata", "update_trtllm_mha_graph_metadata"),
("aiter_unified_attention", "scatter_ragged_to_page_table_kernel"),
("aiter_unified_attention", "scatter_req_to_token_to_page_table_kernel"),
("cache_move", "store_cache_4d"),
("cache_move", "set_kv_buffer_prefix_valid_tiled"),
("cache_move", "copy_all_layer_kv_cache_tiled"),
("mla_buffer", "set_mla_kv_buffer_triton"),
python/sglang/kernels/ops/kvcache/cache_move.py
-182
View File
@@ -122,185 +122,3 @@ def copy_all_layer_kv_cache_func(
num_warps=kv_copy_config["num_warps"],
num_stages=2,
)
# ---------------------------------------------------------------------------
# store_cache_4d — single-launch Triton write into the 4-D page-major envelope
# K/V views. At `PAGE_SIZE = 1` the kernel constexpr-folds to byte-identical
# addresses as the slot-major envelope view; at `PAGE_SIZE > 1` it uses the
# same `(page_id, tok_in_p)` split the attention read kernels use.
# ---------------------------------------------------------------------------
@triton.jit
def store_cache_4d_kernel(
k_view_ptr,
v_view_ptr,
cache_k_ptr,
cache_v_ptr,
loc_ptr,
# Strides in ELEMENTS (not bytes); wrapper passes view.stride(D)
# directly. K and V may have different head_dim → different per-token
# strides, so we carry both.
stride_k_page,
stride_k_tok,
stride_v_page,
stride_v_tok,
stride_src_k_row,
stride_src_v_row,
K_ROW_DIM: tl.constexpr, # head_num * head_dim
V_ROW_DIM: tl.constexpr, # head_num * v_head_dim
PAGE_SIZE: tl.constexpr,
BLOCK: tl.constexpr,
):
"""Token-parallel Triton write into a 4-D envelope-strided K/V view.
Grid: ``(N, ceil(max(K_ROW_DIM, V_ROW_DIM) / BLOCK), 2)`` where:
- axis 0 → one program per token (loc[i])
- axis 1 → blocks within one slot's K (or V) row
- axis 2 → 0 = K, 1 = V (two-tensor write fused into one launch)
For each token i, the kernel writes:
page_id = loc[i] // PAGE_SIZE
tok_in_p = loc[i] % PAGE_SIZE
k_view[page_id, tok_in_p, :, :] = cache_k[i, :, :]
v_view[page_id, tok_in_p, :, :] = cache_v[i, :, :]
Cuda-graph safe: no Python branching on tensor values, no `.item()`,
all shapes/strides known at launch time.
"""
pid_n = tl.program_id(0)
pid_b = tl.program_id(1)
pid_kv = tl.program_id(2)
# 1. Resolve destination slot in the 4-D view.
loc = tl.load(loc_ptr + pid_n).to(tl.int64)
if PAGE_SIZE == 1:
page_id = loc
tok_in_p = tl.zeros([], dtype=tl.int64)
else:
page_id = loc // PAGE_SIZE
tok_in_p = loc % PAGE_SIZE
# 2. Compute per-tensor source/dest pointers.
base_off = pid_b * BLOCK + tl.arange(0, BLOCK)
if pid_kv == 0:
mask = base_off < K_ROW_DIM
# The trailing (head_num, head_dim) axes of `k_view` are
# contiguous: stride[-1]==1, stride[-2]==head_dim. So we can
# treat them as a flat K_ROW_DIM dimension addressed by `base_off`.
# The wrapper asserts this invariant.
src_ptr = cache_k_ptr + pid_n * stride_src_k_row + base_off
dst_ptr = (
k_view_ptr + page_id * stride_k_page + tok_in_p * stride_k_tok + base_off
)
else:
mask = base_off < V_ROW_DIM
src_ptr = cache_v_ptr + pid_n * stride_src_v_row + base_off
dst_ptr = (
v_view_ptr + page_id * stride_v_page + tok_in_p * stride_v_tok + base_off
)
src = tl.load(src_ptr, mask=mask)
tl.store(dst_ptr, src, mask=mask)
def store_cache_4d(
k_view: torch.Tensor,
v_view: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
loc: torch.Tensor,
page_size: int,
) -> None:
"""One-launch Triton write into the 4-D page-major envelope K/V views.
Writes ``cache_k[i]`` and ``cache_v[i]`` to
``k_view[loc[i]//ps, loc[i]%ps, :, :]`` (and analogously for V) for
``i in [0, N)``.
Contract:
- ``k_view``, ``v_view``: 4-D ``(num_pages, page_size, head_num,
head_dim*)``, contiguous in the trailing ``(head_num, head_dim)``
dims (i.e., ``stride[-1] == 1`` and ``stride[-2] == head_dim``).
- ``cache_k``, ``cache_v``: 3-D ``(N, head_num, head_dim*)``,
contiguous in the trailing ``(head_num, head_dim)`` dims.
- ``loc``: 1-D int64 or int32, N elements, values in
``[0, num_pages * page_size)``. The caller is responsible for
clamping any negative entries to ≥ 0.
- At ``page_size == 1`` the kernel produces byte-identical output
to the legacy advanced-indexing path.
Returns nothing; writes in place.
"""
if loc.numel() == 0:
return
assert k_view.is_cuda and v_view.is_cuda, "store_cache_4d: CUDA only"
assert k_view.ndim == 4 and v_view.ndim == 4, (
f"store_cache_4d: k_view/v_view must be 4-D, "
f"got {k_view.ndim}/{v_view.ndim}"
)
assert cache_k.ndim == 3 and cache_v.ndim == 3, (
f"store_cache_4d: cache_k/cache_v must be 3-D, "
f"got {cache_k.ndim}/{cache_v.ndim}"
)
assert cache_k.shape[0] == cache_v.shape[0] == loc.numel(), (
"store_cache_4d: cache_k/cache_v/loc batch dim mismatch: "
f"{cache_k.shape[0]}, {cache_v.shape[0]}, {loc.numel()}"
)
assert k_view.dtype == v_view.dtype == cache_k.dtype == cache_v.dtype, (
"store_cache_4d: dtype mismatch: "
f"k_view={k_view.dtype}, v_view={v_view.dtype}, "
f"cache_k={cache_k.dtype}, cache_v={cache_v.dtype}"
)
# Stride invariants — the kernel addresses (head_num, head_dim) as one
# flat ROW_DIM dimension; this requires the trailing two dims to be
# contiguous. This holds for the page-major envelope views
# (k_stride = (page_bytes/itemsize, k_row_bytes/itemsize, head_dim, 1)) and
# for cache_k/cache_v produced by the model forward.
assert k_view.stride(-1) == 1 and k_view.stride(-2) == k_view.shape[-1], (
f"store_cache_4d: k_view trailing dims must be contiguous; "
f"got stride={k_view.stride()}, shape={tuple(k_view.shape)}"
)
assert v_view.stride(-1) == 1 and v_view.stride(-2) == v_view.shape[-1], (
f"store_cache_4d: v_view trailing dims must be contiguous; "
f"got stride={v_view.stride()}, shape={tuple(v_view.shape)}"
)
assert cache_k.stride(-1) == 1 and cache_k.stride(-2) == cache_k.shape[-1], (
f"store_cache_4d: cache_k trailing dims must be contiguous; "
f"got stride={cache_k.stride()}, shape={tuple(cache_k.shape)}"
)
assert cache_v.stride(-1) == 1 and cache_v.stride(-2) == cache_v.shape[-1], (
f"store_cache_4d: cache_v trailing dims must be contiguous; "
f"got stride={cache_v.stride()}, shape={tuple(cache_v.shape)}"
)
head_num = k_view.shape[2]
head_dim = k_view.shape[3]
v_head_dim = v_view.shape[3]
K_ROW_DIM = head_num * head_dim
V_ROW_DIM = head_num * v_head_dim
BLOCK = 128
N = loc.numel()
row_dim_max = max(K_ROW_DIM, V_ROW_DIM)
grid = (N, triton.cdiv(row_dim_max, BLOCK), 2)
store_cache_4d_kernel[grid](
k_view,
v_view,
cache_k,
cache_v,
loc,
k_view.stride(0),
k_view.stride(1),
v_view.stride(0),
v_view.stride(1),
cache_k.stride(0),
cache_v.stride(0),
K_ROW_DIM=K_ROW_DIM,
V_ROW_DIM=V_ROW_DIM,
PAGE_SIZE=page_size,
BLOCK=BLOCK,
num_warps=4,
)
python/sglang/kernels/ops/kvcache/kv_indices.py
+3 -3
View File
@@ -105,11 +105,11 @@ def create_flashmla_kv_indices_triton(
req_to_token_ptr_stride: tl.constexpr,
kv_indices_ptr_stride: tl.constexpr,
PAGED_SIZE: tl.constexpr = 64,
# Unified-memory dense-view path (page-major envelope shared with the mamba
# Unified-memory per-layer-view path (page-major envelope shared with the mamba
# sub-pool). req_to_token holds VIRTUAL token ids; the block table the MLA
# kernel consumes must hold DENSE page ids. When v2p_ptr is given, map each
# kernel consumes must hold kernel-facing page ids. When v2p_ptr is given, map each
# virtual page through it to the physical page, then scale by PAGE_MULT
# (= num MLA layers) so the entry addresses the layer's dense per-page block
# (= num MLA layers) so the entry addresses the layer's per-page block
# in the (num_pages*L, page_size, kv_cache_dim) reshaped view. Both default
# to the identity (v2p_ptr None, PAGE_MULT 1) for the static pool.
v2p_ptr=None,
python/sglang/srt/arg_groups/kv_cache_hook.py
+31 -6
View File
@@ -9,6 +9,7 @@ from typing import Any
from sglang.srt.arg_groups.overrides import (
attention_backends_of,
declare_resolution,
model_config_of,
resolved_view,
resolving_view,
use_mla_backend,
@@ -223,7 +224,7 @@ def handle_unified_memory_pool(server_args: Any) -> None:
assert cfg.speculative_algorithm in (None, "DSPARK"), (
"--enable-unified-memory only supports --speculative-algorithm "
"DSPARK (chain draft); other speculative algorithms are not yet "
"audited for the unified pool's virtual/dense loc translation. Got "
"audited for the unified pool's virtual/kernel-facing loc translation. Got "
f"--speculative-algorithm={cfg.speculative_algorithm!r}."
)
if cfg.speculative_algorithm == "DSPARK":
@@ -243,7 +244,7 @@ def handle_unified_memory_pool(server_args: Any) -> None:
f"attention backends {sorted(spec_allowed)} for both prefill "
f"and decode; got {sorted(spec_backends)}. flashinfer / fa3 do "
"not translate speculative verify indices to the unified "
"pool's dense space yet."
"pool's kernel-facing space yet."
)
assert not (cfg.enable_hierarchical_cache or cfg.enable_lmcache), (
"--enable-unified-memory is not yet compatible with hierarchical / "
@@ -283,15 +284,39 @@ def handle_page_major_kv_layout(server_args: Any):
)
if not cfg.enable_page_major_kv_layout:
return
assert cfg.enable_unified_memory, (
"--enable-page-major-kv-layout without --enable-unified-memory is "
"temporarily unsupported: the strided MHA K/V views were removed "
"and the static-pool page-major layout awaits its per-layer-view "
"reimplementation. Run with --enable-unified-memory, or drop "
"--enable-page-major-kv-layout."
)
from sglang.srt.mem_cache.unified_memory_pool import (
unified_memory_supported_for_model,
)
model_config = model_config_of(server_args)
assert unified_memory_supported_for_model(
model_config, use_mla_backend=use_mla_backend(server_args)
), (
"--enable-unified-memory requires uniform K/V rows "
"(head_dim == v_head_dim); this model has "
f"head_dim={model_config.head_dim}, "
f"v_head_dim={model_config.v_head_dim}, "
f"swa_head_dim={model_config.swa_head_dim}, "
f"swa_v_head_dim={model_config.swa_v_head_dim}. The unified "
"pool's per-layer views require a uniform row width; run "
"this model without --enable-unified-memory."
)
# Only the Triton attention kernels read the strided 4-D envelope K/V
# views; FA3 / FlashInfer do not. EXCEPTION: the unified-memory MLA pool
# exposes each layer as a DENSE contiguous per-layer view
# (build_dense_mla_views), which the paged MLA kernels consume directly,
# with their kv_indices / block tables remapped to dense ids. Names below
# exposes each layer as a contiguous per-layer view
# (build_mla_views), which the paged MLA kernels consume directly,
# with their kv_indices / block tables remapped to kernel-facing ids. Names below
# are the RESOLVED ids from attention_backends_of: "flashinfer" is
# FlashInferMLAAttnBackend for an MLA model, "trtllm_mla" the trtllm
# decode kernel; "cutedsl_mla" and "tokenspeed_mla" subclass
# TRTLLMMLABackend and inherit its dense read/write path; "fa3" remaps its
# TRTLLMMLABackend and inherit its read/write path; "fa3" remaps its
# page_table (in-kernel for captured decode, one funnel for eager).
# flashmla / cutlass_mla share the create_flashmla block-table path and
# can be added the same way once exercised.
python/sglang/srt/layers/attention/flashattention_backend.py
+4 -4
View File
@@ -193,7 +193,7 @@ class FlashAttentionBackend(AttentionBackend):
self.needs_cpu_seq_lens = False
self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
# Unified pool: req_to_token holds VIRTUAL ids but the MLA per-layer views
# are DENSE, so every page_table needs remapping. MLA-only -- the MHA/SWA
# are kernel-facing, so every page_table needs remapping. MLA-only -- the MHA/SWA
# sub-pools keep the strided envelope layout FA3 cannot read at all.
self._unified_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
self._unified_dense = self._unified_hooks.enabled and self.use_mla
@@ -1101,16 +1101,16 @@ class FlashAttentionBackend(AttentionBackend):
# the remap into normal_decode_set_metadata, which must write in place.
#
# Placed BEFORE the `// page_size` reduction, in token space: since
# dense(t) = phys_page * (ps * L) + t % ps, dense(page_start) // ps is
# kernel_id(t) = phys_page * (ps * L) + t % ps, dense(page_start) // ps is
# phys_page * L, the dense page id the kernel wants. One site then serves
# both page sizes, and it inherits translate_kv_loc_dense's tombstone
# both page sizes, and it inherits translate_kv_loc_for_kernel's tombstone
# clamp so an unwritten req_to_token slot lands in the page-0 sink.
if self._unified_dense and metadata.page_table is not None:
# Flattened: the page_size == 1 translate path uses index_select,
# which rejects a 2-D index.
pt = metadata.page_table
metadata.page_table = (
self._unified_hooks.translate_kv_loc_dense(pt.reshape(-1))
self._unified_hooks.translate_kv_loc_for_kernel(pt.reshape(-1))
.to(torch.int32)
.view(pt.shape)
)
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
+4 -4
View File
@@ -850,7 +850,7 @@ class FlashInferMLAIndicesUpdaterDecode:
# Unified dense MLA pool: VIRTUAL -> DENSE kv_indices (see prefill updater).
self._translate_kv_loc_dense = unified_mla_hooks(
model_runner.token_to_kv_pool_allocator
).translate_kv_loc_dense
).translate_kv_loc_for_kernel
def update(
self,
@@ -918,7 +918,7 @@ class FlashInferMLAIndicesUpdaterDecode:
# [:paged_kernel_lens_sum] prefix the index kernel just filled is
# translated; the stale tail is left alone so it can never index the
# v2p table out of bounds. The int64 translate result narrows back to
# the buffer's int32 on copy_ (flashinfer requires int32; dense ids
# the buffer's int32 on copy_ (flashinfer requires int32; kernel-facing ids
# fit comfortably).
if self._translate_kv_loc_dense is not None:
valid = kv_indices[:paged_kernel_lens_sum]
@@ -989,11 +989,11 @@ class FlashInferMLAIndicesUpdaterPrefill:
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
# Unified dense MLA pool: kv_indices built from req_to_token are VIRTUAL;
# the paged wrapper reads the dense per-layer view, so remap them to DENSE
# the paged wrapper reads the per-layer view, so remap them to kernel-facing
# token ids. None (identity) unless the unified MLA pool is active.
self._translate_kv_loc_dense = unified_mla_hooks(
model_runner.token_to_kv_pool_allocator
).translate_kv_loc_dense
).translate_kv_loc_for_kernel
def update(
self,
python/sglang/srt/layers/attention/triton_backend.py
+2 -2
View File
@@ -210,11 +210,11 @@ class TritonAttnBackend(AttentionBackend):
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
# dense-view MLA pool (translate_kv_loc_dense falls back to the physical
# 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_dense", None
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.num_draft_tokens = get_spec().speculative_num_draft_tokens
self.speculative_num_steps = get_spec().speculative_num_steps
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+26 -22
View File
@@ -285,22 +285,22 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# Tree-mask scratch is fetched from the target backend only.
self.is_draft_runner = model_runner.is_draft_worker
# Unified-memory dense-view hooks (None on the static pool). req_to_token
# holds VIRTUAL token ids; the block table needs DENSE page ids, so the
# Unified-memory per-layer-view hooks (None on the static pool). req_to_token
# holds VIRTUAL token ids; the block table needs kernel-facing page ids, so the
# kv-index kernels gather virtual->physical page through `_v2p_page_table`
# then scale by `_kernel_page_multiplier` (= num MLA layers). See
# build_dense_mla_views / create_flashmla_kv_indices_triton.
# build_mla_views / create_flashmla_kv_indices_triton.
_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
self._v2p_page_table = _hooks.v2p_page_table
self._kernel_page_multiplier = _hooks.kernel_page_multiplier
self._unified_mla = _hooks.enabled
# virtual token id -> DENSE kernel-facing id, for the KV write loc.
self._translate_kv_loc_dense = _hooks.translate_kv_loc_dense
# Per-forward dense write loc ([:n] view of a capture-stable buffer),
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).
self._decode_dense_loc: Optional[torch.Tensor] = None
self.cuda_graph_out_cache_loc_dense: Optional[torch.Tensor] = None
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
# instead of set_mla_kv_buffer + concat_mla_absorb_q). Disabled under
# async asserts: the fused path writes the pool directly and would
@@ -408,7 +408,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# out-of-graph in init_forward_metadata_out_graph so the in-graph
# set_mla_kv_buffer captures no translate.
if self._unified_mla:
self.cuda_graph_out_cache_loc_dense = torch.zeros(
self.cuda_graph_out_cache_loc_kernel = torch.zeros(
max_num_tokens, dtype=torch.int64, device=self.device
)
num_tokens_per_req = max_num_tokens // max_bs
@@ -635,7 +635,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_dense[:n]
dst = self.cuda_graph_out_cache_loc_kernel[:n]
self._translate_kv_loc_dense(out_cache_loc, out=dst)
# Replay-prep receives the RAW (unpadded) out_cache_loc
# (build_replay_fb_view), but the captured write kernel consumes the
@@ -644,16 +644,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# 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_dense[n:].zero_()
self._decode_dense_loc = dst
self.cuda_graph_out_cache_loc_kernel[n:].zero_()
self._decode_kernel_loc = dst
else:
self._decode_dense_loc = None
self._decode_kernel_loc = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass."""
# Eager path: no capture-stable dense write loc; the pool's _full_translate
# 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_dense_loc = None
self._decode_kernel_loc = None
# Delegate to parent for non-decode modes.
if (
forward_batch.forward_mode.is_extend()
@@ -1058,8 +1058,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
if cos_sin_cache is None:
if save_kv_cache and self._fused_set_kv_concat_q_fp8:
loc = (
self._decode_dense_loc
if self._decode_dense_loc is not None
self._decode_kernel_loc
if self._decode_kernel_loc is not None
else (
None if self._unified_mla else forward_batch.out_cache_loc
)
@@ -1099,15 +1099,15 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert (
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_dense_loc is not None:
# cuda-graph path: dense write loc precomputed out-of-graph, so
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.
query = self._set_kv_and_concat_q_fused(
layer=layer,
loc=self._decode_dense_loc,
loc=self._decode_kernel_loc,
k=k,
k_rope=k_rope,
q=q,
@@ -1115,7 +1115,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
)
if query is None:
self.token_to_kv_pool.set_mla_kv_buffer(
layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
layer,
self._decode_kernel_loc,
k,
k_rope,
loc_is_kernel_facing=True,
)
else:
# eager (or static pool): the pool's _full_translate handles it.
@@ -1259,9 +1263,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert (
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_dense_loc is not None:
if self._decode_kernel_loc is not None:
self.token_to_kv_pool.set_mla_kv_buffer(
layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
layer, self._decode_kernel_loc, k, k_rope, loc_is_kernel_facing=True
)
else:
self.token_to_kv_pool.set_mla_kv_buffer(
python/sglang/srt/layers/attention/unified_mem_hooks.py
+7 -5
View File
@@ -37,7 +37,7 @@ class UnifiedMLAHooks(msgspec.Struct, frozen=True):
# Page-level virtual->physical table, gathered through by block-table kernels.
v2p_page_table: Optional[torch.Tensor]
# Virtual token id -> DENSE kernel-facing id (tombstones clamped to the sink).
translate_kv_loc_dense: Optional[Callable[..., torch.Tensor]]
translate_kv_loc_for_kernel: Optional[Callable[..., torch.Tensor]]
# Dense page stride scale (= number of full-attention MLA layers).
kernel_page_multiplier: int
enabled: bool
@@ -45,18 +45,18 @@ class UnifiedMLAHooks(msgspec.Struct, frozen=True):
_STATIC_POOL = UnifiedMLAHooks(
v2p_page_table=None,
translate_kv_loc_dense=None,
translate_kv_loc_for_kernel=None,
kernel_page_multiplier=1,
enabled=False,
)
def unified_mla_hooks(allocator) -> UnifiedMLAHooks:
"""Probe ``allocator`` for the unified-pool dense-view hooks.
"""Probe ``allocator`` for the unified-pool per-layer-view hooks.
Detection keys on the v2p table, NOT on ``kernel_page_multiplier > 1``: a
rank owning exactly ONE full-attention layer has multiplier 1 while its
``req_to_token`` is still virtual. There the dense id collapses onto the
``req_to_token`` is still virtual. There the kernel-facing id collapses onto the
physical id, so the v2p gather alone is the whole translation.
"""
v2p = getattr(allocator, "full_v2p_page_table", None)
@@ -64,7 +64,9 @@ def unified_mla_hooks(allocator) -> UnifiedMLAHooks:
return _STATIC_POOL
return UnifiedMLAHooks(
v2p_page_table=v2p,
translate_kv_loc_dense=getattr(allocator, "translate_kv_loc_dense", None),
translate_kv_loc_for_kernel=getattr(
allocator, "translate_kv_loc_for_kernel", None
),
kernel_page_multiplier=getattr(allocator, "kernel_page_multiplier", 1),
enabled=True,
)
python/sglang/srt/mem_cache/layout/page_major.py
+54 -58
View File
@@ -10,8 +10,8 @@ Across pages the layout is envelope-major (one ``page_bytes`` block per page).
At ``page_size == 1`` a page is a single slot, so the within-page block is the
per-slot ``[L0_K | L0_V | L1_K | L1_V | ...]`` envelope (token-granularity).
These builders produce per-layer strided views into a raw ``uint8`` buffer; they
hold no allocator/ownership state. ``anchor_bytes`` is the byte offset of the
These builders produce per-layer views into a raw ``uint8`` buffer; they hold
no allocator/ownership state. ``anchor_bytes`` is the byte offset of the
pool's region inside the raw buffer (0 for a standalone pool).
"""
@@ -36,7 +36,7 @@ def mha_entry_bytes(
return layer_num * (k_row_bytes + v_row_bytes)
def build_page_major_mha_views(
def build_mha_views(
raw: torch.Tensor,
*,
layer_num: int,
@@ -48,63 +48,59 @@ def build_page_major_mha_views(
num_pages: int,
anchor_bytes: int = 0,
) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
"""Per-layer K/V views over ``raw`` in the page-major layer-major layout.
"""Per-layer K/V views over ``raw`` for uniform-row MHA.
Each returned view is 4-D ``(num_pages, page_size, head_num, head_dim*)``
with constant strides:
The page envelope ``[L0_K*ps | L0_V*ps | L1_K*ps | ...]`` is a uniform
array of ``2*layer_num`` row-blocks when K and V rows are equally wide, so
it is a valid paged pool under
stride[0] = page_bytes / itemsize # next page
stride[1] = k_row_bytes / itemsize # next slot within layer L's K block
stride[2] = head_dim # next head
stride[3] = 1 # next element
kernel_id(t) = (t // ps) * (ps * 2 * layer_num) + t % ps
V is analogous with ``v_row_bytes`` / ``v_head_dim``. A token id ``t`` reads
page ``t // page_size``, slot ``t % page_size``.
with layer ``l``'s K at block ``2l`` and its V at block ``2l+1``. Each view
is a contiguous ``(num_pages * 2 * layer_num * ps, head_num, head_dim)``.
Views overlap by ``ps`` rows per block, safe because an id always resolves
inside its own block; the last view runs ``(2*layer_num - 1) * ps`` rows
past the envelope, so ``raw`` needs ``UnifiedKVPool.view_tail_pad_bytes``.
"""
assert head_dim == v_head_dim, (
f"build_mha_views requires uniform rows (head_dim == v_head_dim); "
f"got head_dim={head_dim}, v_head_dim={v_head_dim}. Asymmetric-KV "
"models cannot use the unified pool (screened out at startup)."
)
itemsize = store_dtype.itemsize
k_row_bytes = head_num * head_dim * itemsize
v_row_bytes = head_num * v_head_dim * itemsize
entry_bytes = layer_num * (k_row_bytes + v_row_bytes)
page_bytes = page_size * entry_bytes
row_elems = head_num * head_dim
row_bytes = row_elems * itemsize
blocks = 2 * layer_num
page_bytes = page_size * blocks * row_bytes
n_rows = num_pages * blocks * page_size
assert anchor_bytes % itemsize == 0
assert k_row_bytes % itemsize == 0
assert v_row_bytes % itemsize == 0
assert page_bytes % itemsize == 0
last_view_end = (
anchor_bytes + (blocks - 1) * page_size * row_bytes + n_rows * row_bytes
)
assert last_view_end <= raw.numel() * raw.itemsize, (
f"build_mha_views: block {blocks - 1}'s view ends at byte "
f"{last_view_end} but the raw buffer holds only "
f"{raw.numel() * raw.itemsize} bytes; allocate the tail pad "
f"(one page envelope = {page_bytes} B) via view_tail_pad_bytes"
)
as_dtype_view = raw.view(store_dtype)
stride_page = page_bytes // itemsize
stride_tok_k = k_row_bytes // itemsize
stride_tok_v = v_row_bytes // itemsize
k_shape = (num_pages, page_size, head_num, head_dim)
v_shape = (num_pages, page_size, head_num, v_head_dim)
k_stride = (stride_page, stride_tok_k, head_dim, 1)
v_stride = (stride_page, stride_tok_v, v_head_dim, 1)
k_buffer: List[torch.Tensor] = []
v_buffer: List[torch.Tensor] = []
for layer in range(layer_num):
# Layer L's K block starts at L * page_size * (k_row + v_row); V follows.
k_base_bytes = anchor_bytes + layer * page_size * (k_row_bytes + v_row_bytes)
v_base_bytes = k_base_bytes + page_size * k_row_bytes
assert k_base_bytes % itemsize == 0
assert v_base_bytes % itemsize == 0
k_buffer.append(
torch.as_strided(
as_dtype_view,
size=k_shape,
stride=k_stride,
storage_offset=k_base_bytes // itemsize,
k_base_bytes = anchor_bytes + (2 * layer) * page_size * row_bytes
v_base_bytes = k_base_bytes + page_size * row_bytes
for base_bytes, out in ((k_base_bytes, k_buffer), (v_base_bytes, v_buffer)):
assert base_bytes % itemsize == 0
out.append(
torch.as_strided(
as_dtype_view,
size=(n_rows, head_num, head_dim),
stride=(row_elems, head_dim, 1),
storage_offset=base_bytes // itemsize,
)
)
)
v_buffer.append(
torch.as_strided(
as_dtype_view,
size=v_shape,
stride=v_stride,
storage_offset=v_base_bytes // itemsize,
)
)
return k_buffer, v_buffer
@@ -113,7 +109,7 @@ def mla_entry_bytes(*, layer_num: int, kv_cache_dim: int, itemsize: int) -> int:
return layer_num * kv_cache_dim * itemsize
def build_dense_mla_views(
def build_mla_views(
raw: torch.Tensor,
*,
layer_num: int,
@@ -123,23 +119,23 @@ def build_dense_mla_views(
num_pages: int,
anchor_bytes: int = 0,
) -> List[torch.Tensor]:
"""Per-layer DENSE views over ``raw`` for MLA in the page-major layout.
"""Per-layer views over ``raw`` for MLA in the page-major layout.
The page envelope is ``[L0_latent * ps | L1_latent * ps | ...]``. Because all
MLA layers share one uniform row size (``kv_cache_dim``), the envelope is
itself a valid dense paged pool under a re-numbered index space: folding the
itself a valid paged pool under a re-numbered index space: folding the
layer offset ``l * ps * kv_cache_dim`` into each view's storage_offset makes
every per-layer view a plain CONTIGUOUS ``(num_pages * layer_num * ps, 1,
kv_cache_dim)`` tensor, addressed by the layer-independent dense id
kv_cache_dim)`` tensor, addressed by the layer-independent kernel-facing id
dense(t) = (t // ps) * (ps * layer_num) + t % ps (t = physical token)
kernel_id(t) = (t // ps) * (ps * layer_num) + t % ps (t = physical token)
so one shared block table (entry = page * layer_num) serves every layer, and
kernels that require ``.view(-1, page_size, kv_cache_dim)`` (trtllm/cutlass/
flashmla) work on the views natively.
The views overlap each other (view ``l+1`` is view ``l`` shifted by ``ps``
rows); that is safe because layer ``l`` is only ever indexed at dense ids,
rows); that is safe because layer ``l`` is only ever indexed at kernel-facing ids,
which always resolve to layer-``l`` bytes relative to view ``l``'s origin.
Layer ``layer_num-1``'s view extends ``(layer_num-1) * ps`` rows past the
last page envelope, so ``raw`` must carry at least one extra page envelope
@@ -148,13 +144,13 @@ def build_dense_mla_views(
itemsize = store_dtype.itemsize
row_bytes = kv_cache_dim * itemsize
page_bytes = page_size * layer_num * row_bytes
n_dense = num_pages * layer_num * page_size
n_rows = num_pages * layer_num * page_size
assert anchor_bytes % itemsize == 0
last_view_end = (
anchor_bytes + (layer_num - 1) * page_size * row_bytes + (n_dense * row_bytes)
anchor_bytes + (layer_num - 1) * page_size * row_bytes + (n_rows * row_bytes)
)
assert last_view_end <= raw.numel() * raw.itemsize, (
f"build_dense_mla_views: layer {layer_num - 1}'s view ends at byte "
f"build_mla_views: layer {layer_num - 1}'s view ends at byte "
f"{last_view_end} but the raw buffer holds only "
f"{raw.numel() * raw.itemsize} bytes; allocate the tail pad "
f"(one page envelope = {page_bytes} B) via view_tail_pad_bytes"
@@ -168,7 +164,7 @@ def build_dense_mla_views(
views.append(
torch.as_strided(
as_dtype_view,
size=(n_dense, 1, kv_cache_dim),
size=(n_rows, 1, kv_cache_dim),
stride=(kv_cache_dim, kv_cache_dim, 1),
storage_offset=base_bytes // itemsize,
)
python/sglang/srt/mem_cache/memory_pool.py
+23 -130
View File
@@ -45,7 +45,6 @@ from sglang.kernels.ops.attention.dsa.quant_k_cache import (
from sglang.kernels.ops.kvcache.cache_move import (
copy_all_layer_kv_cache_func,
set_kv_buffer_prefix_valid_tiled,
store_cache_4d,
)
from sglang.kernels.ops.kvcache.kvcache import can_use_store_cache, store_cache
from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
@@ -62,9 +61,7 @@ from sglang.srt.mem_cache.index_key_cache import IndexKeyCache
from sglang.srt.mem_cache.kv_vmm_backing import KvVmmBufferOwner
from sglang.srt.mem_cache.layout.page_major import (
build_page_major_mamba_views,
build_page_major_mha_views,
mamba_entry_bytes,
mha_entry_bytes,
)
from sglang.srt.mem_cache.utils import (
get_mla_kv_buffer_triton,
@@ -2866,9 +2863,8 @@ class MHATokenToKVPool(KVCache):
self._move_kv_cache_impl(tgt_loc, src_loc)
def _move_kv_cache_impl(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
# Physical move strategy. Override for layouts that change buffer identity
# (e.g. PageMajorMHATokenToKVPool always uses the native move). The 3-D
# per-layer buffers here ignore page_size in move_kv_cache_native.
# Physical move strategy. Override for layouts that change buffer
# identity (e.g. PageMajorMHATokenToKVPool always uses the native move).
if self.use_native_move_kv_cache:
move_kv_cache_native(self.k_buffer, self.v_buffer, tgt_loc, src_loc)
if getattr(self, "k_scale_buffer", None) is not None:
@@ -3184,18 +3180,11 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
class PageMajorMHATokenToKVPool(MHATokenToKVPool):
"""MHA pool with the page-major (layer-major within a page) page-granularity envelope layout.
"""MHA pool with the page-major page-granularity envelope layout.
All layers/slots share one contiguous ``uint8`` ``_raw`` buffer; per-layer K/V
are 4-D strided views ``(num_pages, page_size, head_num, head_dim*)`` built by
``mem_cache/layout/page_major.py``. Token id ``t`` -> page ``t // page_size``,
slot ``t % page_size``; the reserved padding slot 0 lives in page 0. At
``page_size == 1`` a page is a single slot (token-granularity envelope).
Supported: the standard CUDA Triton attention + native move path. The tiled KV
copy kernel, CPU offloading, and the spec-decode prefix-commit kernel all assume
the per-layer contiguous 3-D layout; here they fail loudly rather than silently
mis-indexing the strided views.
NON-CONSTRUCTIBLE: the strided 4-D view builder and its write kernel are
gone, and ServerArgs rejects the static page-major arm at boot. The class
stays as the seat for the per-layer-view reimplementation.
"""
def __init__(
@@ -3221,89 +3210,12 @@ class PageMajorMHATokenToKVPool(MHATokenToKVPool):
)
def _create_buffers(self):
# One contiguous byte buffer holds all layers/slots; per-layer K/V are
# 4-D strided views in the page-granularity envelope layout (see
# mem_cache/layout/page_major.py).
total_slots = self.size + self.page_size
assert total_slots % self.page_size == 0, (
f"page_major_layer_major needs (size + page_size) divisible by "
f"page_size; got size={self.size}, page_size={self.page_size}"
)
num_pages = total_slots // self.page_size
entry_bytes = mha_entry_bytes(
layer_num=self.layer_num,
head_num=self.head_num,
head_dim=self.head_dim,
v_head_dim=self.v_head_dim,
itemsize=self.store_dtype.itemsize,
)
total_bytes = num_pages * self.page_size * entry_bytes
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.enable_custom_mem_pool
else nullcontext()
):
# Unset slots read as zeros (matches the per-layer pool).
self._raw = torch.zeros(
total_bytes, dtype=torch.uint8, device=self.device
)
self.k_buffer, self.v_buffer = build_page_major_mha_views(
self._raw,
layer_num=self.layer_num,
head_num=self.head_num,
head_dim=self.head_dim,
v_head_dim=self.v_head_dim,
store_dtype=self.store_dtype,
page_size=self.page_size,
num_pages=num_pages,
)
# stride(0) * itemsize is the per-page byte stride; for these strided
# views np.prod(shape[1:]) would not equal it, so compute it directly.
self.k_data_ptrs = torch.tensor(
[x.data_ptr() for x in self.k_buffer],
dtype=torch.uint64,
device=self.device,
)
self.v_data_ptrs = torch.tensor(
[x.data_ptr() for x in self.v_buffer],
dtype=torch.uint64,
device=self.device,
)
self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0)
self.data_strides = torch.tensor(
[x.stride(0) * x.dtype.itemsize for x in (self.k_buffer + self.v_buffer)],
device=self.device,
)
def _store_kv_layer(
self,
layer_idx: int,
loc: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
):
# Single-launch Triton write into the 4-D envelope view. The parent's
# view(-1, row_dim) path can't merge the strided 4-D dims.
store_cache_4d(
self.k_buffer[layer_idx],
self.v_buffer[layer_idx],
cache_k,
cache_v,
loc,
page_size=self.page_size,
)
def _move_kv_cache_impl(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
# Strided 4-D views: the tiled copy kernel assumes stride == row bytes, so
# always take the native move (it splits token ids into
# (page_id, slot_in_page) for the 4-D advanced index).
move_kv_cache_native(
self.k_buffer,
self.v_buffer,
tgt_loc,
src_loc,
page_size=self.page_size,
raise NotImplementedError(
"PageMajorMHATokenToKVPool: the strided 4-D envelope views were "
"removed; the static-pool page-major layout is temporarily "
"unsupported (ServerArgs rejects it at startup). "
"--enable-unified-memory provides the page-major layout with "
"per-layer views."
)
# The methods below assume the per-layer contiguous 3-D layout. The 4-D
@@ -3740,9 +3652,9 @@ 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->dense full-KV translate for the model-level MLA entry points
# 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_dense` for the unified pool.
# identity for a static pool, `translate_kv_loc_for_kernel` for the unified pool.
self._full_translate = lambda ids: ids
self.use_mla = use_mla
if full_kv_pool is not None:
@@ -3985,7 +3897,7 @@ class HybridLinearKVPool(KVCache):
)
else:
# Mirror the MHA branch: `full_loc` is the unified pool's
# pre-translated (dense) loc; None for a static pool.
# pre-translated (kernel-facing) loc; None for a static pool.
write_loc = full_loc if full_loc is not None else loc
with self._transfer_id_context(layer):
self.full_kv_pool.set_kv_buffer(
@@ -4029,16 +3941,16 @@ class HybridLinearKVPool(KVCache):
loc: torch.Tensor,
cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor,
loc_is_dense: bool = False,
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 dense id space here.
# pool, so translate to the kernel-facing id space here.
#
# `loc_is_dense`: the caller already translated `loc` (the unified-pool
# `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_dense:
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)
@@ -4631,17 +4543,11 @@ def move_kv_cache_native(
v_buffer: List[torch.Tensor],
tgt_loc: torch.Tensor,
src_loc: torch.Tensor,
page_size: int = 1,
):
"""Move token-granular K/V rows from ``src_loc`` to ``tgt_loc``.
Supports two buffer shapes:
- 3-D ``[max_slots, head_num, head_dim]`` (per-layer pool): direct advanced
indexing on dim 0; ``page_size`` is ignored.
- 4-D ``[num_pages, page_size, head_num, head_dim]`` (envelope layout): split
each token id into ``(page_id, slot_in_page)`` and use 2-D advanced
indexing. PyTorch resolves the strided byte address via the view's strides.
Buffers are the per-layer 3-D ``[max_slots, head_num, head_dim]`` pools;
direct advanced indexing on dim 0.
"""
if tgt_loc.numel() == 0:
return
@@ -4649,21 +4555,8 @@ def move_kv_cache_native(
tgt_loc_flat = tgt_loc.view(-1).long()
src_loc_flat = src_loc.view(-1).long()
for k_cache, v_cache in zip(k_buffer, v_buffer):
if k_cache.ndim == 4:
if page_size == 1:
# Degenerate (num_pages, 1, head, dim): token id == page id.
k_cache[tgt_loc_flat, 0] = k_cache[src_loc_flat, 0]
v_cache[tgt_loc_flat, 0] = v_cache[src_loc_flat, 0]
else:
tgt_page = tgt_loc_flat // page_size
tgt_tok = tgt_loc_flat % page_size
src_page = src_loc_flat // page_size
src_tok = src_loc_flat % page_size
k_cache[tgt_page, tgt_tok] = k_cache[src_page, src_tok]
v_cache[tgt_page, tgt_tok] = v_cache[src_page, src_tok]
else:
k_cache[tgt_loc_flat] = k_cache[src_loc_flat]
v_cache[tgt_loc_flat] = v_cache[src_loc_flat]
k_cache[tgt_loc_flat] = k_cache[src_loc_flat]
v_cache[tgt_loc_flat] = v_cache[src_loc_flat]
@triton.jit
python/sglang/srt/mem_cache/multi_ended_allocator.py
+71 -91
View File
@@ -114,7 +114,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
need_sort: bool = False,
forward_stream: Optional[torch.cuda.Stream] = None,
lazy_compaction: bool = False,
kernel_page_multiplier: int = 1,
kernel_page_multiplier: Optional[int] = None,
):
spec = unified_buffer.spec(sub_pool_name)
max_slots = unified_buffer.max_slots(sub_pool_name)
@@ -134,11 +134,14 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
self.entry_bytes = spec.entry_bytes()
self.min_slot_index = unified_buffer.min_slot_index(sub_pool_name)
self.is_id_owner = is_id_owner
# Dense (kernel-facing) index space scale: the page-major envelope of a
# multi-layer uniform-entry sub-pool (MLA) is a valid dense paged pool
# once page ids are scaled by layer_num — `translate_kv_loc_dense` emits
# that space. 1 for sub-pools whose kernels take real physical ids.
self.kernel_page_multiplier = kernel_page_multiplier
# Kernel-facing page-stride scale, from the spec that owns the layout.
# `kernel_page_multiplier=` overrides it only for tests pinning the
# multiplier-1 collapse.
self.kernel_page_multiplier = (
spec.blocks_per_page()
if kernel_page_multiplier is None
else kernel_page_multiplier
)
# Zero page envelopes on hand-out — see _maybe_zero_pages.
self._zero_pages_on_alloc = isinstance(kvcache, UnifiedMLATokenToKVPool)
# Overlap mode: `free` drops a wait_stream(forward_stream) barrier so its
@@ -641,7 +644,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
return phys_pages
def _maybe_zero_pages(self, phys_pages: torch.Tensor) -> None:
"""Zero the page ENVELOPES on hand-out (MLA-dense full pool only):
"""Zero the page ENVELOPES on hand-out (MLA full pool only):
the MLA kernels arithmetically mask the rows beyond seq_len, so
never-written page bytes must read as finite values. Runs on the
schedule stream, ordered before the consuming forward by the
@@ -713,58 +716,48 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
result = phys_pages * self.page_size + offsets
return torch.clamp_min(result, 0)
def translate_kv_loc_dense(
def translate_kv_loc_for_kernel(
self,
virt_tokens: torch.Tensor,
*,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Translate virtual token ids to DENSE (kernel-facing) ids.
"""Virtual token ids -> kernel-facing ids:
dense(t) = (t // ps) * (ps * kernel_page_multiplier) + t % ps for the
physical token t — i.e. `translate_kv_loc` with the page stride scaled by
`kernel_page_multiplier` (= layer_num for a dense-view MLA sub-pool; see
`build_dense_mla_views`). Internal machinery (compaction, in-flight write
sets) MUST keep using `translate_kv_loc`: dense ids are for kernels only.
kernel_id(t) = (t // ps) * (ps * kernel_page_multiplier) + t % ps
The tombstone clamp routes -1 entries to dense id 0 — inside the page-0
reserved sink for every layer view. Supports ``out=`` like
`translate_kv_loc` for cuda-graph buffer stability.
Internal machinery (compaction, in-flight write sets) MUST keep using
`translate_kv_loc`: kernel-facing ids are for kernels only. Tombstones (-1)
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.
"""
if self.kernel_page_multiplier == 1:
return self.translate_kv_loc(virt_tokens, out=out)
if out is not None:
ps = self.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)
assert out.dtype == torch.int64, (
f"translate_kv_loc_dense: out= dtype must be int64 (matches v2p), "
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_dense: out= shape {tuple(out.shape)} must "
f"translate_kv_loc_for_kernel: out= shape {tuple(out.shape)} must "
f"match virt_tokens shape {tuple(virt_tokens.shape)}"
)
with record_function("MultiEndedAlloc.translate_kv_loc_dense"):
dense_page_stride = self.page_size * self.kernel_page_multiplier
if self.page_size == 1:
# dense = phys * multiplier; tombstone -1 scales negative → clamp 0.
if out is not None:
tmp = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
tmp = torch.clamp_min(tmp * dense_page_stride, 0)
out.copy_(tmp)
return out
result = torch.index_select(self.virtual_to_physical, 0, virt_tokens)
return torch.clamp_min(result * dense_page_stride, 0)
virt_pages = virt_tokens // self.page_size
offsets = virt_tokens % self.page_size
if out is not None:
torch.index_select(self.virtual_to_physical, 0, virt_pages, out=out)
out.mul_(dense_page_stride)
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)
# tombstoned page: -1*dense_page_stride + offset < 0
out.clamp_(min=0)
return out
phys_pages = self.virtual_to_physical[virt_pages]
result = phys_pages * dense_page_stride + offsets
return torch.clamp_min(result, 0)
return out.clamp_(min=0)
# -- alloc --
@@ -1722,7 +1715,6 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
need_sort: bool = False,
forward_stream: Optional[torch.cuda.Stream] = None,
lazy_compaction: bool = False,
full_kernel_page_multiplier: int = 1,
):
full_max = unified_buffer.max_slots("full")
super().__init__(
@@ -1750,7 +1742,6 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
need_sort=need_sort,
forward_stream=forward_stream,
lazy_compaction=lazy_compaction,
kernel_page_multiplier=full_kernel_page_multiplier,
)
self.mamba_allocator = MultiEndedAllocator(
kvcache=kvcache.mamba_pool,
@@ -1911,26 +1902,25 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
"""Page-level virtual->physical table of the full sub-pool. Kernels that
build the MLA block table directly from req_to_token (e.g. trtllm_mla,
flashmla) gather through this to turn a VIRTUAL page into a physical one,
then scale by `kernel_page_multiplier` to reach the dense per-page block.
then scale by `kernel_page_multiplier` to reach the per-page block.
"""
return self.full_attn_allocator.virtual_to_physical
def translate_kv_loc_dense(
def translate_kv_loc_for_kernel(
self,
loc: torch.Tensor,
*,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Full-pool virtual TOKEN ids -> DENSE (kernel-facing) ids. Falls back
to the physical translate when `kernel_page_multiplier == 1` (MHA)."""
return self.full_attn_allocator.translate_kv_loc_dense(loc, out=out)
"""Full-pool virtual TOKEN ids -> kernel-facing ids."""
return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
def translate_kv_indices_for_transfer(
self, kv_indices: torch.Tensor
) -> torch.Tensor:
"""Virtual TOKEN ids -> PHYSICAL token ids for the PD transfer engine.
PHYSICAL, not dense: the transfer registers page ENVELOPES (see
PHYSICAL, not kernel-facing: the transfer registers page ENVELOPES (see
`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`).
"""
return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
@@ -2255,45 +2245,35 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
*,
out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""SWA-layer read path: virtual TOKEN ids -> swa-physical TOKEN ids (int32,
matching the non-shared API). Page math against the swa side's v2p table.
Supports ``out=`` (int32, same shape) for cuda-graph buffer stability.
"""
if out is not None:
assert out.dtype == torch.int32, (
f"translate_loc_from_full_to_swa: out= dtype must be int32 "
f"(matches SWA Triton kernel contract), got {out.dtype}"
)
assert out.shape == kv_indices.shape, (
f"translate_loc_from_full_to_swa: out= shape "
f"{tuple(out.shape)} must match kv_indices shape "
f"{tuple(kv_indices.shape)}"
)
# Tombstone-safety clamp (mirrors the full-side clamp): tombstoned (-1)
# v2p_swa entries must not reach `swa_k_buffer[-1]` (illegal under replay).
# Clamp to 0 routes them to the reserved padding sink (slot 0).
if self.swa_attn_allocator.page_size == 1:
if out is not None:
# Gather into a transient int64, then cast into out (`out.copy_`).
tmp = torch.index_select(
self.swa_attn_allocator.virtual_to_physical, 0, kv_indices
)
tmp = torch.clamp_min(tmp, 0)
out.copy_(tmp.to(torch.int32))
return out
result = self.swa_attn_allocator.virtual_to_physical[kv_indices]
result = torch.clamp_min(result, 0)
return result.to(torch.int32)
ps = self.swa_attn_allocator.page_size
virt_pages = kv_indices // ps
offsets = kv_indices % ps
swa_phys_pages = self.swa_attn_allocator.virtual_to_physical[virt_pages]
result = (swa_phys_pages * ps + offsets).to(torch.int32)
result = torch.clamp_min(result, 0)
if out is not None:
out.copy_(result)
return out
return result
"""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)
@property
def kernel_page_multiplier(self) -> int:
return self.full_attn_allocator.kernel_page_multiplier
@property
def full_v2p_page_table(self) -> torch.Tensor:
"""Page-level virtual->physical table of the full sub-pool."""
return self.full_attn_allocator.virtual_to_physical
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)
@property
def swa_kernel_page_multiplier(self) -> int:
return self.swa_attn_allocator.kernel_page_multiplier
@property
def swa_v2p_page_table(self) -> torch.Tensor:
"""Page-level virtual->physical table of the SWA sub-pool."""
return self.swa_attn_allocator.virtual_to_physical
# -- alloc --
python/sglang/srt/mem_cache/unified_memory_pool.py
+152 -158
View File
@@ -29,17 +29,15 @@ from dataclasses import dataclass
from typing import Dict, List, NamedTuple, Optional, Tuple
import torch
import triton
from torch.profiler import record_function
from sglang.kernels.ops.kvcache.cache_move import store_cache_4d_kernel
from sglang.kernels.ops.kvcache.zero_pages import zero_pages
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
from sglang.srt.mem_cache.layout.page_major import (
build_dense_mla_views,
build_mha_views,
build_mla_views,
build_page_major_mamba_views,
build_page_major_mha_views,
)
from sglang.srt.mem_cache.memory_pool import (
HybridLinearKVPool,
@@ -47,7 +45,6 @@ from sglang.srt.mem_cache.memory_pool import (
MambaPool,
MHATokenToKVPool,
MLATokenToKVPool,
move_kv_cache_native,
unwrap_write_loc,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
@@ -96,6 +93,19 @@ class SubPoolSpec(ABC):
"""Storage dtype (informational). Multi-dtype subclasses return the dominant buffer's."""
raise NotImplementedError
def view_tail_pad_bytes(self, page_size: int) -> int:
"""Bytes this sub-pool's views reach PAST its last page envelope."""
return 0
def blocks_per_page(self) -> int:
"""Row-blocks one page holds in this sub-pool's kernel-facing id space.
The page envelope is a uniform array of equally wide row-blocks, so a
kernel-facing id is the physical page scaled by this count (see
`MultiEndedAllocator.translate_kv_loc_for_kernel`). 1 means the kernel-facing ids are the physical ones.
"""
return 1
@dataclass(frozen=True, kw_only=True)
class MHASubPoolSpec(SubPoolSpec):
@@ -140,6 +150,13 @@ class MHASubPoolSpec(SubPoolSpec):
+ page_size * self.k_row_bytes()
)
def view_tail_pad_bytes(self, page_size: int) -> int:
return page_size * self.entry_bytes()
def blocks_per_page(self) -> int:
"""Row-blocks per page in the kernel-facing id space (one K + one V per layer)."""
return 2 * self.layer_num
def get_dtype(self) -> torch.dtype:
return self.store_dtype
@@ -174,6 +191,14 @@ class MLASubPoolSpec(SubPoolSpec):
def entry_bytes(self) -> int:
return self.layer_num * self.kv_cache_dim * self.store_dtype.itemsize
def view_tail_pad_bytes(self, page_size: int) -> int:
return page_size * self.entry_bytes()
def blocks_per_page(self) -> int:
"""One latent row per layer, so L blocks per page (MHA has 2L: a K
block and a V block per layer)."""
return self.layer_num
def get_dtype(self) -> torch.dtype:
return self.store_dtype
@@ -210,13 +235,28 @@ class MambaSubPoolSpec(SubPoolSpec):
# ---------------------------------------------------------------------------
# UnifiedKVPool — the byte buffer + the strided per-sub-pool views
# UnifiedKVPool — the byte buffer + the per-sub-pool views
# ---------------------------------------------------------------------------
def unified_memory_supported_for_model(model_config, *, use_mla_backend: bool) -> bool:
"""Whether this model's KV geometry can back the unified memory pool."""
return use_mla_backend or not model_config.has_asymmetric_kv
def _assert_kernel_id_bound(*, sub_pool_name: str, n_rows: int) -> None:
"""Check if kernel-facing ids can flow through int32 read-index buffers."""
assert n_rows < 2**31, (
f"sub-pool {sub_pool_name!r}: kernel-facing id space has {n_rows} rows, "
f"exceeding the int32 bound (2^31) that read-index buffers assume. "
"Reduce max_total_num_tokens or the layer count."
)
class UnifiedKVPool:
"""One physical `uint8` byte buffer shared by 2 sub-pools, each exposing
strided per-layer views. Allocators keep byte ranges disjoint; no usage tracking here.
per-layer views over its own byte range (contiguous per layer for KV,
strided for the Mamba state). Allocators keep byte ranges disjoint; no usage tracking here.
"""
def __init__(
@@ -227,7 +267,6 @@ class UnifiedKVPool:
device: str,
enable_memory_saver: bool,
page_size: int = 1,
view_tail_pad_bytes: int = 0,
):
assert page_size >= 1, f"page_size must be >= 1; got {page_size}"
assert len(sub_pool_specs) == 2, (
@@ -253,13 +292,12 @@ class UnifiedKVPool:
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
)
# `view_tail_pad_bytes` extends the ALLOCATION only (dense MLA views are
# per-layer shifted, so the last layer's view reaches past the final page
# envelope); all slot/watermark math stays on the unpadded `total_bytes`.
self.view_tail_pad_bytes = view_tail_pad_bytes
self.view_tail_pad_bytes = max(
spec.view_tail_pad_bytes(page_size) for spec in sub_pool_specs
)
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
self._raw = torch.empty(
total_bytes + view_tail_pad_bytes, dtype=torch.uint8, device=device
total_bytes + self.view_tail_pad_bytes, dtype=torch.uint8, device=device
)
if envs.SGLANG_DEBUG_POISON_POOL.get():
# Debug: bf16-NaN-fill so NaN-unsafe reads of never-written bytes
@@ -275,7 +313,7 @@ class UnifiedKVPool:
self._max_slots: Dict[str, int] = {}
self._anchor_bytes: Dict[str, int] = {}
self._min_slot_index: Dict[str, int] = {}
# MHA: (k_buffer, v_buffer); MLA: [per-layer dense views];
# MHA: (k_buffer, v_buffer); MLA: [per-layer per-layer views];
# Mamba: (conv_state_list, temporal_state)
self._mha_views: Dict[str, Tuple[List[torch.Tensor], List[torch.Tensor]]] = {}
self._mla_views: Dict[str, List[torch.Tensor]] = {}
@@ -405,7 +443,12 @@ class UnifiedKVPool:
max_slots: int,
page_size: int,
) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
return build_page_major_mha_views(
num_pages = max_slots // page_size
_assert_kernel_id_bound(
sub_pool_name=spec.name,
n_rows=num_pages * spec.blocks_per_page() * page_size,
)
return build_mha_views(
self._raw,
layer_num=spec.layer_num,
head_num=spec.head_num,
@@ -413,7 +456,7 @@ class UnifiedKVPool:
v_head_dim=spec.v_head_dim,
store_dtype=spec.store_dtype,
page_size=page_size,
num_pages=max_slots // page_size,
num_pages=num_pages,
anchor_bytes=anchor_bytes,
)
@@ -424,13 +467,18 @@ class UnifiedKVPool:
max_slots: int,
page_size: int,
) -> List[torch.Tensor]:
return build_dense_mla_views(
num_pages = max_slots // page_size
_assert_kernel_id_bound(
sub_pool_name=spec.name,
n_rows=num_pages * spec.blocks_per_page() * page_size,
)
return build_mla_views(
self._raw,
layer_num=spec.layer_num,
kv_cache_dim=spec.kv_cache_dim,
store_dtype=spec.store_dtype,
page_size=page_size,
num_pages=max_slots // page_size,
num_pages=num_pages,
anchor_bytes=anchor_bytes,
)
@@ -450,10 +498,16 @@ class UnifiedKVPool:
class UnifiedMHATokenToKVPool(MHATokenToKVPool):
"""MHA KV pool whose `k_buffer`/`v_buffer` are strided views into a `UnifiedKVPool`.
"""MHA KV pool whose per-layer `k_buffer`/`v_buffer` are `build_mha_views`
views into a `UnifiedKVPool` (requires uniform K/V rows).
Relocation uses the native move (strided views break the tiled Triton kernel that
assumes stride == row bytes). `set_kv_buffer` gets PHYSICAL slot ids; never translates.
Views are contiguous `(n_rows, head_num, head_dim)`; locs are
kernel_id(t) = (t // ps) * (ps * 2 * layer_num) + t % ps
which is layer- and K/V-independent, each view's storage_offset folding in
its block origin (layer l's K at block 2l, V at 2l+1). `move_kv_cache` is
the exception: compaction passes REAL physical token ids.
"""
def __init__(
@@ -467,17 +521,19 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
enable_alt_stream: bool = True,
):
spec = unified_buffer.mha_spec(sub_pool_name)
k_buffer, v_buffer = unified_buffer.mha_views_for(sub_pool_name)
k_views, v_views = unified_buffer.mha_views_for(sub_pool_name)
max_slots = unified_buffer.max_slots(sub_pool_name)
self._unified_buffer = unified_buffer
self._sub_pool_name = sub_pool_name
self._k_views = k_buffer
self._v_views = v_buffer
self._page_size = page_size
self._k_views = k_views
self._v_views = v_views
self._num_pages = max_slots // page_size
self._page_bytes = page_size * spec.entry_bytes()
view_rows = self._num_pages * spec.blocks_per_page() * page_size
super().__init__(
size=max_slots - 1, # -1 for reserved slot 0
size=view_rows - page_size,
page_size=page_size,
dtype=spec.store_dtype,
head_num=spec.head_num,
@@ -489,126 +545,73 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
start_layer=start_layer,
end_layer=end_layer,
enable_alt_stream=enable_alt_stream,
enable_kv_cache_copy=False, # strided views — force native move
enable_kv_cache_copy=False,
kv_cache_layout="page_major",
)
def _create_buffers(self):
self.k_buffer = self._k_views
self.v_buffer = self._v_views
# For external inspectors only; the native move path doesn't consume them.
self.k_data_ptrs = torch.tensor(
[x.data_ptr() for x in self.k_buffer],
dtype=torch.uint64,
device=self.device,
)
self.v_data_ptrs = torch.tensor(
[x.data_ptr() for x in self.v_buffer],
dtype=torch.uint64,
device=self.device,
)
self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0)
self.data_strides = torch.tensor(
[x.stride(0) * x.dtype.itemsize for x in (self.k_buffer + self.v_buffer)],
device=self.device,
)
def _clear_buffers(self):
# Lifetime owned by UnifiedKVPool; do not delete the views.
pass
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
# tgt_loc/src_loc are PHYSICAL slot ids; native move only (strided views).
if tgt_loc.numel() == 0:
return
with record_function("UnifiedMHA.move_kv_cache"):
move_kv_cache_native(
self.k_buffer,
self.v_buffer,
tgt_loc,
src_loc,
page_size=self._page_size,
)
def get_kv_size_bytes(self):
return 0, 0 # UnifiedKVPool logs the total; per-sub-pool would double-count
def set_kv_buffer(
self,
layer,
loc: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
k_scale=None,
v_scale=None,
layer_id_override: Optional[int] = None,
dcp_kv_mask: Optional[torch.Tensor] = None,
):
# Decode context parallel (dcp_kv_mask) unsupported; fail loud.
assert dcp_kv_mask is None, (
"UnifiedMHATokenToKVPool.set_kv_buffer: decode context parallel "
"(dcp_kv_mask) is not supported with --enable-unified-memory."
)
# Bypass super().set_kv_buffer: the parent's `k_cache.view(-1, row_dim)` can't
# merge our 4-D layer-major view (stride[0]=page_bytes) at page_size>1. Call
# store_cache_4d_kernel directly. `loc` is PHYSICAL token ids — no v2p translate.
with record_function("UnifiedMHA.set_kv_buffer"):
if cache_k.dtype != self.dtype:
if k_scale is not None:
cache_k.div_(k_scale)
if v_scale is not None:
cache_v.div_(v_scale)
cache_k = cache_k.to(self.dtype)
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype)
layer_id = (
layer.layer_id if layer_id_override is None else layer_id_override
) - self.start_layer
k_view = self.k_buffer[layer_id]
v_view = self.v_buffer[layer_id]
ps = self._page_size
N = loc.numel()
if N == 0:
return
head_num = k_view.shape[2]
head_dim = k_view.shape[3]
v_head_dim = v_view.shape[3]
K_ROW_DIM = head_num * head_dim
V_ROW_DIM = head_num * v_head_dim
BLOCK = 128
row_dim_max = K_ROW_DIM if K_ROW_DIM > V_ROW_DIM else V_ROW_DIM
store_cache_4d_kernel[(N, triton.cdiv(row_dim_max, BLOCK), 2)](
k_view,
v_view,
cache_k,
cache_v,
loc,
k_view.stride(0),
k_view.stride(1),
v_view.stride(0),
v_view.stride(1),
cache_k.stride(0),
cache_v.stride(0),
K_ROW_DIM=K_ROW_DIM,
V_ROW_DIM=V_ROW_DIM,
PAGE_SIZE=ps,
BLOCK=BLOCK,
num_warps=4,
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
"""Relocate slots by whole page envelope.
`tgt_loc`/`src_loc` are REAL physical token ids, not kernel-facing ids.
"""
if tgt_loc.numel() == 0:
return
# The envelope view below starts at byte 0, so this sub-pool must be
# anchored there; a non-zero anchor moves another sub-pool's bytes, and
# the ids stay in range so nothing downstream notices.
assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0
ps = self.page_size
tgt_pages = tgt_loc.view(-1, ps)[:, 0] // ps
src_pages = src_loc.view(-1, ps)[:, 0] // ps
with record_function("UnifiedMHA.move_kv_cache"):
env = self._unified_buffer._raw[: self._num_pages * self._page_bytes].view(
self._num_pages, self._page_bytes
)
env[tgt_pages] = env[src_pages]
def get_contiguous_buf_infos(self):
raise NotImplementedError(
"unified layout has no per-layer contiguous regions; "
"KV transfer / disaggregation is unsupported."
)
def get_cpu_copy(self, indices, mamba_indices=None):
raise NotImplementedError(
"CPU offloading is unsupported under the unified layout."
)
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
raise NotImplementedError(
"CPU offloading is unsupported under the unified layout."
)
def set_kv_buffer_prefix_valid(self, *args, **kwargs):
raise NotImplementedError(
"prefix-valid commit is unsupported under the unified layout "
"(_set_kv_buffer_prefix_valid_impl assumes token-id indexing)."
)
class UnifiedMLATokenToKVPool(MLATokenToKVPool):
"""MLA KV pool whose per-layer `kv_buffer` entries are DENSE views into a
`UnifiedKVPool` (see `build_dense_mla_views`).
"""MLA KV pool whose per-layer `kv_buffer` entries are kernel-facing views into a
`UnifiedKVPool` (see `build_mla_views`).
Loc-space contract: every loc this pool receives through the KVCache API
(`set_kv_buffer` / `set_mla_kv_buffer` / `get_mla_kv_buffer`, and the
kv_indices consumed by attention kernels reading `get_key_buffer` /
`get_value_buffer`) is a DENSE id — the `translate_kv_loc_dense` output
`get_value_buffer`) is a kernel-facing id — the `translate_kv_loc_for_kernel` output
dense(t) = (t // ps) * (ps * layer_num) + t % ps
kernel_id(t) = (t // ps) * (ps * layer_num) + t % ps
which is layer-independent (the layer offset is folded into each view's
storage_offset), so the stock `MLATokenToKVPool` read/write methods work on
@@ -638,13 +641,12 @@ 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()
# Dense row count per view; also the OOB bound for dense locs.
self._dense_size = self._num_pages * spec.layer_num * page_size
self._view_rows = self._num_pages * spec.layer_num * page_size
super().__init__(
# OOB checks bound locs by `size + page_size`; dense ids run to
# `_dense_size` (page 0 is the reserved padding sink).
size=self._dense_size - page_size,
# OOB checks bound locs by `size + page_size`; kernel-facing ids run to
# `_view_rows` (page 0 is the reserved padding sink).
size=self._view_rows - page_size,
page_size=page_size,
dtype=kv_cache_dtype,
kv_lora_rank=spec.kv_lora_rank,
@@ -669,8 +671,8 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
``raw_ptr + physical_page_id * page_envelope_bytes``.
The transfer item is the whole page envelope (all layers of one page)
rather than a per-layer region, because the per-layer dense views
overlap and index in dense ids. Both sides must therefore build the
rather than a per-layer region, because the per-layer per-layer views
overlap and index in kernel-facing ids. Both sides must therefore build the
pool with identical specs.
"""
# The address formula omits the anchor; a nonzero one would mis-address.
@@ -681,7 +683,7 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
"""Relocate whole page envelopes.
`tgt_loc`/`src_loc` are REAL physical token ids (NOT dense ids): both
`tgt_loc`/`src_loc` are REAL physical token ids (NOT kernel-facing ids): both
compaction paths expand page ids into page-major-ordered token runs
(`pages[:, None] * ps + offsets`), relied on here to recover the page
lists. One contiguous envelope copy replaces the per-layer strided moves.
@@ -1164,16 +1166,12 @@ def init_unified_mamba_pools(
max_total_num_tokens * full_spec.entry_bytes()
+ max_mamba_cache_size * mamba_spec.entry_bytes()
)
# Dense MLA views are per-layer shifted, so the last layer's view reaches one
# page envelope past the final page — allocation-only tail pad (~page bytes).
view_tail_pad_bytes = page_size * full_spec.entry_bytes() if use_mla_backend else 0
shared_pool = UnifiedKVPool(
total_bytes=total_bytes,
sub_pool_specs=[full_spec, mamba_spec],
device=device,
enable_memory_saver=enable_memory_saver,
page_size=page_size,
view_tail_pad_bytes=view_tail_pad_bytes,
)
req_to_token_pool = UnifiedHybridReqToTokenPool(
unified_buffer=shared_pool,
@@ -1233,9 +1231,6 @@ def init_unified_mamba_pools(
need_sort=need_sort,
forward_stream=forward_stream,
lazy_compaction=lazy_compaction,
full_kernel_page_multiplier=(
len(full_attention_layer_ids) if use_mla_backend else 1
),
)
# Wrap the composite's mamba MultiEndedAllocator in a slot allocator (PHYSICAL view).
@@ -1249,9 +1244,9 @@ def init_unified_mamba_pools(
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 dense space internally
# 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_dense
token_to_kv_pool._full_translate = allocator.translate_kv_loc_for_kernel
logger.info(
"[unified-memory-pool] ============================================================"
@@ -1263,7 +1258,7 @@ def init_unified_mamba_pools(
if use_mla_backend:
logger.info(
"[unified-memory-pool] full_layers=%d, mamba_layers=%d, kv_lora_rank=%d, "
"qk_rope_head_dim=%d, page_size=%d (dense views, kernel_page_multiplier=%d, "
"qk_rope_head_dim=%d, page_size=%d (per-layer views, kernel_page_multiplier=%d, "
"view_tail_pad=%d B)",
len(full_attention_layer_ids),
len(mamba_layer_ids),
@@ -1271,18 +1266,20 @@ def init_unified_mamba_pools(
qk_rope_head_dim,
page_size,
len(full_attention_layer_ids),
view_tail_pad_bytes,
shared_pool.view_tail_pad_bytes,
)
else:
logger.info(
"[unified-memory-pool] full_layers=%d, mamba_layers=%d, head_num=%d, head_dim=%d, "
"page_size=%d, is_draft_worker=%s",
"page_size=%d, is_draft_worker=%s (%s)",
len(full_attention_layer_ids),
len(mamba_layer_ids),
head_num,
head_dim,
page_size,
is_draft_worker,
"per-layer views, kernel_page_multiplier=%d, view_tail_pad=%d B"
% (full_spec.blocks_per_page(), shared_pool.view_tail_pad_bytes),
)
logger.info(
"[unified-memory-pool] total_bytes=%d, max_total_num_tokens=%d, max_mamba_cache_size=%d, "
@@ -1418,24 +1415,12 @@ class UnifiedSWAKVPool(SWAKVPool):
return # no-op in shared mode (the swa-side v2p IS the mapping)
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
"""Virtual token ids -> swa-physical token ids (int32)."""
"""Virtual token ids -> swa kernel-facing ids (int64)."""
assert self._swa_allocator is not None, (
"UnifiedSWAKVPool.translate_loc_from_full_to_swa called before "
"attach_allocators"
)
ps = self._swa_allocator.page_size
# Tombstone-safety clamp, matching MultiEndedAllocator.translate_kv_loc:
# a tombstoned v2p entry (-1) must not reach the caller as a negative
# loc. Clamp to 0 routes it to the reserved padding sink instead.
if ps == 1:
swa_locs = self._swa_allocator.virtual_to_physical[kv_indices]
else:
virt_pages = kv_indices // ps
offsets = kv_indices % ps
swa_phys_pages = self._swa_allocator.virtual_to_physical[virt_pages]
# Tombstoned page: -1 * ps + offset lands in [-ps, -1].
swa_locs = swa_phys_pages * ps + offsets
return swa_locs.clamp(min=0).to(torch.int32)
return self._swa_allocator.translate_kv_loc_for_kernel(kv_indices)
def get_state_buf_infos(self):
return self.swa_kv_pool.get_contiguous_buf_infos()
@@ -1664,6 +1649,15 @@ def init_unified_swa_pools(
"[unified-memory-pool] ============================================================"
)
logger.info("[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=SWA hybrid")
logger.info(
"[unified-memory-pool] %s",
"per-layer views, kernel_page_multiplier full=%d swa=%d, view_tail_pad=%d B"
% (
full_spec.blocks_per_page(),
swa_spec.blocks_per_page(),
shared_pool.view_tail_pad_bytes,
),
)
logger.info(
"[unified-memory-pool] full_layers=%d, swa_layers=%d, head_num=%d, head_dim=%d, "
"v_head_dim=%d, swa_head_num=%d, swa_head_dim=%d, swa_v_head_dim=%d, "