diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md new file mode 100644 index 000000000..2e81de2ab --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md @@ -0,0 +1,542 @@ +# NSA Prefill CP HiCache Per-Layer Backup Design Plan + +Date: 2026-05-27 + +Branch context: `cjy-cp-refactor` after `367dff06f3` (`Keep CP HiCache draft KV invisible until joint readiness`). + +This note records the current code reading and the design constraints for turning +HiCache backup into a per-layer data-plane operation under NSA prefill CP shared +KV, HiCache, and EAGLE/MTP draft KV. In this document, **backup** means the +D2H target/draft KV copy into host HiCache. It is not the CLI +`--hicache-write-policy write_back` eviction policy. + +## Scope + +In scope: + +- prefill-side CP shared KV (`--enable-nsa-prefill-cp-shared-kv`), +- HiCache L2 host backup/load/host eviction, +- target KV plus EAGLE/MTP draft KV as one logical cache object, +- write-through / write-through-selective style backup, +- per-layer D2H backup after each layer has produced immutable KV, +- low-synchronization CP rank behavior. + +Out of scope for this stage: + +- decode-side cache lifecycle, +- TBO/two-batch overlap scheduling, +- L3/storage write-back or prefetch, +- fully async emergency demotion for `--hicache-write-policy write_back`, +- cross-instance cache sharing. + +## Current Code Facts + +### Radix node state + +`TreeNode` currently stores both device and host residency fields in +`python/sglang/srt/mem_cache/radix_cache.py`: + +- `value`: device KV locs, full logical loc tensor in CP shared-KV mode. +- `host_value`: legacy non-CP physical host locs. +- `host_len`: logical host-backed length in CP HiCache mode. +- `cp_hicache`: CP host metadata for owned positions, page owner pattern, + target host slots, and optional draft host slots. +- `evicted` means `value is None`. +- `backuped` only reflects `host_value is not None`, so CP code must use + `HiRadixCache._node_backuped()` instead of the plain property. + +Current CP host-valid predicate is in +`python/sglang/srt/mem_cache/hiradix_cache.py`: + +```text +_node_backuped(node): + if CP HiCache: + host_len > 0 + required draft metadata exists + node is not in ongoing_write_through +``` + +This is already close to the desired request-visible rule: in-flight backup is +not host-visible. + +### Current backup path + +Current CP backup path: + +1. `HiRadixCache.write_backup(node)` calls `cache_controller.write(node.value)`. +2. `HiCacheController._write_cp()` treats `node.value` as full logical locs. +3. `_write_cp()` derives local owned positions through + `CpSharedKVLayout.owned_by_this_rank()`. +4. `_write_cp()` converts local logical locs to physical locs. +5. `_write_cp()` allocates target host slots and, when draft HiCache is + attached, draft host slots for the same local physical count. +6. Draft host allocation failure frees target host slots and returns failure. +7. Zero-owned ranks enqueue a completed no-op write ack. +8. Non-zero-owned ranks enqueue target and draft `CacheOperation`s. +9. `HiCacheController.start_writing()` merges queued ops and calls + `backup_from_device_all_layer()` for target, then draft, on `write_stream`. +10. One `HiCacheAck` is appended for the merged logical node ids. +11. `HiRadixCache.write_backup()` stores `host_len`, `cp_hicache`, inserts the + node into `ongoing_write_through`, and locks the device source when this is + write-through rather than eviction-time write-back. +12. `writing_check()` polls completed acks, uses `all_reduce(MIN)` on completed + ack prefix length, releases write locks, and makes the node host-visible by + removing it from `ongoing_write_through`. + +Important current limitation: backup is all-layer D2H. Load is already +per-layer H2D, but backup waits until it is submitted as one all-layer transfer. + +### Current load path + +Current CP load path: + +1. `match_prefix()` may return a host-hit node only when `_node_backuped()` is + true. +2. `schedule_policy.add_one_req()` calls `init_load_back()` when + `host_hit_length > 0`. +3. `load_back()` walks host-backed evicted nodes and calls + `HiCacheController.load_cp()`. +4. `load_cp()` replays `node.cp_hicache.page_owners` with + `alloc_pages_with_owners()` so the new logical device allocation matches the + write-time owner pattern. +5. It queues target and draft H2D using the same owned physical device slots. +6. `start_loading()` copies per layer and completes `LayerDoneCounter` events. +7. The batch carries `hicache_consumer_index`; KV pool reads wait for the layer + event before using that layer. + +This establishes the precedent that transfer can be layer-wise while radix +state remains scheduler-owned. + +### Current host eviction path + +Host eviction currently only targets host-only logical leaves: + +- `_update_host_leaf_status()` removes a node from `evictable_host_leaves` when + `not node.evicted` or `node.lock_ref > 0`. +- `_evict_host_for_physical_slots()` selects from `evictable_host_leaves`, skips + malformed/non-backed/pinned/host-referenced nodes, frees target and draft + host slots through `evict_cp_host()`, clears `host_len`, `cp_hicache`, and + `host_value`, then removes the logical host leaf. + +Therefore, in the current code, **device-valid + host-valid nodes are normally +not host-evictable**. Host eviction is intended for +`DEVICE_EVICTED + HOST_VALID` nodes. + +### Device-hit but host-missing is valid + +The design must not assume `device hit => host exists`. The following states +are valid today and must remain valid: + +- a freshly inserted node before write-through threshold is reached, +- a node whose backup failed because host was full, +- a node under in-flight backup, +- a node whose host copy was never made because the policy is selective, +- a node recomputed back into device after an earlier host-only node was evicted + from host. + +In these cases, service uses device KV directly. Host cache is only an optional +backup copy. + +## Desired State Model + +Per-node host/device state should be explicit enough to separate slot ownership, +transfer progress, and request visibility: + +```text +DEVICE_VALID_HOST_NONE +DEVICE_VALID_HOST_RESERVED +DEVICE_VALID_HOST_BACKUP_IN_FLIGHT +DEVICE_VALID_HOST_VALID +DEVICE_EVICTED_HOST_VALID +DEVICE_LOAD_PENDING_HOST_VALID +HOST_INVALID +``` + +Required visibility rules: + +- `HOST_RESERVED` is not a cache hit. +- `HOST_BACKUP_IN_FLIGHT` is not a cache hit. +- `HOST_VALID` is the first state where `match_prefix()` may return a host hit. +- `HOST_RESERVED` / `HOST_BACKUP_IN_FLIGHT` must not enter + `evictable_host_leaves`. +- Device source KV must stay protected until backup commit or rollback. +- Host destination slots must stay protected until backup commit or rollback. + +The current `ongoing_write_through` map partially models this, but it mixes too +many meanings: reserved host slots, submitted D2H, pending host visibility, and +device source protection. Per-layer backup should introduce an explicit pending +backup object or equivalent fields. + +## Host-Full Reservation Flow + +Host-full handling should be reserve-first, evict-on-failure, retry-once: + +```text +try reserve target+draft host slots +if success: + mark HOST_RESERVED and submit per-layer backup +else: + evict host-only HOST_VALID victims for required local physical slots + retry reserve target+draft host slots + if success: + mark HOST_RESERVED and submit per-layer backup + else: + skip backup; keep DEVICE_VALID_HOST_NONE +``` + +Rules: + +1. Host eviction victims must be `DEVICE_EVICTED_HOST_VALID`, unlocked, + unpinned, not a load source, and not a backup destination. +2. Reservation has target+draft atomic semantics. Target allocation success and + draft allocation failure must rollback target allocation. +3. The retry loop must be bounded. Repeated scheduler hot retries under host + pressure are worse than skipping backup. +4. Backup failure under write-through/selective is not a correctness failure. + The node remains device-resident and host-missing. +5. Zero-owned ranks must still create the same logical pending backup op with a + no-op event. + +## Low-Synchronization Policy + +Do not introduce per-layer collectives. + +Per-layer D2H should be tracked by local CUDA events or layer completion +counters only. CP rank synchronization should happen at coarse logical gates: + +- reserve/submit failure slow path: one global success/abort check after local + retry, +- final backup visibility commit: one final logical ack check and one batched contiguous-op prefix decision, +- debug/invariant validation: optional and off the hot path. + +Fast path target: + +```text +reserve local target/draft host slots +submit local per-layer D2H or zero-owned no-op +record local completion +commit completed logical prefix in deterministic order +``` + +Avoid: + +```text +all_reduce per layer +all_reduce per node victim in a tight host-eviction loop +all_reduce every scheduler tick when no completion prefix advanced +``` + +## Per-Layer Backup Data Plane + +The existing host pool exposes `load_to_device_per_layer()` but not a matching +backup method. Per-layer backup needs a new mirror API for every relevant host +pool type: + +```text +backup_from_device_per_layer(device_pool, host_indices, device_indices, layer_id, io_backend) +``` + +For NSA host pools this must copy both payloads for the same layer: + +- MLA/NSA KV payload, +- NSA indexer K/scale payload. + +The existing all-layer NSA backup calls `super().backup_from_device_all_layer()` +and `_backup_indexer_from_device_all_layer()`. The per-layer API must mirror +that shape with `_backup_indexer_from_device_per_layer()`. + +Ordering requirement: + +```text +layer i KV store complete + -> backup layer i target D2H + -> backup layer i draft D2H when draft is attached +``` + +The backup hook should be a **layer-end hook**: it fires after layer `i` has +issued and ordered its device KV store. Candidate integration points are the +attention layer forward path after `save_kv_cache` writes have been issued, or a +model runner/controller hook that receives the current layer id after the +layer's attention backend completes. A next-layer-start hook is only an +implementation fallback because it is indirect and needs a separate final-layer +hook. The design requirement is that backup must be ordered after the KV store +for that layer and before final host visibility. + +## Radix Tree Safety + +The radix tree should not be globally locked while per-layer backup runs. +Instead, pending state should be node-scoped: + +- mark node as backup-pending, +- protect device source and host destination, +- block host-hit visibility, +- block host eviction of pending destination, +- block device eviction of pending source. + +### Split policy + +First implementation should use **defer-split-on-pending-backup** for nodes +with pending backup state. + +If `_split_node()` would split a node in `HOST_RESERVED` or +`HOST_BACKUP_IN_FLIGHT`, do not split the in-flight backup op. The request that +needs the partial-prefix split should be blocked from this scheduling attempt +and moved to a pending/waiting path, while the existing backup continues toward +its normal final ack. After the backup commits or rolls back, the request can be +scheduled again and `_split_node()` can operate on stable node metadata. + +This policy is deliberately compatible with future `bs > 1`: one request's +partial-prefix split must not mutate another request's in-flight backup state or +force all ranks to split pending target/draft host metadata. It may delay the +new request behind D2H completion, but preserves deterministic radix structure, +owner-lane metadata, and target/draft strong sync. + +Longer-term alternatives are op-aware split, pending-segment aliases, or +synchronous drain-before-split, but they are higher-risk and should not be part +of the first per-layer backup pass. + +## Target/Draft Strong-Sync Contract + +Draft/MTP host and device payloads are a shadow of the target radix node. +Per-layer backup must preserve this invariant: + +```text +target HOST_VALID <=> draft HOST_VALID +target DEVICE_VALID <=> draft DEVICE_VALID +``` + +Consequences: + +- target and draft host slots are reserved together, +- target and draft layer D2H are submitted under one logical backup op, +- a target layer cannot make the node partially host-visible, +- draft failure rolls back target reservation, +- final commit requires all target and draft layer events complete, +- host eviction frees target and draft host slots together. + +## Owner-Pattern Contract + +CP shared KV uses logical page ownership. Host backup metadata must preserve +write-time `page_owners`, because load-back allocates new logical device pages +and then uses `owned_positions` to select physical shards. + +Per-layer backup does not change the owner pattern. Reservation and metadata +must still be node/page based, not layer based: + +```text +one CpHiCacheNodeMetadata per logical node +same owned_positions for every layer +same page_owners for load replay +same target/draft host indices for every layer payload +``` + +Layer granularity is a transfer scheduling property only, not a radix metadata +shape change. + +## Implementation Plan + +### Phase 0: Tests for current invariants + +Add or extend unit tests before behavior changes: + +- device-valid host-valid nodes are not host-evictable, +- device-valid host-missing nodes remain valid device hits, +- in-flight backup nodes are not host hits, +- in-flight backup nodes cannot be host-evicted, +- host reservation failure leaves node device-only, +- draft reservation failure rolls back target reservation, +- zero-owned ranks still produce a logical completion op, +- `_split_node()` on pending backup does not split the node; the triggering + request is deferred/requeued in the chosen first-pass behavior. + +### Phase 1: Explicit pending backup state + +Introduce a pending backup object owned by `HiRadixCache` or `HiCacheController`: + +```text +node_id +node reference +logical_len +metadata candidate +state: RESERVED | SUBMITTED | LOCAL_DONE | COMMITTED | ABORTED +target host indices +draft host indices +per-layer completion bitmap/events +``` + +Keep `_node_backuped()` false while the node has a pending backup object. + +### Phase 2: Reservation API split + +Split current `_write_cp()` into reserve and submit phases: + +```text +reserve_write_cp(device_indices, node_id) -> metadata candidate or failure +submit_write_cp_all_layer(pending) # compatibility baseline +submit_write_cp_layer(pending, layer_id) # new per-layer path +``` + +The all-layer path should remain as a guarded fallback until per-layer tests are +passing, but unexpected CP+MTP invariant failures must raise rather than silently +falling back. + +### Phase 3: Per-layer host-pool backup API + +Add `backup_from_device_per_layer()` to host pool implementations used by this +path. NSA implementation must include indexer K/scale per layer. + +Initial tests can use fake host pools to verify call ordering and target/draft +pairing without requiring GPU kernels. + +### Phase 4: Per-layer controller submission + +Add a controller-side per-layer backup producer that: + +- owns one transfer op per logical pending backup, +- submits target and draft copies for each layer, +- records per-layer completion events, +- exposes one final completion event / status for the logical node. + +No per-layer all-reduce is allowed. + +### Phase 5: Radix final commit and rollback + +Move host visibility from reservation time to final commit: + +- final commit writes `host_len` and `cp_hicache` or marks them visible, +- commit removes the pending op, +- commit releases the device source lock, +- rollback frees target/draft host slots and releases the device source lock, +- host eviction can only see committed `HOST_VALID` nodes. + +### Phase 6: Host-full slow path + +Refactor host full handling so reserve failure returns required local physical +slots, host eviction frees only valid host-only victims, and reservation retries +once. Add logs/counters for backup skipped due to host pressure. + +### Phase 7: Split safety + +Implement defer/requeue-on-pending-split for nodes with pending backup. Add +tests around partial prefix matches during in-flight backup, including a future +`bs > 1` style case where one request would split a node being backed up for +another request. + +### Phase 8: Verification on remote ETE + +Remote validation should run in the container, not local pytest for CUDA-heavy +paths: + +- CP shared KV + HiCache + EAGLE/MTP short prompts, +- host-hit scenario with accept length stability, +- host pressure scenario that forces host eviction before reservation retry, +- zero-owned rank coverage, +- cache hit after split-like prefixes, +- no leak in target or draft host allocators after abort/rollback. + +## Open Questions / Decisions + +1. **Layer completion hook.** The chosen semantic hook is layer-end after that + layer's KV store is ordered. The exact code location still needs + implementation-time confirmation. A next-layer-start hook is a fallback only + if it preserves the same ordering and adds an explicit final-layer trigger. +2. **Pending split behavior.** The chosen first pass is defer/requeue the + request that would split a node with pending backup. Do not split the + in-flight backup op, including for future `bs > 1`. +3. **Ack batching threshold.** Current `writing_check()` can all-reduce on + every progress poll. The first per-layer implementation should keep one + final logical ack per node and check it at final visibility time, not per + layer; batching threshold for final commit remains a later performance pass. +4. **Failure policy under reserve mismatch.** Capacity pressure should skip + backup; malformed target/draft metadata should fail fast. If a local rank + fails reservation after deterministic host eviction while another succeeds, + the slow path should globally abort the backup op and rollback local slots. + +## Implementation Notes: P4-P8 Pass + +Implemented in the first P4-P8 pass: + +- host pools now expose `backup_from_device_per_layer()`; +- controller has `submit_write_cp_layer()` with one final logical ack after all + target/draft layers; +- `HiRadixCache.write_backup()` now reserves host slots first, keeps the node + host-invisible in `pending_host_backups`, retries reservation once after + host-only eviction, and commits host visibility only when the write ack + completes; +- pending backup split attempts are deferred instead of mutating the radix tree; +- request admission marks deferred prefix matches and keeps those requests in + the waiting path for a later retry; +- KV pools expose a no-op-safe layer-store notifier. NSA defers notification + until `set_index_k_scale_buffer()` so the indexer payload is ordered after the + KV payload. + +Important limitation: the current `sgl_kernel.kvcacheio` surface has H2D +per-layer page-first kernels and D2H all-layer page-first kernels, but not a +D2H per-layer LF->PF page-first kernel. The per-layer host-pool method +therefore raises for `kernel + page_first` per-layer backup instead of silently +falling back to a slow torch copy. The production backup path remains on +`submit_write_cp_all_layer()` as a compatibility baseline until that kernel is +added; the per-layer controller path is covered by fake/direct unit tests and is +ready for kernel integration. + +## Implementation Notes: tai-kernel LF->PF pass + +The missing production kernel is now supplied by `tai-kernel`: + +```text +tai_kernel.nsa_prefill.transfer_kv_per_layer_mla_lf_pf( + src_layer_buffer, + dst_page_first_host_buffer, + src_indices, + dst_indices, + layer_id, + item_size, + dst_layout_dim, +) +``` + +This covers the GLM5/NSA production shape: + +- MLA/NSA target KV page-first host backup; +- NSA indexer page-first host backup; +- target and draft backup through the same logical reservation/ack path. + +MHA page-first remains fail-fast until it has a dedicated kernel. + +The current radix insertion path creates host reservations after request KV has +already been materialized. Therefore `submit_write_cp_per_layer()` performs a +safe catch-up submission of all layers through the per-layer API for current +call sites. The layer-store notifier path is also wired: + +```text +submit_write_cp_per_layer(..., catch_up_all_layers=False) + -> on_layer_kv_stored(layer_id, source="target" | "draft") + -> submit_write_cp_layer(reservation, layer_id, submit_target=..., submit_draft=...) + -> final ack only after all target/draft layers are submitted +``` + +This keeps the public state model unchanged: reserved or partially copied host +payloads remain invisible until the final write ack commits +`pending_host_backups`. Target and draft layer readiness are tracked +separately so an early target hook cannot copy draft KV before draft has stored +the same layer. + +## Summary + +Per-layer backup should be implemented as a data-plane refinement, not a radix +metadata shape change: + +```text +radix state synchronous +host reservation explicit and pending +D2H backup per layer and local-event driven +host visibility only after final logical commit +target/draft atomic +owner pattern preserved +host-full path evicts host-only valid victims then retries reservation once +``` + +The main correctness boundary is request visibility. A reserved or partially +copied host payload is not a cache hit. The main performance boundary is +synchronization. Per-layer backup must not add per-layer CP collectives; it +should reduce transfer tail while keeping final commit batched and deterministic. diff --git a/docs/superpowers/plans/2026-05-27-cp-hicache-per-layer-backup.md b/docs/superpowers/plans/2026-05-27-cp-hicache-per-layer-backup.md new file mode 100644 index 000000000..605bf62cd --- /dev/null +++ b/docs/superpowers/plans/2026-05-27-cp-hicache-per-layer-backup.md @@ -0,0 +1,336 @@ +# CP HiCache Per-Layer Backup Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Convert CP HiCache backup from all-layer D2H submission to host-reserved, per-layer D2H backup with final host-visible commit, while preserving target/draft strong sync and low CP-rank synchronization. + +**Architecture:** Keep radix/control-plane state synchronous and request-visible only after final commit. Split backup into reservation, per-layer data transfer, and final commit/rollback. Target and EAGLE/MTP draft payloads remain one logical cache object sharing target owner pattern metadata. + +**Tech Stack:** Python, SGLang SRT, `HiRadixCache`, `HiCacheController`, `HostKVCache` / `NSATokenToKVPoolHost`, `CpHiCacheNodeMetadata`, CUDA streams/events, CP shared KV owner-pattern allocator. + +--- + +## File map + +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` + - Add explicit pending backup visibility state. + - Keep `_node_backuped()` false for reserved/in-flight backup. + - Implement final commit/rollback and host-full reserve retry semantics. + - Add split safety for pending backup nodes by deferring/requeueing the request that would split them. +- Modify: `python/sglang/srt/managers/cache_controller.py` + - Split CP write reservation from transfer submission. + - Add per-layer backup submission bookkeeping. + - Keep target/draft operations atomic under one logical node id. +- Modify: `python/sglang/srt/mem_cache/memory_pool_host.py` + - Add `backup_from_device_per_layer()` mirror APIs. + - Add NSA indexer per-layer backup support. +- Modify as needed after locating the lowest-risk layer completion hook: + - likely attention/model executor files under `python/sglang/srt/layers/` or `python/sglang/srt/models/`. +- Modify tests: + - `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + - `test/registered/unit/managers/test_hicache_controller_cp.py` + - add new focused tests only if existing files become too large. + +## Task 1: Lock current request-visible invariants with tests + +**Files:** +- Modify: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` +- Modify: `test/registered/unit/managers/test_hicache_controller_cp.py` + +- [x] **Step 1: Add a unit test showing device-valid host-missing remains a device hit** + +Use a fake CP HiCache `HiRadixCache` node with `value` set, `host_len=0`, and `cp_hicache=None`. Assert `match_prefix()` returns device indices and `host_hit_length == 0`. + +- [x] **Step 2: Add a unit test showing in-flight backup is not host-visible** + +Create a node with valid `host_len`/`cp_hicache` and put `node.id` in `ongoing_write_through`. Assert `_node_backuped(node)` is false and `match_prefix()` does not return it as `last_host_node`. + +- [x] **Step 3: Add a unit test showing host eviction skips device-valid nodes** + +Create a CP host-backed node with `value` not `None`, call `_update_host_leaf_status(node)`, and assert the node is not in `evictable_host_leaves`. + +- [x] **Step 4: Add controller tests for target/draft reservation rollback** + +Extend the existing fake host pool tests so target allocation succeeds and draft allocation fails. Assert target host indices are freed and no write op is queued. + +- [x] **Step 5: Run targeted unit tests** + +Run: + +```bash +python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q +``` + +Expected: all existing tests plus the new invariant tests pass. + +Actual P1-P3 verification: ran in `g0034` / `sglang-glm5-dev-2` container with `cd /sgl-workspace/sglang-tai && python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q`; result `95 passed, 3 warnings in 22.06s`. + +## Task 2: Introduce explicit pending backup state + +**Files:** +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [x] **Step 1: Add a small pending-backup record** + +Add a dataclass near `CpHiCacheNodeMetadata` or a private internal class near `HiRadixCache`: + +```python +@dataclass +class PendingHiCacheBackup: + node: TreeNode + metadata: CpHiCacheNodeMetadata + logical_len: int + submitted: bool = False + local_done: bool = False +``` + +If import cycles make `TreeNode` typing awkward, type the field as `object` and keep construction internal to `HiRadixCache`. + +- [x] **Step 2: Add a pending map** + +Initialize: + +```python +self.pending_host_backups: Dict[int, PendingHiCacheBackup] = {} +``` + +Keep `ongoing_write_through` until the controller ack path is migrated; treat the new map as the source of host visibility. + +- [x] **Step 3: Update `_node_backuped()`** + +For CP HiCache, return false when `node.id` is in `pending_host_backups` or `ongoing_write_through`. Preserve existing draft metadata fail-fast behavior. + +- [x] **Step 4: Add commit and rollback helpers** + +Add internal helpers: + +```python +def _commit_pending_backup(self, node_id: int) -> None: + pending = self.pending_host_backups.pop(node_id) + node = pending.node + node.host_len = pending.logical_len + node.cp_hicache = pending.metadata + self.dec_node_lock_ref(node) + +def _rollback_pending_backup(self, node_id: int) -> None: + pending = self.pending_host_backups.pop(node_id) + self.cache_controller.evict_cp_host(pending.metadata) + self.dec_node_lock_ref(pending.node) +``` + +Adjust exact freeing helper if metadata is reserved before becoming committed. + +- [x] **Step 5: Test host visibility** + +Add a test that pending backup with metadata is not host-visible before `_commit_pending_backup()` and becomes visible after commit. + +## Task 3: Split CP write into reserve and submit + +**Files:** +- Modify: `python/sglang/srt/managers/cache_controller.py` +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Test: `test/registered/unit/managers/test_hicache_controller_cp.py` + +- [x] **Step 1: Extract reservation logic from `_write_cp()`** + +Create a method with the current allocation and metadata construction logic: + +```python +def reserve_write_cp(self, device_indices: torch.Tensor, node_id: int = -1) -> HiCacheWriteResult | HiCacheWriteFailure: + ... +``` + +It must compute `owned_positions`, `page_owners`, local physical device indices, target host indices, and draft host indices, but not submit D2H. + +- [x] **Step 2: Represent reserved physical indices** + +The pending op needs local physical device indices in addition to metadata. Store these in a controller-side pending write op, not in request-visible `CpHiCacheNodeMetadata` if possible. + +- [x] **Step 3: Make `_write_cp()` call reserve + submit-all-layer** + +Keep existing behavior working by implementing `_write_cp()` as: + +```python +reserved = self.reserve_write_cp(device_indices, node_id=node_id) +if failure: + return failure +self.submit_write_cp_all_layer(reserved, node_id=node_id) +return reserved result +``` + +- [x] **Step 4: Update tests to assert reserve alone queues no transfer** + +Fake host pools should show allocations happened but no `backup_from_device_all_layer()` call until submit. + +- [x] **Step 5: Keep zero-owned no-op semantics** + +A zero-owned rank still creates a logical no-op ack on submit, not during reserve, so logical op cardinality remains aligned. + +## Task 4: Add host-pool per-layer backup APIs + +**Files:** +- Modify: `python/sglang/srt/mem_cache/memory_pool_host.py` +- Test: `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py` or a new focused unit file + +- [x] **Step 1: Add abstract/base method** + +Add to the base host pool class: + +```python +def backup_from_device_per_layer(self, device_pool, host_indices, device_indices, layer_id, io_backend): + raise NotImplementedError +``` + +- [x] **Step 2: Implement MHA/MLA per-layer backup mirror** + +Mirror `load_to_device_per_layer()` direction with D2H transfer kernels/direct calls. Use the same index shapes and layout restrictions as the all-layer backup path. + +Note: current `sgl_kernel.kvcacheio` exposes H2D per-layer page-first kernels and D2H all-layer page-first kernels, but not D2H per-layer LF->PF page-first kernels. The first implementation therefore fail-fast raises for production `kernel + page_first` per-layer backup instead of silently using a slow torch fallback. The runtime write path still uses the all-layer compatibility baseline until the kernel exists. + +- [x] **Step 3: Implement NSA indexer per-layer backup** + +Add `_backup_indexer_from_device_per_layer()` mirroring `_load_indexer_to_device_per_layer()`, then call it from `NSATokenToKVPoolHost.backup_from_device_per_layer()` after the KV payload backup. + +- [x] **Step 4: Unit test fake/direct path call shape** + +Use fake pools or monkeypatched transfer functions to assert that backing up layer `i` calls the target layer and indexer layer exactly once with matching host/device indices. + +## Task 5: Submit and finalize per-layer backup + +**Files:** +- Modify: `python/sglang/srt/managers/cache_controller.py` +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Test: `test/registered/unit/managers/test_hicache_controller_cp.py` + +- [x] **Step 1: Add per-layer submission state** + +Represent one logical pending backup op with target/draft host indices, target/draft physical device indices, node id, and a per-layer event list. + +- [x] **Step 2: Add submit-layer method** + +Add: + +```python +def submit_write_cp_layer(self, pending, layer_id: int) -> None: + ... +``` + +It submits target and draft D2H for `layer_id` on `write_stream` and records completion for that layer. + +- [x] **Step 3: Add final ack only after all layers** + +When all layer events are complete, append one `HiCacheAck` for the logical node id. Do not append one ack per layer. + +- [x] **Step 4: Wire final ack handling to commit pending backups** + +Do not poll/commit per layer. When the final logical ack for all target/draft layers completes, call `_commit_pending_backup(node_id)` instead of only removing from `ongoing_write_through`. Keep any all-reduce at this final visibility boundary. + +- [x] **Step 5: Preserve all-layer fallback only as a compatibility baseline** + +Until the layer hook is connected, keep submit-all-layer available. If CP+MTP metadata is malformed, raise instead of silently falling back. + +## Task 6: Connect the layer completion hook + +**Files:** +- Investigate and modify the lowest-risk file after code inspection: + - `python/sglang/srt/layers/radix_attention.py`, or + - the relevant model attention forward method, or + - a model runner/controller hook carrying `layer_id`. +- Test: unit/fake test plus remote ETE. + +- [x] **Step 1: Locate the point after layer KV store is ordered** + +The hook must run at layer end, after `save_kv_cache` has issued and ordered the layer's KV store. It must not run before attention writes the cache. A next-layer-start hook is only acceptable as a fallback if it waits for the previous layer store and includes an explicit final-layer trigger. + +- [x] **Step 2: Add a no-op-safe callback on forward batch or controller** + +Expose a callback such as: + +```python +forward_batch.hicache_backup_notifier.on_layer_kv_stored(layer_id) +``` + +When no pending backup exists, it returns immediately. + +- [x] **Step 3: Submit target/draft layer backup through the controller** + +The callback should enqueue the current layer for every pending backup op whose source includes the current batch's new KV locs. + +- [x] **Step 4: Add tests with a fake notifier** + +Use a fake forward batch or attention layer wrapper to assert layer ids are reported in order and no callback fires before the store point. + +## Task 7: Host-full retry and rollback + +**Files:** +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Modify: `python/sglang/srt/managers/cache_controller.py` +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [x] **Step 1: Make reserve failure return required local physical host slots** + +Keep `HiCacheWriteFailure(required_host_slots=...)` as the capacity signal. + +- [x] **Step 2: Evict host-only valid victims before retry** + +Use `_evict_host_for_physical_slots(required_slots)` and ensure it skips pending backups and device-valid nodes. + +- [x] **Step 3: Retry reservation once** + +If the second reservation fails, rollback any partial reservation and leave the node `DEVICE_VALID_HOST_NONE`. + +- [x] **Step 4: Add tests for retry success and retry failure** + +Test that retry success creates pending backup state. Test that retry failure does not set `host_len`, `cp_hicache`, or pending state. + +## Task 8: Split safety + +**Files:** +- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py` +- Modify: scheduler/request admission path that handles prefix-match retry +- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py` + +- [x] **Step 1: Add pending split guard** + +Before splitting a CP HiCache node, detect pending backup state. + +- [x] **Step 2: First-pass behavior: defer/requeue request before split** + +If a new request would split a node with pending backup, do not split that node and do not drain the in-flight op synchronously. Mark this scheduling attempt as temporarily blocked and return the request to a pending/waiting path. Once the backup final ack commits or rolls back, normal scheduling can retry and split stable metadata if still needed. + +- [x] **Step 3: Keep the rule valid for future `bs > 1`** + +The guard must be request-local: one request trying to split a pending node must not mutate another request's in-flight backup state. The batch builder should either skip that request for the current batch or return the whole candidate to the waiting queue if the current code cannot safely skip within a batch. + +- [x] **Step 4: Add tests** + +Construct a node with pending backup and trigger partial match/split. Assert no split happens while pending, no lock leak occurs, and the request is reported as deferred/retryable. Add a multi-request shaped test that simulates future `bs > 1`: one request owns the pending backup and another request attempts the partial split. + +## Task 9: Remote verification + +**Files:** +- No source changes unless failures are found. + +- [ ] **Step 1: Sync code to remote as ubuntu-owned files** + +Use the established tar/scp workflow and ensure target ownership remains `ubuntu:ubuntu`. + +- [ ] **Step 2: Run targeted unit tests in the remote container** + +Run the CP HiCache metadata/controller unit tests inside `sglang-glm5-dev-2`. + +Actual P4-P8 verification: ran in `g0034` / `sglang-glm5-dev-2` container with +`cd /sgl-workspace/sglang-tai && python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q`; result `101 passed, 3 warnings in 21.96s`. + +- [ ] **Step 3: Run ETE CP HiCache + EAGLE/MTP smoke** + +Use the existing GLM5 command shape with CP shared KV, HiCache, and EAGLE/MTP enabled. Confirm startup, cache hits, and stable accept length. + +- [ ] **Step 4: Run host-pressure scenario** + +Use reduced `--hicache-size` to force host eviction before reservation retry. Confirm no hang, no memory leak, and explicit backup-skip logs only when no host victims exist. + +- [ ] **Step 5: Record results in the design doc or a follow-up session note** + +Document exact command, log path, observed hit rate, accept length, and any failure. diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index a531626d1..e5b95b2d0 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -16,9 +16,9 @@ limitations under the License. import logging import threading import time -from dataclasses import dataclass +from dataclasses import dataclass, field from queue import Empty, Full, Queue -from typing import TYPE_CHECKING, List, NamedTuple, Optional +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set import torch @@ -266,6 +266,32 @@ class HiCacheWriteFailure: metadata: object = None +@dataclass +class HiCacheWriteReservation: + metadata: object + host_indices: torch.Tensor + physical_device_indices: torch.Tensor + node_id: int = -1 + priority: Optional[int] = None + draft_host_indices: Optional[torch.Tensor] = None + required_host_slots: int = 0 + + +@dataclass +class HiCacheLayerWriteState: + reservation: HiCacheWriteReservation + total_layers: int + start_event: object + finish_event: object + layer_events: List[object] = field(default_factory=list) + completed_target_layers: Set[int] = field(default_factory=set) + completed_draft_layers: Set[int] = field(default_factory=set) + host_indices: Optional[torch.Tensor] = None + physical_device_indices: Optional[torch.Tensor] = None + draft_host_indices: Optional[torch.Tensor] = None + ack_appended: bool = False + + class HiCacheController: def __init__( @@ -333,6 +359,10 @@ class HiCacheController: self.layer_num = self.mem_pool_device.layer_num self.layer_done_counter = LayerDoneCounter(self.layer_num) self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter) + if hasattr(self.mem_pool_device, "register_layer_backup_notifier"): + self.mem_pool_device.register_layer_backup_notifier( + lambda layer_id: self.on_layer_kv_stored(layer_id, source="target") + ) self.draft_mem_pool_host = None self.draft_mem_pool_device = None @@ -351,6 +381,7 @@ class HiCacheController: self.draft_write_queue: List[CacheOperation] = [] self.ack_load_queue: List[HiCacheAck] = [] self.ack_write_queue: List[HiCacheAck] = [] + self.pending_layer_writes: Dict[int, HiCacheLayerWriteState] = {} if draft_mem_pool_host is not None or draft_mem_pool_device is not None: self.attach_draft_pool(draft_mem_pool_device, draft_mem_pool_host) @@ -398,6 +429,28 @@ class HiCacheController: draft_mem_pool_device.register_layer_transfer_counter( self.layer_done_counter ) + if hasattr(draft_mem_pool_device, "register_layer_backup_notifier"): + draft_mem_pool_device.register_layer_backup_notifier( + lambda layer_id: self.on_layer_kv_stored(layer_id, source="draft") + ) + + def on_layer_kv_stored(self, layer_id: int, source: str = "target") -> None: + if not self.pending_layer_writes: + return + if source not in ("target", "draft"): + raise ValueError(f"Unknown CP HiCache layer backup source: {source}") + reservations = [ + state.reservation for state in list(self.pending_layer_writes.values()) + ] + for reservation in reservations: + if reservation.node_id not in self.pending_layer_writes: + continue + self.submit_write_cp_layer( + reservation, + layer_id, + submit_target=(source == "target"), + submit_draft=(source == "draft"), + ) def _start_storage_threads(self): """Start storage prefetch/backup threads and their queues. @@ -703,6 +756,7 @@ class HiCacheController: self.load_buffer.clear() self.ack_write_queue.clear() self.ack_load_queue.clear() + self.pending_layer_writes.clear() if self.enable_storage: self.prefetch_thread.join() self.backup_thread.join() @@ -859,12 +913,12 @@ class HiCacheController: device_indices, self.page_size, "physical_device_indices" ) - def _write_cp( + def reserve_write_cp( self, device_indices: torch.Tensor, priority: Optional[int] = None, node_id: int = -1, - ) -> HiCacheWriteResult | HiCacheWriteFailure: + ) -> HiCacheWriteReservation | HiCacheWriteFailure: from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata layout = self.cp_shared_kv_layout @@ -882,32 +936,31 @@ class HiCacheController: f"logical_len={logical_len} page_size={page_size}" ) page_first_locs = device_indices[::page_size] - logical_pages = torch.div( - page_first_locs, page_size, rounding_mode="floor" - ) + 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" ) if owned_positions.numel() == 0: - self._append_completed_write_ack(node_id) logger.info( - "[CacheCtrl-write] _write_cp zero-owned rank: node_id=%d logical_len=%d", + "[CacheCtrl-write] reserve_write_cp zero-owned rank: node_id=%d logical_len=%d", node_id, logical_len, ) - return HiCacheWriteResult( + empty = torch.empty((0,), dtype=torch.int64) + return HiCacheWriteReservation( metadata=CpHiCacheNodeMetadata( logical_len=logical_len, owned_positions=owned_positions, - host_indices=torch.empty((0,), dtype=torch.int64), + host_indices=empty, page_owners=page_owners, page_size=page_size, - draft_host_indices=( - torch.empty((0,), dtype=torch.int64) - if self.has_draft_hicache - else None - ), - ) + draft_host_indices=(empty.clone() if self.has_draft_hicache else None), + ), + host_indices=empty, + physical_device_indices=empty, + node_id=node_id, + priority=priority, + draft_host_indices=(empty.clone() if self.has_draft_hicache else None), ) owned_logical_indices = device_indices[owned_mask] @@ -915,12 +968,13 @@ class HiCacheController: host_indices = self.mem_pool_host.alloc(len(physical_device_indices)) if host_indices is None: logger.info( - "[CacheCtrl-write] _write_cp FAILED (host full): node_id=%d logical_len=%d owned=%d", + "[CacheCtrl-write] reserve_write_cp FAILED (host full): node_id=%d logical_len=%d owned=%d", node_id, logical_len, owned_positions.numel(), ) return HiCacheWriteFailure(required_host_slots=len(physical_device_indices)) + draft_host_indices = None if self.has_draft_hicache: draft_host_indices = self.draft_mem_pool_host.alloc( @@ -929,7 +983,7 @@ class HiCacheController: if draft_host_indices is None: self.mem_pool_host.free(host_indices) logger.info( - "[CacheCtrl-write] _write_cp FAILED (draft host full): node_id=%d logical_len=%d owned=%d", + "[CacheCtrl-write] reserve_write_cp FAILED (draft host full): node_id=%d logical_len=%d owned=%d", node_id, logical_len, owned_positions.numel(), @@ -952,25 +1006,7 @@ class HiCacheController: self.draft_mem_pool_host.free(draft_host_indices) raise - self.write_queue.append( - CacheOperation(host_indices, physical_device_indices, node_id, priority) - ) - if draft_host_indices is not None: - self.draft_write_queue.append( - CacheOperation( - draft_host_indices, physical_device_indices, node_id, priority - ) - ) - self.start_writing() - logger.info( - "[CacheCtrl-write] _write_cp submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s", - node_id, - logical_len, - owned_positions.numel(), - len(physical_device_indices), - draft_host_indices is not None, - ) - return HiCacheWriteResult( + return HiCacheWriteReservation( metadata=CpHiCacheNodeMetadata( logical_len=logical_len, owned_positions=owned_positions, @@ -980,8 +1016,255 @@ class HiCacheController: draft_host_indices=( draft_host_indices.cpu() if draft_host_indices is not None else None ), + ), + host_indices=host_indices, + physical_device_indices=physical_device_indices, + node_id=node_id, + priority=priority, + draft_host_indices=draft_host_indices, + ) + + def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None: + logger.warning( + "[CacheCtrl-write] CP HiCache all-layer backup fallback: node_id=%d logical_len=%d owned=%d physical=%d draft=%s", + reservation.node_id, + reservation.metadata.logical_len, + reservation.metadata.owned_positions.numel(), + len(reservation.physical_device_indices), + reservation.draft_host_indices is not None, + ) + if len(reservation.physical_device_indices) == 0: + self._append_completed_write_ack(reservation.node_id) + logger.info( + "[CacheCtrl-write] submit_write_cp_all_layer zero-owned ack: node_id=%d logical_len=%d", + reservation.node_id, + reservation.metadata.logical_len, + ) + return + + self.write_queue.append( + CacheOperation( + reservation.host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, ) ) + if reservation.draft_host_indices is not None: + self.draft_write_queue.append( + CacheOperation( + reservation.draft_host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, + ) + ) + self.start_writing() + logger.info( + "[CacheCtrl-write] submit_write_cp_all_layer submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s", + reservation.node_id, + reservation.metadata.logical_len, + reservation.metadata.owned_positions.numel(), + len(reservation.physical_device_indices), + reservation.draft_host_indices is not None, + ) + + def _get_or_create_layer_write_state( + self, reservation: HiCacheWriteReservation + ) -> HiCacheLayerWriteState: + state = self.pending_layer_writes.get(reservation.node_id) + if state is not None: + if state.reservation is not reservation: + raise RuntimeError( + f"Conflicting CP HiCache per-layer backup reservation for " + f"node_id={reservation.node_id}" + ) + return state + + draft_layer_num = ( + self.draft_mem_pool_device.layer_num + if reservation.draft_host_indices is not None + else 0 + ) + total_layers = max(self.layer_num, draft_layer_num) + if total_layers <= 0: + raise RuntimeError( + f"Invalid CP HiCache per-layer backup layer count: {total_layers}" + ) + + state = HiCacheLayerWriteState( + reservation=reservation, + total_layers=total_layers, + start_event=device_module.Event(), + finish_event=device_module.Event(), + ) + if len(reservation.physical_device_indices) > 0: + op = CacheOperation( + reservation.host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, + ) + state.host_indices, state.physical_device_indices = self.move_indices( + op, self.mem_pool_host + ) + if reservation.draft_host_indices is not None: + draft_op = CacheOperation( + reservation.draft_host_indices, + reservation.physical_device_indices, + reservation.node_id, + reservation.priority, + ) + state.draft_host_indices, _ = self.move_indices( + draft_op, self.draft_mem_pool_host + ) + self.pending_layer_writes[reservation.node_id] = state + state.start_event.record() + return state + + def submit_write_cp_layer( + self, + reservation: HiCacheWriteReservation, + layer_id: int, + *, + submit_target: bool = True, + submit_draft: bool = True, + ) -> None: + """Submit one layer of a reserved CP host backup. + + Target and draft D2H are still one logical operation: this method only + queues a final write ack when all target/draft layers have been submitted. + """ + + state = self._get_or_create_layer_write_state(reservation) + if layer_id < 0 or layer_id >= state.total_layers: + raise ValueError( + f"layer_id={layer_id} is outside CP HiCache backup layer range " + f"[0, {state.total_layers}) for node_id={reservation.node_id}" + ) + needs_target = ( + submit_target + and layer_id < self.layer_num + and layer_id not in state.completed_target_layers + ) + needs_draft = ( + submit_draft + and state.draft_host_indices is not None + and layer_id < self.draft_mem_pool_device.layer_num + and layer_id not in state.completed_draft_layers + ) + if not needs_target and not needs_draft: + return + + layer_event = device_module.Event() + layer_event.record() + state.layer_events.append(layer_event) + + if len(reservation.physical_device_indices) > 0: + with device_module.stream(self.write_stream): + state.start_event.wait(self.write_stream) + layer_event.wait(self.write_stream) + if needs_target: + self.mem_pool_host.backup_from_device_per_layer( + self.mem_pool_device, + state.host_indices, + state.physical_device_indices, + layer_id, + self.io_backend, + ) + if needs_draft: + self.draft_mem_pool_host.backup_from_device_per_layer( + self.draft_mem_pool_device, + state.draft_host_indices, + state.physical_device_indices, + layer_id, + self.io_backend, + ) + + if needs_target: + state.completed_target_layers.add(layer_id) + if needs_draft: + state.completed_draft_layers.add(layer_id) + + target_done = len(state.completed_target_layers) >= self.layer_num + draft_done = ( + state.draft_host_indices is None + or len(state.completed_draft_layers) + >= self.draft_mem_pool_device.layer_num + ) + if not target_done or not draft_done: + return + if state.ack_appended: + return + + with device_module.stream(self.write_stream): + state.finish_event.record() + if state.host_indices is not None and state.host_indices.is_cuda: + state.host_indices.record_stream(self.write_stream) + if ( + state.physical_device_indices is not None + and state.physical_device_indices.is_cuda + ): + state.physical_device_indices.record_stream(self.write_stream) + if ( + state.draft_host_indices is not None + and state.draft_host_indices.is_cuda + ): + state.draft_host_indices.record_stream(self.write_stream) + + state.ack_appended = True + self.pending_layer_writes.pop(reservation.node_id, None) + self.ack_write_queue.append( + HiCacheAck(state.start_event, state.finish_event, [reservation.node_id]) + ) + logger.info( + "[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s", + reservation.node_id, + reservation.metadata.logical_len, + state.total_layers, + reservation.draft_host_indices is not None, + ) + + def submit_write_cp_per_layer( + self, + reservation: HiCacheWriteReservation, + *, + catch_up_all_layers: bool = True, + ) -> None: + """Register a CP write reservation for per-layer host backup. + + Current radix insertion calls write_backup after the request KV has already + been materialized. For that path, catch up by submitting all layers + immediately through the per-layer API. Future early reservations can use + catch_up_all_layers=False and rely on on_layer_kv_stored(). + """ + + state = self._get_or_create_layer_write_state(reservation) + if not catch_up_all_layers: + logger.info( + "[CacheCtrl-write] submit_write_cp_per_layer registered: node_id=%d logical_len=%d layers=%d draft=%s", + reservation.node_id, + reservation.metadata.logical_len, + state.total_layers, + reservation.draft_host_indices is not None, + ) + return + + total_layers = state.total_layers + for layer_id in range(total_layers): + self.submit_write_cp_layer(reservation, layer_id) + + def _write_cp( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + ) -> HiCacheWriteResult | HiCacheWriteFailure: + reservation = self.reserve_write_cp(device_indices, priority, node_id) + if isinstance(reservation, HiCacheWriteFailure): + return reservation + self.submit_write_cp_per_layer(reservation) + return HiCacheWriteResult(metadata=reservation.metadata) def set_draft_kv_pool(self, draft_device_pool, draft_host_pool) -> None: """Register draft KV pools so L2 ops piggyback draft transfers.""" diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c058169bf..ec52db19b 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -643,6 +643,7 @@ class Req(ReqDllmMixin): self.last_host_node: Any = None self.last_host_backup_node: Any = None self.host_hit_length = 0 + self.prefix_match_deferred_by_pending_backup = False # Tokens loaded from storage backend (L3) during prefetch for this request self.storage_hit_length = 0 # The node to lock until for swa radix tree lock ref @@ -930,6 +931,9 @@ class Req(ReqDllmMixin): self.cache_protected_len = match_result.cache_protected_len else: self.cache_protected_len = len(self.prefix_indices) + self.prefix_match_deferred_by_pending_backup = ( + match_result.pending_backup_deferred_node is not None + ) if self.is_dllm(): self._update_block_offset_for_dllm() diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 2abbbb667..241e28e65 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -212,6 +212,9 @@ class SchedulePolicy: match_result.last_host_node, match_result.host_hit_length, ) + r.prefix_match_deferred_by_pending_backup = ( + match_result.pending_backup_deferred_node is not None + ) # NOTE(sang): This logic is for in-batch prefix caching; # If there are more than 1 request that have small matching prefix from @@ -748,6 +751,9 @@ class PrefillAdder: def add_one_req( self, req: Req, has_chunked_req: bool, truncation_align_size: Optional[int] ): + if getattr(req, "prefix_match_deferred_by_pending_backup", False): + return AddReqResult.OTHER + if (self.prefill_delayer_single_pass is not None) and ( not self.prefill_delayer_single_pass.negotiate_should_allow_prefill( local_prefillable=True, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 7ab2d5bb0..563566cca 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -144,6 +144,7 @@ class MatchResult(NamedTuple): host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None + pending_backup_deferred_node: Any = None class BasePrefixCache(ABC, PrefixCacheTrait): diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e89112c02..93f258f75 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -14,7 +14,11 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch from sglang.srt.environ import envs -from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation +from sglang.srt.managers.cache_controller import ( + HiCacheController, + HiCacheWriteFailure, + PrefetchOperation, +) from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, @@ -232,6 +236,24 @@ class CpHiCacheNodeMetadata: ) +@dataclass +class PendingHiCacheBackup: + node: TreeNode + metadata: CpHiCacheNodeMetadata + logical_len: int + submitted: bool = False + local_done: bool = False + locked: bool = True + + +class HiCachePendingBackupSplit(Exception): + def __init__(self, node: TreeNode): + self.node = node + super().__init__( + f"Cannot split node_id={getattr(node, 'id', None)} while CP HiCache backup is pending" + ) + + class HiRadixCache(RadixCache): def _create_token_to_kv_pool_host( @@ -433,6 +455,8 @@ class HiRadixCache(RadixCache): # record the nodes with ongoing write through self.ongoing_write_through = {} + # record CP host reservations/backups that are not request-visible yet. + self.pending_host_backups: Dict[int, PendingHiCacheBackup] = {} # record the node segments with ongoing load back self.ongoing_load_back = {} # record the ongoing prefetch requests @@ -886,6 +910,8 @@ class HiRadixCache(RadixCache): self.cache_controller.clear_draft_host_pool() # Clear per-request tracking dicts self.prefetch_loaded_tokens_by_reqid.clear() + if hasattr(self, "pending_host_backups"): + self.pending_host_backups.clear() self.evictable_host_leaves.clear() self.pinned_size_ = 0 super().reset() @@ -945,7 +971,10 @@ class HiRadixCache(RadixCache): return True def _node_host_write_pending(self, node: TreeNode) -> bool: - return getattr(node, "id", None) in getattr(self, "ongoing_write_through", {}) + node_id = getattr(node, "id", None) + return node_id in getattr(self, "ongoing_write_through", {}) or node_id in getattr( + self, "pending_host_backups", {} + ) def _node_host_write_ready(self, node: TreeNode) -> bool: return not self._node_host_write_pending(node) @@ -960,6 +989,27 @@ class HiRadixCache(RadixCache): ) return node.host_value is not None + def _commit_pending_backup(self, node_id: int) -> TreeNode: + pending = self.pending_host_backups.pop(node_id) + node = pending.node + node.host_len = pending.logical_len + node.cp_hicache = pending.metadata + node.host_value = None + if pending.locked: + self.dec_node_lock_ref(node) + return node + + def _rollback_pending_backup(self, node_id: int) -> TreeNode: + pending = self.pending_host_backups.pop(node_id) + node = pending.node + self.cache_controller.evict_cp_host(pending.metadata) + if node.cp_hicache is pending.metadata: + node.cp_hicache = None + node.host_len = 0 + if pending.locked: + self.dec_node_lock_ref(node) + return node + def _node_host_len(self, node: TreeNode) -> int: if self._uses_cp_hicache: return node.host_len @@ -979,30 +1029,26 @@ class HiRadixCache(RadixCache): write_back, ) if self._uses_cp_hicache: - result = self.cache_controller.write( + result = self.cache_controller.reserve_write_cp( device_indices=node.value, node_id=node.id, ) - metadata = getattr(result, "metadata", None) - required_host_slots = 0 - if metadata is None: + if isinstance(result, HiCacheWriteFailure): required_host_slots = result.required_host_slots - self._evict_host_for_physical_slots( - required_host_slots, - synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, - ) - if metadata is None: + self._evict_host_for_physical_slots( + required_host_slots, + synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, + ) logger.info( "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d", node.id, required_host_slots, ) - result = self.cache_controller.write( + result = self.cache_controller.reserve_write_cp( device_indices=node.value, node_id=node.id, ) - metadata = getattr(result, "metadata", None) - if metadata is None: + if isinstance(result, HiCacheWriteFailure): logger.info( "[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d", node.id, @@ -1010,18 +1056,29 @@ class HiRadixCache(RadixCache): ) return 0 - node.host_len = len(node.value) - node.cp_hicache = metadata - node.host_value = None self.ongoing_write_through[node.id] = node if not write_back: self.inc_node_lock_ref(node) + self.pending_host_backups[node.id] = PendingHiCacheBackup( + node=node, + metadata=result.metadata, + logical_len=len(node.value), + submitted=True, + locked=not write_back, + ) + try: + self.cache_controller.submit_write_cp_per_layer(result) + except Exception: + self.ongoing_write_through.pop(node.id, None) + self._rollback_pending_backup(node.id) + raise logger.info( - "[HiCache-write] write_backup CP SUCCESS: node_id=%d host_len=%d owned_positions=%d ongoing_writes=%d", + "[HiCache-write] write_backup CP SUBMITTED: node_id=%d logical_len=%d owned_positions=%d ongoing_writes=%d pending_backups=%d", node.id, - node.host_len, - node.cp_hicache.owned_positions.numel(), + len(node.value), + result.metadata.owned_positions.numel(), len(self.ongoing_write_through), + len(self.pending_host_backups), ) return len(node.value) @@ -1102,13 +1159,22 @@ class HiRadixCache(RadixCache): self.write_backup(node) def writing_check(self, write_back=False): + def complete_write_ack_node(ack_id: int) -> TreeNode: + if ack_id in self.pending_host_backups: + backuped_node = self._commit_pending_backup(ack_id) + self.ongoing_write_through.pop(ack_id, None) + return backuped_node + backuped_node = self.ongoing_write_through.pop(ack_id) + self.dec_node_lock_ref(backuped_node) + return backuped_node + if write_back: # blocking till all write back complete while len(self.ongoing_write_through) > 0: for _, finish_event, ack_list in self.cache_controller.ack_write_queue: finish_event.synchronize() for ack_id in ack_list: - backuped_node = self.ongoing_write_through.pop(ack_id) + backuped_node = complete_write_ack_node(ack_id) if self.enable_storage: self.write_backup_storage(backuped_node) self.cache_controller.ack_write_queue.clear() @@ -1148,9 +1214,8 @@ class HiRadixCache(RadixCache): _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0) finish_event.synchronize() for ack_id in ack_list: - backuped_node = self.ongoing_write_through.pop(ack_id) + backuped_node = complete_write_ack_node(ack_id) released_nodes.append(ack_id) - self.dec_node_lock_ref(backuped_node) if self.enable_storage: self.write_backup_storage(backuped_node) finish_count -= 1 @@ -1985,7 +2050,13 @@ class HiRadixCache(RadixCache): host_hit_length=0, ) - value, last_node = self._match_prefix_helper(self.root_node, key) + deferred_node = None + try: + value, last_node = self._match_prefix_helper(self.root_node, key) + except HiCachePendingBackupSplit as exc: + value = [] + last_node = exc.node.parent if exc.node.parent is not None else self.root_node + deferred_node = exc.node if value: value = torch.cat(value) else: @@ -2018,6 +2089,7 @@ class HiRadixCache(RadixCache): last_device_node=last_node, last_host_node=last_host_node, host_hit_length=host_hit_length, + pending_backup_deferred_node=deferred_node, ) def prefetch_from_storage( @@ -2121,6 +2193,11 @@ class HiRadixCache(RadixCache): child.pin_expiry = time.monotonic() + child.pin_ttl prefix_len = self.key_match_fn(child.key, key) if prefix_len < len(child.key): + if ( + self._uses_cp_hicache + and child.id in getattr(self, "pending_host_backups", {}) + ): + raise HiCachePendingBackupSplit(child) new_node = self._split_node(child.key, child, prefix_len) if not new_node.evicted: value.append(new_node.value) @@ -2138,6 +2215,11 @@ class HiRadixCache(RadixCache): return value, node def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): + if ( + self._uses_cp_hicache + and child.id in getattr(self, "pending_host_backups", {}) + ): + raise HiCachePendingBackupSplit(child) # child node split into new_node -> child new_node = TreeNode(priority=child.priority) new_node.children = {self.get_child_key_fn(key[split_len:]): child} diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 86179c82f..8f6302ed0 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -694,6 +694,8 @@ class KVCache(abc.ABC): # default state for optional layer-wise transfer control self.layer_transfer_counter = None + self.layer_backup_notifiers = [] + self.layer_backup_notify_after_indexer = False # for disagg with nvlink self.enable_custom_mem_pool, self.custom_mem_pool, _ = ( @@ -745,6 +747,15 @@ class KVCache(abc.ABC): def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter): self.layer_transfer_counter = layer_transfer_counter + def register_layer_backup_notifier(self, notifier): + self.layer_backup_notifiers.append(notifier) + + def _notify_layer_kv_stored(self, layer_id: int, source: str = "kv"): + if self.layer_backup_notify_after_indexer and source != "indexer": + return + for notifier in self.layer_backup_notifiers: + notifier(layer_id) + def get_cpu_copy(self, indices): raise NotImplementedError() @@ -1036,6 +1047,7 @@ class MHATokenToKVPool(KVCache): alt_stream=self.alt_stream, same_kv_dim=self.same_kv_dim, ) + self._notify_layer_kv_stored(layer_id - self.start_layer) def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get(): @@ -1583,6 +1595,7 @@ class MLATokenToKVPool(KVCache): ) else: self.kv_buffer[layer_id - self.start_layer][loc] = cache_k + self._notify_layer_kv_stored(layer_id - self.start_layer) def set_mla_kv_buffer( self, @@ -1624,6 +1637,7 @@ class MLATokenToKVPool(KVCache): cache_k_nope, cache_k_rope, ) + self._notify_layer_kv_stored(layer_id - self.start_layer) def get_mla_kv_buffer( self, @@ -1845,6 +1859,7 @@ class NSATokenToKVPool(MLATokenToKVPool): use_nsa=True, override_kv_cache_dim=override_dim, ) + self.layer_backup_notify_after_indexer = True # self.index_k_dtype = torch.float8_e4m3fn # self.index_k_scale_dtype = torch.float32 self.index_head_dim = index_head_dim @@ -1951,6 +1966,7 @@ class NSATokenToKVPool(MLATokenToKVPool): index_buf_accessor.SetKAndS.execute( pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale ) + self._notify_layer_kv_stored(layer_id - self.start_layer, source="indexer") def get_state_buf_infos(self): data_ptrs = [ diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 72c466c17..c7ba5b557 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -2,7 +2,7 @@ import abc import logging import threading from collections import defaultdict -from functools import wraps +from functools import lru_cache, wraps from typing import Optional import psutil @@ -52,6 +52,20 @@ if _is_npu: logger = logging.getLogger(__name__) +@lru_cache(maxsize=1) +def _load_tai_transfer_kv_per_layer_mla_lf_pf(): + try: + from tai_kernel.nsa_prefill import transfer_kv_per_layer_mla_lf_pf + + return transfer_kv_per_layer_mla_lf_pf + except Exception as exc: + raise RuntimeError( + "Per-layer D2H backup for page_first MLA/NSA HiCache requires " + "tai_kernel.nsa_prefill.transfer_kv_per_layer_mla_lf_pf. " + "Build/sync tai-kernel with the LF->PF per-layer kvcacheio op." + ) from exc + + def synchronized(func): @wraps(func) def wrapper(self, *args, **kwargs): @@ -221,6 +235,15 @@ class HostKVCache(abc.ABC): """ raise NotImplementedError() + @abc.abstractmethod + def backup_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ) -> None: + """ + Backup KV data from the device memory pool to the host memory pool for a specific layer. + """ + raise NotImplementedError() + @abc.abstractmethod def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend @@ -577,6 +600,77 @@ class MHATokenToKVPoolHost(HostKVCache): else: raise ValueError(f"Unsupported IO backend: {io_backend}") + def backup_from_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + ): + if io_backend == "kernel": + if self.layout == "layer_first": + transfer_kv_per_layer( + src_k=device_pool.k_buffer[layer_id], + dst_k=self.k_buffer[layer_id], + src_v=device_pool.v_buffer[layer_id], + dst_v=self.v_buffer[layer_id], + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + ) + elif self.layout == "page_first": + raise NotImplementedError( + "Per-layer D2H backup for page_first MHA kernel layout requires " + "a dedicated LF->PF per-layer kernel; use all-layer backup until " + "that kernel exists." + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + dst_layers=[self.k_buffer[layer_id], self.v_buffer[layer_id]], + src_indices=device_indices, + dst_indices=host_indices, + page_size=self.page_size, + ) + elif self.layout == "page_first_direct": + for host_index, device_index in zip( + host_indices.cpu().tolist(), device_indices.cpu().tolist() + ): + host_page = host_index // self.page_size + host_offset = host_index % self.page_size + self.k_buffer[host_page, layer_id, host_offset] = ( + device_pool.k_buffer[layer_id][device_index] + ) + self.v_buffer[host_page, layer_id, host_offset] = ( + device_pool.v_buffer[layer_id][device_index] + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "kernel_ascend": + if self.layout == "page_first_direct": + if layer_id == 0: + transfer_kv_dim_exchange( + device_indices=device_indices, + host_indices=host_indices, + device_k=device_pool.k_buffer, + host_k=self.k_buffer, + device_v=device_pool.v_buffer, + host_v=self.v_buffer, + page_size=self.page_size, + direction=TransferDirection.D2H, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: if self.layout == "layer_first": data_page = self.kv_buffer[:, :, index : index + self.page_size, :, :] @@ -993,6 +1087,75 @@ class MLATokenToKVPoolHost(HostKVCache): else: raise ValueError(f"Unsupported IO backend: {io_backend}") + def backup_from_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + ): + if io_backend == "kernel": + if self.layout == "layer_first": + transfer_kv_per_layer_mla( + src=device_pool.kv_buffer[layer_id], + dst=self.kv_buffer[layer_id], + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + ) + elif self.layout == "page_first": + _load_tai_transfer_kv_per_layer_mla_lf_pf()( + device_pool.kv_buffer[layer_id], + self.kv_buffer, + device_indices, + host_indices, + layer_id=layer_id, + item_size=self.token_stride_size, + dst_layout_dim=self.layout_dim, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=[device_pool.kv_buffer[layer_id]], + dst_layers=[self.kv_buffer[layer_id]], + src_indices=device_indices, + dst_indices=host_indices, + page_size=self.page_size, + ) + elif self.layout == "page_first_direct": + for host_index, device_index in zip( + host_indices.cpu().tolist(), device_indices.cpu().tolist() + ): + host_page = host_index // self.page_size + host_offset = host_index % self.page_size + self.kv_buffer[host_page, layer_id, host_offset] = ( + device_pool.kv_buffer[layer_id][device_index] + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "kernel_ascend": + if self.layout == "page_first_kv_split": + if layer_id == 0: + transfer_kv_dim_exchange( + device_indices=device_indices, + host_indices=host_indices, + device_k=device_pool.k_buffer, + host_k=self.k_buffer, + device_v=device_pool.v_buffer, + host_v=self.v_buffer, + device_index_k=device_pool.index_k_buffer, + host_index_k=self.index_k_buffer, + page_size=self.page_size, + direction=TransferDirection.D2H, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: if self.layout == "layer_first": data_page = self.kv_buffer[:, index : index + self.page_size, :, :] @@ -1299,6 +1462,56 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): else: raise ValueError(f"Unsupported IO backend: {io_backend}") + def _backup_indexer_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + host_page_indices, device_page_indices = self._get_indexer_page_indices( + host_indices, device_indices + ) + use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0 + if use_kernel: + if self.layout == "layer_first": + transfer_kv_per_layer_mla( + src=device_pool.index_k_with_scale_buffer[layer_id], + dst=self.index_k_with_scale_buffer[layer_id], + src_indices=device_page_indices, + dst_indices=host_page_indices, + item_size=self.indexer_page_stride_size, + ) + elif self.layout == "page_first": + _load_tai_transfer_kv_per_layer_mla_lf_pf()( + device_pool.index_k_with_scale_buffer[layer_id], + self.index_k_with_scale_buffer, + device_page_indices, + host_page_indices, + layer_id=layer_id, + item_size=self.indexer_page_stride_size, + dst_layout_dim=self.indexer_layout_dim, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=[device_pool.index_k_with_scale_buffer[layer_id]], + dst_layers=[self.index_k_with_scale_buffer[layer_id]], + src_indices=device_page_indices, + dst_indices=host_page_indices, + page_size=1, + ) + elif self.layout == "page_first_direct": + for host_page, device_page in zip( + host_page_indices.cpu().tolist(), + device_page_indices.cpu().tolist(), + ): + self.index_k_with_scale_buffer[host_page, layer_id, 0] = ( + device_pool.index_k_with_scale_buffer[layer_id][device_page] + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + def load_to_device_per_layer( self, device_pool, @@ -1323,3 +1536,13 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): self._backup_indexer_from_device_all_layer( device_pool, host_indices, device_indices, io_backend ) + + def backup_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + super().backup_from_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) + self._backup_indexer_from_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 473b38dee..c13e4f02d 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -108,6 +108,10 @@ for _schema in ( from sglang.srt.managers.cache_controller import HiCacheController from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata +from sglang.srt.mem_cache.memory_pool_host import ( + MLATokenToKVPoolHost, + NSATokenToKVPoolHost, +) from sglang.srt.mem_cache.radix_cache import TreeNode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -120,6 +124,7 @@ class FakeHostPool: self.alloc_result = alloc_result self.alloc_calls = [] self.backups = [] + self.layer_backups = [] self.loads = [] self.frees = [] self.page_size = 4 @@ -136,6 +141,13 @@ class FakeHostPool: ): self.backups.append((host_indices.clone(), device_indices.clone(), device_pool)) + def backup_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + self.layer_backups.append( + (host_indices.clone(), device_indices.clone(), layer_id, device_pool) + ) + def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): @@ -155,10 +167,105 @@ class FakeDevicePool: def __init__(self, name="target", layer_num=1): self.name = name self.layer_num = layer_num + self.layer_backup_notifiers = [] def register_layer_transfer_counter(self, counter): self.counter = counter + def register_layer_backup_notifier(self, notifier): + self.layer_backup_notifiers.append(notifier) + + +class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase): + def test_mla_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self): + calls = [] + + def fake_kernel(src, dst, src_indices, dst_indices, **kwargs): + calls.append((src, dst, src_indices.clone(), dst_indices.clone(), kwargs)) + + host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) + host_pool.layout = "page_first" + host_pool.token_stride_size = 16 + host_pool.layout_dim = 64 + host_pool.kv_buffer = torch.empty((32, 4, 1, 16), dtype=torch.uint8) + device_pool = type("DevicePool", (), {})() + device_pool.kv_buffer = torch.empty((4, 32, 1, 16), dtype=torch.uint8) + host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64) + device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64) + + with patch( + "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_mla_lf_pf", + return_value=fake_kernel, + ): + host_pool.backup_from_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id=2, + io_backend="kernel", + ) + + self.assertEqual(len(calls), 1) + src, dst, src_indices, dst_indices, kwargs = calls[0] + expected_src = device_pool.kv_buffer[2] + self.assertEqual(src.data_ptr(), expected_src.data_ptr()) + self.assertEqual(src.shape, expected_src.shape) + self.assertEqual(src.stride(), expected_src.stride()) + self.assertIs(dst, host_pool.kv_buffer) + self.assertEqual(src_indices.tolist(), [12, 13, 14, 15]) + self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7]) + self.assertEqual(kwargs["layer_id"], 2) + self.assertEqual(kwargs["item_size"], 16) + self.assertEqual(kwargs["dst_layout_dim"], 64) + + def test_nsa_indexer_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self): + calls = [] + + def fake_kernel(src, dst, src_indices, dst_indices, **kwargs): + calls.append((src, dst, src_indices.clone(), dst_indices.clone(), kwargs)) + + host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost) + host_pool.layout = "page_first" + host_pool.page_size = 4 + host_pool.indexer_page_stride_size = 32 + host_pool.indexer_layout_dim = 96 + host_pool.index_k_with_scale_buffer = torch.empty( + (16, 3, 1, 32), dtype=torch.uint8 + ) + device_pool = type("DevicePool", (), {})() + device_pool.index_k_with_scale_buffer = torch.empty( + (3, 16, 32), dtype=torch.uint8 + ) + host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64) + device_indices = torch.tensor( + [12, 13, 14, 15, 28, 29, 30, 31], dtype=torch.int64 + ) + + with patch( + "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_mla_lf_pf", + return_value=fake_kernel, + ): + host_pool._backup_indexer_from_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id=1, + io_backend="kernel", + ) + + self.assertEqual(len(calls), 1) + src, dst, src_indices, dst_indices, kwargs = calls[0] + expected_src = device_pool.index_k_with_scale_buffer[1] + self.assertEqual(src.data_ptr(), expected_src.data_ptr()) + self.assertEqual(src.shape, expected_src.shape) + self.assertEqual(src.stride(), expected_src.stride()) + self.assertIs(dst, host_pool.index_k_with_scale_buffer) + self.assertEqual(src_indices.tolist(), [3, 7]) + self.assertEqual(dst_indices.tolist(), [2, 5]) + self.assertEqual(kwargs["layer_id"], 1) + self.assertEqual(kwargs["item_size"], 32) + self.assertEqual(kwargs["dst_layout_dim"], 96) + class FakeAllocator: def __init__(self, alloc_result=None): @@ -320,7 +427,8 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(result.metadata.owned_positions.tolist(), [4, 5, 6, 7]) self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103]) self.assertEqual(host_pool.alloc_calls, [4]) - self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) + 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): host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) @@ -413,9 +521,11 @@ class TestHiCacheControllerCPWrite(CustomTestCase): ) self.assertEqual(host_pool.alloc_calls, [4]) self.assertEqual(draft_host_pool.alloc_calls, [4]) - self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) - self.assertEqual(draft_host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) - self.assertIs(draft_host_pool.backups[0][2], draft_device_pool) + self.assertEqual(host_pool.backups, []) + self.assertEqual(draft_host_pool.backups, []) + self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertIs(draft_host_pool.layer_backups[0][3], draft_device_pool) self.assertEqual(len(controller.ack_write_queue), 1) self.assertEqual(controller.ack_write_queue[0].node_ids, [77]) @@ -438,6 +548,178 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(host_pool.backups, []) self.assertEqual(draft_host_pool.backups, []) + def test_cp_reserve_write_queues_no_transfer_until_submit(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.arange(4, 20, dtype=torch.int64) + + reservation = controller.reserve_write_cp(logical_locs, node_id=79) + + self.assertEqual(reservation.metadata.logical_len, 16) + self.assertEqual(host_pool.alloc_calls, [4]) + self.assertEqual(host_pool.backups, []) + self.assertEqual(controller.write_queue, []) + self.assertEqual(controller.ack_write_queue, []) + + with self.assertLogs( + "sglang.srt.managers.cache_controller", level="WARNING" + ) as logs: + controller.submit_write_cp_all_layer(reservation) + + self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [79]) + self.assertIn("all-layer backup fallback", "\n".join(logs.output)) + + def test_cp_reserve_zero_owned_queues_no_ack_until_submit(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + controller = self.make_controller(host_pool, cp_rank=3) + logical_locs = torch.arange(4, 8, dtype=torch.int64) + + reservation = controller.reserve_write_cp(logical_locs, node_id=80) + + self.assertEqual(reservation.metadata.host_indices.tolist(), []) + self.assertEqual(host_pool.alloc_calls, []) + self.assertEqual(controller.ack_write_queue, []) + + with self.assertLogs( + "sglang.srt.managers.cache_controller", level="WARNING" + ) as logs: + controller.submit_write_cp_all_layer(reservation) + + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [80]) + self.assertEqual(host_pool.backups, []) + self.assertIn("all-layer backup fallback", "\n".join(logs.output)) + + def test_cp_reserve_draft_allocation_failure_rolls_back_without_transfer(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool(None) + controller = self.make_controller( + host_pool, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=FakeDevicePool("draft"), + ) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + + result = controller.reserve_write_cp(logical_locs, node_id=81) + + self.assertIsNone(result.metadata) + self.assertEqual(result.required_host_slots, 4) + self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103]) + self.assertEqual(host_pool.backups, []) + self.assertEqual(draft_host_pool.backups, []) + self.assertEqual(controller.write_queue, []) + self.assertEqual(controller.draft_write_queue, []) + + def test_cp_submit_write_cp_layer_pairs_target_and_draft_with_single_final_ack(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool( + torch.tensor([200, 201, 202, 203], dtype=torch.int64) + ) + allocator = FakeAllocator() + allocator.device_pool = FakeDevicePool("target", layer_num=2) + draft_device_pool = FakeDevicePool("draft", layer_num=2) + controller = self.make_controller( + host_pool, + allocator=allocator, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=draft_device_pool, + ) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + reservation = controller.reserve_write_cp(logical_locs, node_id=82) + + controller.submit_write_cp_layer(reservation, 0) + + self.assertEqual(controller.ack_write_queue, []) + self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(host_pool.layer_backups[0][2], 0) + self.assertEqual(draft_host_pool.layer_backups[0][2], 0) + + controller.submit_write_cp_layer(reservation, 1) + + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [82]) + self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1]) + self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0, 1]) + + def test_cp_submit_write_cp_layer_zero_owned_final_ack_once(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + allocator = FakeAllocator() + allocator.device_pool = FakeDevicePool("target", layer_num=2) + controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3) + logical_locs = torch.arange(4, 8, dtype=torch.int64) + reservation = controller.reserve_write_cp(logical_locs, node_id=83) + + controller.submit_write_cp_layer(reservation, 0) + self.assertEqual(controller.ack_write_queue, []) + + controller.submit_write_cp_layer(reservation, 1) + + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [83]) + self.assertEqual(host_pool.layer_backups, []) + + def test_cp_layer_hook_submits_registered_write_without_all_layer_backup(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + allocator = FakeAllocator() + allocator.device_pool = FakeDevicePool("target", layer_num=2) + controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + reservation = controller.reserve_write_cp(logical_locs, node_id=84) + + controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False) + + self.assertEqual(host_pool.backups, []) + self.assertEqual(host_pool.layer_backups, []) + self.assertEqual(controller.ack_write_queue, []) + + controller.on_layer_kv_stored(0) + + self.assertEqual(host_pool.layer_backups[0][2], 0) + self.assertEqual(controller.ack_write_queue, []) + + controller.on_layer_kv_stored(1) + + self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1]) + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [84]) + self.assertEqual(host_pool.backups, []) + + def test_cp_layer_hook_waits_for_draft_source_before_final_ack(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + draft_host_pool = FakeHostPool( + torch.tensor([200, 201, 202, 203], dtype=torch.int64) + ) + allocator = FakeAllocator() + allocator.device_pool = FakeDevicePool("target", layer_num=1) + draft_device_pool = FakeDevicePool("draft", layer_num=1) + controller = self.make_controller( + host_pool, + allocator=allocator, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=draft_device_pool, + ) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + reservation = controller.reserve_write_cp(logical_locs, node_id=85) + + controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False) + controller.on_layer_kv_stored(0, source="target") + + self.assertEqual([x[2] for x in host_pool.layer_backups], [0]) + self.assertEqual(draft_host_pool.layer_backups, []) + self.assertEqual(controller.ack_write_queue, []) + + controller.on_layer_kv_stored(0, source="draft") + + self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0]) + self.assertEqual(len(controller.ack_write_queue), 1) + self.assertEqual(controller.ack_write_queue[0].node_ids, [85]) + def test_generate_storage_config_constructs_config_at_runtime(self): controller = HiCacheController.__new__(HiCacheController) controller.mem_pool_device = FakeDevicePool() 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 4ca621fa2..89a4624f0 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -90,10 +90,15 @@ for _schema in ( if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): raise +from sglang.srt.managers.cache_controller import ( + HiCacheWriteFailure, + HiCacheWriteReservation, +) from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams from sglang.srt.mem_cache.hiradix_cache import ( CpHiCacheNodeMetadata, HiRadixCache, + PendingHiCacheBackup, _compute_shared_hicache_token_capacities, ) from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode @@ -531,6 +536,49 @@ class FakeZeroOwnedWriteController: ) +class FakeReserveWriteController: + write_policy = "write_through" + has_draft_hicache = False + + def __init__(self, results): + self.results = list(results) + self.reservations = [] + self.submitted = [] + self.evicted_host_indices = [] + + def reserve_write_cp(self, device_indices, priority=None, node_id=-1): + self.reservations.append((device_indices.clone(), node_id)) + result = self.results.pop(0) + return result(device_indices) if callable(result) else result + + def submit_write_cp_all_layer(self, reservation): + self.submitted.append(reservation) + + def submit_write_cp_per_layer(self, reservation): + self.submitted.append(reservation) + + def evict_cp_host(self, metadata): + self.evicted_host_indices.append(metadata.host_indices.clone()) + return len(metadata.host_indices) + + +def make_write_reservation(device_indices, node_id=0, host_start=90): + host_indices = torch.arange(host_start, host_start + len(device_indices)) + metadata = CpHiCacheNodeMetadata( + logical_len=len(device_indices), + owned_positions=torch.arange(len(device_indices), dtype=torch.int64), + host_indices=host_indices, + page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8), + page_size=1, + ) + return HiCacheWriteReservation( + metadata=metadata, + host_indices=host_indices, + physical_device_indices=device_indices.clone(), + node_id=node_id, + ) + + class FakeEvictionStrategy: def get_priority(self, node): return 0 @@ -643,6 +691,65 @@ class TestHiRadixCacheCPBackup(CustomTestCase): cache.ongoing_write_through.clear() self.assertTrue(cache._node_backuped(node)) + def test_node_backuped_excludes_explicit_pending_cp_backup_until_commit(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + dec_locked = [] + cache.dec_node_lock_ref = lambda node: dec_locked.append(node) + node = TreeNode() + node.id = 128 + node.host_len = 64 + metadata = CpHiCacheNodeMetadata( + logical_len=64, + owned_positions=torch.empty((0,), dtype=torch.int64), + host_indices=torch.empty((0,), dtype=torch.int64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + node.cp_hicache = metadata + cache.pending_host_backups[node.id] = PendingHiCacheBackup( + node=node, metadata=metadata, logical_len=64 + ) + + self.assertFalse(cache._node_backuped(node)) + + cache._commit_pending_backup(node.id) + + self.assertTrue(cache._node_backuped(node)) + self.assertEqual(dec_locked, [node]) + + def test_rollback_pending_cp_backup_frees_reserved_host_slots(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + evicted = [] + cache.cache_controller = types.SimpleNamespace( + evict_cp_host=lambda metadata: evicted.append(metadata) or 2 + ) + dec_locked = [] + cache.dec_node_lock_ref = lambda node: dec_locked.append(node) + cache.pending_host_backups = {} + node = TreeNode() + node.id = 129 + metadata = CpHiCacheNodeMetadata( + logical_len=64, + owned_positions=torch.tensor([0, 1], dtype=torch.int64), + host_indices=torch.tensor([10, 11], dtype=torch.int64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + cache.pending_host_backups[node.id] = PendingHiCacheBackup( + node=node, metadata=metadata, logical_len=64 + ) + + cache._rollback_pending_backup(node.id) + + self.assertEqual(evicted, [metadata]) + self.assertEqual(dec_locked, [node]) + self.assertNotIn(node.id, cache.pending_host_backups) + def test_single_node_write_lock_updates_device_evictable_leaf_set(self): cache = HiRadixCache.__new__(HiRadixCache) cache.disable = False @@ -726,12 +833,18 @@ class TestHiRadixCacheCPBackup(CustomTestCase): def test_write_backup_retries_by_required_physical_slots(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True - cache.cache_controller = FakeWriteController(required_host_slots=1) + reservation_factory = lambda device_indices: make_write_reservation( + device_indices, node_id=122 + ) + cache.cache_controller = FakeReserveWriteController( + [HiCacheWriteFailure(required_host_slots=1), reservation_factory] + ) cache.evictable_host_leaves = set() cache.eviction_strategy = FakeEvictionStrategy() cache.get_child_key_fn = lambda key: key.token_ids[0] cache._record_remove_event = lambda node: None cache.ongoing_write_through = {} + cache.pending_host_backups = {} cache.inc_node_lock_ref = lambda node: None root = TreeNode() @@ -754,6 +867,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase): cache.evictable_host_leaves.add(evictable_node) node = TreeNode() + node.id = 122 node.value = torch.arange(16, dtype=torch.int64) cache.write_backup(node) @@ -764,13 +878,57 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(evictable_node.host_len, 0) self.assertIsNone(evictable_node.cp_hicache) self.assertNotIn(1, root.children) - self.assertEqual(node.host_len, 16) + self.assertEqual(node.host_len, 0) + self.assertIsNone(node.cp_hicache) + self.assertIn(node.id, cache.pending_host_backups) + self.assertEqual(len(cache.cache_controller.submitted), 1) + + def test_write_backup_cp_retry_failure_leaves_node_device_only(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = FakeReserveWriteController( + [ + HiCacheWriteFailure(required_host_slots=2), + HiCacheWriteFailure(required_host_slots=2), + ] + ) + cache.evictable_host_leaves = set() + cache.eviction_strategy = FakeEvictionStrategy() + cache._record_remove_event = lambda node: None + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + + node = TreeNode() + node.id = 124 + node.value = torch.arange(16, dtype=torch.int64) + + backed_len = cache.write_backup(node) + + self.assertEqual(backed_len, 0) + self.assertEqual(node.host_len, 0) + self.assertIsNone(node.cp_hicache) + self.assertEqual(cache.pending_host_backups, {}) + self.assertEqual(cache.cache_controller.submitted, []) def test_write_backup_cp_success_returns_logical_length_for_zero_owned_rank(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True - cache.cache_controller = FakeZeroOwnedWriteController() + cache.cache_controller = FakeReserveWriteController( + [lambda device_indices: HiCacheWriteReservation( + metadata=CpHiCacheNodeMetadata( + logical_len=len(device_indices), + owned_positions=torch.empty((0,), dtype=torch.int64), + host_indices=torch.empty((0,), dtype=torch.int64), + page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8), + page_size=1, + ), + host_indices=torch.empty((0,), dtype=torch.int64), + physical_device_indices=torch.empty((0,), dtype=torch.int64), + node_id=123, + )] + ) cache.ongoing_write_through = {} + cache.pending_host_backups = {} cache.inc_node_lock_ref = lambda node: None node = TreeNode() node.id = 123 @@ -779,8 +937,10 @@ class TestHiRadixCacheCPBackup(CustomTestCase): backed_len = cache.write_backup(node) self.assertEqual(backed_len, 16) - self.assertEqual(node.host_len, 16) - self.assertEqual(node.cp_hicache.host_indices.tolist(), []) + self.assertEqual(node.host_len, 0) + self.assertIsNone(node.cp_hicache) + self.assertIn(node.id, cache.pending_host_backups) + self.assertEqual(len(cache.cache_controller.submitted), 1) def test_attach_storage_backend_rejects_cp_hicache_without_controller_call(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -875,6 +1035,26 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): self.assertEqual(child.cp_hicache.owned_positions.tolist(), [0, 4]) self.assertEqual(child.cp_hicache.host_indices.tolist(), [22, 23]) + def test_cp_host_leaf_status_skips_device_valid_host_backed_node(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.evictable_host_leaves = set() + node = TreeNode() + node.value = torch.arange(4, dtype=torch.int64) + node.host_len = 4 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.tensor([0], dtype=torch.int64), + host_indices=torch.tensor([55], dtype=torch.int64), + page_owners=torch.zeros(max(4, 0), dtype=torch.int8), + page_size=1, + ) + cache.evictable_host_leaves.add(node) + + cache._update_host_leaf_status(node) + + self.assertNotIn(node, cache.evictable_host_leaves) + def test_cp_host_eviction_uses_physical_freed_slots_for_progress(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -1268,6 +1448,38 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertIs(result.last_host_node, node) self.assertIs(result.last_device_node, root) + def test_cp_match_prefix_device_valid_host_missing_remains_device_hit(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.device = "cpu" + cache.disable = False + cache.page_size = 1 + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.key_match_fn = lambda child_key, key: sum( + 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs + ) + 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.parent = root + node.key = RadixKey(list(range(8))) + node.value = torch.arange(8, dtype=torch.int64) + node.host_value = None + node.host_len = 0 + node.cp_hicache = None + root.children[0] = node + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8))))) + + self.assertEqual(result.device_indices.tolist(), list(range(8))) + self.assertEqual(result.host_hit_length, 0) + self.assertIs(result.last_device_node, node) + self.assertIs(result.last_host_node, root) + def test_cp_match_prefix_does_not_admit_inflight_write_as_host_hit(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -1311,6 +1523,51 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertIs(result.last_device_node, node) self.assertIs(result.last_host_node, root) + def test_cp_match_prefix_defers_split_when_child_backup_pending(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.device = "cpu" + cache.disable = False + cache.page_size = 1 + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.key_match_fn = lambda child_key, key: sum( + 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs + ) + 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 + child = TreeNode() + child.id = 130 + child.parent = root + child.key = RadixKey(list(range(8))) + child.value = torch.arange(8, dtype=torch.int64) + child.host_len = 0 + metadata = CpHiCacheNodeMetadata( + logical_len=8, + owned_positions=torch.arange(8, dtype=torch.int64), + host_indices=torch.arange(50, 58, dtype=torch.int64), + page_owners=torch.zeros(8, dtype=torch.int8), + page_size=1, + ) + cache.pending_host_backups[child.id] = PendingHiCacheBackup( + node=child, metadata=metadata, logical_len=8 + ) + root.children[0] = child + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 9]))) + + self.assertIs(result.pending_backup_deferred_node, child) + self.assertIs(root.children[0], child) + self.assertEqual(child.key.token_ids, list(range(8))) + self.assertEqual(result.device_indices.tolist(), []) + self.assertIs(result.last_device_node, root) + def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True