From 7cfc3c1324064e2dcbe34caa8c92b93dd442e753 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 05:28:53 +0800 Subject: [PATCH] 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 --- ..._prefill_cp_page_aligned_cache_contract.md | 70 +++++++++ .../sglang/srt/managers/cache_controller.py | 79 +++++++--- .../srt/mem_cache/cp_shared_kv_layout.py | 47 ++++++ python/sglang/srt/mem_cache/hiradix_cache.py | 96 ++++++++---- python/sglang/srt/mem_cache/radix_cache.py | 11 +- .../managers/test_hicache_controller_cp.py | 38 ++++- .../test_cp_hicache_load_back_owner_lanes.py | 20 +++ .../mem_cache/test_cp_hicache_metadata.py | 139 +++++++++++++++++- 8 files changed, 440 insertions(+), 60 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 3e36b4b0a..b58e543fc 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/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: diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 58d4337ef..81f1e012c 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/cp_shared_kv_layout.py b/python/sglang/srt/mem_cache/cp_shared_kv_layout.py index 7a038fe84..788944d56 100644 --- a/python/sglang/srt/mem_cache/cp_shared_kv_layout.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 3c9bd7831..b2812816b 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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 ( diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index d440257d2..6a55209ce 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index ea3af2f86..e0816b672 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/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)) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py index cf9c625ca..ad7f0bdec 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +++ b/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) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 478eb0565..eeaec8700 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/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