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

[kv-shard 1/4] Logical-page placement with UnifiedRadixCache (#38356)

Browse Source
This commit is contained in:
Shunkangz
2026-09-10 22:06:42 -07:00
committed by GitHub
parent 94ce940ff8
commit 40a84d6dfc
11 changed files with 1551 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+10
View File
@@ -1131,6 +1131,15 @@ class Req(ReqDllmMixin):
self.lock_receipt: DecLockRefParams = DecLockRefParams()
# Whether the prefill-time SWA tree lock has been released early
self.swa_prefix_lock_released: bool = False
# Logical-page KV sharding: rotation base of the chain this request
# extends (owner of position-page P is (base + P) % shard_size).
# Refreshed at every sharded alloc — read through last_node, or drawn
# least-full for a new chain — and consumed by the radix insert to
# stamp new tree nodes. Allocation itself must NOT read it back when
# a tree node is available (the cache_unfinished_req dedup rebind
# would make it stale); the only allocation-time reader is the
# ChunkCache fallback, which has no tree nodes and no rebind.
self.kv_rotation_base: Optional[int] = None
# Whether or not if it is chunked. It increments whenever
# it is chunked, and decrement whenever chunked request is
@@ -1821,6 +1830,7 @@ class Req(ReqDllmMixin):
self.indexer_topk = None
self.last_node = None
self.kv.cache_protected_len = 0
self.kv_rotation_base = None
self.num_matched_prefix_tokens = 0
self.lock_receipt = DecLockRefParams()
self.swa_prefix_lock_released = False
python/sglang/srt/mem_cache/allocation.py
+63 -3
View File
@@ -17,6 +17,7 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
maybe_write_dsv4_decode,
maybe_write_dsv4_extend,
)
from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
from sglang.srt.mem_cache.common import (
MAMBA_STATE_PER_REQ_NO_CACHE,
@@ -181,13 +182,24 @@ def alloc_paged_token_slots_extend(
req_pool_indices: Optional[torch.Tensor] = None,
batch=None,
):
# Over estimate the number of tokens: assume each request needs a new page.
# Over estimate the number of tokens: assume each request needs a new page
# (one page per CLASS per request under sharding — the min-class
# availability floor must cover every request's ceil(K_i/N) rounding,
# matching the N*ps per-request admission reserve).
allocator = tree_cache.token_to_kv_pool_allocator
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
num_tokens = extend_num_tokens + len(seq_lens_cpu) * (
allocator.page_size * page_interleave_shard_size(allocator)
)
evict_from_tree_cache(tree_cache, num_tokens)
is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c128_attn_allocator")
extra_alloc_kwargs = {}
kv_shard_rotation_bases = None
if page_interleave_shard_size(allocator) > 1:
kv_shard_rotation_bases = _kv_shard_rotation_bases(
tree_cache=tree_cache, batch=batch, prefix_lens_cpu=prefix_lens_cpu
)
extra_alloc_kwargs["rotation_bases"] = kv_shard_rotation_bases
if is_dsv4:
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
# Per-call per-req table for the C128 KV last_loc lookup.
@@ -223,9 +235,55 @@ def alloc_paged_token_slots_extend(
tree_cache.pretty_print()
raise RuntimeError(error_msg)
if kv_shard_rotation_bases is not None:
# The allocator resolved None entries (new chains) in place from the
# least-full class at each request's turn; record the bases for the
# radix insert to stamp onto new tree nodes (UnifiedTreeNode.rotation_base).
for req, base in zip(batch.reqs, kv_shard_rotation_bases):
req.kv_rotation_base = base
return out_cache_loc
def _kv_shard_rotation_bases(
tree_cache: BasePrefixCache, batch: ScheduleBatch, prefix_lens_cpu: torch.Tensor
) -> list:
"""Per-request rotation bases ``b_i`` of the batch's chains, host-only.
The owner class of position-page P is ``(b_i + P) % shard_size``. Rules:
- Read through ``req.last_node`` at alloc time, never a value cached on
the request: ``cache_unfinished_req`` can rebind a chunked request onto
another chain's canonical locs between chunks, changing the base. The
read goes through ``tree_cache.rotation_base_of`` because the node
handle is tree-specific (a NodeId on the unified tree).
- A request without a cached prefix starts a new chain: None here — the
allocator draws from the least-full class at that request's turn (so
the draw sees earlier requests' pops in the same batch) and resolves
the entry in place.
- ChunkCache has no tree nodes (``last_node`` is None); its chunked
continuations fall back to the base recorded on the request at the
previous chunk's alloc (no cross-request reuse, no rebind there).
"""
assert batch is not None and len(batch.reqs) == len(prefix_lens_cpu)
bases = []
for i, req in enumerate(batch.reqs):
if int(prefix_lens_cpu[i]) == 0:
bases.append(None)
continue
node_base = tree_cache.rotation_base_of(req.last_node)
if node_base is not None:
bases.append(node_base)
else:
assert req.kv_rotation_base is not None, (
"sharded extend with a cached prefix but no rotation base: "
"req.last_node carries none and the request recorded none "
"at a previous alloc"
)
bases.append(req.kv_rotation_base)
return bases
def alloc_req_slots(
req_to_token_pool: ReqToTokenPool,
reqs: list[Req],
@@ -273,7 +331,9 @@ def alloc_req_slots(
def _alloc_page_size(batch: ScheduleBatch) -> int:
# DCP swaps in an allocator whose page_size is the configured page_size *
# dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on
# the real allocator's page_size there. Elsewhere the two are equal.
# the real allocator's page_size there. Elsewhere the two are equal --
# including under KV sharding, which widens the index space but keeps the
# allocator page at the physical page.
if (_is_hip or _is_cuda) and get_parallel().dcp_enabled:
return batch.tree_cache.token_to_kv_pool_allocator.page_size
return batch.tree_cache.page_size
python/sglang/srt/mem_cache/allocator/page_interleave.py
+406
View File
@@ -0,0 +1,406 @@
# 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.
# ==============================================================================
"""Allocator for logical-page KV sharding (see mem_cache/page_interleave.py).
Rotated owner-classed page allocation:
physical pages are handed out one at a time from N mirrored per-rank class
free lists. Logical page ``l = loc // ps`` is owned by rank ``l % N`` and is
that rank's local physical page ``l // N``, so class ``r``'s free list IS
rank ``r``'s free pages. Ownership is derived from ``loc`` everywhere (write
filter, gather plan, scratch translation, P/D send filter), which makes any
class choice correct; balance is pure policy:
- ``class(position-page P of a chain) = (b + P) % N`` where ``b`` is the
chain's rotation base — a new chain draws ``b`` from the least-full class,
an extension continues from the phase recorded on ``req.last_node``
(``TreeNode.rotation_base``). Within one cached prefix the owners are
exactly cyclic, so per-rank page counts differ by <= 1 and the prefix
gather stays a regular padded allgather.
Every list, cursor decision, and pop is a pure function of the mirrored
alloc/free stream, so the state is byte-identical across shard-group ranks
by construction (SPMD, no consensus protocol). A freed page is immediately
reusable: nothing strands a page until its whole logical group is free.
``available_size`` reports the MIN-CLASS capacity floor
(``N * min_r free_pages(r) * ps``): in-flight P/D transfers lock tree nodes,
so an aggregate gate could admit a request whose tight class has nothing
evictable — fail-loud where min-class admission defers. Rotation plus
least-full seeding keeps the classes near-balanced, so the floor tracks the
aggregate within the bounded skew.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import KVCache
class PageInterleavePoolAllocator(PagedTokenToKVPoolAllocator):
"""Classed paged allocator over the shard-widened logical index space.
The inherited ``page_size`` is the PHYSICAL page — the working quantum of
every seam that reads it (radix-tree match, chunk flooring, admission
reserve arithmetic, free alignment, wire sampling). The widening is pure
index space: ``size`` logical slots = ``shard_size`` x one rank's
physical slots; per-rank HBM is unchanged.
"""
def __init__(
self,
size: int,
physical_page_size: int,
shard_size: int,
dtype: torch.dtype,
device: str,
kvcache: KVCache,
need_sort: bool,
shard_spec=None,
):
assert shard_size > 1, "PageInterleavePoolAllocator requires shard_size > 1"
# Set before super().__init__: the base constructor calls clear(),
# which builds the class lists from these.
self.physical_page_size = physical_page_size
self.shard_size = shard_size
# PageShardSpec (None only in unit tests): carries the assembly
# scratch capacities the PrefillAdder gates batch admission on.
self.shard_spec = shard_spec
super().__init__(
size * shard_size,
page_size=physical_page_size,
dtype=dtype,
device=device,
kvcache=kvcache,
need_sort=need_sort,
)
def clear(self):
# Logical pages l ∈ [0, shard_size) are reserved — physical page 0 of
# every rank (loc ∈ [0, shard_size * ps) is the padded/trash range).
# Class r's free list holds the allocatable logical pages with
# l % shard_size == r, ascending — exactly rank r's free physical
# pages 1..pages_per_rank at local page l // shard_size.
pages_per_rank = self.num_pages // self.shard_size
self.class_free_pages: List[torch.Tensor] = [
torch.arange(1, pages_per_rank + 1, dtype=torch.int64, device=self.device)
* self.shard_size
+ r
for r in range(self.shard_size)
]
self.class_release_pages: List[torch.Tensor] = [
torch.empty((0,), dtype=torch.int64, device=self.device)
for _ in range(self.shard_size)
]
# None: free right away. A list: hold frees until free_group_end().
self.free_group = None
# free_segment() routes back to free(), so this stays empty; the paged
# base's free_group_end() reads it unconditionally.
self.free_page_ids_group = []
# Neutralize the base single-list attributes: every consumer of this
# allocator must go through the classed API (or fail loud on None),
# never a stale flat free list.
self.free_pages = None
self.release_pages = None
# ---- classed accounting ---------------------------------------------------
def class_free_page_counts(self) -> List[int]:
"""Free pages per class (free + release) — the per-class watermarks
exported at the scheduler invariant-checker seam."""
return [
len(free) + len(release)
for free, release in zip(self.class_free_pages, self.class_release_pages)
]
def available_size(self) -> int:
# Min-class capacity floor: a K-page request needs up to
# ceil(K / N) pages of EACH class (cyclic draws), so admission must
# gate on the tightest class, not the aggregate — the aggregate can
# be large while one class is fully protected by locked chains.
return self.shard_size * self.page_size * min(self.class_free_page_counts())
def aggregate_free_size(self) -> int:
"""Total free logical slots across all classes — the accounting
identity's `available` term (the invariant checker); NOT an admission
gate (see available_size)."""
return self.page_size * sum(self.class_free_page_counts())
def least_full_class(self) -> int:
"""Rotation base for a new chain: the class with the most free pages,
ties broken by the lowest class id. A pure function of the mirrored
class fills (SPMD-safe). Least-full seeding self-corrects the
per-chain <= 1-page remainders instead of letting them drift."""
counts = self.class_free_page_counts()
return max(range(self.shard_size), key=lambda r: (counts[r], -r))
def merge_and_sort_free(self):
for r in range(self.shard_size):
if len(self.class_release_pages[r]) > 0:
merged = torch.cat(
(self.class_free_pages[r], self.class_release_pages[r])
)
self.class_free_pages[r], _ = torch.sort(merged)
self.class_release_pages[r] = torch.empty(
(0,), dtype=torch.int64, device=self.device
)
# ---- alloc / free -----------------------------------------------------------
def alloc(self, need_size: int):
raise NotImplementedError(
"PageInterleavePoolAllocator allocates through alloc_extend only "
"(class draws follow the chain's rotation; a bare alloc has no "
"rotation base)"
)
@staticmethod
def _class_counts(start_class: int, new_pages: int, shard_size: int) -> List[int]:
"""Pages drawn from each class by a cyclic run of ``new_pages`` pages
starting at ``start_class`` — closed-form, host-only."""
return [
new_pages // shard_size
+ (1 if (c - start_class) % shard_size < new_pages % shard_size else 0)
for c in range(shard_size)
]
def alloc_extend(
self,
prefix_lens: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor,
extend_num_tokens: int,
num_new_pages: int = None,
rotation_bases: Optional[List[Optional[int]]] = None,
):
"""Pop one page per new position-page, class ``(b_i + P) % N`` per
request.
``rotation_bases`` carries one entry per request: the chain's
rotation base, or None for a new chain. None entries are resolved IN
PLACE from the least-full class at that request's turn — a draw sees
the pops of earlier requests in the same batch, keeping the balance
policy exact at bs > 1 — and the caller reads the resolved values
back to stamp the requests for the radix insert.
Sync-free: per-class pop counts are closed-form from the host
lengths and the host rotation bases (the reason the base is host
metadata on radix nodes rather than derived from device locs).
Returns None when some needed class cannot supply its pages — the
caller's admission gate (min-class available_size + the N*ps
per-request reserve) makes that unreachable outside true OOM.
"""
ps, shard_size = self.page_size, self.shard_size
bs = len(prefix_lens_cpu)
assert rotation_bases is not None and len(rotation_bases) == bs, (
"sharded alloc_extend needs one rotation base slot per request "
"(alloc_paged_token_slots_extend derives them from req.last_node)"
)
# Pass 1 (host, mirrored): resolve draw-at-turn bases and check the
# per-class supply against simulated fills, so the batch either
# commits whole or defers whole.
sim_counts = self.class_free_page_counts()
start_classes: List[int] = []
new_pages_list: List[int] = []
for i in range(bs):
prefix_len = int(prefix_lens_cpu[i])
seq_len = int(seq_lens_cpu[i])
assert prefix_len % ps == 0, (
f"sharded extends start page-aligned (the radix-tree match "
f"quantum and the chunk flooring guarantee it), got "
f"prefix_len={prefix_len}, page_size={ps}"
)
new_pages = -(-seq_len // ps) - prefix_len // ps
assert new_pages > 0
if rotation_bases[i] is None:
# New chain: least-full over the simulated fills (mirrored).
rotation_bases[i] = max(
range(shard_size), key=lambda r: (sim_counts[r], -r)
)
start_class = (rotation_bases[i] + prefix_len // ps) % shard_size
start_classes.append(start_class)
new_pages_list.append(new_pages)
for c, need in enumerate(
self._class_counts(start_class, new_pages, shard_size)
):
sim_counts[c] -= need
if self.debug_mode and prefix_len > 0:
# Owner congruence of the prefix end: the host-tracked phase
# must agree with the loc-derived owner (one D2H sync, debug
# only). A mismatch means a stale rotation base — the pool's
# plan would translate into the wrong rank's scratch block.
actual = int(last_loc[i].item()) // ps % shard_size
expected = (start_class - 1) % shard_size
assert actual == expected, (
f"rotation base out of sync with the prefix locs (req "
f"{i}): owner of the last prefix page is {actual}, host "
f"phase says {expected}"
)
assert extend_num_tokens == sum(
int(seq_lens_cpu[i]) - int(prefix_lens_cpu[i]) for i in range(bs)
)
if min(sim_counts) < 0:
return None
# Under need_sort the pops slice class_free_pages only; merge the
# release lists in whenever any class's free list alone is shorter
# than its total need (= total count - simulated remainder).
needs = [
total - remaining
for total, remaining in zip(self.class_free_page_counts(), sim_counts)
]
if self.need_sort and any(
needs[c] > len(self.class_free_pages[c]) for c in range(shard_size)
):
self.merge_and_sort_free()
# Pass 2: commit the pops, one position-ordered page vector per
# request, concatenated in batch order (= out_cache_loc order).
out_parts = []
for i in range(bs):
new_pages = new_pages_list[i]
start_class = start_classes[i]
counts = self._class_counts(start_class, new_pages, shard_size)
# Interleave the class pops into one position-ordered page
# vector: position-page j draws class (start_class + j) % N, so
# class c fills plan slots (c - start_class) % N, +N, +2N, ...
pages = torch.empty((new_pages,), dtype=torch.int64, device=self.device)
for c in range(shard_size):
if counts[c] == 0:
continue
pages[
torch.arange(
(c - start_class) % shard_size,
new_pages,
shard_size,
device=self.device,
)
] = self.class_free_pages[c][: counts[c]]
self.class_free_pages[c] = self.class_free_pages[c][counts[c] :]
extend_len = int(seq_lens_cpu[i]) - int(prefix_lens_cpu[i])
offsets = torch.arange(extend_len, dtype=torch.int64, device=self.device)
out_parts.append(pages[offsets // ps] * ps + offsets % ps)
out_indices = torch.cat(out_parts) if len(out_parts) > 1 else out_parts[0]
if self.debug_mode:
assert len(torch.unique(out_indices)) == len(out_indices)
return out_indices
def alloc_decode(
self,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor,
):
raise NotImplementedError(
"PageInterleavePoolAllocator does not support decode allocation "
"(logical-page KV sharding runs on prefill nodes only)"
)
def free(self, free_index: torch.Tensor):
if free_index.numel() == 0:
return
if self.free_group is not None:
# Match the base allocator's ownership contract: callers may
# overwrite req_to_token views before free_group_end consumes the
# deferred indices.
self.free_group.append(self._copy_for_free_group(free_index))
return
# The tree quantum is the physical page, so every free covers whole
# pages (any partial in-page coverage means the rest of that page
# belongs to the same free). Split by owner class; a freed page is
# immediately reusable — nothing strands.
pages = torch.unique(free_index.long() // self.page_size)
owners = pages % self.shard_size
for r in range(self.shard_size):
freed = pages[owners == r]
if freed.numel() == 0:
continue
if self.need_sort:
self.class_release_pages[r] = torch.cat(
(freed, self.class_release_pages[r])
)
else:
self.class_free_pages[r] = torch.cat((freed, self.class_free_pages[r]))
if self.debug_mode:
self.debug_check_classes()
def free_segment(self, free_index: torch.Tensor, *, start_pos: int):
# Back to the base contract (plain free): the paged override derives
# page representatives by striding and hands them to _release_page_ids,
# which writes the stock free_pages list this allocator does not use.
# start_pos buys nothing here — free() keys pages by owner class, which
# needs the page ids anyway.
self.free(free_index)
def _release_page_ids(self, *page_ids: torch.Tensor):
raise NotImplementedError(
"PageInterleavePoolAllocator keeps free pages in per-owner class "
"lists; a caller reaching the stock free_pages list would leak them"
)
def _debug_check_no_duplicate_pages(self):
# The base sweep concatenates the neutralized flat lists; the classed
# census is the equivalent check here.
self.debug_check_classes()
# ---- state / debug ----------------------------------------------------------
def resize(self, config) -> None:
raise NotImplementedError(
"post-capture KV resizing is not supported under logical-page KV sharding"
)
def get_all_free_pages(self) -> torch.Tensor:
"""All free logical page ids across classes (free + release) — for
the scheduler invariant checker's use-after-free / double-free sweep
(page unit = the physical page = self.page_size). The paged base
reads the flat lists this allocator neutralizes, so the classed
census is the override."""
return torch.cat(self.class_free_pages + self.class_release_pages)
def debug_check_classes(self):
for r in range(self.shard_size):
merged = torch.cat((self.class_free_pages[r], self.class_release_pages[r]))
assert bool((merged % self.shard_size == r).all()), (
f"page of another owner leaked into class {r}"
)
assert len(torch.unique(merged)) == len(merged), (
f"double free: duplicate pages in class {r}"
)
def page_interleave_shard_size(allocator: BaseTokenToKVPoolAllocator) -> int:
"""Shard-group size of a widened allocator, or 1 for stock allocators.
The scheduler-side seam predicate: the sites that must treat the
index space as shard-widened (capacity reporting, admission reserve,
rotation-base plumbing, invariant slack) branch on this, exactly
parallel to their existing DCP conditions.
"""
if isinstance(allocator, PageInterleavePoolAllocator):
return allocator.shard_size
return 1
python/sglang/srt/mem_cache/base_prefix_cache.py
+23
View File
@@ -80,6 +80,11 @@ class InsertParams:
priority: int = 0
track_adopted_ranges: bool = False
# Logical-page KV sharding: rotation base of the chain the inserted
# values belong to (stamped onto new tree nodes; None when sharding is
# off). See UnifiedTreeNode.rotation_base.
rotation_base: Optional[int] = None
@dataclasses.dataclass
class InsertResult:
@@ -90,6 +95,13 @@ class InsertResult:
last_device_node: Any = None
mamba_exist: bool = False
swa_branch_inserted: bool = False
# Logical-page KV sharding: the un-matched tail was NOT inserted because
# its rotation base disagrees with the matched chain's (a cross-chain
# graft would break the cyclic-owner gather contract). The tail's pages
# stay owned by the inserting request; callers must not dedup/rebind
# past prefix_len.
rotation_tail_declined: bool = False
inserted_host_node: Any = None
host_insert_dropped: bool = False
adopted_ranges: Optional[dict[ComponentType, list[tuple[int, int]]]] = None
@@ -380,6 +392,17 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
"""The hash chain of the node's ancestors, in root-to-parent order."""
return node.get_prefix_hash_values(node.parent)
def rotation_base_of(self, node: Any) -> Optional[int]:
"""Logical-page KV sharding: the rotation base stamped on ``node``.
``node`` is whatever this cache stores in ``req.last_node`` (a NodeId
for the unified tree, None for caches without tree nodes). None means
"no base available here", which sends the alloc path to the base the
request recorded at its previous alloc. Tree caches that keep the
per-chain base override this. See UnifiedTreeNode.rotation_base.
"""
return None
@abstractmethod
def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
pass
python/sglang/srt/mem_cache/common.py
+33 -2
View File
@@ -10,6 +10,7 @@ from sglang.kernels.ops.memory.common import (
_get_last_loc_safe_kernel as _get_last_loc_safe_kernel,
)
from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel
from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
from sglang.srt.mem_cache.hicache_storage import PoolTransfer
@@ -176,6 +177,32 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
tree_cache.evict_for_alloc(
EvictParams(num_tokens=num_tokens - available_size)
)
_evict_until_allocatable(tree_cache, allocator, num_tokens)
def _evict_until_allocatable(
tree_cache: BasePrefixCache, allocator, num_tokens: int
) -> None:
"""Keep evicting the shortfall until `num_tokens` are allocatable.
Under classed page sharding available_size() reports the MIN-CLASS
capacity floor, so a single evict() sized in tokens can raise that floor by
less than the number of tokens it freed: the evicted pages spread across
all owner classes. Looping is deterministic, so it stays mirrored across
the ranks of a shard group. Stock allocators need no extra pass.
"""
if page_interleave_shard_size(allocator) <= 1:
return
while True:
available_size = allocator.available_size()
if available_size >= num_tokens:
return
shortfall = num_tokens - available_size
result = tree_cache.evict(
EvictParams(num_tokens=max(shortfall, allocator.page_size))
)
if result.num_tokens_evicted == 0:
return
def retraction_backup(
@@ -294,12 +321,16 @@ def _release_overallocated_kv_indices(
f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}"
)
# Align to the ALLOCATOR's page, which under DCP is wider than the kernel
# page: paged free() releases the whole page containing any freed index, so
# a boundary aligned only to the kernel page could free a widened page whose
# head rows are still live.
if page_size > 1:
start_p = ceil_align(start_p, page_size)
if start_p < end_p:
# start_p is aligned to the allocator's physical page size above, so it
# never shares a page with cache_finished_req's tail free in this group.
# start_p is aligned to the allocator's page above, so it never shares a
# page with cache_finished_req's tail free in this group.
tree_cache.free_kv_row(req.kv, [(start_p, end_p)])
python/sglang/srt/mem_cache/page_interleave.py
+86
View File
@@ -0,0 +1,86 @@
# 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.
# ==============================================================================
"""Logical-page KV sharding: placement arithmetic and shard-group resolution.
A *logical page* is ``shard_size`` consecutive, ``shard_size``-aligned physical
pages — one per rank of the shard group. Everything above the memory pool
(radix tree, allocator free list, ``req_to_token``, scheduler budgets) sees
only logical token slots, identical on every rank (SPMD); everything below
sees only its own physical pages; the boundary is the pure bijection below.
The shard group is the group across which KV storage is replicated today and
therefore can be striped without extra compute-time communication:
- GQA/MHA models: the **attention CP group** — prefill CP already allgathers
the full chunk's K/V to every CP rank (``cp_allgather_and_save_kv_cache``).
- MLA models: the **attention TP group** — the latent KV projection is
``ReplicatedLinear``, so every attn-TP rank computes identical latent KV.
"""
from __future__ import annotations
import logging
import msgspec
import torch
logger = logging.getLogger(__name__)
class PageShardSpec(msgspec.Struct, frozen=True):
"""Everything both sides of any transfer need to reproduce the layout."""
shard_rank: int
shard_size: int
page_size: int # physical page size (kernel-visible)
max_prefix_tokens: int # scratch prefix-region capacity, granule-aligned
chunk_tokens: int # scratch chunk-region capacity, granule-aligned
@property
def logical_page_size(self) -> int:
"""The N*page_size span of one logical page. It bounds the assembly
scratch and rounds chunked_prefill_size; the allocator and the tree
both keep drawing and matching at the physical ``page_size``."""
return self.shard_size * self.page_size
class PageInterleavePlacement:
"""``loc = Q*(N*ps) + r*ps + o`` -> owner ``r``, local physical row ``Q*ps + o``.
Pure, stateless and invertible, so nothing has to be stored or kept
coherent to translate. Because the allocator is
mirrored, logical group ``Q`` resolves to local rows ``[Q*ps, (Q+1)*ps)``
on every rank, so a reader computes a peer's source offset from arithmetic
alone. Owned tokens form ``page_size``-long contiguous runs in both the
logical and the local space.
"""
def __init__(self, spec: PageShardSpec):
self.spec = spec
def owner_of(self, loc: torch.Tensor) -> torch.Tensor:
ps, n = self.spec.page_size, self.spec.shard_size
return (loc % (n * ps)) // ps
def local_index(self, loc: torch.Tensor) -> torch.Tensor:
ps, n = self.spec.page_size, self.spec.shard_size
return (loc // (n * ps)) * ps + loc % ps
def local_mask(self, loc: torch.Tensor, rank: int) -> torch.Tensor:
return self.owner_of(loc) == rank
def filter_local(self, loc: torch.Tensor, rank: int) -> torch.Tensor:
"""Logical slots -> this rank's physical pool rows, order-preserving."""
return self.local_index(loc[self.local_mask(loc, rank)])
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+89 -1
View File
@@ -141,6 +141,14 @@ class UnifiedTreeNode:
# Anchor NodeId of an in-flight H->D load-back reading this node's
# host slots; such host copies must not be reclaimed until the ack.
self.load_back_pending_id: Optional[int] = None
# Logical-page KV sharding: rotation base of the chain this node's Full
# KV pages belong to — the owner rank of position-page P along the chain
# is (rotation_base + P) % shard_size. A host-side mirror of the
# loc-derived owners (the Full value is a device tensor; reading it
# would put a D2H sync on the alloc path): set at insert from the
# inserting request's host base, copied on split, read through
# req.last_node at alloc time. None when sharding is off.
self.rotation_base: Optional[int] = None
def component(self, component_type: ComponentType) -> ComponentData:
return self.component_data[component_type]
@@ -515,6 +523,13 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
"""Whether the node is the tree root."""
return self.node_by_id(node_id) is self.root_node
supports_rotation_base = True
def rotation_base_of(self, node_id: NodeId) -> Optional[int]:
"""Logical-page KV sharding: the node's chain rotation base, or None
when sharding is off (and on the root, which starts no chain)."""
return self.node_by_id(node_id).rotation_base
def get_last_hash_value(self, node_id: NodeId) -> Optional[str]:
"""The node's last page hash, or None when it was never hashed."""
return self.node_by_id(node_id).get_last_hash_value()
@@ -999,6 +1014,20 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
),
)
if params.rotation_base is not None:
conflict = self._rotation_conflict(key, params.rotation_base)
if conflict is not None:
total_prefix_length, node = conflict
return InsertStepResult(
actions=[],
result=InsertResult(
prefix_len=total_prefix_length,
last_device_node=node.id,
rotation_tail_declined=True,
adopted_ranges={} if params.track_adopted_ranges else None,
),
)
self._ongoing_insert_walk_state = _InsertWalkState(
phase=_InsertPhase.WALK,
node=self.root_node,
@@ -1013,6 +1042,52 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
)
return self._advance_insert()
def _rotation_conflict(
self, key: RadixKey, rotation_base: int
) -> Optional[tuple[int, UnifiedTreeNode]]:
"""Logical-page KV sharding: refuse an insert that would splice two
rotation runs into one cached path.
Read-only mirror of ``_insert_walk_step``'s matching. Returns
``(matched_len, deepest_matched_node)`` when the chain this key lands on
was allocated under a different rotation base than the inserting
request's pages, else None.
Grafting across a base discontinuity would create a cached path whose
page owners are not one cyclic run — the padded-allgather / ``k // N``
translation contract — so later readers would crash on a negative pad or
silently read the wrong rank's scratch rows. Such a discontinuity means
two requests sharing a prefix were planned in pipelined batches before
either's insert landed, or the match was capped below the cached prefix.
Unlike the flat radix tree, the unified insert transfers page ownership
at three points, not just the tail graft: the tail leaf
(``_add_new_node``), an evicted node restored from the request's fresh
pages (``_unevict_node_on_insert``), and a component re-pointing a
matched node's Full value at the request's pages
(``update_component_on_insert_overlap``). All three are downstream of
this same predicate, so it is evaluated once, up front, and declines the
whole insert — which also keeps the walk's duplicate frees from running,
as the declined request stays on its own pages.
"""
node = self.root_node
total_prefix_length = 0
while len(key) > 0:
child_key = key.child_key(self.page_size)
if child_key not in node.children:
break
child = node.children[child_key]
prefix_len = child.key.match(key, page_size=self.page_size)
node = child
total_prefix_length += prefix_len
key = key[prefix_len:]
if prefix_len < len(child.key):
# The walk would split here; the fragment inherits this base.
break
if node is self.root_node or node.rotation_base == rotation_base:
return None
return total_prefix_length, node
def resume_insert(self) -> InsertStepResult:
"""Continue the suspended insert after its step actions were executed."""
if self._ongoing_insert_walk_state is None:
@@ -1149,7 +1224,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
state.total_prefix_length + len(state.key),
)
state.target_node = self._add_new_node(
state.node, state.key, state.value, priority=state.priority
state.node,
state.key,
state.value,
priority=state.priority,
rotation_base=state.params.rotation_base,
)
state.is_new_leaf = True
else:
@@ -1224,6 +1303,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
new_node.creation_time = child.creation_time
# Split fragments stay on the anchor's root path for the ack's walk.
new_node.load_back_pending_id = child.load_back_pending_id
# The rotation base is constant along a chain (position-page P keeps
# owner (b + P) % N on both sides of the split).
new_node.rotation_base = child.rotation_base
self._for_each_component_lru(child, UnifiedLRUList.remove_node)
@@ -1272,10 +1354,16 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
key: RadixKey,
value: torch.Tensor,
priority: int = 0,
rotation_base: Optional[int] = None,
) -> UnifiedTreeNode:
new_node = self._new_node(priority=priority)
new_node.parent = parent
new_node.key = key
# Chain-constant under sharding: the pre-flight decline in
# begin_insert() guarantees this tail continues the matched prefix's
# rotation, so stamping the inserting request's base keeps every node
# on a root path carrying the same base.
new_node.rotation_base = rotation_base
new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone()
parent.children[key.child_key(self.page_size)] = new_node
self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value)
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+16
View File
@@ -185,6 +185,22 @@ class UnifiedTreeCoreInterface(ABC):
"""Whether the node is the tree root."""
...
# Logical-page KV sharding: whether this core stamps and honors
# UnifiedTreeNode.rotation_base. A core that does not cannot serve a
# sharded allocator (it would never decline a cross-base graft), and
# UnifiedRadixCache.__init__ rejects that pairing at construction.
supports_rotation_base: bool = False
def rotation_base_of(self, node_id: NodeId) -> Optional[int]:
"""Logical-page KV sharding: the node's chain rotation base, or None
when sharding is off (and on the root, which starts no chain).
Concrete, not abstract: a core that does not track rotation bases
stays constructible, and its None means "sharding is off" -- never
"sharding is on but unknown", which the constructor gate rules out.
"""
return None
@abstractmethod
def get_last_hash_value(self, node_id: NodeId) -> Optional[str]:
"""The node's last page hash, or None when it was never hashed."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+54 -2
View File
@@ -13,6 +13,9 @@ import torch
from sglang.srt.distributed.communication_tags import P2PTag
from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import CacheOperation
from sglang.srt.mem_cache.allocator.page_interleave import (
page_interleave_shard_size,
)
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
@@ -200,6 +203,22 @@ class UnifiedRadixCache(BasePrefixCache):
for component in self.components.values():
component.tree_core = self.tree_core
if (
page_interleave_shard_size(params.token_to_kv_pool_allocator) > 1
and not self.tree_core.supports_rotation_base
):
# A core that does not model rotation_base would never decline a
# cross-base graft, and the resulting cached path's page owners are
# not one cyclic run: later readers take a negative allgather pad or
# silently read another rank's scratch rows. Fail at construction
# rather than corrupt reads at serve time.
raise ValueError(
"logical-page KV sharding requires a tree core that tracks "
"rotation bases; "
f"SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND={self._tree_core_backend!r} "
"does not."
)
# Session ref tracking (--enable-session-radix-cache).
self.session_refs = UnifiedSessionRefTracker(
components=self._components_tuple,
@@ -939,6 +958,7 @@ class UnifiedRadixCache(BasePrefixCache):
insert_params = InsertParams(
prev_prefix_len=req.kv.cache_protected_len,
priority=getattr(req, "priority", 0) or 0,
rotation_base=req.kv_rotation_base,
)
# components prepare insert data + return effective cache_len
@@ -975,8 +995,17 @@ class UnifiedRadixCache(BasePrefixCache):
insert_params.value = values
result = self.insert(insert_params)
# Free unaligned tail (+ deferred truncation tail)
ranges = [(page_aligned_len, len(kv_indices))]
# Free unaligned tail (+ deferred truncation tail). A rotation
# decline inserted nothing, so the whole span past the protected
# prefix stayed request-owned and is released here instead.
free_from = (
# min(): the protected prefix can already run past a truncated
# cache_len, and free_kv_row takes ascending ranges only.
min(req.kv.cache_protected_len, len(kv_indices))
if result.rotation_tail_declined
else page_aligned_len
)
ranges = [(free_from, len(kv_indices))]
if tail_free_start is not None:
ranges.append((tail_free_start, len(kv_indices_full)))
self.free_kv_row(req.kv, ranges)
@@ -1026,6 +1055,7 @@ class UnifiedRadixCache(BasePrefixCache):
prev_prefix_len=req.kv.cache_protected_len,
chunked=chunked,
priority=getattr(req, "priority", 0) or 0,
rotation_base=req.kv_rotation_base,
)
effective_cache_len = len(token_ids)
for comp in self._components_tuple:
@@ -1072,6 +1102,23 @@ class UnifiedRadixCache(BasePrefixCache):
insert_params.value = values
result = self.insert(insert_params)
if result.rotation_tail_declined:
# Rotation-base discontinuity with the matched chain (pipelined
# batches raced this request's insert against another chain over
# the same prefix). Adopting the canonical locs would leave this
# request's row mixing two rotation runs, which the cyclic-owner
# gather contract forbids -- keep the request entirely on its own
# pages: no dedup free, no rebind, no protection change. The insert
# declined before its walk, so nothing was freed underneath us. The
# final cache_finished_req releases everything past the protected
# prefix.
req.prefix_indices = kv_indices_orig.to(dtype=torch.int64, copy=True)
for comp in self._components_tuple:
comp.cleanup_after_caching_req(
req, is_finished=False, insert_params=insert_params
)
return
# Match prefix. SWA insertion retains one extra window before the
# page-aligned boundary, so the normal match remains safe to repoint.
match_result = self.match_prefix(MatchPrefixParams(key=radix_key, req=req))
@@ -1733,6 +1780,11 @@ class UnifiedRadixCache(BasePrefixCache):
def get_prefix_hash_values(self, node_id: NodeId) -> list[str]:
return self.tree_core.get_prefix_hash_values(node_id)
def rotation_base_of(self, node_id: Optional[NodeId]) -> Optional[int]:
if node_id is None:
return None
return self.tree_core.rotation_base_of(node_id)
def query_storage_hit_length(
self,
last_host_node_id: NodeId,
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+1
View File
@@ -94,6 +94,7 @@ class MockReq:
cache_protected_len=cache_protected_len,
swa_evicted_seqlen=0,
)
self.kv_rotation_base = None
def get_fill_ids(self):
return self.full_untruncated_fill_ids[: self.extend_range.end]
test/registered/unit/mem_cache/test_page_interleave_shard.py
+770
View File
@@ -0,0 +1,770 @@
# 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.
# ==============================================================================
"""Unit tests for logical-page KV cache sharding (CPU only).
Pins the pure arithmetic that rotated owner-classed allocation hangs on:
1. The placement bijection ``loc = Q*(N*ps) + r*ps + o`` — owner / local-row
round-trip, disjoint equal partition across ranks.
2. ``PageInterleavePoolAllocator`` — N mirrored class free lists, rotated
class draws (owners exactly cyclic along a chain), least-full root
seeding, min-class admission accounting, zero stranding (a freed page is
immediately reusable).
3. The host rotation base on ``UnifiedTreeNode`` — stamped at insert, copied
on split, read through ``last_node``, and the pre-flight that declines an
insert whose pages carry a different base than the chain it would join.
"""
import unittest
import unittest.mock
from array import array
from types import SimpleNamespace
import torch
from sglang.srt.mem_cache.allocator.page_interleave import (
PageInterleavePoolAllocator,
page_interleave_shard_size,
)
from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
from sglang.srt.mem_cache.page_interleave import (
PageInterleavePlacement,
PageShardSpec,
)
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
N = 4 # shard size
PS = 16 # physical page size
GS = N * PS # full-group span (N physical pages)
def _make_spec(shard_rank=0, max_prefix_groups=64, chunk_pages=32):
return PageShardSpec(
shard_rank=shard_rank,
shard_size=N,
page_size=PS,
max_prefix_tokens=max_prefix_groups * GS,
chunk_tokens=chunk_pages * PS,
)
def _make_allocator(pages_per_rank=32, need_sort=False):
return PageInterleavePoolAllocator(
size=pages_per_rank * PS, # physical token slots of one rank
physical_page_size=PS,
shard_size=N,
dtype=torch.bfloat16,
device="cpu",
kvcache=None,
need_sort=need_sort,
)
def _alloc_extend_batch(alloc, prefix_lens, seq_lens, rotation_bases, last_locs=None):
"""Drive alloc_extend for a batch; rotation_bases is resolved in place."""
if last_locs is None:
last_locs = [-1] * len(prefix_lens)
return alloc.alloc_extend(
prefix_lens=torch.tensor(prefix_lens, dtype=torch.int64),
prefix_lens_cpu=torch.tensor(prefix_lens, dtype=torch.int64),
seq_lens=torch.tensor(seq_lens, dtype=torch.int64),
seq_lens_cpu=torch.tensor(seq_lens, dtype=torch.int64),
last_loc=torch.tensor(last_locs, dtype=torch.int64),
extend_num_tokens=sum(s - p for p, s in zip(prefix_lens, seq_lens)),
rotation_bases=rotation_bases,
)
def _alloc_extend(alloc, prefix_len, seq_len, rotation_base, last_loc=-1):
return _alloc_extend_batch(
alloc, [prefix_len], [seq_len], [rotation_base], [last_loc]
)
class TestPlacement(CustomTestCase):
def test_owner_local_round_trip(self):
pl = PageInterleavePlacement(_make_spec())
loc = torch.arange(0, 37 * GS + 5)
owner = pl.owner_of(loc)
local = pl.local_index(loc)
# Reconstruct loc from (group, owner, in-page offset): the bijection.
group = loc // GS
self.assertTrue(torch.equal(group * GS + owner * PS + loc % PS, loc))
# Local rows are group-major: [Q*ps, (Q+1)*ps) — identical on every
# rank (symmetric allocation); owner only selects WHICH rank stores.
self.assertTrue(torch.equal(local, group * PS + loc % PS))
def test_filter_local_partitions_disjoint_and_equal(self):
pl = PageInterleavePlacement(_make_spec())
loc = torch.arange(0, 10 * GS)
parts = [pl.filter_local(loc, r) for r in range(N)]
self.assertEqual(sum(p.numel() for p in parts), loc.numel())
# Equal shares of whole groups.
self.assertEqual(len({p.numel() for p in parts}), 1)
# Every rank's local rows for a full range are the same integers
# (each rank stores its own stripe at the SAME rows).
for p in parts[1:]:
self.assertTrue(torch.equal(p, parts[0]))
def test_owned_tokens_form_page_runs(self):
pl = PageInterleavePlacement(_make_spec(shard_rank=2))
loc = torch.arange(0, 3 * GS)
mask = pl.local_mask(loc, 2)
# Owner-2 tokens are exactly [2*ps, 3*ps) of every group.
expect = (loc % GS >= 2 * PS) & (loc % GS < 3 * PS)
self.assertTrue(torch.equal(mask, expect))
class TestClassedAllocator(CustomTestCase):
def test_index_space_widened_classes_mirror_ranks(self):
alloc = _make_allocator(pages_per_rank=32)
self.assertEqual(alloc.size, 32 * PS * N) # logical slots
self.assertEqual(alloc.page_size, PS) # the PHYSICAL page quantum
self.assertEqual(page_interleave_shard_size(alloc), N)
# Class r holds exactly rank r's allocatable pages: l % N == r,
# local pages 1..32 (page 0 reserved on every rank).
self.assertEqual(alloc.class_free_page_counts(), [32] * N)
for r in range(N):
pages = alloc.class_free_pages[r]
self.assertTrue(bool((pages % N == r).all()))
self.assertTrue(torch.equal(pages // N, torch.arange(1, 33)))
def test_rotation_worked_example_zero_stranding(self):
"""Two-turn worked example at ps=16: turn 1 allocates cyclic owners
from the root base; the turn-boundary free returns its page whole and
immediately reusable; turn 2 continues the rotation and reuses the
freed page before any fresh one."""
alloc = _make_allocator()
total = alloc.available_size()
# Turn 1: 122 tokens = 8 position-pages, root base 0.
base = alloc.least_full_class()
self.assertEqual(base, 0) # all classes equal -> lowest id
out = _alloc_extend(alloc, 0, 122, base)
pages = out[::PS] // PS
# Owners exactly cyclic from the base; in-page offsets positional.
self.assertTrue(torch.equal(pages % N, torch.arange(8) % N))
self.assertTrue(torch.equal(out % PS, torch.arange(122) % PS))
# Boundary: cache 112 (7 pages), free the sub-ps tail's page whole.
alloc.free(out[112:122])
# Page 7's owner is (0 + 7) % 4 = 3: back on class 3, reusable now.
self.assertEqual(alloc.class_free_page_counts(), [30, 30, 30, 31])
# Turn 2: prefix 112, extend to 244 (9 new pages P7..P15).
out2 = _alloc_extend(alloc, 112, 244, base, last_loc=int(out[111]))
pages2 = out2[::PS] // PS
self.assertTrue(torch.equal(pages2 % N, (7 + torch.arange(9)) % N))
# The freed page is the class-3 head: reused before any fresh page.
self.assertEqual(int(pages2[0]), int(pages[7]))
# Nothing stranded: freeing the chain restores full capacity.
alloc.free(out[:112])
alloc.free(out2)
self.assertEqual(alloc.available_size(), total)
self.assertEqual(alloc.class_free_page_counts(), [32] * N)
def test_min_class_admission_accounting(self):
"""available_size is the MIN-CLASS floor: draining one class must
zero the admission budget even while the aggregate stays large —
an aggregate gate would over-admit into the alloc path's fail-loud
RuntimeError when the tight class is protected."""
alloc = _make_allocator(pages_per_rank=4)
outs = [_alloc_extend(alloc, 0, PS, rotation_base=3) for _ in range(4)]
self.assertEqual(alloc.class_free_page_counts(), [4, 4, 4, 0])
self.assertEqual(alloc.available_size(), 0)
self.assertEqual(alloc.aggregate_free_size(), 12 * PS)
# A draw needing the empty class defers (None), never raises.
self.assertIsNone(_alloc_extend(alloc, 0, N * PS, rotation_base=0))
# A free of one class-3 page lifts the floor by one page per class.
alloc.free(outs[0])
self.assertEqual(alloc.available_size(), N * PS)
def test_least_full_root_seeding(self):
"""Roots draw from the class with the most free pages (ties: lowest
id). Uniform 1-page roots therefore spread with skew <= 1."""
alloc = _make_allocator(pages_per_rank=32)
for i in range(2 * N + 1):
base = alloc.least_full_class()
_alloc_extend(alloc, 0, PS, rotation_base=base)
counts = alloc.class_free_page_counts()
self.assertLessEqual(max(counts) - min(counts), 1, counts)
# 9 single-page roots at N=4: classes filled 3,2,2,2.
self.assertEqual(alloc.class_free_page_counts(), [29, 30, 30, 30])
def test_chain_rotation_run_property(self):
"""Within one chain (root + arbitrary ps-aligned extensions) the
owners are exactly cyclic, so per-rank owned page counts differ by
<= 1 — the padded-allgather block contract ceil(K/N). Guards the
class-interleave scatter in alloc_extend."""
for shard_size in (2, 4, 8):
alloc = PageInterleavePoolAllocator(
size=256 * PS,
physical_page_size=PS,
shard_size=shard_size,
dtype=torch.bfloat16,
device="cpu",
kvcache=None,
need_sort=False,
)
lens = [3 * PS, 5 * PS, PS, 7 * PS] # chunked extensions
base = alloc.least_full_class()
chain = []
prefix = 0
for ext in lens:
out = _alloc_extend(alloc, prefix, prefix + ext, base)
chain.append(out)
prefix += ext
locs = torch.cat(chain)
pages = locs[::PS] // PS
owners = pages % shard_size
expect = torch.arange(pages.numel()) % shard_size
self.assertTrue(torch.equal(owners, (int(owners[0]) + expect) % shard_size))
per_rank = torch.bincount(owners, minlength=shard_size)
self.assertLessEqual(int(per_rank.max() - per_rank.min()), 1)
def test_free_splits_by_owner_class(self):
alloc = _make_allocator()
out = _alloc_extend(alloc, 0, 6 * PS, rotation_base=1)
before = alloc.class_free_page_counts()
# Free pages 2 and 3 of the chain (owners 3 and 0) in one call, via
# the free-group batching path the scheduler uses.
alloc.free_group_begin()
alloc.free(out[2 * PS : 3 * PS])
alloc.free(out[3 * PS : 4 * PS])
alloc.free_group_end()
after = alloc.class_free_page_counts()
deltas = [a - b for a, b in zip(after, before)]
self.assertEqual(deltas, [1, 0, 0, 1]) # classes (1+2)%4=3 and (1+3)%4=0
def test_grouped_free_owns_indices_before_caller_mutation(self):
"""Deferred frees must snapshot req_to_token views: the scheduler may
overwrite the backing row before free_group_end consumes them."""
alloc = _make_allocator()
out = _alloc_extend(alloc, 0, 2 * PS, rotation_base=2)
first_page = out[:PS]
owner = int(first_page[0] // PS % N)
before = alloc.class_free_page_counts()
alloc.free_group_begin()
alloc.free(first_page)
first_page.zero_()
alloc.free_group_end()
after = alloc.class_free_page_counts()
self.assertEqual(after[owner], before[owner] + 1)
self.assertEqual(
[after[r] - before[r] for r in range(N)],
[1 if r == owner else 0 for r in range(N)],
)
def test_free_segment_returns_pages_to_their_classes(self):
# The radix cache frees through free_segment/free_segments. The paged
# base routes those to the stock free_pages list, which this allocator
# never reads, so the override must land them in the class lists.
alloc = _make_allocator()
total = alloc.available_size()
out = _alloc_extend(alloc, 0, 3 * PS, rotation_base=2)
self.assertLess(alloc.available_size(), total)
alloc.free_segment(out, start_pos=0)
self.assertEqual(alloc.available_size(), total)
self.assertEqual(alloc.class_free_page_counts(), [32] * N)
def test_free_segments_splits_at_a_page_boundary(self):
alloc = _make_allocator()
total = alloc.available_size()
out = _alloc_extend(alloc, 0, 4 * PS, rotation_base=0)
# Two disjoint ascending segments of one request's kv row.
alloc.free_segments([(out[: 2 * PS], 0), (out[2 * PS :], 2 * PS)])
self.assertEqual(alloc.available_size(), total)
self.assertEqual(alloc.class_free_page_counts(), [32] * N)
def test_need_sort_merges_per_class(self):
alloc = _make_allocator(pages_per_rank=4, need_sort=True)
out = _alloc_extend(alloc, 0, 4 * N * PS, rotation_base=0) # everything
self.assertEqual(alloc.available_size(), 0)
alloc.free(out) # lands in the per-class release lists
self.assertEqual(alloc.available_size(), 4 * N * PS)
# A fresh draw forces the per-class merge+sort and succeeds.
out2 = _alloc_extend(alloc, 0, N * PS, rotation_base=0)
self.assertIsNotNone(out2)
pages = out2[::PS] // PS
self.assertTrue(torch.equal(pages % N, torch.arange(N) % N))
def test_unsupported_paths_fail_loud(self):
alloc = _make_allocator()
with self.assertRaises(NotImplementedError):
alloc.alloc(GS)
with self.assertRaises(NotImplementedError):
alloc.alloc_decode(
torch.tensor([PS + 1]), torch.tensor([PS + 1]), torch.tensor([PS - 1])
)
def test_batch_alloc_per_request_rotation(self):
"""bs > 1: each request draws its own cyclic run; out_cache_loc is
the batch-order concatenation (write_cache_indices' contract), and a
None base is resolved from the least-full class AT THAT REQUEST'S
TURN — the draw must see earlier requests' pops in the same batch,
or uniform short batches would all pile onto one class."""
alloc = _make_allocator()
# req0: extension of a base-1 chain with a 2-page prefix;
# req1 and req2: new chains (drawn in place).
bases = [1, None, None]
out = _alloc_extend_batch(
alloc,
prefix_lens=[2 * PS, 0, 0],
seq_lens=[5 * PS, 3 * PS, PS],
rotation_bases=bases,
# req0's last prefix page must carry owner (1 + 1) % 4 = 2.
last_locs=[(5 * N + 2) * PS + PS - 1, -1, -1],
)
self.assertEqual(out.numel(), 3 * PS + 3 * PS + PS)
# Batch-order concatenation, per-request cyclic owners.
pages = out[::PS] // PS
self.assertTrue(
torch.equal(pages[:3] % N, (1 + 2 + torch.arange(3)) % N) # req0
)
b1, b2 = bases[1], bases[2]
self.assertIsNotNone(b1)
self.assertIsNotNone(b2)
self.assertTrue(torch.equal(pages[3:6] % N, (b1 + torch.arange(3)) % N))
self.assertEqual(int(pages[6]) % N, b2)
# req1's draw saw req0's pops (classes 3,0,1 used once each -> class
# 2 is fullest... all equal except used {3,0,1} -> least-full = 2);
# req2's draw saw req1's pops on top.
self.assertEqual(b1, 2)
self.assertEqual(b2, 1) # after req1 used {2,3,0}: class 1 fullest
# The whole batch is one no-duplicate allocation.
self.assertEqual(len(torch.unique(out)), out.numel())
def test_batch_alloc_defers_whole_when_a_class_is_short(self):
"""A batch either commits whole or returns None (mirrored decision):
partial commits would desync the free lists from the retry."""
alloc = _make_allocator(pages_per_rank=2)
counts_before = alloc.class_free_page_counts()
out = _alloc_extend_batch(
alloc,
prefix_lens=[0, 0],
seq_lens=[4 * PS, 5 * PS], # 9 pages: class need exceeds 2 somewhere
rotation_bases=[0, 0],
)
self.assertIsNone(out)
self.assertEqual(alloc.class_free_page_counts(), counts_before)
class TestEvictUntilAllocatable(CustomTestCase):
"""The evict-then-allocate contract under min-class accounting: one
evict() sized in tokens can raise the tight class by less than the
tokens it freed (evicted pages spread across classes), so the alloc
path iterates. Guards the two termination conditions of
_evict_until_allocatable."""
def _allocator_with_tight_class(self):
alloc = _make_allocator(pages_per_rank=4)
# Four 1-page chains, all in class 3: the tight class.
outs = [_alloc_extend(alloc, 0, PS, rotation_base=3) for _ in range(4)]
assert alloc.available_size() == 0
return alloc, outs
def _tree_stub(self, alloc, frees):
from sglang.srt.mem_cache.base_prefix_cache import EvictResult
stub = SimpleNamespace(calls=0)
def evict(params):
stub.calls += 1
if not frees:
return EvictResult(num_tokens_evicted=0)
head = frees.pop(0)
alloc.free(head)
return EvictResult(num_tokens_evicted=head.numel())
stub.evict = evict
return stub
def test_iterates_until_min_class_covers(self):
from sglang.srt.mem_cache.common import _evict_until_allocatable
alloc, outs = self._allocator_with_tight_class()
# Each round frees ONE class-3 page (a whole 1-page chain): reaching
# a min-class floor of 2 pages takes 2 rounds.
tree = self._tree_stub(alloc, list(outs))
_evict_until_allocatable(tree, alloc, 2 * N * PS)
self.assertGreaterEqual(alloc.available_size(), 2 * N * PS)
self.assertEqual(tree.calls, 2)
def test_terminates_when_tree_dry(self):
from sglang.srt.mem_cache.common import _evict_until_allocatable
alloc, _ = self._allocator_with_tight_class()
tree = self._tree_stub(alloc, []) # nothing evictable
_evict_until_allocatable(tree, alloc, PS)
self.assertEqual(alloc.available_size(), 0) # need unmet, but no hang
self.assertEqual(tree.calls, 1)
def _unified_tree(page_size=4, pool_size=256, disable=False):
"""A CPU-only UnifiedRadixCache with only the Full component."""
dtype = torch.float16
kv_pool = MHATokenToKVPool(
size=pool_size,
page_size=page_size,
dtype=dtype,
head_num=2,
head_dim=8,
layer_num=1,
device="cpu",
enable_memory_saver=False,
)
allocator = PagedTokenToKVPoolAllocator(
size=pool_size,
page_size=page_size,
dtype=dtype,
device="cpu",
kvcache=kv_pool,
need_sort=False,
)
req_pool = ReqToTokenPool(
size=8,
max_context_len=128,
device="cpu",
enable_memory_saver=False,
)
return UnifiedRadixCache(
CacheInitParams(
disable=disable,
req_to_token_pool=req_pool,
token_to_kv_pool_allocator=allocator,
page_size=page_size,
eviction_policy="lru",
tree_components=(ComponentType.FULL,),
)
)
def _insert(tree, tokens, rotation_base=None, value=None):
if value is None:
value = tree.token_to_kv_pool_allocator.alloc(len(tokens))
return tree.insert(
InsertParams(
key=RadixKey(array("q", tokens)),
value=value.to(dtype=torch.int64),
rotation_base=rotation_base,
)
)
def _node(tree, node_id):
return tree.tree_core.node_by_id(node_id)
def _match_len(tree, tokens):
res = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
return len(res.device_indices)
class TestUnifiedRotationBase(CustomTestCase):
"""The host rotation base on UnifiedTreeNode: the one new piece of
metadata. The Full component's value is a device tensor, so the base must
survive inserts and splits purely host-side or the alloc path gains a D2H
sync."""
def test_insert_stamps_split_copies(self):
tree = _unified_tree()
_insert(tree, list(range(12)), rotation_base=2)
# A shorter lookup splits the node at the match boundary: BOTH halves
# keep the chain's base (position-page P keeps owner (b+P)%N on both
# sides of any split).
probe = list(range(8)) + [99, 98, 97, 96]
_insert(tree, probe, rotation_base=2)
res = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", probe))))
tail = _node(tree, res.last_device_node)
self.assertEqual(tail.rotation_base, 2)
parent = tail.parent
self.assertEqual(parent.rotation_base, 2)
for child in parent.children.values():
self.assertEqual(child.rotation_base, 2)
def test_new_chain_gets_its_own_base(self):
tree = _unified_tree()
_insert(tree, list(range(8)), rotation_base=1)
_insert(tree, list(range(100, 108)), rotation_base=3)
r1 = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", range(8)))))
r2 = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", range(100, 108))))
)
self.assertEqual(_node(tree, r1.last_device_node).rotation_base, 1)
self.assertEqual(_node(tree, r2.last_device_node).rotation_base, 3)
def test_extension_tail_node_stamped_from_request(self):
tree = _unified_tree()
_insert(tree, list(range(8)), rotation_base=1)
# A longer insert of the same chain dedups the prefix and stamps the
# tail node with the (same, chain-constant) base.
_insert(tree, list(range(16)), rotation_base=1)
res = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", range(16)))))
self.assertEqual(_node(tree, res.last_device_node).rotation_base, 1)
def test_unsharded_inserts_keep_none(self):
tree = _unified_tree()
_insert(tree, list(range(8)))
res = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", range(8)))))
self.assertIsNone(_node(tree, res.last_device_node).rotation_base)
def test_rotation_base_of_reads_through_the_cache_boundary(self):
"""The alloc path holds a NodeId, not a node: the base must be
readable through the tree-cache API (BasePrefixCache.rotation_base_of
defaults to None, so an unsharded cache sends the alloc path to the
request's recorded base)."""
tree = _unified_tree()
_insert(tree, list(range(8)), rotation_base=2)
res = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", range(8)))))
self.assertEqual(tree.rotation_base_of(res.last_device_node), 2)
self.assertIsNone(tree.rotation_base_of(None))
self.assertIsNone(tree.rotation_base_of(tree.tree_core.root_node_handle()))
class TestShardedCoreGate(CustomTestCase):
"""A tree core that does not model rotation_base would never decline a
cross-base graft. Pairing one with a sharded allocator must fail at
construction, not produce wrong-owner gathers at serve time."""
def test_python_core_supports_rotation_base(self):
tree = _unified_tree()
self.assertTrue(tree.tree_core.supports_rotation_base)
def test_sharded_allocator_rejects_a_core_without_rotation_base(self):
tree = _unified_tree()
params = CacheInitParams(
disable=False,
req_to_token_pool=tree.req_to_token_pool,
token_to_kv_pool_allocator=_make_allocator(pages_per_rank=8),
page_size=PS,
eviction_policy="lru",
tree_components=(ComponentType.FULL,),
)
with unittest.mock.patch.object(
UnifiedTreeCore, "supports_rotation_base", False
):
with self.assertRaisesRegex(ValueError, "rotation bases"):
UnifiedRadixCache(params)
def test_unsharded_allocator_accepts_any_core(self):
params = CacheInitParams(
disable=False,
req_to_token_pool=ReqToTokenPool(
size=8, max_context_len=128, device="cpu", enable_memory_saver=False
),
token_to_kv_pool_allocator=_unified_tree().token_to_kv_pool_allocator,
page_size=4,
eviction_policy="lru",
tree_components=(ComponentType.FULL,),
)
with unittest.mock.patch.object(
UnifiedTreeCore, "supports_rotation_base", False
):
UnifiedRadixCache(params) # no raise: sharding is off
class _GraftReq:
"""Minimal Req stand-in for cache_unfinished/finished_req."""
def __init__(self, fill_ids, req_pool_idx=0):
self.fill_ids = list(fill_ids)
self.origin_input_ids = array("q", fill_ids)
self.output_ids = array("q", [])
self.kv = SimpleNamespace(
req_pool_idx=req_pool_idx,
cache_protected_len=0,
swa_evicted_seqlen=0,
)
self.extra_key = None
self.cache_salt = None
self.prefix_indices = torch.empty(0, dtype=torch.int64)
self.last_node = None
self.priority = 0
self.kv_rotation_base = None
self.lock_receipt = DecLockRefParams()
self.swa_prefix_lock_released = False
self.finished_reason = None
self.session = None
def get_fill_ids(self):
return array("q", self.fill_ids)
class TestRotationGraftDecline(CustomTestCase):
"""The overlap disagg-prefill loop plans batch t+1 before batch t's radix
insert lands, so two requests sharing a prefix can allocate under
different rotation bases. Grafting the second one's tail under the first
chain leaves the cached path's page owners not one cyclic run, so a later
reader either crashes on a negative allgather pad or silently reads the
wrong rank's scratch rows. Inserts must refuse the graft.
Unlike the flat radix tree, the unified insert also transfers page
ownership inside the MATCHED region (an evicted node restored from the
request's fresh pages, a component re-pointing a matched node's Full value
at them). The decline therefore runs before the walk and refuses the whole
insert, which additionally keeps the walk's duplicate frees from running
under a request that is about to keep its own pages.
"""
PS = 4 # tree quantum for these tests
def _tree_with_spy(self):
"""Record every KV row release. The unified tree frees through two
seams: the caller's free_kv_row -> free_segments, and the insert
walk's FreeDeviceKV -> free_segment."""
tree = _unified_tree(page_size=self.PS)
allocator = tree.token_to_kv_pool_allocator
freed = []
real_free_segments = allocator.free_segments
real_free_segment = allocator.free_segment
def spy_segments(segments):
freed.extend(torch.as_tensor(seg).clone() for seg, _start in segments)
return real_free_segments(segments)
def spy_segment(free_index, *, start_pos):
freed.append(torch.as_tensor(free_index).clone())
return real_free_segment(free_index, start_pos=start_pos)
allocator.free_segments = spy_segments
allocator.free_segment = spy_segment
return tree, freed
def _seed_chain(self, tree, tokens, base):
_insert(tree, tokens, rotation_base=base)
def _own_row(self, tree, req, n_tokens):
"""Give the request its own allocated KV row and return the locs."""
locs = tree.token_to_kv_pool_allocator.alloc(n_tokens).to(dtype=torch.int64)
tree.req_to_token_pool.req_to_token[req.kv.req_pool_idx, :n_tokens] = locs
return locs
def test_foreign_base_tail_declined(self):
tree = _unified_tree(page_size=self.PS)
self._seed_chain(tree, list(range(12)), base=1)
# Same 8-token prefix, different suffix, allocated under base 3.
key = list(range(8)) + [90, 91, 92, 93]
res = _insert(tree, key, rotation_base=3)
self.assertTrue(res.rotation_tail_declined)
self.assertEqual(res.prefix_len, 8)
# The suffix is NOT cached: a full-key match stops at the seam.
self.assertEqual(_match_len(tree, key), 8)
def test_empty_page_aligned_key_insert(self):
"""A finished request with fewer cached tokens than one tree page
inserts an EMPTY page-aligned key; the empty-key early return must
precede the rotation pre-flight."""
tree = _unified_tree(page_size=self.PS)
res = _insert(tree, [1, 2], rotation_base=1)
self.assertEqual(res.prefix_len, 0)
self.assertFalse(res.rotation_tail_declined)
def test_same_base_tail_attaches(self):
tree = _unified_tree(page_size=self.PS)
self._seed_chain(tree, list(range(12)), base=1)
key = list(range(8)) + [90, 91, 92, 93]
res = _insert(tree, key, rotation_base=1)
self.assertFalse(res.rotation_tail_declined)
self.assertEqual(_match_len(tree, key), 12)
def test_no_matched_chain_never_declines(self):
"""A request that matches nothing starts its own chain: the guard
only fires against an EXISTING chain's base."""
tree = _unified_tree(page_size=self.PS)
self._seed_chain(tree, list(range(12)), base=1)
res = _insert(tree, list(range(50, 62)), rotation_base=3)
self.assertFalse(res.rotation_tail_declined)
self.assertEqual(_match_len(tree, list(range(50, 62))), 12)
def test_unsharded_insert_onto_a_based_chain_never_declines(self):
"""rotation_base=None means sharding is off for this insert: the
pre-flight must not fire, or every unsharded path would stop caching."""
tree = _unified_tree(page_size=self.PS)
self._seed_chain(tree, list(range(8)), base=1)
key = list(range(8)) + [90, 91, 92, 93]
res = _insert(tree, key, rotation_base=None)
self.assertFalse(res.rotation_tail_declined)
self.assertEqual(_match_len(tree, key), 12)
def test_full_match_under_a_foreign_base_declines(self):
"""No tail to graft, but the unified insert would still hand the
request's pages to the matched chain (unevict-on-insert, component
Full re-point). The pre-flight declines that too."""
tree = _unified_tree(page_size=self.PS)
self._seed_chain(tree, list(range(12)), base=1)
res = _insert(tree, list(range(12)), rotation_base=3)
self.assertTrue(res.rotation_tail_declined)
def test_cache_unfinished_decline_keeps_request_on_own_pages(self):
tree, freed = self._tree_with_spy()
self._seed_chain(tree, list(range(8)), base=1)
req = _GraftReq(list(range(8)) + [90, 91, 92, 93])
req.kv_rotation_base = 3
own_locs = self._own_row(tree, req, 12)
tree.cache_unfinished_req(req)
# No dedup free, no rebind: the request keeps its own locs whole.
self.assertEqual([t.tolist() for t in freed], [])
self.assertTrue(torch.equal(req.prefix_indices, own_locs))
self.assertEqual(req.kv.cache_protected_len, 0)
self.assertTrue(
torch.equal(tree.req_to_token_pool.req_to_token[0, :12], own_locs)
)
def test_cache_finished_decline_frees_duplicates_and_suffix(self):
tree, freed = self._tree_with_spy()
self._seed_chain(tree, list(range(8)), base=1)
req = _GraftReq(list(range(8)) + [90, 91, 92, 93])
req.kv_rotation_base = 3
own_locs = self._own_row(tree, req, 12)
tree.cache_finished_req(req, kv_len_to_handle=12)
released = torch.cat(freed)
# Everything past the protected prefix is released: the duplicates of
# the matched region AND the declined tail (nothing leaks, nothing is
# grafted).
self.assertEqual(set(released.tolist()), set(own_locs.tolist()))
self.assertEqual(_match_len(tree, req.fill_ids), 8)
def test_cache_finished_same_base_keeps_the_tail_cached(self):
"""Control for the decline test: with an agreeing base the tail is
grafted and only the matched duplicates are freed."""
tree, freed = self._tree_with_spy()
self._seed_chain(tree, list(range(8)), base=1)
req = _GraftReq(list(range(8)) + [90, 91, 92, 93])
req.kv_rotation_base = 1
own_locs = self._own_row(tree, req, 12)
tree.cache_finished_req(req, kv_len_to_handle=12)
self.assertEqual(_match_len(tree, req.fill_ids), 12)
released = torch.cat(freed) if freed else torch.empty(0, dtype=torch.int64)
# Only the 8 duplicate rows go back; the tail stays live in the tree.
self.assertEqual(set(released.tolist()), set(own_locs[:8].tolist()))
if __name__ == "__main__":
unittest.main()