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

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

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

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

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

• Step 5: Run targeted unit tests

Run:

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

• Step 1: Add a small pending-backup record

Add a dataclass near CpHiCacheNodeMetadata or a private internal class near HiRadixCache:

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

• Step 2: Add a pending map

Initialize:

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.

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

• Step 4: Add commit and rollback helpers

Add internal helpers:

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.

• 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

• Step 1: Extract reservation logic from _write_cp()

Create a method with the current allocation and metadata construction logic:

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.

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

• Step 3: Make _write_cp() call reserve + submit-all-layer

Keep existing behavior working by implementing _write_cp() as:

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

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

• 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

• Step 1: Add abstract/base method

Add to the base host pool class:

def backup_from_device_per_layer(self, device_pool, host_indices, device_indices, layer_id, io_backend):
    raise NotImplementedError

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

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

• 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

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

• Step 2: Add submit-layer method

Add:

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.

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

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

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

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

• Step 2: Add a no-op-safe callback on forward batch or controller

Expose a callback such as:

forward_batch.hicache_backup_notifier.on_layer_kv_stored(layer_id)

When no pending backup exists, it returns immediately.

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

• 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

• Step 1: Make reserve failure return required local physical host slots

Keep HiCacheWriteFailure(required_host_slots=...) as the capacity signal.

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

• Step 3: Retry reservation once

If the second reservation fails, rollback any partial reservation and leave the node DEVICE_VALID_HOST_NONE.

• 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

• Step 1: Add pending split guard

Before splitting a CP HiCache node, detect pending backup state.

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

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

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