Preserve CP HiCache valid tails while padding physical pages

CP HiCache now keeps radix and scheduler-visible lengths as valid tokens while host/device transfers reserve and replay the padded physical page span. Exact valid-tail write, insertion, and match paths no longer fall back to page-flooring; the physical owner-lane contract still uses padded page metadata.

Constraint: Scheduler prefix indices must never include padded tail locs.
Constraint: Host/device transfer and owner-lane admission remain page-based.
Rejected: Pad to cp_size or 2*cp_size pages | wastes KV and recreates short-tail fallback behavior.
Rejected: Expose padded locs through load_cp return | would leak fake tokens into req.prefix_indices.
Confidence: medium
Scope-risk: moderate
Directive: Do not implement split-inside-tail by duplicating page_owners without a page-sharing/refcount design.
Tested: local py_compile for touched CP HiCache/radix/controller files and tests.
Tested: remote g0034 CP HiCache impacted suites: 143 passed, 5 warnings.
Tested: remote g0034 CP shared KV C1-C5 suite: 122 passed, 5 warnings.
Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson/starlette.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>

docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md

@@ -549,6 +549,26 @@ Tests:
 - Host hit reports `100` to scheduler but reserves/loads `128` physical slots.
 - `page_owners` length is `2`.
 
+Implementation findings recorded before editing:
+
+- `init_load_back()` appends the tensor returned by `load_back()` directly into
+  `req.prefix_indices`.  Therefore a CP `load_cp()` implementation must not
+  expose padded tail locations through the returned tensor; padded locations may
+  be allocated and used for transfer, but the scheduler-visible return must be
+  valid-token length only.
+- `load_back()` currently slices each reloaded node from a single flat
+  `device_indices` tensor with `node.host_len`.  Once `host_len` becomes
+  valid-token length, the flat physical allocation offset must advance by
+  `metadata.padded_len`, not by `host_len`.
+- `load_back()` also increments device-cache residency accounting from the
+  tensor returned to the scheduler.  If `load_cp()` returns only valid locs,
+  residency accounting must still use the padded physical allocation length so
+  GPU cache pressure is not undercounted.
+- `_cp_load_back_node_owner_page_counts()` currently rejects non-page-aligned
+  device `node.value`.  That is acceptable for the existing page-aligned write
+  path, but it is a separate C7/C8 follow-up once radix nodes can store
+  valid-length tails backed by padded physical pages.
+
 ### C7. Radix page alignment currently floors partial tails
 
 Current state:
@@ -577,6 +597,56 @@ Tests:
 - Split a node around the tail page and verify valid/padded lengths remain
   consistent.
 
+Implementation findings recorded before editing:
+
+- `prepare_write_backup_for_req()` floors `req.fill_ids` through
+  `page_align_keys()` before reserving CP host backup.  This is the path that
+  drops short tails from write-through HiCache.
+- `_probe_existing_radix_prefix_len_no_split()` also floors its probe key.  If
+  write preparation keeps valid tails, this probe must use the same valid-key
+  contract or it can reserve duplicate tail spans.
+- `match_prefix()` floors lookup keys before matching.  Even if a valid-tail
+  radix node exists, current lookup will not report that tail as a device/host
+  hit.
+- `_split_node()` delegates CP HiCache metadata split to
+  `CpHiCacheNodeMetadata.split()`, which still requires page-boundary splits.
+  Supporting arbitrary valid-tail splits is therefore not just a metadata
+  change; it needs explicit tail-page ownership semantics or a delayed split.
+- `pin_prefix()` also floors keys.  Pinning is not on the hot cache-hit path,
+  but it will remain page-granular until the radix valid-tail contract is fully
+  audited.
+- `HiRadixCache` inherits `RadixCache.cache_finished_req()` and
+  `cache_unfinished_req()`.  Both floor keys before `insert()`, so prepared CP
+  backups for valid tails would be rolled back as `insert_miss` unless the
+  insertion path also keeps valid-tail keys.
+- `_key_match_paged()` assumes all keys were page-aligned; when two equal
+  partial-tail keys are compared, it can advance by a full `page_size` and
+  return a prefix length larger than `min(len(key0), len(key1))`.  Valid-tail
+  radix keys require the final compare step to advance only by the remaining
+  chunk length.
+
+Implemented C7 slice:
+
+- CP HiCache write preparation now keeps the valid tail instead of flooring
+  `fill_ids` before reservation.
+- CP `cache_finished_req()` / `cache_unfinished_req()` insertion keeps valid
+  tails when `_uses_cp_hicache` is true; non-CP radix behavior still floors.
+- CP `match_prefix()` keeps valid-tail lookup keys and reports the valid
+  `host_hit_length`.
+- CP write reservation pads only to the current tail page and stores
+  `metadata.logical_len == valid_len` with `metadata.padded_len` as the physical
+  span.
+- `_key_match_paged()` now returns the true valid prefix length for a partial
+  final page.
+
+Remaining C7 limitation:
+
+- Splitting an already-backed CP HiCache node inside a padded tail page still
+  needs a deliberate design.  The current safe slice supports exact valid-tail
+  hits and writes; divergent requests that force a split inside a physical tail
+  page are still a follow-up because splitting one physical host page across two
+  radix nodes would otherwise double-count or lose ownership metadata.
+
 ### C8. Owner-lane capacity must be padded-page based end to end
 
 Current state:

python/sglang/srt/managers/cache_controller.py

@@ -38,7 +38,10 @@ from sglang.srt.layers.dp_attention import (
     get_attention_tp_size,
     is_dp_attention_enabled,
 )
-from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+from sglang.srt.mem_cache.cp_shared_kv_layout import (
+    CpSharedKVLayout,
+    pad_token_locs_to_page_boundary,
+)
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool
 from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices
 from sglang.srt.utils import get_device_module
@@ -931,21 +934,22 @@ class HiCacheController:
     ) -> HiCacheWriteReservation | HiCacheWriteFailure:
         from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
 
+        page_size = self.page_size
         layout = self.cp_shared_kv_layout
-        owned_mask = layout.owned_by_this_rank(device_indices)
-        owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu()
         logical_len = len(device_indices)
+        padded_device_indices = pad_token_locs_to_page_boundary(
+            device_indices,
+            page_size,
+            name="CP HiCache write device_indices",
+        )
+        padded_len = int(padded_device_indices.numel())
+        owned_mask = layout.owned_by_this_rank(padded_device_indices)
+        owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu()
         # Capture the global owner pattern (one int8 per logical PAGE; identical
         # on all CP ranks since it's a pure function of the logical page ids).
         # load_cp will replay this pattern via alloc_pages_with_owners so the
         # saved owned_positions correctly index the new allocation.
-        page_size = self.page_size
-        if logical_len % page_size != 0:
-            raise ValueError(
-                f"_write_cp expects page-aligned device_indices, got "
-                f"logical_len={logical_len} page_size={page_size}"
-            )
-        page_first_locs = device_indices[::page_size]
+        page_first_locs = padded_device_indices[::page_size]
         logical_pages = torch.div(page_first_locs, page_size, rounding_mode="floor")
         page_owners = layout.owner_for_logical_pages(logical_pages).to(
             dtype=torch.int8, device="cpu"
@@ -960,6 +964,7 @@ class HiCacheController:
             return HiCacheWriteReservation(
                 metadata=CpHiCacheNodeMetadata(
                     logical_len=logical_len,
+                    padded_len=padded_len,
                     owned_positions=owned_positions,
                     host_indices=empty,
                     page_owners=page_owners,
@@ -973,7 +978,7 @@ class HiCacheController:
                 draft_host_indices=(empty.clone() if self.has_draft_hicache else None),
             )
 
-        owned_logical_indices = device_indices[owned_mask]
+        owned_logical_indices = padded_device_indices[owned_mask]
         physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices)
         host_indices = self.mem_pool_host.alloc(len(physical_device_indices))
         if host_indices is None:
@@ -1019,6 +1024,7 @@ class HiCacheController:
         return HiCacheWriteReservation(
             metadata=CpHiCacheNodeMetadata(
                 logical_len=logical_len,
+                padded_len=padded_len,
                 owned_positions=owned_positions,
                 host_indices=host_indices.cpu(),
                 page_owners=page_owners,
@@ -1350,24 +1356,47 @@ class HiCacheController:
         if device_indices is None:
             return None
 
-        logical_len_expected = sum(node.host_len for node in nodes_to_load)
-        if device_indices.numel() != logical_len_expected:
+        padded_len_expected = sum(
+            int(getattr(node.cp_hicache, "padded_len", node.host_len))
+            for node in nodes_to_load
+        )
+        if device_indices.numel() != padded_len_expected:
             self.mem_pool_device_allocator.free(device_indices)
             raise RuntimeError(
                 "alloc_pages_with_owners returned unexpected length: "
-                f"got {device_indices.numel()}, expected {logical_len_expected}"
+                f"got {device_indices.numel()}, expected {padded_len_expected}"
             )
 
         host_chunks = []
         draft_host_chunks = []
         physical_chunks = []
+        visible_chunks = []
         offset = 0
         for node in nodes_to_load:
-            node_device_indices = device_indices[offset : offset + node.host_len]
-            offset += node.host_len
+            meta = node.cp_hicache
+            valid_len = int(getattr(meta, "valid_len", node.host_len))
+            padded_len = int(getattr(meta, "padded_len", node.host_len))
+            if valid_len != int(node.host_len):
+                self.mem_pool_device_allocator.free(device_indices)
+                raise RuntimeError(
+                    "CP HiCache load metadata valid length mismatch: "
+                    f"node_id={getattr(node, 'id', '?')} "
+                    f"host_len={node.host_len} valid_len={valid_len}"
+                )
+            if padded_len < valid_len:
+                self.mem_pool_device_allocator.free(device_indices)
+                raise RuntimeError(
+                    "CP HiCache load metadata padded length is shorter than "
+                    "valid length: "
+                    f"node_id={getattr(node, 'id', '?')} "
+                    f"valid_len={valid_len} padded_len={padded_len}"
+                )
+            node_device_indices = device_indices[offset : offset + padded_len]
+            offset += padded_len
+            visible_chunks.append(node_device_indices[:valid_len])
             if self.has_draft_hicache:
                 draft_host_indices = getattr(
-                    node.cp_hicache, "draft_host_indices", None
+                    meta, "draft_host_indices", None
                 )
                 if draft_host_indices is None:
                     self.mem_pool_device_allocator.free(device_indices)
@@ -1378,7 +1407,7 @@ class HiCacheController:
             else:
                 draft_host_indices = None
 
-            owned_positions = node.cp_hicache.owned_positions.to(device_indices.device)
+            owned_positions = meta.owned_positions.to(device_indices.device)
             if owned_positions.numel() == 0:
                 continue
             selected_logical_locs = node_device_indices[owned_positions]
@@ -1397,10 +1426,18 @@ class HiCacheController:
             physical_chunks.append(
                 self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
             )
-            host_chunks.append(node.cp_hicache.host_indices)
+            host_chunks.append(meta.host_indices)
             if draft_host_indices is not None:
                 draft_host_chunks.append(draft_host_indices)
 
+        visible_device_indices = (
+            torch.cat(visible_chunks)
+            if visible_chunks
+            else torch.empty(
+                (0,), dtype=device_indices.dtype, device=device_indices.device
+            )
+        )
+
         if not host_chunks:
             # Keep CP load ACK rows identical across ranks.  A zero-owned rank
             # still queues a zero-length op with the logical node id so a later
@@ -1414,7 +1451,7 @@ class HiCacheController:
                     node_id,
                 )
             )
-            return device_indices
+            return visible_device_indices
 
         host_indices = torch.cat(host_chunks)
         physical_device_indices = torch.cat(physical_chunks)
@@ -1450,7 +1487,7 @@ class HiCacheController:
                 )
             )
 
-        return device_indices
+        return visible_device_indices
 
     def move_indices(self, op: CacheOperation, mem_pool_host=None):
         mem_pool_host = mem_pool_host or self.mem_pool_host

python/sglang/srt/mem_cache/cp_shared_kv_layout.py

@@ -6,6 +6,53 @@ import numpy as np
 import torch
 
 
+def pad_token_locs_to_page_boundary(
+    token_locs: torch.Tensor,
+    page_size: int,
+    *,
+    name: str = "token_locs",
+) -> torch.Tensor:
+    """Pad a valid token-loc span to the end of its physical tail page.
+
+    The returned tensor keeps the original valid locs first and appends only
+    tail-page slack locs.  It does not pad to CP size and it does not create
+    fake request tokens.
+    """
+
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+    valid_len = int(token_locs.numel())
+    if valid_len == 0:
+        return token_locs
+
+    padding_len = (-valid_len) % page_size
+    if padding_len == 0:
+        return token_locs
+
+    if not token_locs.is_cuda:
+        first_offset = int(torch.remainder(token_locs[0], page_size).item())
+        if first_offset != 0:
+            raise ValueError(
+                f"{name} must start at a page boundary before padding, got "
+                f"first_loc={int(token_locs[0].item())} page_size={page_size}"
+            )
+        last_offset = int(torch.remainder(token_locs[-1], page_size).item())
+        if last_offset + padding_len != page_size - 1:
+            raise ValueError(
+                f"{name} does not end inside the expected tail page: "
+                f"last_loc={int(token_locs[-1].item())} valid_len={valid_len} "
+                f"padding_len={padding_len} page_size={page_size}"
+            )
+
+    pad_offsets = torch.arange(
+        1,
+        padding_len + 1,
+        dtype=token_locs.dtype,
+        device=token_locs.device,
+    )
+    return torch.cat([token_locs, token_locs[-1:] + pad_offsets])
+
+
 @dataclass(frozen=True)
 class CpSharedKVLayout:
     """Map CP-group logical KV locations to per-rank physical KV locations.

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -33,7 +33,10 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
     MatchResult,
 )
-from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+from sglang.srt.mem_cache.cp_shared_kv_layout import (
+    CpSharedKVLayout,
+    pad_token_locs_to_page_boundary,
+)
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
     MLATokenToKVPool,
@@ -44,7 +47,6 @@ from sglang.srt.mem_cache.radix_cache import (
     RadixKey,
     TreeNode,
     compute_node_hash_values,
-    page_align_keys,
     split_node_hash_value,
 )
 from sglang.srt.mem_cache.utils import convert_to_bigram_key
@@ -124,6 +126,8 @@ def _compute_shared_hicache_token_capacities(
 
 @dataclass
 class CpHiCacheNodeMetadata:
+    # Legacy name kept for existing call sites.  Under the page-aligned cache
+    # contract this is the radix/scheduler-visible valid token length.
     logical_len: int
     owned_positions: torch.Tensor
     host_indices: torch.Tensor
@@ -138,6 +142,9 @@ class CpHiCacheNodeMetadata:
     # NEW: needed at split() time to convert split_len into a page index
     # without plumbing it from the caller.
     page_size: int
+    # Physical host/device span.  When omitted, legacy callers keep the old
+    # page-aligned invariant by using logical_len as the physical length.
+    padded_len: Optional[int] = None
     draft_host_indices: Optional[torch.Tensor] = None
 
     def __post_init__(self):
@@ -145,9 +152,19 @@ class CpHiCacheNodeMetadata:
             raise ValueError(f"logical_len must be non-negative, got {self.logical_len}")
         if self.page_size <= 0:
             raise ValueError(f"page_size must be positive, got {self.page_size}")
-        if self.logical_len % self.page_size != 0:
+        if self.padded_len is None:
+            self.padded_len = self.logical_len
+        self.padded_len = int(self.padded_len)
+        if self.padded_len < 0:
+            raise ValueError(f"padded_len must be non-negative, got {self.padded_len}")
+        if self.padded_len < self.logical_len:
             raise ValueError(
-                f"logical_len ({self.logical_len}) must be a multiple of "
+                f"padded_len ({self.padded_len}) must be >= logical_len "
+                f"({self.logical_len})"
+            )
+        if self.padded_len % self.page_size != 0:
+            raise ValueError(
+                f"padded_len ({self.padded_len}) must be a multiple of "
                 f"page_size ({self.page_size})"
             )
         self.owned_positions = self.owned_positions.to(
@@ -163,11 +180,11 @@ class CpHiCacheNodeMetadata:
             self.draft_host_indices = self.draft_host_indices.to(
                 device="cpu", dtype=torch.int64
             ).detach().clone()
-        expected_num_pages = self.logical_len // self.page_size
+        expected_num_pages = self.padded_len // self.page_size
         if self.page_owners.numel() != expected_num_pages:
             raise ValueError(
                 f"page_owners length ({self.page_owners.numel()}) must equal "
-                f"logical_len/page_size ({expected_num_pages})"
+                f"padded_len/page_size ({expected_num_pages})"
             )
         if expected_num_pages > 0 and bool((self.page_owners < 0).any()):
             raise ValueError("page_owners entries must be non-negative")
@@ -186,14 +203,21 @@ class CpHiCacheNodeMetadata:
             )
         if self.owned_positions.numel() > 0:
             if torch.any(self.owned_positions < 0) or torch.any(
-                self.owned_positions >= self.logical_len
+                self.owned_positions >= self.padded_len
             ):
-                raise ValueError("owned_positions must be in [0, logical_len)")
+                raise ValueError(
+                    "owned_positions must be in [0, padded_len) "
+                    "(legacy [0, logical_len) for unpadded metadata)"
+                )
             if torch.any(self.owned_positions[1:] <= self.owned_positions[:-1]):
                 raise ValueError(
                     "owned_positions must be sorted and strictly increasing"
                 )
 
+    @property
+    def valid_len(self) -> int:
+        return self.logical_len
+
     def split(
         self, split_len: int
     ) -> tuple["CpHiCacheNodeMetadata", "CpHiCacheNodeMetadata"]:
@@ -216,6 +240,7 @@ class CpHiCacheNodeMetadata:
                 host_indices=self.host_indices[parent_mask],
                 page_owners=self.page_owners[:split_pages],
                 page_size=self.page_size,
+                padded_len=split_len,
                 draft_host_indices=(
                     self.draft_host_indices[parent_mask]
                     if self.draft_host_indices is not None
@@ -228,6 +253,7 @@ class CpHiCacheNodeMetadata:
                 host_indices=self.host_indices[child_mask],
                 page_owners=self.page_owners[split_pages:],
                 page_size=self.page_size,
+                padded_len=self.padded_len - split_len,
                 draft_host_indices=(
                     self.draft_host_indices[child_mask]
                     if self.draft_host_indices is not None
@@ -780,13 +806,13 @@ class HiRadixCache(RadixCache):
         logical_len = int(device_indices.numel())
         if logical_len == 0:
             return self._cp_zero_counts(layout.cp_size)
-        if logical_len % self.page_size != 0:
-            raise RuntimeError(
-                "CP HiCache capacity planning requires page-aligned device indices: "
-                f"logical_len={logical_len} page_size={self.page_size}"
-            )
+        padded_device_indices = pad_token_locs_to_page_boundary(
+            device_indices,
+            self.page_size,
+            name="CP HiCache capacity device_indices",
+        )
         logical_pages = torch.div(
-            device_indices[:: self.page_size],
+            padded_device_indices[:: self.page_size],
             self.page_size,
             rounding_mode="floor",
         )
@@ -904,6 +930,7 @@ class HiRadixCache(RadixCache):
 
         page_owners: List[int] = []
         host_hit_len = 0
+        physical_hit_len = 0
         for node in nodes_to_load:
             metadata = getattr(node, "cp_hicache", None)
             if metadata is None:
@@ -912,26 +939,35 @@ class HiRadixCache(RadixCache):
                     f"load_node_id={getattr(node, 'id', '?')} root_node_id={node_id}"
                 )
             node_host_len = int(self._node_host_len(node))
-            if node_host_len % self.page_size != 0:
+            node_padded_len = int(getattr(metadata, "padded_len", node_host_len))
+            if node_padded_len % self.page_size != 0:
                 raise RuntimeError(
-                    "CP HiCache load-back requires page-aligned host lengths: "
+                    "CP HiCache load-back requires page-aligned physical lengths: "
                     f"load_node_id={getattr(node, 'id', '?')} "
-                    f"host_len={node_host_len} page_size={self.page_size}"
+                    f"padded_len={node_padded_len} page_size={self.page_size}"
                 )
-            if int(getattr(metadata, "logical_len", node_host_len)) != node_host_len:
+            if int(getattr(metadata, "valid_len", node_host_len)) != node_host_len:
                 raise RuntimeError(
                     "CP HiCache load-back metadata length mismatch: "
                     f"load_node_id={getattr(node, 'id', '?')} "
-                    f"host_len={node_host_len} metadata_len={metadata.logical_len}"
+                    f"host_len={node_host_len} metadata_len={metadata.valid_len}"
+                )
+            if node_padded_len < node_host_len:
+                raise RuntimeError(
+                    "CP HiCache load-back physical length is shorter than "
+                    "valid host length: "
+                    f"load_node_id={getattr(node, 'id', '?')} "
+                    f"host_len={node_host_len} padded_len={node_padded_len}"
                 )
             page_owners.extend(int(owner) for owner in metadata.page_owners.tolist())
             host_hit_len += node_host_len
+            physical_hit_len += node_padded_len
 
-        expected_pages = host_hit_len // self.page_size
+        expected_pages = physical_hit_len // self.page_size
         if len(page_owners) != expected_pages:
             raise RuntimeError(
                 "CP HiCache load-back page owner count mismatch: "
-                f"node_id={node_id} host_hit_len={host_hit_len} "
+                f"node_id={node_id} physical_hit_len={physical_hit_len} "
                 f"page_size={self.page_size} page_owners={len(page_owners)}"
             )
 
@@ -1956,9 +1992,6 @@ class HiRadixCache(RadixCache):
             return 0
 
         key, _ = self.maybe_bigram_convert(key)
-        if self.page_size != 1:
-            page_aligned_len = len(key) // self.page_size * self.page_size
-            key = key[:page_aligned_len]
 
         if len(key) == 0:
             return 0
@@ -2016,7 +2049,6 @@ class HiRadixCache(RadixCache):
 
         token_ids = req.fill_ids
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
-        keys = page_align_keys(keys, self.page_size)
         existing_prefix_len = self._probe_existing_radix_prefix_len_no_split(
             RadixKey(
                 keys,
@@ -2820,11 +2852,16 @@ class HiRadixCache(RadixCache):
 
             self.ongoing_load_back[last_hit_node.id] = last_hit_node
             offset = 0
+            physical_loaded_len = 0
             for loaded_node in nodes_to_load:
                 host_len = self._node_host_len(loaded_node)
                 loaded_node.value = device_indices[offset : offset + host_len].clone()
                 offset += host_len
-            self.evictable_size_ += len(device_indices)
+                metadata = getattr(loaded_node, "cp_hicache", None)
+                physical_loaded_len += int(
+                    getattr(metadata, "padded_len", host_len)
+                )
+            self.evictable_size_ += physical_loaded_len
             self.inc_lock_ref(last_hit_node)
 
             if self.metrics_collector is not None:
@@ -2836,9 +2873,10 @@ class HiRadixCache(RadixCache):
                 )
 
             logger.info(
-                "[HiCache-load] load_back CP SUCCESS: node_id=%d loaded_tokens=%d",
+                "[HiCache-load] load_back CP SUCCESS: node_id=%d loaded_tokens=%d physical_tokens=%d",
                 last_hit_node.id,
                 len(device_indices),
+                physical_loaded_len,
             )
             return device_indices
 
@@ -3150,7 +3188,7 @@ class HiRadixCache(RadixCache):
             )
 
         page_aligned_len = len(key)
-        if self.page_size != 1:
+        if self.page_size != 1 and not self._uses_cp_hicache:
             page_aligned_len = len(key) // self.page_size * self.page_size
             key = key[:page_aligned_len]
 
@@ -3177,7 +3215,7 @@ class HiRadixCache(RadixCache):
         host_hit_length = 0
         last_host_node = last_node
         if self._uses_cp_hicache:
-            while last_node.evicted:
+            while last_node != self.root_node and last_node.evicted:
                 host_hit_length += self._node_host_len(last_node)
                 last_node = last_node.parent
             while (

python/sglang/srt/mem_cache/radix_cache.py

@@ -209,9 +209,10 @@ def _key_match_paged(key0: RadixKey, key1: RadixKey, page_size: int):
 
     i = 0
     while i < min_len:
-        if key0.token_ids[i : i + page_size] != key1.token_ids[i : i + page_size]:
+        step = min(page_size, min_len - i)
+        if key0.token_ids[i : i + step] != key1.token_ids[i : i + step]:
             break
-        i += page_size
+        i += step
 
     return i
 
@@ -486,7 +487,8 @@ class RadixCache(BasePrefixCache):
 
         # Maybe convert to bigram keys for EAGLE
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
-        keys = page_align_keys(keys, self.page_size)
+        if not getattr(self, "_uses_cp_hicache", False):
+            keys = page_align_keys(keys, self.page_size)
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
         prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
@@ -542,7 +544,8 @@ class RadixCache(BasePrefixCache):
 
         # Maybe convert to bigram keys for EAGLE
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
-        keys = page_align_keys(keys, self.page_size)
+        if not getattr(self, "_uses_cp_hicache", False):
+            keys = page_align_keys(keys, self.page_size)
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
         prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)

test/registered/unit/managers/test_hicache_controller_cp.py

@@ -701,15 +701,20 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
         self.assertEqual(host_pool.backups, [])
         self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
 
-    def test_cp_write_rejects_incomplete_owned_physical_page(self):
+    def test_cp_write_accepts_valid_tail_and_pads_owned_physical_page(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         controller = self.make_controller(host_pool, cp_rank=1)
         logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
 
-        with self.assertRaisesRegex(
-            ValueError, "_write_cp expects page-aligned device_indices"
-        ):
-            controller.write(logical_locs, node_id=21)
+        result = controller.write(logical_locs, node_id=21)
+
+        self.assertEqual(result.metadata.logical_len, 3)
+        self.assertEqual(result.metadata.valid_len, 3)
+        self.assertEqual(result.metadata.padded_len, 4)
+        self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3])
+        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
+        self.assertEqual(host_pool.alloc_calls, [4])
+        self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
 
     def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
@@ -1052,6 +1057,29 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
         self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
 
+    def test_cp_load_returns_valid_locs_while_transferring_padded_tail_page(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 72, dtype=torch.int64))
+        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
+        node = TreeNode()
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            padded_len=8,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            page_owners=torch.tensor([3, 0], dtype=torch.int8),
+            page_size=4,
+        )
+
+        device_indices = controller.load_cp([node], node_id=112)
+        controller.start_loading()
+
+        self.assertEqual(device_indices.tolist(), list(range(64, 70)))
+        self.assertEqual(allocator.owner_alloc_calls, [[3, 0]])
+        self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
+        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
+
     def test_cp_load_frees_unexpected_owner_allocator_length(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))

test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py

@@ -232,6 +232,26 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase):
         self.assertEqual(plan.deficit_by_owner, [0, 1, 0, 0])
         self.assertEqual(plan.host_hit_len, 16)
 
+    def test_load_back_plan_keeps_valid_hit_len_while_planning_padded_pages(self):
+        allocator = _make_allocator(page_size=4, cp_size=4)
+        cache = _make_cache(allocator)
+        node = TreeNode(id=12)
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            padded_len=8,
+            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
+            host_indices=torch.arange(4, dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=4,
+        )
+
+        plan = cache._build_cp_load_back_plan([node], node_id=node.id)
+
+        self.assertEqual(plan.host_hit_len, 6)
+        self.assertEqual(plan.page_owners, [0, 1])
+        self.assertEqual(plan.required_by_owner, [1, 1, 0, 0])
+
     def test_load_back_plan_fails_closed_without_cp_metadata(self):
         allocator = _make_allocator()
         cache = _make_cache(allocator)

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -109,7 +109,7 @@ from sglang.srt.mem_cache.hiradix_cache import (
     PreparedCpHiCacheBackup,
     _compute_shared_hicache_token_capacities,
 )
-from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
+from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, _key_match_paged
 from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -199,6 +199,11 @@ class TestHiRadixCacheCPDraftHostPool(CustomTestCase):
 
 
 class TestCpHiCacheNodeMetadata(CustomTestCase):
+    def test_paged_key_match_returns_valid_tail_length_not_next_page(self):
+        key = RadixKey(list(range(6)))
+
+        self.assertEqual(_key_match_paged(key, key, page_size=4), 6)
+
     def test_split_zero_len_moves_all_positions_to_child(self):
         metadata = CpHiCacheNodeMetadata(
             logical_len=8,
@@ -268,6 +273,21 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
         self.assertEqual(metadata.owned_positions.dtype, torch.int64)
         self.assertEqual(metadata.host_indices.dtype, torch.int64)
 
+    def test_valid_length_can_be_shorter_than_physical_padded_length(self):
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=100,
+            padded_len=128,
+            owned_positions=torch.tensor([0, 63, 100, 127], dtype=torch.int64),
+            host_indices=torch.tensor([10, 11, 12, 13], dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=64,
+        )
+
+        self.assertEqual(metadata.logical_len, 100)
+        self.assertEqual(metadata.valid_len, 100)
+        self.assertEqual(metadata.padded_len, 128)
+        self.assertEqual(metadata.page_owners.tolist(), [0, 1])
+
     def test_non_int64_inputs_are_converted(self):
         metadata = CpHiCacheNodeMetadata(
             logical_len=4,
@@ -1162,6 +1182,80 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
             [{"catch_up_all_layers": False}],
         )
 
+    def test_prepare_write_backup_for_req_keeps_valid_tail_length(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.disable = False
+        cache._uses_cp_hicache = True
+        cache.is_eagle = False
+        cache.page_size = 4
+        cache.req_to_token_pool = types.SimpleNamespace(
+            req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
+        )
+        cache.cache_controller = FakeReserveWriteController(
+            [
+                lambda device_indices, node_id: make_write_reservation(
+                    device_indices, node_id=node_id, host_start=170
+                )
+            ]
+        )
+        cache.maybe_bigram_convert = lambda key: (key, None)
+        cache.root_node = TreeNode()
+        cache.root_node.children = {}
+        cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
+
+        req = types.SimpleNamespace(
+            rid="rid-tail-prepare",
+            fill_ids=list(range(6)),
+            cache_protected_len=0,
+            req_pool_idx=0,
+            is_chunked=0,
+            cp_hicache_prepared_backup=None,
+        )
+
+        cache.prepare_write_backup_for_req(req)
+
+        self.assertIsNotNone(req.cp_hicache_prepared_backup)
+        self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 6)
+        self.assertEqual(cache.cache_controller.reservations[0][0].tolist(), list(range(6)))
+
+    def test_cache_finished_req_keeps_cp_valid_tail_insert_key(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.disable_finished_insert = False
+        cache.disable = False
+        cache._uses_cp_hicache = True
+        cache.is_eagle = False
+        cache.page_size = 4
+        cache.req_to_token_pool = types.SimpleNamespace(
+            req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
+        )
+        freed = []
+        cache.token_to_kv_pool_allocator = types.SimpleNamespace(
+            free=lambda indices: freed.append(indices.clone())
+        )
+        inserted = []
+        cache.insert = lambda params: inserted.append(params) or types.SimpleNamespace(
+            prefix_len=0
+        )
+        cache.dec_lock_ref = lambda node: None
+
+        req = types.SimpleNamespace(
+            origin_input_ids=list(range(6)),
+            output_ids=[],
+            req_pool_idx=0,
+            extra_key=None,
+            cache_protected_len=0,
+            last_node=TreeNode(),
+            cp_hicache_prepared_backup=None,
+            pop_committed_kv_cache=lambda: 6,
+        )
+
+        cache.cache_finished_req(req)
+
+        self.assertEqual(len(inserted), 1)
+        self.assertEqual(inserted[0].key.token_ids, list(range(6)))
+        self.assertEqual(inserted[0].value.tolist(), list(range(6)))
+        self.assertEqual([indices.tolist() for indices in freed], [[], []])
+
     def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache.disable = False
@@ -2136,6 +2230,49 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         self.assertIs(result.last_host_node, cache.root_node)
         self.assertEqual(result.host_hit_length, 0)
 
+    def test_cp_match_prefix_reports_valid_tail_host_hit(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache._uses_cp_hicache = True
+        cache.device = "cpu"
+        cache.disable = False
+        cache.page_size = 4
+        cache.ongoing_write_through = {}
+        cache.pending_host_backups = {}
+        cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
+        cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
+        cache.key_match_fn = lambda child_key, key: _key_match_paged(
+            child_key, key, page_size=4
+        )
+        cache.maybe_bigram_convert = lambda key: (key, None)
+        root = TreeNode()
+        root.key = RadixKey([])
+        root.value = torch.empty((0,), dtype=torch.int64)
+        root.host_len = 0
+        cache.root_node = root
+        node = TreeNode()
+        node.id = 140
+        node.parent = root
+        node.key = RadixKey(list(range(6)))
+        node.value = None
+        node.host_value = None
+        node.host_len = 6
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=6,
+            padded_len=8,
+            owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
+            host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=4,
+        )
+        root.children[(0, 1, 2, 3)] = node
+
+        result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6)))))
+
+        self.assertEqual(result.device_indices.tolist(), [])
+        self.assertEqual(result.host_hit_length, 6)
+        self.assertIs(result.last_device_node, root)
+        self.assertIs(result.last_host_node, node)
+
     def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = False