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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user