[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
-        # never shares a page with cache_finished_req's tail free in this group.
+        # start_p is aligned to the allocator's page above, so it never shares a
+        # page with cache_finished_req's tail free in this group.
         tree_cache.free_kv_row(req.kv, [(start_p, end_p)])
 
 

python/sglang/srt/mem_cache/page_interleave.py

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