chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Make CP HiCache backup admission deterministic

Browse Source 
CP HiCache write-through under shared KV was still using rank-wide collectives to decide host reservation eviction, and per-layer backup registration could be bypassed before the final forward boundary. This moves backup registration to the final run_batch pre-forward boundary, forwards it through SessionAwareCache, exposes fallback paths as explicit warnings, and introduces deterministic owner-lane capacity planning for CP host reservation.

Constraint: CP shared-KV ranks must keep target and draft host reservations owner-lane consistent without adding hot-path collective synchronization

Constraint: Remote CUDA validation must run in the g0034 container, not locally

Rejected: Keep reserve_slots_max all_reduce as the default admission path | observed reserve collectives reaching double-digit and occasional 100ms+ latency

Rejected: Silent post-forward catch-up backup | hides when per-layer forward-overlap backup is not actually active

Confidence: medium

Scope-risk: broad

Directive: Do not reintroduce CP HiCache hot-path collectives without a measured mismatch case and explicit fallback warning

Tested: py_compile for modified Python modules and CP HiCache metadata test file in remote g0034 container

Tested: python3 -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q in remote g0034 container (75 passed, 5 warnings)

Tested: git diff --check HEAD~1..HEAD

Not-tested: Local pytest blocked by missing pybase64 in the local environment

Not-tested: Full CP HiCache + MTP E2E after the no-collective reservation change

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-27 23:11:23 +08:00
co-authored by OmX
parent f355fdd39e
commit 40a8de5fd1
11 changed files with 1476 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md
+349
View File
@@ -0,0 +1,349 @@
# CP HiCache No-Collective Capacity Plan
> **Goal:** remove CP HiCache hot-path all-reduce usage by making host capacity,
> victim selection, and async visibility deterministic-by-construction.
>
> **Scope:** prefill-side CP shared KV + HiCache L2 host cache + EAGLE/MTP draft
> shadow KV. Excludes decode, TBO, L3/storage prefetch, and cross-instance cache.
## Current Code Facts
The current hot collectives are in `python/sglang/srt/mem_cache/hiradix_cache.py`:
1. `reserve_slots_max`
- `_sync_cp_write_required_host_slots()`
- Called by `_reserve_write_cp_indices_collectively()`.
- Purpose: force all CP ranks into the same reserve/evict/retry branch when
any rank's local host pool is full.
2. `host_evict_done_min`
- `_evict_host_for_physical_slots(..., synchronize_across_ranks=True)`.
- Purpose: keep all ranks in the host-eviction loop until every rank has
freed enough local host slots.
3. `writing_check_min`
- `writing_check()`.
- Purpose: commit the same prefix of completed write acks on all ranks so
radix host visibility does not diverge.
Storage/prefetch collectives (`storage_queue_min`, `prefetch_terminate_max`,
`prefetch_completed_min`) are L3/storage scope and are intentionally out of this
stage.
Relevant controller facts in `python/sglang/srt/managers/cache_controller.py`:
- `reserve_write_cp()` already derives the global page owner pattern from
logical pages and stores it in `CpHiCacheNodeMetadata.page_owners`.
- The local physical host allocation only stores pages owned by this rank.
- Draft/MTP host allocation follows target ownership and is rolled back if
target/draft cannot be reserved atomically.
- `load_cp()` replays `page_owners` through `alloc_pages_with_owners()` and then
filters to this rank's local physical shard.
## Root Issue
The current code uses collectives as an admission oracle because the control
plane does not have a replicated view of rank-local capacity. That works for
correctness, but it turns ordinary backup/eviction/check ticks into CPU
collective traffic.
The replacement model is:
```text
same radix tree
same logical nodes
same page_owners
same target/draft atomic state
same deterministic victim order
=> every rank can independently derive every rank's required/free host pages
```
The local allocator should only execute an already-decided logical plan. If it
fails after the deterministic plan says it must succeed, that is an invariant
violation or a slow-path abort, not a normal retry path.
## Invariants
1. Radix tree mutations stay synchronous and deterministic.
2. Transfer work stays asynchronous and rank-local.
3. Target and draft KV are one logical cache object.
4. Host capacity is accounted by logical CP owner, not by asking peers.
5. Host eviction victims are chosen from a deterministic logical order.
6. Host-resident pending writes are not host-evictable.
7. No new hot-path collectives are allowed.
## Design
### 1. Replicated CP host capacity ledger
Add a small ledger owned by `HiRadixCache`:
```text
target_host_capacity_pages_by_rank[cp_rank]
draft_host_capacity_pages_by_rank[cp_rank]
target_host_used_pages_by_rank[cp_rank]
draft_host_used_pages_by_rank[cp_rank]
pending_backup_pages_by_rank[cp_rank]
```
Counts are derived from `CpHiCacheNodeMetadata.page_owners`:
```text
pages_by_owner = bincount(page_owners, minlength=cp_size)
```
For a local rank, `pages_by_owner[cp_rank] * page_size` must match
`len(metadata.host_indices)`. For remote ranks, the ledger still knows the
logical count even though it does not know remote physical host indices.
### 2. Deterministic host admission
Before calling `reserve_write_cp()`:
```text
required = pages_by_owner(new_node.page_owners)
if ledger lacks target or draft capacity:
evict deterministic host victims until enough capacity exists
if capacity still insufficient:
skip backup for every rank before any physical reservation
else:
reserve local target/draft host slots
```
This removes `reserve_slots_max`: every rank reaches the same success/skip
decision without voting.
### 3. Deterministic host eviction
Replace `_evict_host_for_physical_slots(required_host_slots, sync=True)` with a
logical planner:
```text
victims = sorted(evictable_host_leaves, key=deterministic_host_evict_key)
for victim in victims:
if victim is host-valid, device-evicted, unlocked, unpinned, not pending:
subtract pages_by_owner(victim.page_owners) from required deficits
add victim to victim_plan
stop when all deficits are satisfied
```
Then each rank applies the same `victim_plan` and frees only its local
`metadata.host_indices` and `metadata.draft_host_indices`.
This removes `host_evict_done_min`. The stop condition is logical deficits, not
local freed-token count.
The eviction key must not depend on unordered set iteration or wall-clock tie
breaks. Use a deterministic tuple such as:
```text
(node.priority, node.logical_last_access_seq, node.id)
```
If a logical access sequence is not available in the first pass, use
`(node.priority, node.id)` for CP host eviction and accept lower cache-quality
temporarily rather than risking rank-divergent victim selection.
### 4. Writing visibility without completion votes
`writing_check_min` is harder than capacity admission. It exists because
`writing_check()` currently mutates request-visible radix state when local CUDA
events complete. If one rank commits earlier than another, the next request can
see different host-hit state.
The no-collective replacement is to separate logical visibility from data
readiness:
```text
HOST_PENDING metadata is inserted deterministically on all ranks
HOST_PENDING is matchable but not host-evictable
load from HOST_PENDING waits on the local write event before H2D
local write completion only releases local physical lifetimes
no background completion path mutates radix topology independently
```
This makes all ranks choose the same host-hit path. Slower ranks pay local
event wait before reading host bytes; faster ranks do not need to ask peers.
First implementation may use the final write event as the dependency. A later
optimization can store per-layer write events so H2D load for layer `i` only
waits for D2H write of layer `i`.
## Implementation Plan
## Implementation Status
- **P1 implemented in observer mode.**
- Added CP owner-lane token counting from `CpHiCacheNodeMetadata.page_owners`.
- Added metadata invariant checks for local `owned_positions`, target
`host_indices`, and draft `draft_host_indices`.
- Added a recomputed host capacity snapshot that separates committed and
pending target/draft host residency without changing runtime behavior.
- **P2 implemented in observer mode.**
- Added deterministic host victim planning keyed by logical node metadata
rather than by `set` iteration order.
- Planner skips device-valid nodes, pending backup nodes, host-protected nodes,
pinned nodes, malformed metadata, and non-host-backed leaves.
- Current eviction behavior is unchanged.
- **Online debug option added:** `SGLANG_CP_HICACHE_CAPACITY_DEBUG=1`.
- When enabled, write admission computes the deterministic pre-reserve plan and
compares its eviction-needed decision with the current collective
`reserve_slots_max` result.
- Logs use `[HiCache-capacity-debug] write_admission_compare ...`.
- This is observer-only and does not remove or add collectives.
### P1: Observer-only ledger
Files:
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Add helper methods:
- `_cp_owner_page_counts(page_owners) -> list[int]`
- `_cp_metadata_page_counts(metadata) -> list[int]`
- `_cp_host_capacity_snapshot()`
- `_cp_assert_metadata_counts(metadata)`
2. Initialize ledger capacities from target/draft host pool token capacity divided
by `page_size`.
3. Update ledger on committed host metadata, rollback, and host eviction in
observer mode only.
4. Add tests for:
- zero-owned rank metadata,
- target+draft counts match page owner counts,
- rollback subtracts pending counts exactly once,
- host eviction subtracts target and draft together.
No behavior change in P1.
### P2: Deterministic host victim planner
Files:
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Add `_cp_host_evict_key(node)`.
2. Add `_plan_cp_host_evictions(required_pages_by_rank)`.
3. Ensure planner skips:
- non-evicted device nodes,
- pending backups,
- host-ref protected nodes,
- pinned nodes before expiry,
- malformed host metadata.
4. Unit-test that two different `evictable_host_leaves` insertion orders produce
the same victim node ids.
No collective removal yet; compare the planned victim count against current
eviction in debug logs/assertions.
### P3: Replace reserve MAX with planned admission
Files:
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Replace `_reserve_write_cp_indices_collectively()` with
`_reserve_write_cp_indices_planned()`.
2. Compute required page counts before physical reserve.
3. If ledger says insufficient, run the deterministic host eviction plan.
4. If still insufficient, skip backup before calling `reserve_write_cp()`.
5. If `reserve_write_cp()` fails after planned admission, raise an invariant
error with:
- node id,
- page owner counts,
- ledger free counts,
- local host allocator availability,
- draft allocator availability.
This removes `reserve_slots_max`.
### P4: Remove synchronized host eviction loop
Files:
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
Steps:
1. Stop passing `synchronize_across_ranks=True`.
2. Delete the `all_ranks_done()` all-reduce loop.
3. Apply the planned victim list deterministically.
4. Keep a fail-fast assertion if applying the victim list does not free the
ledger-predicted local target/draft slots.
This removes `host_evict_done_min`.
### P5: Remove write completion voting
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`
- Test: `test/registered/unit/managers/test_hicache_controller_cp.py`
Steps:
1. Extend CP host metadata/state with a write dependency:
- final write event first,
- per-layer write event later.
2. Commit logical `HOST_PENDING` metadata at deterministic insert/attach time.
3. Change match/load predicates:
- pending host metadata can be selected for host hit,
- pending host metadata cannot be host-evicted,
- load waits on local pending write event before reading host data.
4. Convert `writing_check()` into local cleanup only:
- no radix-visible state mutation,
- no CP all-reduce,
- local source lifetime release after local event completion.
5. Add tests:
- pending host hit is selected consistently,
- host eviction skips pending node,
- load from pending node waits on dependency,
- target+draft pending states are atomic.
This removes `writing_check_min`.
### P6: Verification
Local CPU-only checks:
```bash
python3 -m py_compile python/sglang/srt/mem_cache/hiradix_cache.py python/sglang/srt/managers/cache_controller.py
python3 -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q
```
Remote CUDA/ETE checks only in the container:
```bash
ssh g0034 'docker exec sglang-glm5-dev-2 bash -lc "cd /sgl-workspace/sglang-tai && python3 -m py_compile python/sglang/srt/mem_cache/hiradix_cache.py python/sglang/srt/managers/cache_controller.py"'
```
ETE acceptance criteria:
- no `[HiCache-collective] tag=reserve_slots_max` in CP HiCache run,
- no `[HiCache-collective] tag=host_evict_done_min`,
- after P5, no `[HiCache-collective] tag=writing_check_min`,
- cache hit stays non-zero on repeated prompt,
- speculative accept length does not collapse after cache hit,
- no target/draft host allocator leak after abort/rollback,
- no local physical reserve failure after ledger admission.
## Risk Notes
1. Removing `reserve_slots_max` before deterministic eviction is unsafe.
2. Removing `writing_check_min` without pending-host dependencies is unsafe.
3. Host victim selection must not rely on wall-clock priority alone.
4. Draft/MTP must remain a shadow of target metadata; no independent draft radix
state should be introduced.
5. Any new collective introduced during this work must be explicitly reported and
justified before merging.
docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md
+82
View File
@@ -571,6 +571,88 @@ insertion key without splitting the tree, and only reserves the suffix beyond
`max(cache_protected_len, existing_insertion_prefix_len)`. This preserves the
cold-request behavior while avoiding duplicate prepared backups on exact repeats.
## 0SM Backup Direction: Direct Page-First-Direct
If the goal is to make backup overlap forward compute without stealing SMs, the
preferred direction is **not** to further optimize the `kernel + page_first`
backup kernels. Those paths are CUDA kernels by construction, including the
current `tai-kernel` per-layer LF->PF implementation, and therefore consume SM
resources.
The 0SM-oriented path should be:
```text
--hicache-io-backend direct
--hicache-mem-layout page_first_direct
```
In this mode, all-layer LF->PF backup already has a direct-copy implementation
through `sgl_kernel.kvcacheio.transfer_kv_all_layer_direct_lf_pf()`, which uses
`cudaMemcpyBatchAsync` when the CUDA runtime/driver supports it and falls back to
non-blocking copy slices otherwise. This is the copy-engine style path and is
the right base for 0SM backup.
The current gap is per-layer D2H backup. `sgl_kernel.kvcacheio` exposes
`transfer_kv_per_layer_direct_pf_lf()` for page-first-direct host-to-device load,
and `transfer_kv_all_layer_direct_lf_pf()` for all-layer device-to-host backup,
but it does not expose the per-layer device-to-host sibling:
```text
transfer_kv_per_layer_direct_lf_pf(
src_ptrs, # one target/draft device layer, or K/V layer pair
dst_ptrs, # page_first_direct host pool
src_indices, # local physical device token/page spans
dst_indices, # reserved host token/page spans
layer_id,
page_size,
)
```
That operator can reuse the existing internal page-first-direct direct-copy
implementation with `IsLf2Pf=true` and `start_layer_id=layer_id`. It should
require page-aligned contiguous spans and fail fast if the layout/backend does
not match. After this operator exists, the current Python-loop fallbacks in
`memory_pool_host.py` for `direct + page_first_direct` per-layer backup should
route to it for:
- MLA/NSA target KV,
- EAGLE/MTP draft KV,
- NSA indexer state.
### Backend split status
The current controller has one `io_backend` field, and uses it for both load and
backup:
```text
start_loading() -> load_to_device_per_layer(..., self.io_backend)
start_writing() -> backup_from_device_all_layer(..., self.io_backend)
submit_write_cp_layer() -> backup_from_device_per_layer(..., self.io_backend)
```
It also uses the same backend in `move_indices()` to decide whether indices live
on CUDA tensors (`kernel`) or CPU tensors (`direct`). Therefore, the current
runtime supports `direct` as a whole HiCache CPU/GPU I/O mode, but does not
support a pure configuration knob for "backup direct, load kernel".
Splitting load and backup backends is possible, but should be treated as a
second step rather than the first 0SM backup step. The reason is layout
compatibility: `direct` backup wants `page_first_direct`, while the existing
`kernel` load path supports `page_first`, not `page_first_direct`. A split mode
would therefore also need either a `kernel + page_first_direct` load path or a
separate layout conversion path, which adds more correctness surface than first
making full direct mode work well.
Recommended sequence:
1. Implement `transfer_kv_per_layer_direct_lf_pf()` in `sgl-kernel`.
2. Route `direct + page_first_direct` per-layer target/draft/indexer backup to
that operator.
3. Validate full direct mode for CP shared KV + HiCache + MTP/EAGLE.
4. Only if direct load is a bottleneck, introduce separate
`backup_io_backend`/`load_io_backend` fields and the extra layout support
required by a mixed direct-backup/kernel-load mode.
## Summary
Per-layer backup should be implemented as a data-plane refinement, not a radix