Keep CP HiCache host visibility behind per-layer target and draft backup

CP shared KV needs HiCache backup to overlap with layer execution without exposing partially copied host state. Split CP backup into reservation, pending radix state, per-layer target/draft D2H submission, and one final ack-driven visibility commit. The all-layer path remains available only as an explicit fallback and now logs a warning when used.

Constraint: CP shared KV owner-lane metadata and draft/MTP KV must stay strongly synchronized with target KV.
Constraint: Local CUDA tests are disallowed; CUDA verification was run only in the g0034 container.
Rejected: Let target layer hooks copy draft KV too | draft may not have stored that layer yet, which can corrupt MTP accept behavior.
Rejected: Silent all-layer fallback | it hides performance regressions and makes ETE logs ambiguous.
Confidence: medium
Scope-risk: broad
Directive: Reserved or partially copied host payloads must remain invisible until final ack commits pending_host_backups.
Tested: g0034 docker /sgl-workspace/sglang-tai PYTHONPATH=/mnt/beegfs/cjy/tai-kernel/python:python python -m pytest test/registered/unit/managers/test_hicache_controller_cp.py -q -> 49 passed.
Tested: g0034 docker /sgl-workspace/sglang-tai PYTHONPATH=/mnt/beegfs/cjy/tai-kernel/python:python python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q -> 58 passed.
Tested: g0034 docker /mnt/beegfs/cjy/tai-kernel PYTHONPATH=python python -m pytest tests/nsa_prefill/test_kvcacheio_lf_pf.py -q -> 7 passed.
Not-tested: Long-running GLM5 CP+HiCache+MTP ETE throughput and host-pressure soak.
This commit is contained in:
laoyao0822
2026-05-27 07:45:16 +08:00
parent 367dff06f3
commit 03529319a1
11 changed files with 2104 additions and 72 deletions
@@ -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.
@@ -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.
+321 -38
View File
@@ -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."""
@@ -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()
@@ -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,
@@ -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):
+106 -24
View File
@@ -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}
@@ -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 = [
+224 -1
View File
@@ -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
)
@@ -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()
@@ -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