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

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

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

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

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

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
+        # start_p is aligned to the allocator's page above, so it never shares a
-        # never shares a page with cache_finished_req's tail free in this group.
+        # 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

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

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

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

@@ -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)
+            # Free unaligned tail (+ deferred truncation tail). A rotation
-            ranges = [(page_aligned_len, len(kv_indices))]
+            # 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

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

@@ -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()