Overlap CP HiCache backup without exposing partial host state

CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss.

Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary.

Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary.

Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile.

Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work.

Confidence: medium

Scope-risk: moderate

Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven.

Tested: local git diff --check

Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files

Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings)

Not-tested: full GLM5 ETE server rerun after this commit
This commit is contained in:
laoyao0822
2026-05-27 09:50:47 +08:00
parent 03529319a1
commit f355fdd39e
11 changed files with 747 additions and 57 deletions

View File

@@ -226,6 +226,18 @@ all_reduce per node victim in a tight host-eviction loop
all_reduce every scheduler tick when no completion prefix advanced
```
Current correctness note: CP host reservation now synchronizes
`required_host_slots` with `all_reduce(MAX)` before the host-eviction retry.
This forces every rank into the same reserve/evict/retry branch and avoids
collective mismatches when one rank is host-full and another rank reserves
successfully. It is intentionally a coarse slow-path collective, not a
per-layer collective, but it can become a performance cost when host pressure
is frequent because every reservation failure pays at least one rank-wide MAX
sync and retry failures pay a second one. Treat this as a correctness guard to
be amortized later with batched reservation epochs, deterministic host
watermarks, or less frequent proactive host eviction; do not move it into the
per-layer data path.
## Per-Layer Backup Data Plane
The existing host pool exposes `load_to_device_per_layer()` but not a matching
@@ -415,6 +427,13 @@ 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.
The current first implementation uses rank-wide
`MAX(required_host_slots)` synchronization before eviction/retry so all CP ranks
observe the same capacity-pressure branch. This fixes correctness under
rank-skewed host pressure, but it is a known performance risk if host-full
becomes common. Follow-up optimization should reduce how often this slow path
is reached rather than adding more collectives around it.
### Phase 7: Split safety
Implement defer/requeue-on-pending-split for nodes with pending backup. Add
@@ -436,10 +455,14 @@ paths:
## 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.
1. **Layer completion hook.** The chosen semantic hook is explicit model
layer-end after that layer's target/draft KV and NSA indexer writes have
been enqueued. Store-side notifiers are intentionally not the ownership
boundary because MLA, NSA indexer, TAI fused store, Triton store, and
zero-local paths can bypass one another. The first implementation hooks
`DeepseekV2DecoderLayer.forward()` and the TBO `op_comm_postprocess_layer()`
path; TBO children only notify from subbatch 1 so a layer is not backed up
after only the first child has written its KV.
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`.
@@ -466,9 +489,8 @@ Implemented in the first P4-P8 pass:
- 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.
- KV pools expose a no-op-safe explicit layer-end notifier. Store methods and
fused store kernels do not drive backup progress.
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
@@ -503,14 +525,16 @@ This covers the GLM5/NSA production shape:
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:
The first production forward-overlap path reserves host slots after
`prepare_for_extend()` and before the model forward, then calls
`submit_write_cp_per_layer(..., catch_up_all_layers=False)`. Radix insertion
only attaches that prepared reservation; it does not replay an all-layer
catch-up copy. The model layer-end hook is the forward progress driver:
```text
submit_write_cp_per_layer(..., catch_up_all_layers=False)
-> on_layer_kv_stored(layer_id, source="target" | "draft")
-> DeepseekV2 layer-end calls token_to_kv_pool.notify_layer_end_for_backup(layer_id)
-> HiCacheController.on_layer_end(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
```
@@ -521,6 +545,32 @@ payloads remain invisible until the final write ack commits
separately so an early target hook cannot copy draft KV before draft has stored
the same layer.
The previous store-notifier attempt was rejected as the primary correctness
boundary. It was too easy for fused paths to skip it: NSA fused indexer store
bypassed `set_index_k_scale_buffer()`, MLA can use `SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1`,
and zero-local paths do not necessarily touch the same setter. Missing one
store notifier left the prepared backup in `pending_layer_writes` forever: the
first request could finish prefill without a final write ack, and a second
identical request could hit the pending node and wedge the scheduler. The
layer-end hook makes backup progress independent of the chosen store kernel;
store-side hooks should remain absent or debug-only.
One more correctness/performance edge was observed with repeated EAGLE requests.
Scheduler prefix matching and final radix insertion can have different effective
lengths because matching is based on the schedulable prefix while insertion uses
the final page-aligned key. With EAGLE bigram keys this can leave
`req.cache_protected_len` shorter than the prefix that insertion will find in the
tree. If early backup reserves `[cache_protected_len, insertion_key_len)` blindly,
the system copies a duplicate page, then final insertion discovers that no new
node was created and rolls the prepared backup back with `reason=insert_miss`.
This is safe but wastes host reservation, per-layer callbacks, and write-stream
work on repeated prompts.
The early backup path now probes the existing radix prefix for the final
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.
## Summary
Per-layer backup should be implemented as a data-plane refinement, not a radix