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>
This commit is contained in:
laoyao0822
2026-05-29 05:28:53 +08:00
co-authored by OmX
parent c551623ca8
commit 7cfc3c1324
8 changed files with 440 additions and 60 deletions
@@ -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: