Prevent unattached CP HiCache write acks from crashing chunked prefill
Prepared per-layer CP HiCache backups can enqueue their final write ack before the radix node is attached. Chunked prefill exposed this window when a separate catch-up backup made writing_check drain the ack queue and pop an unattached node id. Keep ready but unattached prepared acks queued until insert either attaches the prepared backup or rollback removes the orphan ack. Also document the reactive host free-room eviction plan separately from this state-machine fix. Constraint: CP HiCache prepared backup transfer can complete before radix insertion attaches node state Rejected: Drop unknown ack ids | would orphan a later successful prepared attach and leak write state Rejected: Chunked-only guard | the invalid assumption is in the generic CP write ack state machine Confidence: high Scope-risk: narrow Directive: Do not drain CP write acks unless every ack id is registered in pending_host_backups or ongoing_write_through Tested: Remote red-green test_cp_hicache_metadata.py::TestHiRadixCacheCPBackup::test_writing_check_defers_unattached_prepared_ack Tested: Remote PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py (116 passed) Tested: python -m py_compile python/sglang/srt/mem_cache/hiradix_cache.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py Not-tested: Full chunked-prefill ETE replay after this commit Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
@@ -0,0 +1,205 @@
|
||||
# NSA Prefill CP HiCache Reactive Host Free-Room Plan
|
||||
|
||||
Date: 2026-06-02
|
||||
|
||||
## Goal
|
||||
|
||||
Keep CP HiCache host/L2 from running at the exact-full boundary. When a host
|
||||
write reservation discovers that capacity is insufficient, evict enough host
|
||||
cache so the reservation can complete and a configured free-room watermark
|
||||
remains after the reservation.
|
||||
|
||||
This is intentionally **reactive**. It must not resurrect the old proactive
|
||||
per-scheduler-tick host eviction loop.
|
||||
|
||||
## Current eviction behavior
|
||||
|
||||
### L1/device cache eviction
|
||||
|
||||
Entry points:
|
||||
|
||||
- `python/sglang/srt/mem_cache/common.py::evict_from_tree_cache`
|
||||
- `python/sglang/srt/mem_cache/common.py::_evict_for_compute_owner_lanes`
|
||||
- `python/sglang/srt/mem_cache/hiradix_cache.py::HiRadixCache.evict`
|
||||
|
||||
Current behavior:
|
||||
|
||||
- Standard device allocator path checks `available < num_tokens`.
|
||||
- If insufficient, it calls `tree_cache.evict(EvictParams(num_tokens=num_tokens))`.
|
||||
- CP owner-lane path computes physical owner-lane page deficits and evicts
|
||||
`deficit_pages * page_size`, without the old `* cp_size` over-eviction.
|
||||
- `HiRadixCache.evict` picks victims from `evictable_leaves` by priority:
|
||||
- host-backed nodes are demoted by freeing device slots only;
|
||||
- unbacked nodes are either backed up first under `write_back` or deleted;
|
||||
- eviction stops after `num_evicted >= params.num_tokens`.
|
||||
|
||||
There is no stable L1 free-room watermark in the current branch.
|
||||
|
||||
### L2/host cache eviction
|
||||
|
||||
Entry points:
|
||||
|
||||
- `python/sglang/srt/mem_cache/hiradix_cache.py::HiRadixCache.evict_host`
|
||||
- `python/sglang/srt/mem_cache/hiradix_cache.py::_evict_host_for_physical_slots`
|
||||
|
||||
Current behavior:
|
||||
|
||||
- CP HiCache host eviction frees host-only, backed, unpinned, unreferenced host
|
||||
leaves.
|
||||
- CP target and draft host slots are freed together through
|
||||
`cache_controller.evict_cp_host(...)`.
|
||||
- Eviction stops after the exact requested physical slot count is freed.
|
||||
|
||||
There is no host/L2 free-room watermark in the current branch.
|
||||
|
||||
### CP write-admission host reservation
|
||||
|
||||
Entry points:
|
||||
|
||||
- `python/sglang/srt/mem_cache/hiradix_cache.py::_cp_build_write_admission`
|
||||
- `python/sglang/srt/mem_cache/hiradix_cache.py::_plan_cp_host_evictions`
|
||||
- `python/sglang/srt/mem_cache/hiradix_cache.py::_reserve_write_cp_indices_no_collective`
|
||||
|
||||
Current behavior:
|
||||
|
||||
```text
|
||||
required_by_owner = required host tokens/pages for this write
|
||||
target_deficit = max(0, required_by_owner - target_available_by_owner)
|
||||
draft_deficit = max(0, required_by_owner - draft_available_by_owner)
|
||||
deficit = max(target_deficit, draft_deficit)
|
||||
```
|
||||
|
||||
The eviction plan only covers the exact per-owner deficit. After a successful
|
||||
reservation, the affected owner lanes can be left near zero free host slots,
|
||||
which makes later writes repeatedly enter the eviction path.
|
||||
|
||||
## Historical implementation to avoid
|
||||
|
||||
The old free-room work around commit `579fe3d41` maintained host room with a
|
||||
proactive "Round 2.5" scheduler-tick eviction path. That direction is not the
|
||||
desired fix here because it can:
|
||||
|
||||
- scan and evict on every scheduler tick;
|
||||
- introduce hot-path synchronization pressure;
|
||||
- evict repeatedly under sustained host pressure even if the next reservation
|
||||
would still fit;
|
||||
- recreate the earlier performance-regression shape.
|
||||
|
||||
The new design should not add a per-tick maintenance loop.
|
||||
|
||||
## Proposed design
|
||||
|
||||
Add a **reactive host free-room watermark** to CP HiCache write admission.
|
||||
|
||||
When a reservation would fit, do nothing:
|
||||
|
||||
```text
|
||||
if required <= available:
|
||||
deficit = 0
|
||||
```
|
||||
|
||||
When a reservation would not fit, evict enough for the write and for post-write
|
||||
free room:
|
||||
|
||||
```text
|
||||
if required > available:
|
||||
deficit = required + free_room_target - available
|
||||
```
|
||||
|
||||
Because eviction happens before reservation, reserving `required` after this
|
||||
eviction should leave approximately `free_room_target` host slots available on
|
||||
the affected owner lanes.
|
||||
|
||||
## Scope for the first implementation
|
||||
|
||||
In scope:
|
||||
|
||||
1. CP HiCache host/L2 write-admission path.
|
||||
2. Per CP owner-lane free-room target.
|
||||
3. Target and draft host pools must both satisfy the same free-room contract.
|
||||
4. Explicit configuration for the free-room ratio.
|
||||
5. Unit tests for admission planning and target/draft symmetry.
|
||||
|
||||
Out of scope:
|
||||
|
||||
1. L1/device free-room changes.
|
||||
2. Load-back owner-lane device eviction changes.
|
||||
3. Per-tick proactive host eviction.
|
||||
4. Allocator contiguity changes.
|
||||
5. Evict victim merge/sort policy changes.
|
||||
6. New collectives/all-reduce.
|
||||
|
||||
## Configuration
|
||||
|
||||
Preferred first version:
|
||||
|
||||
- Add `--hicache-host-free-room-ratio`.
|
||||
- Default: `0.0` to preserve current behavior.
|
||||
- Remote experiments can explicitly set a small value such as `0.03` or `0.05`.
|
||||
|
||||
The target should be rounded to page units so CP host accounting remains
|
||||
page-granular.
|
||||
|
||||
## Implementation notes
|
||||
|
||||
Likely code points:
|
||||
|
||||
1. Add `hicache_host_free_room_ratio` to `ServerArgs`.
|
||||
2. Store the ratio in `HiRadixCache`.
|
||||
3. Add helper to compute per-owner host free-room targets from
|
||||
`_cp_host_capacity_snapshot()`.
|
||||
4. Change `_cp_build_write_admission()` so host deficits become reactive
|
||||
low-watermark deficits instead of exact deficits.
|
||||
5. Keep `_plan_cp_host_evictions()` unchanged initially: it already accepts a
|
||||
per-owner deficit vector and picks deterministic host victims.
|
||||
6. Keep `_reserve_write_cp_indices_no_collective()` unchanged except for using
|
||||
the updated admission result.
|
||||
|
||||
Important detail:
|
||||
|
||||
The free-room target should only be applied when that pool/lane is already
|
||||
short for the current write. Do not evict merely because `available` is below
|
||||
`free_room_target` if the current reservation still fits.
|
||||
|
||||
## Expected effect
|
||||
|
||||
Compared with exact-deficit eviction:
|
||||
|
||||
- fewer repeated host evictions when host memory is near full;
|
||||
- less scheduler/control-path churn under sustained write pressure;
|
||||
- better chance for later reservations to complete without immediate eviction;
|
||||
- no extra collective cost.
|
||||
|
||||
Compared with old proactive free-room maintenance:
|
||||
|
||||
- lower steady-state overhead;
|
||||
- eviction work only happens on real reservation pressure;
|
||||
- no recurring scheduler-tick eviction loop.
|
||||
|
||||
## Test plan
|
||||
|
||||
Unit tests:
|
||||
|
||||
1. `ratio=0.0` preserves exact-deficit behavior.
|
||||
2. If `required <= available`, no free-room eviction is planned.
|
||||
3. If `required > available`, planned deficit includes `required + room - available`.
|
||||
4. Target/draft symmetry: deficit uses the max of target and draft low-watermark
|
||||
requirements.
|
||||
5. Rounding: free-room target is page-aligned.
|
||||
|
||||
Runtime validation:
|
||||
|
||||
1. Start with default ratio `0.0` to confirm no behavior change.
|
||||
2. Run remote ETE with `--hicache-host-free-room-ratio 0.03`.
|
||||
3. Check host eviction frequency and host free-slot floor in logs.
|
||||
4. Verify no new CP HiCache collective logs are introduced.
|
||||
|
||||
## Risks
|
||||
|
||||
1. Too large a ratio can evict useful host cache and lower hit rate.
|
||||
2. Evicting large nodes may overshoot the target; this is acceptable initially
|
||||
because host eviction is node/page based.
|
||||
3. If host pressure is caused by pinned or in-flight backup nodes, free-room
|
||||
eviction may still fail; failures must stay explicit and warning-level.
|
||||
4. The first version does not improve physical page contiguity; it only reduces
|
||||
eviction frequency.
|
||||
@@ -5036,3 +5036,55 @@ Required correction:
|
||||
- Make stale-tail split/prune atomic from the radix tree perspective: preflight pending/unprunable state before `_split_node()` and drain completed write acks once on the conflict path.
|
||||
- If the stale tail is still pending, do not mutate the tree and do not crash the scheduler. Defer/skip this cache insertion attempt while keeping request KV ownership intact for transfer/release.
|
||||
- Do not delete the stale tail or split an in-flight backup node. Per-layer backup ownership remains node/page based and cannot be repartitioned mid-flight.
|
||||
|
||||
### C119 — 2026-06-02 Prepared per-layer ack can arrive before radix attach under chunked prefill
|
||||
|
||||
Finding:
|
||||
|
||||
- Remote crash in `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260601_205025.log`
|
||||
was `KeyError: 54` in `HiRadixCache.writing_check()` rather than the earlier
|
||||
direct `HiCachePendingBackupSplit` exception.
|
||||
- Sequence observed in the log:
|
||||
1. A chunked request skipped pre-forward backup with
|
||||
`[CP_HICACHE_FALLBACK][prepare_write_backup_skipped] reason=chunked_req`.
|
||||
2. A later chunk for the same request registered prepared per-layer backup for
|
||||
`node_id=54` and enqueued the final ack.
|
||||
3. A separate catch-up all-layer backup for `node_id=53` made
|
||||
`ongoing_write_through` non-empty.
|
||||
4. The stale-tail pending-split path called
|
||||
`_cp_drain_write_acks_before_pending_split() -> writing_check()`.
|
||||
5. `writing_check()` tried to pop ack id `54`, but `54` was not yet in either
|
||||
`pending_host_backups` or `ongoing_write_through` because the prepared
|
||||
backup had not been attached to a radix node yet.
|
||||
|
||||
Root cause:
|
||||
|
||||
- The old `writing_check()` contract assumed every ready ack in
|
||||
`ack_write_queue` already had corresponding radix state in
|
||||
`pending_host_backups` or `ongoing_write_through`.
|
||||
- That assumption is false for prepared per-layer backups: the transfer can
|
||||
complete and enqueue its ack before `insert()` attaches the prepared backup to
|
||||
the newly inserted radix node.
|
||||
- The bug is exposed by chunked prefill because chunked/catch-up writes can make
|
||||
`ongoing_write_through` non-empty while an unattached prepared ack is at the
|
||||
head of the queue.
|
||||
|
||||
Correction:
|
||||
|
||||
- `writing_check()` must only drain a ready ack entry when all `ack.node_ids` are
|
||||
registered in radix write state (`pending_host_backups` or
|
||||
`ongoing_write_through`).
|
||||
- If a ready ack contains an unattached prepared node id, keep the ack queued and
|
||||
stop this drain pass. Later `insert()` will either attach the prepared backup,
|
||||
making the ack committable, or rollback the unattached prepared backup and
|
||||
remove the orphan ack.
|
||||
- Do not silently drop unknown ack ids; dropping would make a later attach wait
|
||||
forever for an ack that no longer exists.
|
||||
|
||||
Regression test:
|
||||
|
||||
- `test_cp_hicache_metadata.py::TestHiRadixCacheCPBackup::test_writing_check_defers_unattached_prepared_ack`
|
||||
constructs an ack queue with an unattached prepared ack `[54]` ahead of an
|
||||
attached catch-up ack `[53]`. `writing_check()` must not raise `KeyError`, must
|
||||
preserve both ack entries, and must not release the attached node behind the
|
||||
unattached head entry.
|
||||
|
||||
@@ -2583,10 +2583,31 @@ class HiRadixCache(RadixCache):
|
||||
if ack_queue_len == 0:
|
||||
return
|
||||
|
||||
def ack_registered_for_radix_state(ack_id: int) -> bool:
|
||||
return (
|
||||
ack_id in self.pending_host_backups
|
||||
or ack_id in self.ongoing_write_through
|
||||
)
|
||||
|
||||
finish_count = 0
|
||||
for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
|
||||
if not finish_event.query():
|
||||
break
|
||||
if not all(ack_registered_for_radix_state(ack_id) for ack_id in ack_list):
|
||||
# Prepared per-layer CP backups can finish and enqueue their final
|
||||
# ack before the corresponding radix node is attached. A later
|
||||
# catch-up/write-back node may make ongoing_write_through non-empty
|
||||
# and enter writing_check while that prepared ack is still
|
||||
# unattached. Keep the ack queued until insert either attaches the
|
||||
# prepared backup or rolls it back.
|
||||
logger.debug(
|
||||
"[HiCache-write] writing_check defers unattached ack ids: "
|
||||
"ack_ids=%s ongoing=%s pending=%s",
|
||||
ack_list,
|
||||
list(self.ongoing_write_through.keys()),
|
||||
list(self.pending_host_backups.keys()),
|
||||
)
|
||||
break
|
||||
finish_count += 1
|
||||
local_finish_count = finish_count
|
||||
queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
|
||||
|
||||
@@ -1778,6 +1778,41 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
self.assertEqual(cache.cache_controller.ack_write_queue, [])
|
||||
self.assertEqual(evicted, [reservation.metadata])
|
||||
|
||||
def test_writing_check_defers_unattached_prepared_ack(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache._uses_cp_hicache = True
|
||||
cache.tp_world_size = 1
|
||||
cache.enable_storage = False
|
||||
cache.pending_host_backups = {}
|
||||
|
||||
class ReadyEvent:
|
||||
def query(self):
|
||||
return True
|
||||
|
||||
def synchronize(self):
|
||||
pass
|
||||
|
||||
attached_node = TreeNode()
|
||||
attached_node.id = 53
|
||||
dec_locked = []
|
||||
cache.ongoing_write_through = {53: attached_node}
|
||||
cache.dec_node_lock_ref = lambda node: dec_locked.append(node.id)
|
||||
cache.cache_controller = types.SimpleNamespace(
|
||||
ack_write_queue=[
|
||||
HiCacheAck(ReadyEvent(), ReadyEvent(), [54]),
|
||||
HiCacheAck(ReadyEvent(), ReadyEvent(), [53]),
|
||||
],
|
||||
)
|
||||
|
||||
cache.writing_check()
|
||||
|
||||
self.assertEqual(
|
||||
[ack.node_ids for ack in cache.cache_controller.ack_write_queue],
|
||||
[[54], [53]],
|
||||
)
|
||||
self.assertEqual(cache.ongoing_write_through, {53: attached_node})
|
||||
self.assertEqual(dec_locked, [])
|
||||
|
||||
def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache._uses_cp_hicache = True
|
||||
|
||||
Reference in New Issue
Block a user