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/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