sglang/docs/superpowers/plans/2026-05-27-cp-hicache-per-layer-backup.md

# 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.