diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md new file mode 100644 index 000000000..06219758a --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_hicache_reactive_host_free_room_plan.md @@ -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. diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 785d987a0..28f679bf8 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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. diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e340f971e..1e3e65a88 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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") diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index b4565abda..8f8bc9b47 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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