diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4290ad32b..efde185a4 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 5e109a4c8..75d33454e 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -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 diff --git a/python/sglang/srt/mem_cache/allocator/page_interleave.py b/python/sglang/srt/mem_cache/allocator/page_interleave.py new file mode 100644 index 000000000..0fd38a35e --- /dev/null +++ b/python/sglang/srt/mem_cache/allocator/page_interleave.py @@ -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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index f1f774d07..5034ad6a5 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 8ac5ca73a..9757fe57a 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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)]) diff --git a/python/sglang/srt/mem_cache/page_interleave.py b/python/sglang/srt/mem_cache/page_interleave.py new file mode 100644 index 000000000..9092b4ed0 --- /dev/null +++ b/python/sglang/srt/mem_cache/page_interleave.py @@ -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)]) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 83206f4b2..84e77362e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -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) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 94fc9e5ec..64c36443a 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -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.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f602d2c52..19bed9e53 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index c37ae1927..601175449 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -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] diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py new file mode 100644 index 000000000..09e467046 --- /dev/null +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -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()