docs: add CP HiCache host plan

.gitignore

@@ -272,3 +272,6 @@ sgl-kernel/csrc/**/*_musa/
 *.glb
 *.ply
 *.npz
+
+# Git worktrees
+.worktrees/

docs/superpowers/specs/2026-05-07-cp-hicache-host-design.md

@@ -0,0 +1,429 @@
+# CP Shared KV + HiCache Host Integration Design
+
+**Goal:** Deliver a first reliable integration between NSA CP shared KV and HiCache host caching for GLM 5 inference.
+
+**Architecture:** Keep scheduler, radix cache, and req-to-token in the CP-group logical address space. Keep device and host KV pools physical and per-rank sharded. Add a CP-aware HiCache controller adapter plus explicit per-node host metadata so only owned physical shards are transferred while radix state remains logically complete.
+
+**Tech Stack:** Python, SGLang server runtime, NSA `NSATokenToKVPool`, `CPSharedPagedTokenToKVPoolAllocator`, `CpSharedKVLayout`, `HiRadixCache`, `HiCacheController`.
+
+---
+
+## Scope
+
+This design covers the first implementation stage only:
+
+- Support `--enable-nsa-prefill-cp-shared-kv` with HiCache host write/load/evict.
+- Keep HiCache storage backends unsupported in this mode.
+- Keep radix tree, scheduler admission, and req-to-token state in full logical KV loc space.
+- Keep each CP rank's device KV pool and host KV pool sharded by physical owner pages.
+
+This design does not cover:
+
+- HiCache storage backend backup or prefetch.
+- Storage key changes for File, NIXL, Mooncake, EIC, HF3FS, or other backends.
+- Decode-side CP/shared KV.
+- HiSparse NSA KV pools.
+- A full replacement of HiRadixCache with upstream UnifiedRadixTree or HybridCacheController.
+
+If CP shared KV and `hicache_storage_backend` are enabled together, startup must fail fast with a clear error. Silent fallback is not acceptable because the existing storage path assumes page-hash sequences and contiguous host indices that do not describe CP-owned sparse shards.
+
+---
+
+## Existing Code Facts
+
+The current CP shared KV branch already separates logical and physical KV locs:
+
+- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py` provides owner masks and logical-to-physical conversions for torch and NumPy paths.
+- `python/sglang/srt/mem_cache/allocator.py::CPSharedPagedTokenToKVPoolAllocator` tracks logical allocator size while storing physical pool size separately.
+- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py` profiles physical token capacity and expands logical capacity by CP size.
+- Runtime persistent writes already use `CpSharedKVLayout` to filter owned logical locs and convert to physical locs before touching local KV buffers.
+
+The current HiCache path is not CP-aware:
+
+- `HiRadixCache.write_backup()` passes `node.value` directly to `HiCacheController.write()`.
+- `HiCacheController.write()` allocates host indices for the full tensor length and transfers those indices as if they were physical device indices.
+- `HiCacheController.load()` allocates device indices for the full host tensor length and transfers those indices as if they were physical device indices.
+- `HiRadixCache.load_back()` uses `len(node.host_value)` to compute host hit length and to slice restored device indices.
+- `HiRadixCache._split_node()` slices `node.host_value` when radix nodes split.
+
+Under CP shared KV, `node.value` is a full logical loc tensor. It cannot be passed directly into physical host/device transfer functions.
+
+---
+
+## Core Invariants
+
+The implementation must preserve these invariants:
+
+1. `TreeNode.value` remains a full logical loc tensor.
+2. `req_to_token` remains full logical loc state.
+3. Scheduler admission and radix matching remain based on full logical token lengths.
+4. `HostKVCache` remains physical-only and does not import or depend on `CpSharedKVLayout`.
+5. CP filtering and logical-to-physical conversion happen in `HiCacheController` or a small adapter owned by it.
+6. `TreeNode.host_value` keeps the legacy non-CP meaning: physical host indices tensor.
+7. CP mode does not store CP-owned host subsets in `TreeNode.host_value`.
+8. CP mode uses explicit metadata to represent host backup state.
+9. Host eviction and load-back accounting report logical token counts, even though local physical transfers are sparse.
+
+---
+
+## Metadata Model
+
+Add explicit CP host metadata for HiRadixCache nodes.
+
+Recommended structure:
+
+```python
+@dataclass
+class CpHiCacheNodeMetadata:
+    logical_len: int
+    owned_positions: torch.Tensor
+    host_indices: torch.Tensor
+
+    def split(self, split_len: int) -> tuple["CpHiCacheNodeMetadata", "CpHiCacheNodeMetadata"]:
+        if split_len < 0 or split_len > self.logical_len:
+            raise ValueError(
+                f"split_len must be in [0, {self.logical_len}], got {split_len}"
+            )
+        parent_mask = self.owned_positions < split_len
+        child_mask = ~parent_mask
+        parent = CpHiCacheNodeMetadata(
+            logical_len=split_len,
+            owned_positions=self.owned_positions[parent_mask],
+            host_indices=self.host_indices[parent_mask],
+        )
+        child = CpHiCacheNodeMetadata(
+            logical_len=self.logical_len - split_len,
+            owned_positions=self.owned_positions[child_mask] - split_len,
+            host_indices=self.host_indices[child_mask],
+        )
+        return parent, child
+```
+
+Field semantics:
+
+- `logical_len` is the full logical token count represented by this node.
+- `owned_positions` are node-local logical token positions owned by this CP rank.
+- `host_indices` are local physical host slots storing the owned token KV data.
+- `len(owned_positions) == len(host_indices)`.
+- `owned_positions` can be empty on a rank that owns no page in this node.
+- `owned_positions` are node-local positions, not old logical loc values. This avoids depending on evicted logical locs remaining reusable after load-back.
+- `owned_positions` and `host_indices` are stored as CPU `torch.int64` tensors. Radix tree operations are CPU-side bookkeeping, and keeping metadata on CPU avoids retaining GPU tensors in long-lived tree nodes.
+
+Add fields to `TreeNode` or attach equivalent metadata from HiRadixCache:
+
+- `host_len: int = 0`
+- `cp_hicache: Optional[CpHiCacheNodeMetadata] = None`
+
+For non-CP HiCache, existing `host_value` behavior remains unchanged.
+
+For CP HiCache, backup state is represented by `host_len > 0` and `cp_hicache is not None`.
+
+Radix backup checks must become cache-type aware. A helper is preferable to changing the global `TreeNode.backuped` property for all cache classes:
+
+```python
+def _node_backuped(self, node: TreeNode) -> bool:
+    if self._uses_cp_hicache:
+        return node.host_len > 0 and node.cp_hicache is not None
+    return node.host_value is not None
+```
+
+Use this helper inside `HiRadixCache` paths that currently read `node.backuped`. Do not change SWA, Mamba, or plain Radix semantics as part of this integration unless a touched path requires it.
+
+---
+
+## Controller Adapter
+
+`HiCacheController` is the transfer chokepoint and should own CP-aware translation.
+
+CP HiCache mode must be detected from the allocator or explicit layout, not from the KV pool alone. Use `isinstance(token_to_kv_pool_allocator, CPSharedPagedTokenToKVPoolAllocator)` or an explicitly passed `CpSharedKVLayout` to enable the CP adapter. After CP mode is detected, validate that the underlying KV pool is `NSATokenToKVPool` for this stage.
+
+Add initialization state:
+
+- `uses_cp_hicache: bool`
+- `cp_shared_kv_layout: Optional[CpSharedKVLayout]`
+- `page_size: int`
+
+The controller can derive CP mode from the allocator by checking for CP-specific fields such as `cp_size` and `cp_rank`, or it can receive an explicit layout from `HiRadixCache`. Prefer explicit construction in `HiRadixCache` and passing the layout into `HiCacheController` so the adapter has one source of truth.
+
+### CP Write
+
+External call shape stays close to the current API:
+
+```python
+metadata = cache_controller.write(device_indices=node.value, node_id=node.id)
+```
+
+CP behavior:
+
+1. Treat `device_indices` as full logical locs.
+2. Build `owned_mask = layout.owned_by_this_rank(device_indices)`.
+3. Build `owned_positions = owned_mask.nonzero(as_tuple=True)[0]`.
+4. Build `physical_device_indices = layout.logical_locs_to_physical(device_indices[owned_mask])`.
+5. Allocate `host_indices = mem_pool_host.alloc(len(physical_device_indices))`.
+6. If allocation fails, return a failure that exposes `required_host_slots=len(physical_device_indices)` so `HiRadixCache` can evict physical host slots and retry.
+7. If the owned count is zero, skip host allocation and skip physical transfer, but still return metadata with `logical_len=len(device_indices)`.
+8. Enqueue or run the existing host transfer using `host_indices` and `physical_device_indices`.
+9. Store returned metadata tensors on CPU: `owned_positions.cpu()` and `host_indices.cpu()`.
+10. Return `CpHiCacheNodeMetadata(logical_len, owned_positions, host_indices)`.
+
+Every logical write-back must produce the same acknowledgement cardinality on every TP rank. If a rank owns zero physical slots for the node, the controller must still create a completed no-op write acknowledgement when `HiRadixCache` adds the node to `ongoing_write_through`. This preserves the existing `writing_check()` invariant that all ranks pop the same logical ack ids after `all_reduce(MIN)`.
+
+The non-CP path continues returning the existing host indices tensor.
+
+### CP Load
+
+The load path must allocate full logical device locs but transfer only owned physical shards.
+
+Recommended API:
+
+```python
+device_indices = cache_controller.load_cp(nodes_to_load)
+```
+
+CP behavior:
+
+1. Compute `logical_len = sum(node.host_len for node in nodes_to_load)`.
+2. Allocate `device_indices = mem_pool_device_allocator.alloc(logical_len)`.
+3. For each node, take its slice of `device_indices`.
+4. Use `node.cp_hicache.owned_positions` to select the newly allocated logical locs owned by this rank.
+5. Convert selected logical locs to physical locs with `layout.logical_locs_to_physical(selected_logical_locs)`.
+6. Concatenate all owned host indices and owned physical device indices.
+7. If the owned count is zero, skip physical transfer and still return the full logical `device_indices`.
+8. Enqueue or run the existing host-to-device transfer for the owned physical pairs.
+9. Return the full logical `device_indices` to HiRadixCache.
+
+CP load-back has the same acknowledgement cardinality requirement as CP write-back. A zero-owned rank must still produce a completed no-op load acknowledgement or equivalent completed producer event so lock release and producer ids remain consistent across ranks.
+
+The non-CP path continues using the existing `load(host_indices=host_indices)` API shape.
+
+---
+
+## HiRadixCache Changes
+
+### Initialization
+
+`HiRadixCache.__init__()` must detect CP shared KV mode and construct the controller adapter.
+
+It should also enforce:
+
+- CP shared KV + HiCache storage backend is unsupported and must fail during server argument validation or cache initialization.
+- CP shared KV + HiCache is valid only for `NSATokenToKVPool` in this stage.
+- CP HiCache mode is enabled from `CPSharedPagedTokenToKVPoolAllocator` or an explicit `CpSharedKVLayout`, not merely because the device pool is `NSATokenToKVPool`.
+
+### Backup State
+
+Introduce helpers inside `HiRadixCache`:
+
+```python
+def _node_backuped(self, node: TreeNode) -> bool:
+    if self._uses_cp_hicache:
+        return node.host_len > 0 and node.cp_hicache is not None
+    return node.host_value is not None
+
+
+def _node_host_len(self, node: TreeNode) -> int:
+    if self._uses_cp_hicache:
+        return node.host_len
+    return len(node.host_value)
+
+
+def _node_host_evict_indices(self, node: TreeNode) -> torch.Tensor:
+    if self._uses_cp_hicache:
+        return node.cp_hicache.host_indices
+    return node.host_value
+```
+
+Use these helpers in paths currently reading `node.backuped`, `len(node.host_value)`, or `node.host_value` for host eviction/load-back.
+
+### Write Backup
+
+Non-CP path remains unchanged.
+
+CP path:
+
+1. Call controller CP write with full logical `node.value`.
+2. If write reports insufficient host slots, evict host entries until physical freed slots are at least the requested `required_host_slots`, then retry.
+3. On success, set `node.host_len = len(node.value)` and `node.cp_hicache = metadata`.
+4. Keep `node.host_value is None` in CP mode.
+5. Add the node to `ongoing_write_through` if write-through acknowledgement is required.
+6. Locking and metrics count logical tokens, not owned physical tokens.
+
+`_inc_hit_count()` must use `_node_backuped(node)` instead of `node.backuped`, otherwise CP-backed nodes would appear unbacked because `node.host_value` remains `None` in CP mode and the write-through path would repeatedly call `write_backup()`.
+
+### Split Node
+
+`_split_node()` must split CP metadata when a backed-up node is split:
+
+- Parent metadata keeps positions `< split_len`.
+- Child metadata keeps positions `>= split_len`, rebased by subtracting `split_len`.
+- Parent `host_len = split_len`.
+- Child `host_len = old_host_len - split_len`.
+- Host indices are split according to the owned positions mask, not by slicing first `split_len` entries.
+
+Non-CP `host_value` slicing remains unchanged.
+
+### Load Back
+
+Non-CP path remains unchanged.
+
+CP path:
+
+1. Collect `nodes_to_load` while walking evicted ancestors.
+2. Compute `host_hit_len = sum(node.host_len for node in nodes_to_load)`.
+3. Apply `load_back_threshold` and `mem_quota` checks to `host_hit_len`.
+4. Call controller CP load with `nodes_to_load`.
+5. If allocation fails, evict device logical tokens and retry.
+6. Assign full logical slices from returned `device_indices` to each `node.value` using `node.host_len`.
+7. Update `evictable_size_`, locks, and metrics by full logical lengths.
+
+The `load_back_threshold` check intentionally uses logical length in this first stage. Using local physical owned-token count would let different CP ranks make different load-back decisions. A physical-transfer threshold would require a CP-group reduction and is deferred.
+
+### Host Eviction
+
+Non-CP path remains unchanged.
+
+CP path:
+
+1. Host eviction candidate selection still uses radix logical nodes.
+2. Free only `node.cp_hicache.host_indices` from the local host pool.
+3. Track two counts separately: physical freed slots for allocator retry progress, and logical evicted tokens for metrics/debug output.
+4. Clear `node.host_len` and `node.cp_hicache`.
+5. Remove the node from the radix tree as the existing host eviction path does.
+
+Host eviction for CP write retry must progress by physical freed slots, not by logical `host_len`. A node can have a large logical `host_len` but zero local `host_indices` on a rank. Counting that logical length as eviction progress could stop the eviction loop without freeing the physical host slots needed by the retry.
+
+---
+
+## Capacity Policy
+
+Host memory remains per-rank physical capacity based.
+
+The existing `HostKVCache` default sizing uses `device_pool.size * hicache_ratio` when `hicache_size <= 0`. Under CP shared KV, `device_pool.size` is already physical per-rank size, so this is the desired behavior.
+
+Do not change host pool sizing to logical capacity. That would multiply host memory by CP size per rank and defeat the CP memory-saving goal.
+
+Do not add a new CP-specific host ratio parameter in the first stage. Existing `hicache_ratio` and `hicache_size` remain the user-facing controls.
+
+---
+
+## dev-hicache Fix Integration
+
+Do not blindly cherry-pick the `dev-hicache` commits onto `dev-cp-refactor`. The CP branch already modified allocator/common/model runner code, and direct cherry-picks risk overwriting CP shared KV logic.
+
+Manually port fixes in this order:
+
+1. Correctness fixes before CP HiCache integration:
+   - Load-back `mem_quota` scheduling budget from `a666ab552`.
+   - Paged allocator extend OOM accounting from `d6ca2fd0c`, merged carefully with compute-owner lane eviction in `common.py`.
+   - `defaultdict(TreeNode)` ghost-node fix from `e008d2cf3` for radix-like nodes.
+2. CP host HiCache integration from this design.
+3. HiCache performance fixes from `25b6a957c`:
+   - Incremental eviction heap.
+   - Cached timestamp in hot loops.
+   - Persistent all-reduce tensors.
+   - Async write acknowledgement thread.
+
+The short bigram page lookup fix from upstream commit `7be96c8ce` is a separate small correctness task and should be manually ported with a unit test.
+
+The upstream SWA HiCache storage type registration is out of scope for GLM 5 NSA host-only integration.
+
+---
+
+## Error Handling
+
+Startup or initialization must fail with clear messages for unsupported combinations:
+
+- CP shared KV + HiCache storage backend.
+- CP shared KV + HiSparse NSA pool.
+- CP shared KV + non-NSA KV pool in HiCache host integration.
+
+Runtime allocation failure handling follows existing HiCache behavior:
+
+- Host write allocation failure returns `None`; HiRadixCache evicts host entries and retries.
+- Device load allocation failure returns `None`; HiRadixCache evicts device entries and retries.
+- If retry still fails, `load_back()` returns `None` and the request recomputes the prefix as today.
+
+For CP metadata consistency failures, fail loudly with assertions in unit-testable helper methods. Examples:
+
+- `metadata.logical_len` must match the node logical length.
+- `len(metadata.owned_positions) == len(metadata.host_indices)`.
+- split lengths must be within `[0, logical_len]`.
+- `owned_positions` must be sorted and in range.
+- `owned_positions` and `host_indices` must be CPU tensors before being stored on `TreeNode`.
+
+---
+
+## Testing Strategy
+
+Use SGLang unit-test conventions for new tests:
+
+- Place tests under `test/registered/unit/`.
+- Inherit from `CustomTestCase`.
+- Register CPU tests with `register_cpu_ci` where possible.
+- Avoid launching a server for metadata and controller mapping tests.
+
+Required unit coverage:
+
+1. `CpHiCacheNodeMetadata.split()`:
+   - page-aligned split.
+   - `split_len == 0`.
+   - split with no owned positions on this rank.
+   - split where owned positions are on both sides.
+   - invalid split values raise errors.
+2. Controller CP write mapping:
+   - full logical locs produce only owned host allocation.
+   - physical device indices passed to host transfer match `CpSharedKVLayout`.
+   - zero-owned node returns metadata without allocating host slots.
+   - zero-owned node still produces one completed logical write acknowledgement if the node enters `ongoing_write_through`.
+3. Controller CP load mapping:
+   - full logical device locs are allocated.
+   - owned positions map to newly allocated logical locs, then to physical locs.
+   - zero-owned metadata skips transfer but returns full logical locs.
+   - zero-owned metadata still produces a completed logical load acknowledgement or producer event.
+4. HiRadixCache metadata behavior:
+   - `_node_backuped()` returns true for CP nodes with `host_len` and metadata.
+   - `_inc_hit_count()` uses `_node_backuped()` and does not repeatedly write back already CP-backed nodes.
+   - `_split_node()` splits `host_len` and CP metadata correctly.
+   - `load_back()` uses `host_len` rather than `len(host_value)` in CP mode.
+   - host eviction frees owned host indices and uses physical freed slots as retry progress.
+5. Server args/config validation:
+   - CP shared KV + `hicache_storage_backend` fails fast.
+
+Required regression checks:
+
+- CP shared KV layout tests.
+- Radix cache unit tests.
+- Scheduler prefill-adder tests covering `mem_quota`.
+- New CP HiCache metadata/controller tests.
+
+Optional manual GPU smoke:
+
+- Launch GLM5/NSA CP shared KV with HiCache host-only.
+- Use repeated-prefix requests to trigger backup, device eviction, and host load-back.
+- Confirm no logical/physical index assertion failures and no host/storage unsupported path is entered.
+
+---
+
+## Implementation Order
+
+1. Add tests for existing correctness fixes that are currently missing from the CP branch.
+2. Manually port load-back `mem_quota`, allocator OOM accounting, and `defaultdict(TreeNode)` fixes.
+3. Add CP+storage fail-fast validation.
+4. Add `CpHiCacheNodeMetadata` and unit tests.
+5. Add CP adapter methods in `HiCacheController` with mock unit tests.
+6. Wire `HiRadixCache` backup, split, load-back, and host eviction to CP metadata helpers.
+7. Run CPU/unit regression tests.
+8. Manually smoke test GPU CP+HiCache host-only if hardware/model access is available.
+9. Port performance fixes from `25b6a957c` as separate patches with regression tests.
+
+---
+
+## Open Follow-Up Work
+
+These items are deliberately deferred:
+
+- CP-aware HiCache storage key and prefetch design.
+- Support for File/NIXL/Mooncake/EIC/HF3FS storage backends under CP HiCache.
+- Independent CP-specific host ratio configuration.
+- Decode-side CP/shared KV HiCache support.
+- Upstream UnifiedRadixTree or HybridCacheController migration.