diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md new file mode 100644 index 000000000..b933d7f3f --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md @@ -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. diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md index 5652fbfbd..4682f75f3 100644 --- a/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md @@ -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 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3a7f612fa..d29199e73 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -215,6 +215,7 @@ class Envs: SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1) SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False) SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False) + SGLANG_CP_HICACHE_CAPACITY_DEBUG = EnvBool(False) SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 69c00c05f..7f79b33d5 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -38,8 +38,8 @@ def _prefetch_log(message: str, *args) -> None: def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None: - logger.info( - "CP shared KV index prefetch fallback (%s): " + message, + logger.warning( + "[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message, reason, *args, ) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index d98b1d45a..36d506d5e 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -123,8 +123,8 @@ def _log_cp_shared_kv_index_prefetch_fallback( dedupe here or later real fallbacks become invisible during profiling. """ - logger.info( - "CP shared KV index prefetch fallback (%s): " + message, + logger.warning( + "[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message, reason, *args, ) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index d496fc85d..76e8c89e8 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -58,8 +58,8 @@ def log_cp_shared_kv_direct_write_fallback( Warmup can hit the same fallback reason as a later real request, so de-duplicating by reason hides correctness/performance issues after startup. """ - logger.info( - "CP shared KV direct-write fallback (%s): " + message, + logger.warning( + "[CP_SHARED_KV_FALLBACK][direct_write] reason=%s " + message, reason, *args, ) @@ -821,8 +821,8 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor( def _log_tai_in_seq_rerange_fallback(reason: str, message: str, *args) -> None: - logger.info( - "CP NSA in-seq all-gather rerange tai-kernel fallback (%s): " + message, + logger.warning( + "[CP_SHARED_KV_FALLBACK][tai_in_seq_rerange] reason=%s " + message, reason, *args, ) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 440bebd1b..c39864d5a 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -1026,7 +1026,9 @@ class HiCacheController: def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None: logger.warning( - "[CacheCtrl-write] CP HiCache all-layer backup fallback: node_id=%d logical_len=%d owned=%d physical=%d draft=%s", + "[CP_HICACHE_FALLBACK][submit_write_cp_all_layer] " + "CP HiCache all-layer backup fallback. " + "node_id=%d logical_len=%d owned=%d physical=%d draft=%s", reservation.node_id, reservation.metadata.logical_len, reservation.metadata.owned_positions.numel(), @@ -1250,6 +1252,16 @@ class HiCacheController: ) return + logger.warning( + "[CP_HICACHE_FALLBACK][catch_up_all_layers] " + "CP HiCache write is submitting all layers after KV materialization; " + "forward-overlap per-layer backup was not active. " + "node_id=%d logical_len=%d layers=%d draft=%s", + reservation.node_id, + reservation.metadata.logical_len, + state.total_layers, + reservation.draft_host_indices is not None, + ) total_layers = state.total_layers for layer_id in range(total_layers): self.submit_write_cp_layer(reservation, layer_id) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 9984c961c..847ce4f89 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2457,11 +2457,6 @@ class Scheduler( ) new_batch.prepare_for_extend() - if self.enable_hierarchical_cache and hasattr( - self.tree_cache, "prepare_write_backup_for_req" - ): - for req in can_run_list: - self.tree_cache.prepare_write_backup_for_req(req) # Record prefill stats for logging after forward new_batch.prefill_stats = PrefillStats.from_adder( @@ -2579,6 +2574,38 @@ class Scheduler( self.batch_record_ct = (self.batch_record_ct + 1) % 2 self.batch_record_buf[self.batch_record_ct] = model_worker_batch + def _prepare_hicache_write_backups_before_forward( + self, batch: ScheduleBatch + ) -> None: + """Reserve/register CP HiCache write-through backups immediately before forward. + + CP shared-KV HiCache per-layer D2H backup must be registered after + `prepare_for_extend()` has allocated output KV slots, but before the + model starts emitting per-layer KV. Keeping this at the final + `run_batch` boundary prevents alternate prefill event loops from + bypassing the early-backup path and falling back to post-forward + catch-up copies. + """ + + if not self.enable_hierarchical_cache: + return + if batch.forward_mode not in ( + ForwardMode.EXTEND, + ForwardMode.SPLIT_PREFILL, + ForwardMode.DLLM_EXTEND, + ): + return + + prepare_fn = getattr(self.tree_cache, "prepare_write_backup_for_req", None) + if prepare_fn is None: + inner_cache = getattr(self.tree_cache, "inner", None) + prepare_fn = getattr(inner_cache, "prepare_write_backup_for_req", None) + if prepare_fn is None: + return + + for req in batch.reqs: + prepare_fn(req) + def run_batch( self, batch: ScheduleBatch, @@ -2601,6 +2628,8 @@ class Scheduler( if batch.forward_mode.is_prebuilt(): return self._run_batch_prebuilt(batch) + self._prepare_hicache_write_backups_before_forward(batch) + # Run forward if self.is_generation: if self.spec_algorithm.is_none() or self.enable_overlap: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 0629ea8ac..0ff09acac 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -257,6 +257,37 @@ class PreparedCpHiCacheBackup: attached: bool = False +@dataclass(frozen=True) +class CpHiCacheCapacitySnapshot: + target_capacity: Tuple[int, ...] + draft_capacity: Optional[Tuple[int, ...]] + committed_target: Tuple[int, ...] + committed_draft: Tuple[int, ...] + pending_target: Tuple[int, ...] + pending_draft: Tuple[int, ...] + + @property + def target_used(self) -> Tuple[int, ...]: + return tuple( + committed + pending + for committed, pending in zip(self.committed_target, self.pending_target) + ) + + @property + def draft_used(self) -> Tuple[int, ...]: + return tuple( + committed + pending + for committed, pending in zip(self.committed_draft, self.pending_draft) + ) + + +@dataclass(frozen=True) +class CpHiCacheEvictionPlan: + victims: Tuple[TreeNode, ...] + planned_freed: Tuple[int, ...] + remaining_deficit: Tuple[int, ...] + + class HiCachePendingBackupSplit(Exception): def __init__(self, node: TreeNode): self.node = node @@ -473,6 +504,14 @@ class HiRadixCache(RadixCache): # record the ongoing prefetch requests self.ongoing_prefetch = {} self.ongoing_backup = {} + # Temporary aggregate accounting for CP HiCache collectives. These + # calls are on latency-sensitive scheduler paths; keep logging coarse + # enough to avoid per-request spam while still exposing hot collectives. + self._cp_hicache_collective_stats = {} + self._cp_hicache_capacity_debug = ( + envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get() + ) + self._cp_hicache_capacity_debug_stats = {} # track per-request tokens loaded from storage (L3 hits) # key: request_id, value: number of tokens actually loaded from storage self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {} @@ -501,6 +540,441 @@ class HiRadixCache(RadixCache): super().__init__(params=params) + def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None: + start = time.perf_counter() + torch.distributed.all_reduce(tensor, op=op, group=group) + elapsed_ms = (time.perf_counter() - start) * 1000 + + if not hasattr(self, "_cp_hicache_collective_stats"): + self._cp_hicache_collective_stats = {} + stats = self._cp_hicache_collective_stats.setdefault( + tag, + { + "count": 0, + "total_ms": 0.0, + "max_ms": 0.0, + }, + ) + stats["count"] += 1 + stats["total_ms"] += elapsed_ms + stats["max_ms"] = max(stats["max_ms"], elapsed_ms) + + count = stats["count"] + if count == 1 or count % 128 == 0 or elapsed_ms >= 10.0: + logger.info( + "[HiCache-collective] tag=%s count=%d total_ms=%.3f avg_ms=%.3f max_ms=%.3f last_ms=%.3f", + tag, + count, + stats["total_ms"], + stats["total_ms"] / count, + stats["max_ms"], + elapsed_ms, + ) + + def _cp_hicache_cp_size(self, page_owners: Optional[torch.Tensor] = None) -> int: + layout = getattr( + getattr(self, "cache_controller", None), "cp_shared_kv_layout", None + ) + cp_size = getattr(layout, "cp_size", None) + if cp_size is not None: + return int(cp_size) + tp_world_size = int(getattr(self, "tp_world_size", 1) or 1) + if tp_world_size > 1: + return tp_world_size + if page_owners is not None and page_owners.numel() > 0: + non_negative = page_owners.to(device="cpu", dtype=torch.int64) + non_negative = non_negative[non_negative >= 0] + if non_negative.numel() > 0: + return int(non_negative.max().item()) + 1 + return 1 + + def _cp_hicache_cp_rank(self) -> int: + layout = getattr( + getattr(self, "cache_controller", None), "cp_shared_kv_layout", None + ) + return int(getattr(layout, "cp_rank", getattr(self, "_tp_group_rank", 0)) or 0) + + def _cp_zero_counts(self, cp_size: Optional[int] = None) -> Tuple[int, ...]: + cp_size = self._cp_hicache_cp_size() if cp_size is None else int(cp_size) + return tuple(0 for _ in range(cp_size)) + + @staticmethod + def _cp_add_counts(lhs: Tuple[int, ...], rhs: Tuple[int, ...]) -> Tuple[int, ...]: + if len(lhs) != len(rhs): + raise RuntimeError( + f"CP HiCache count length mismatch: {len(lhs)} vs {len(rhs)}" + ) + return tuple(int(a) + int(b) for a, b in zip(lhs, rhs)) + + @staticmethod + def _cp_sub_counts_floor_zero( + lhs: Tuple[int, ...], rhs: Tuple[int, ...] + ) -> Tuple[int, ...]: + if len(lhs) != len(rhs): + raise RuntimeError( + f"CP HiCache count length mismatch: {len(lhs)} vs {len(rhs)}" + ) + return tuple(max(0, int(a) - int(b)) for a, b in zip(lhs, rhs)) + + def _cp_owner_token_counts( + self, + page_owners: torch.Tensor, + *, + page_size: int, + cp_size: Optional[int] = None, + ) -> Tuple[int, ...]: + cp_size = ( + self._cp_hicache_cp_size(page_owners) + if cp_size is None + else int(cp_size) + ) + counts = [0 for _ in range(cp_size)] + owners = page_owners.to(device="cpu", dtype=torch.int64) + for owner in owners.tolist(): + owner = int(owner) + if owner < 0: + continue + if owner >= cp_size: + raise RuntimeError( + "CP HiCache metadata owner exceeds CP size: " + f"owner={owner} cp_size={cp_size}" + ) + counts[owner] += int(page_size) + return tuple(counts) + + def _cp_metadata_token_counts( + self, metadata: CpHiCacheNodeMetadata + ) -> Tuple[int, ...]: + return self._cp_owner_token_counts( + metadata.page_owners, + page_size=metadata.page_size, + ) + + def _cp_assert_metadata_counts( + self, metadata: CpHiCacheNodeMetadata, *, context: str + ) -> Tuple[int, ...]: + counts = self._cp_metadata_token_counts(metadata) + cp_rank = self._cp_hicache_cp_rank() + if cp_rank >= len(counts): + raise RuntimeError( + "CP HiCache owner lane mismatch: " + f"context={context} cp_rank={cp_rank} cp_size={len(counts)}" + ) + expected_local = counts[cp_rank] + actual_local = int(metadata.owned_positions.numel()) + if actual_local != expected_local: + raise RuntimeError( + "CP HiCache owner lane mismatch: " + f"context={context} cp_rank={cp_rank} expected_local={expected_local} " + f"owned_positions={actual_local}" + ) + if int(metadata.host_indices.numel()) != expected_local: + raise RuntimeError( + "CP HiCache host index owner lane mismatch: " + f"context={context} cp_rank={cp_rank} expected_local={expected_local} " + f"host_indices={metadata.host_indices.numel()}" + ) + draft_host_indices = getattr(metadata, "draft_host_indices", None) + if draft_host_indices is not None and int(draft_host_indices.numel()) != expected_local: + raise RuntimeError( + "CP HiCache draft host index owner lane mismatch: " + f"context={context} cp_rank={cp_rank} expected_local={expected_local} " + f"draft_host_indices={draft_host_indices.numel()}" + ) + return counts + + def _cp_walk_radix_nodes(self) -> List[TreeNode]: + root = getattr(self, "root_node", None) + if root is None: + return [] + nodes: List[TreeNode] = [] + stack = list(getattr(root, "children", {}).values()) + while stack: + node = stack.pop() + nodes.append(node) + stack.extend(getattr(node, "children", {}).values()) + return nodes + + def _cp_host_capacity_snapshot(self) -> CpHiCacheCapacitySnapshot: + cp_size = self._cp_hicache_cp_size() + target_capacity = tuple( + int(getattr(getattr(self, "token_to_kv_pool_host", None), "size", 0)) + for _ in range(cp_size) + ) + controller = getattr(self, "cache_controller", None) + draft_pool = getattr(controller, "draft_mem_pool_host", None) + has_draft = bool(getattr(controller, "has_draft_hicache", False)) + draft_capacity = ( + tuple(int(getattr(draft_pool, "size", 0)) for _ in range(cp_size)) + if has_draft + else None + ) + committed_target = self._cp_zero_counts(cp_size) + committed_draft = self._cp_zero_counts(cp_size) + pending_target = self._cp_zero_counts(cp_size) + pending_draft = self._cp_zero_counts(cp_size) + + pending = getattr(self, "pending_host_backups", {}) + pending_node_ids = set(pending.keys()) + for node in self._cp_walk_radix_nodes(): + metadata = getattr(node, "cp_hicache", None) + if metadata is None or getattr(node, "host_len", 0) <= 0: + continue + if getattr(node, "id", None) in pending_node_ids: + continue + counts = self._cp_assert_metadata_counts( + metadata, context=f"snapshot_committed node_id={getattr(node, 'id', '?')}" + ) + committed_target = self._cp_add_counts(committed_target, counts) + if getattr(metadata, "draft_host_indices", None) is not None: + committed_draft = self._cp_add_counts(committed_draft, counts) + + for node_id, pending_backup in pending.items(): + metadata = pending_backup.metadata + counts = self._cp_assert_metadata_counts( + metadata, context=f"snapshot_pending node_id={node_id}" + ) + pending_target = self._cp_add_counts(pending_target, counts) + if getattr(metadata, "draft_host_indices", None) is not None: + pending_draft = self._cp_add_counts(pending_draft, counts) + + return CpHiCacheCapacitySnapshot( + target_capacity=target_capacity, + draft_capacity=draft_capacity, + committed_target=committed_target, + committed_draft=committed_draft, + pending_target=pending_target, + pending_draft=pending_draft, + ) + + def _cp_required_host_tokens_by_rank( + self, device_indices: torch.Tensor + ) -> Tuple[int, ...]: + layout = getattr( + getattr(self, "cache_controller", None), "cp_shared_kv_layout", None + ) + if layout is None: + return self._cp_zero_counts() + logical_len = int(device_indices.numel()) + if logical_len == 0: + return self._cp_zero_counts(layout.cp_size) + if logical_len % self.page_size != 0: + raise RuntimeError( + "CP HiCache capacity planning requires page-aligned device indices: " + f"logical_len={logical_len} page_size={self.page_size}" + ) + logical_pages = torch.div( + device_indices[:: self.page_size], + self.page_size, + rounding_mode="floor", + ) + page_owners = layout.owner_for_logical_pages(logical_pages).to( + dtype=torch.int8, device="cpu" + ) + return self._cp_owner_token_counts( + page_owners, + page_size=self.page_size, + cp_size=layout.cp_size, + ) + + def _cp_host_available_tokens_by_rank( + self, snapshot: CpHiCacheCapacitySnapshot, *, draft: bool = False + ) -> Tuple[int, ...]: + if draft: + if snapshot.draft_capacity is None: + return tuple(2**62 for _ in snapshot.target_capacity) + return self._cp_sub_counts_floor_zero( + snapshot.draft_capacity, snapshot.draft_used + ) + return self._cp_sub_counts_floor_zero( + snapshot.target_capacity, snapshot.target_used + ) + + def _cp_host_evict_key(self, node: TreeNode): + return ( + int(getattr(node, "priority", 0)), + int(getattr(node, "id", 0)), + ) + + def _cp_host_leaf_is_plannable_victim(self, node: TreeNode) -> bool: + if node == getattr(self, "root_node", None): + return False + if not getattr(node, "evicted", False): + return False + parent = getattr(node, "parent", None) + if parent is None: + return False + try: + parent_key = self.get_child_key_fn(node.key) + except Exception: + return False + if getattr(parent, "children", {}).get(parent_key) is not node: + return False + pin_expiry = float(getattr(node, "pin_expiry", 0.0) or 0.0) + pin_active = pin_expiry > 0 and time.monotonic() <= pin_expiry + if pin_active: + return False + effective_host_refs = int(getattr(node, "host_ref_counter", 0) or 0) + if pin_expiry > 0: + # Current eviction clears an expired pin before checking host refs. + # Mirror that decision without mutating the tree in observer mode. + effective_host_refs = max(0, effective_host_refs - 1) + if effective_host_refs > 0: + return False + if self._node_host_write_pending(node): + return False + try: + return self._node_backuped(node) + except RuntimeError: + return False + + def _plan_cp_host_evictions( + self, required_tokens_by_rank: Tuple[int, ...] + ) -> CpHiCacheEvictionPlan: + deficits = [max(0, int(v)) for v in required_tokens_by_rank] + cp_size = len(deficits) + planned_freed = [0 for _ in range(cp_size)] + victims: List[TreeNode] = [] + if all(v <= 0 for v in deficits): + return CpHiCacheEvictionPlan((), tuple(planned_freed), tuple(deficits)) + + leaves = sorted( + list(getattr(self, "evictable_host_leaves", set())), + key=self._cp_host_evict_key, + ) + for node in leaves: + if not self._cp_host_leaf_is_plannable_victim(node): + continue + metadata = getattr(node, "cp_hicache", None) + if metadata is None: + continue + counts = self._cp_assert_metadata_counts( + metadata, context=f"eviction_plan node_id={getattr(node, 'id', '?')}" + ) + if len(counts) != cp_size: + raise RuntimeError( + "CP HiCache eviction plan CP size mismatch: " + f"required={cp_size} metadata={len(counts)}" + ) + victims.append(node) + for rank, count in enumerate(counts): + planned_freed[rank] += int(count) + deficits[rank] = max(0, deficits[rank] - int(count)) + if all(v <= 0 for v in deficits): + break + + return CpHiCacheEvictionPlan( + victims=tuple(victims), + planned_freed=tuple(planned_freed), + remaining_deficit=tuple(deficits), + ) + + def _cp_capacity_debug_enabled(self) -> bool: + return bool( + getattr(self, "_cp_hicache_capacity_debug", False) + or envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get() + ) + + def _cp_build_write_admission_plan( + self, device_indices: torch.Tensor, *, node_id: int, phase: str + ) -> dict: + required = self._cp_required_host_tokens_by_rank(device_indices) + snapshot = self._cp_host_capacity_snapshot() + target_available = self._cp_host_available_tokens_by_rank(snapshot) + draft_available = self._cp_host_available_tokens_by_rank(snapshot, draft=True) + target_deficit = tuple( + max(0, req - avail) for req, avail in zip(required, target_available) + ) + draft_deficit = tuple( + max(0, req - avail) for req, avail in zip(required, draft_available) + ) + deficit = tuple(max(a, b) for a, b in zip(target_deficit, draft_deficit)) + current_compatible_need = tuple( + req if missing > 0 else 0 for req, missing in zip(required, deficit) + ) + eviction_plan = self._plan_cp_host_evictions(deficit) + return { + "node_id": node_id, + "phase": phase, + "required": required, + "target_available": target_available, + "draft_available": draft_available, + "deficit": deficit, + "current_compatible_need": current_compatible_need, + "eviction_plan": eviction_plan, + } + + def _cp_debug_build_write_admission_plan( + self, device_indices: torch.Tensor, *, node_id: int, phase: str + ) -> Optional[dict]: + if not self._cp_capacity_debug_enabled(): + return None + try: + return self._cp_build_write_admission_plan( + device_indices, node_id=node_id, phase=phase + ) + except Exception as exc: + logger.warning( + "[HiCache-capacity-debug] failed to build write admission plan: node_id=%d phase=%s error=%s", + node_id, + phase, + exc, + ) + return None + + def _cp_debug_compare_write_admission( + self, + plan: Optional[dict], + *, + result, + planned_required_slots: int, + ) -> None: + if plan is None or not self._cp_capacity_debug_enabled(): + return + stats = getattr(self, "_cp_hicache_capacity_debug_stats", None) + if stats is None: + stats = self._cp_hicache_capacity_debug_stats = {} + tag = "write_admission" + entry = stats.setdefault(tag, {"count": 0, "mismatch": 0}) + entry["count"] += 1 + + planned_required_slots = max(plan["current_compatible_need"], default=0) + planned_needs_eviction = planned_required_slots > 0 + reservation_failed = isinstance(result, HiCacheWriteFailure) + mismatch = ( + reservation_failed + and not planned_needs_eviction + and all(v <= 0 for v in plan["eviction_plan"].remaining_deficit) + ) + if mismatch: + entry["mismatch"] += 1 + should_log = mismatch or entry["count"] == 1 or entry["count"] % 128 == 0 + if not should_log: + return + + eviction_plan = plan["eviction_plan"] + log_fn = logger.warning if mismatch else logger.info + log_fn( + "[HiCache-capacity-debug] write_admission_compare node_id=%d phase=%s " + "count=%d mismatch=%d planned_required=%d local_failed=%s " + "planned_current_required=%d required=%s target_avail=%s draft_avail=%s " + "deficit=%s planned_freed=%s remaining_deficit=%s victims=%s", + plan["node_id"], + plan["phase"], + entry["count"], + entry["mismatch"], + int(planned_required_slots), + isinstance(result, HiCacheWriteFailure), + int(planned_required_slots), + plan["required"], + plan["target_available"], + plan["draft_available"], + plan["deficit"], + eviction_plan.planned_freed, + eviction_plan.remaining_deficit, + [getattr(node, "id", None) for node in eviction_plan.victims], + ) + def shutdown(self): """Best-effort auto-detach of storage backend on process shutdown. @@ -1021,84 +1495,130 @@ class HiRadixCache(RadixCache): self.dec_node_lock_ref(node) return node - def _sync_cp_write_required_host_slots(self, result) -> int: - """Return the max host slots any CP rank needs before write backup. + def _evict_cp_host_plan( + self, plan: dict, *, node_id: int, phase: str + ) -> bool: + eviction_plan = plan["eviction_plan"] + if any(v > 0 for v in eviction_plan.remaining_deficit): + logger.warning( + "[CP_HICACHE_FALLBACK][cp_host_reservation_plan_insufficient] " + "reason=insufficient_evictable_host_slots node_id=%d phase=%s " + "required=%s target_avail=%s draft_avail=%s deficit=%s " + "planned_freed=%s remaining_deficit=%s victims=%s", + node_id, + phase, + plan["required"], + plan["target_available"], + plan["draft_available"], + plan["deficit"], + eviction_plan.planned_freed, + eviction_plan.remaining_deficit, + [getattr(node, "id", None) for node in eviction_plan.victims], + ) + return False - CP ranks must make the reserve/evict/retry decision collectively. A - local successful reservation is not enough: another CP rank may be full - and enter host eviction, whose implementation contains collective - all_reduce calls. If successful ranks skip that branch they can instead - enter unrelated collectives (for example disagg polling), causing a - process-group mismatch. - """ + if len(eviction_plan.victims) == 0: + return True - required_host_slots = ( - int(result.required_host_slots) - if isinstance(result, HiCacheWriteFailure) - else 0 - ) - if ( - getattr(self, "tp_group", None) is None - or getattr(self, "tp_world_size", 1) <= 1 - ): - return required_host_slots + local_freed = 0 + victim_ids = [] + for victim in eviction_plan.victims: + victim_id = getattr(victim, "id", None) + if not self._cp_host_leaf_is_plannable_victim(victim): + raise RuntimeError( + "CP HiCache deterministic eviction plan diverged before " + f"reservation: node_id={node_id} phase={phase} victim_id={victim_id}" + ) - required = torch.tensor(required_host_slots, dtype=torch.int64, device="cpu") - torch.distributed.all_reduce( - required, op=torch.distributed.ReduceOp.MAX, group=self.tp_group - ) - return int(required.item()) + victim_ids.append(victim_id) + self._record_remove_event(victim) + local_freed += self.cache_controller.evict_cp_host(victim.cp_hicache) + victim.host_len = 0 + victim.cp_hicache = None + victim.host_value = None + self._remove_host_leaf(victim) - def _release_cp_write_reservation_for_retry(self, result, node_id: int) -> None: - if isinstance(result, HiCacheWriteFailure): - return logger.info( - "[HiCache-write] release local CP reservation before collective retry: node_id=%d owned_positions=%d", + "[HiCache-evict] deterministic CP host eviction before write: " + "node_id=%d phase=%s victims=%s local_freed=%d planned_freed=%s", node_id, - result.metadata.owned_positions.numel(), + phase, + victim_ids, + local_freed, + eviction_plan.planned_freed, ) - self.cache_controller.evict_cp_host(result.metadata) + return True - def _reserve_write_cp_indices_collectively( + def _reserve_write_cp_indices_no_collective( self, device_indices: torch.Tensor, node_id: int ): + plan = self._cp_build_write_admission_plan( + device_indices, node_id=node_id, phase="initial" + ) + planned_required_slots = max(plan["current_compatible_need"], default=0) + if planned_required_slots > 0: + if not self._evict_cp_host_plan(plan, node_id=node_id, phase="initial"): + return HiCacheWriteFailure(required_host_slots=planned_required_slots) + result = self.cache_controller.reserve_write_cp( device_indices=device_indices, node_id=node_id, ) - required_host_slots = self._sync_cp_write_required_host_slots(result) - if required_host_slots <= 0: + debug_plan = plan if self._cp_capacity_debug_enabled() else None + self._cp_debug_compare_write_admission( + debug_plan, + result=result, + planned_required_slots=planned_required_slots, + ) + if not isinstance(result, HiCacheWriteFailure): return result - self._release_cp_write_reservation_for_retry(result, node_id) - self._evict_host_for_physical_slots( - required_host_slots, - synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, + retry_plan = self._cp_build_write_admission_plan( + device_indices, node_id=node_id, phase="retry_after_local_failure" ) + retry_required_slots = max( + retry_plan["current_compatible_need"], default=int(result.required_host_slots) + ) + if retry_required_slots <= 0: + raise RuntimeError( + "CP HiCache host reservation failed although deterministic " + "capacity planner predicted no eviction need: " + f"node_id={node_id} required_host_slots={result.required_host_slots}" + ) + if not self._evict_cp_host_plan( + retry_plan, node_id=node_id, phase="retry_after_local_failure" + ): + return HiCacheWriteFailure(required_host_slots=retry_required_slots) + logger.info( - "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d", + "[HiCache-write] write_backup CP retry after deterministic host eviction: " + "node_id=%d needed_slots=%d", node_id, - required_host_slots, + retry_required_slots, ) result = self.cache_controller.reserve_write_cp( device_indices=device_indices, node_id=node_id, ) - required_host_slots = self._sync_cp_write_required_host_slots(result) - if required_host_slots <= 0: + self._cp_debug_compare_write_admission( + retry_plan if self._cp_capacity_debug_enabled() else None, + result=result, + planned_required_slots=retry_required_slots, + ) + if not isinstance(result, HiCacheWriteFailure): return result - self._release_cp_write_reservation_for_retry(result, node_id) logger.info( - "[HiCache-write] write_backup CP FAILED after collective retry: node_id=%d len=%d needed_slots=%d", + "[HiCache-write] write_backup CP FAILED after deterministic retry: " + "node_id=%d len=%d needed_slots=%d", node_id, len(device_indices) if device_indices is not None else 0, - required_host_slots, + int(result.required_host_slots), ) - return HiCacheWriteFailure(required_host_slots=required_host_slots) + return result - def _reserve_write_cp_collectively(self, node: TreeNode): - return self._reserve_write_cp_indices_collectively(node.value, node.id) + def _reserve_write_cp(self, node: TreeNode): + return self._reserve_write_cp_indices_no_collective(node.value, node.id) def _attach_prepared_cp_backup( self, node: TreeNode, prepared: PreparedCpHiCacheBackup @@ -1165,6 +1685,17 @@ class HiRadixCache(RadixCache): ) self.cache_controller.evict_cp_host(prepared.metadata) + def _warn_cp_hicache_fallback(self, fallback_name: str, reason: str, **kwargs): + details = " ".join(f"{key}={value}" for key, value in kwargs.items()) + if details: + details = " " + details + logger.warning( + "[CP_HICACHE_FALLBACK][%s] reason=%s%s", + fallback_name, + reason, + details, + ) + def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int: """Return the already-present radix prefix length without mutating the tree. @@ -1212,13 +1743,31 @@ class HiRadixCache(RadixCache): return total_prefix_len def prepare_write_backup_for_req(self, req) -> None: - if ( - self.disable - or not self._uses_cp_hicache - or self.cache_controller.write_policy == "write_back" - or getattr(req, "is_chunked", 0) > 0 - or getattr(req, "cp_hicache_prepared_backup", None) is not None - ): + if self.disable or not self._uses_cp_hicache: + return + if self.cache_controller.write_policy == "write_back": + self._warn_cp_hicache_fallback( + "prepare_write_backup_skipped", + "write_back_policy", + rid=getattr(req, "rid", ""), + ) + return + if getattr(req, "is_chunked", 0) > 0: + self._warn_cp_hicache_fallback( + "prepare_write_backup_skipped", + "chunked_req", + rid=getattr(req, "rid", ""), + is_chunked=getattr(req, "is_chunked", 0), + ) + return + if getattr(req, "cp_hicache_prepared_backup", None) is not None: + prepared = getattr(req, "cp_hicache_prepared_backup", None) + self._warn_cp_hicache_fallback( + "prepare_write_backup_skipped", + "already_prepared", + rid=getattr(req, "rid", ""), + node_id=getattr(prepared, "node_id", None), + ) return token_ids = req.fill_ids @@ -1239,17 +1788,26 @@ class HiRadixCache(RadixCache): req.req_pool_idx, start : len(keys) ].to(dtype=torch.int64, copy=True) if len(kv_indices) == 0: + self._warn_cp_hicache_fallback( + "prepare_write_backup_skipped", + "empty_kv_indices", + rid=getattr(req, "rid", ""), + start=start, + key_len=len(keys), + ) return node_id = TreeNode.counter TreeNode.counter += 1 - result = self._reserve_write_cp_indices_collectively(kv_indices, node_id) + result = self._reserve_write_cp_indices_no_collective(kv_indices, node_id) if isinstance(result, HiCacheWriteFailure): - logger.info( - "[HiCache-write] prepare CP backup skipped after host reservation failure: node_id=%d rid=%s len=%d", - node_id, - getattr(req, "rid", ""), - len(kv_indices), + self._warn_cp_hicache_fallback( + "prepare_write_backup_reservation_failed", + "host_reservation_failed", + node_id=node_id, + rid=getattr(req, "rid", ""), + logical_len=len(kv_indices), + required_host_slots=result.required_host_slots, ) return @@ -1294,12 +1852,21 @@ class HiRadixCache(RadixCache): write_back, ) if self._uses_cp_hicache: - result = self._reserve_write_cp_collectively(node) + if not write_back: + self._warn_cp_hicache_fallback( + "post_forward_catch_up_backup", + "write_backup_called_without_prepared_cp_backup", + node_id=node.id, + logical_len=len(node.value) if node.value is not None else 0, + ) + result = self._reserve_write_cp(node) if isinstance(result, HiCacheWriteFailure): - logger.info( - "[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d", - node.id, - len(node.value) if node.value is not None else 0, + self._warn_cp_hicache_fallback( + "post_forward_catch_up_backup_failed", + "host_reservation_failed", + node_id=node.id, + logical_len=len(node.value) if node.value is not None else 0, + required_host_slots=result.required_host_slots, ) return 0 @@ -1442,10 +2009,11 @@ class HiRadixCache(RadixCache): queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") if self.tp_world_size > 1: # synchronize TP workers to make the same update to radix cache - torch.distributed.all_reduce( + self._cp_hicache_all_reduce( queue_size, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, + tag="writing_check_min", ) finish_count = int(queue_size.item()) @@ -1783,8 +2351,11 @@ class HiRadixCache(RadixCache): if not synchronize_across_ranks: return bool(local_done) done = torch.tensor(local_done, dtype=torch.int, device="cpu") - torch.distributed.all_reduce( - done, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + self._cp_hicache_all_reduce( + done, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, + tag="host_evict_done_min", ) return bool(done.item()) @@ -2133,8 +2704,11 @@ class HiRadixCache(RadixCache): dtype=torch.int, ) if self.tp_world_size > 1: - torch.distributed.all_reduce( - qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + self._cp_hicache_all_reduce( + qsizes, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, + tag="storage_queue_min", ) n_revoke, n_backup, n_release = map(int, qsizes.tolist()) @@ -2181,10 +2755,11 @@ class HiRadixCache(RadixCache): [1 - int(can_terminate), int(operation_terminated)], dtype=torch.int, ) - torch.distributed.all_reduce( + self._cp_hicache_all_reduce( states, op=torch.distributed.ReduceOp.MAX, group=self.tp_group, + tag="prefetch_terminate_max", ) can_terminate = states[0].item() == 0 operation_terminated = states[1].item() == 1 @@ -2222,10 +2797,11 @@ class HiRadixCache(RadixCache): completed_tokens_tensor = torch.tensor( min_completed_tokens, dtype=torch.int ) - torch.distributed.all_reduce( + self._cp_hicache_all_reduce( completed_tokens_tensor, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, + tag="prefetch_completed_min", ) min_completed_tokens = completed_tokens_tensor.item() fetched_token_ids = token_ids[:min_completed_tokens] diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 0298a3a5a..edc22775d 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -331,6 +331,12 @@ class SessionAwareCache(BasePrefixCache): def flush_write_through_acks(self) -> None: return self.inner.flush_write_through_acks() + def prepare_write_backup_for_req(self, req) -> None: + prepare_fn = getattr(self.inner, "prepare_write_backup_for_req", None) + if prepare_fn is not None: + return prepare_fn(req) + return None + def check_hicache_events(self): return self.inner.check_hicache_events() 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 34d66a066..d3592c2fa 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -108,6 +108,7 @@ from sglang.srt.mem_cache.hiradix_cache import ( _compute_shared_hicache_token_capacities, ) from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode +from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -597,6 +598,26 @@ class FakeEvictionStrategy: return 0 +class FakeCpLayout: + def __init__(self, cp_size=4, cp_rank=0, page_size=1): + self.cp_size = cp_size + self.cp_rank = cp_rank + self.page_size = page_size + + def owner_for_logical_pages(self, logical_pages): + owners = torch.remainder(logical_pages - 1, self.cp_size) + return torch.where(logical_pages <= 0, torch.full_like(owners, -1), owners) + + def logical_locs_to_physical(self, logical_locs): + return logical_locs + + def owned_by_this_rank(self, logical_locs): + logical_pages = torch.div( + logical_locs, self.page_size, rounding_mode="floor" + ) + return self.owner_for_logical_pages(logical_pages) == self.cp_rank + + class FakeEvictDeviceController: write_policy = "write_through" @@ -610,6 +631,20 @@ class FakeTokenAllocator: class TestHiRadixCacheCPBackup(CustomTestCase): + def test_session_aware_cache_forwards_cp_hicache_prepare(self): + calls = [] + + class Inner: + def prepare_write_backup_for_req(self, req): + calls.append(req) + + wrapper = SessionAwareCache(Inner()) + req = object() + + wrapper.prepare_write_backup_for_req(req) + + self.assertEqual(calls, [req]) + def test_node_backuped_uses_cp_metadata(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -843,19 +878,21 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(node.hit_count, 1) - def test_write_backup_retries_by_required_physical_slots(self): + def test_write_backup_uses_deterministic_host_eviction_before_reserve(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True + cache.page_size = 1 reservation_factory = lambda device_indices: make_write_reservation( device_indices, node_id=122 ) - cache.cache_controller = FakeReserveWriteController( - [HiCacheWriteFailure(required_host_slots=1), reservation_factory] - ) + cache.cache_controller = FakeReserveWriteController([reservation_factory]) + cache.cache_controller.cp_shared_kv_layout = FakeCpLayout(cp_size=1, cp_rank=0) + cache.token_to_kv_pool_host = types.SimpleNamespace(size=16) cache.evictable_host_leaves = set() cache.eviction_strategy = FakeEvictionStrategy() cache.get_child_key_fn = lambda key: key.token_ids[0] cache._record_remove_event = lambda node: None + cache._update_host_leaf_status = lambda node: None cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.inc_node_lock_ref = lambda node: None @@ -871,8 +908,8 @@ class TestHiRadixCacheCPBackup(CustomTestCase): evictable_node.host_len = 4 evictable_node.cp_hicache = CpHiCacheNodeMetadata( logical_len=4, - owned_positions=torch.tensor([0], dtype=torch.int64), - host_indices=torch.tensor([55], dtype=torch.int64), + owned_positions=torch.arange(4, dtype=torch.int64), + host_indices=torch.arange(55, 59, dtype=torch.int64), page_owners=torch.zeros(max(4, 0), dtype=torch.int8), page_size=1, ) @@ -886,7 +923,8 @@ class TestHiRadixCacheCPBackup(CustomTestCase): cache.write_backup(node) self.assertEqual( - cache.cache_controller.evicted_host_indices[0].tolist(), [55] + cache.cache_controller.evicted_host_indices[0].tolist(), + list(range(55, 59)), ) self.assertEqual(evictable_node.host_len, 0) self.assertIsNone(evictable_node.cp_hicache) @@ -896,51 +934,124 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertIn(node.id, cache.pending_host_backups) self.assertEqual(len(cache.cache_controller.submitted), 1) - def test_write_backup_rolls_back_local_success_when_peer_needs_host_eviction(self): + def test_write_backup_cp_post_forward_path_warns_fallback(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True - cache.tp_world_size = 2 - cache.tp_group = object() cache.cache_controller = FakeReserveWriteController( [ - lambda device_indices: make_write_reservation( - device_indices, node_id=130, host_start=90 - ), - lambda device_indices: make_write_reservation( - device_indices, node_id=130, host_start=100 - ), + lambda device_indices, node_id: make_write_reservation( + device_indices, node_id=node_id + ) ] ) - evictions = [] - cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: ( - evictions.append((required, synchronize_across_ranks)) or required - ) cache.ongoing_write_through = {} cache.pending_host_backups = {} cache.inc_node_lock_ref = lambda node: None - reduce_values = iter([4, 0]) + node = TreeNode() + node.id = 133 + node.value = torch.arange(16, dtype=torch.int64) - def fake_all_reduce(tensor, op=None, group=None): - tensor.fill_(next(reduce_values)) + with self.assertLogs( + "sglang.srt.mem_cache.hiradix_cache", level="WARNING" + ) as captured: + cache.write_backup(node) + + self.assertTrue( + any( + "[CP_HICACHE_FALLBACK][post_forward_catch_up_backup]" in message + for message in captured.output + ) + ) + + def test_prepare_write_backup_for_req_chunked_skip_warns_fallback(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.cache_controller = FakeReserveWriteController([]) + + req = types.SimpleNamespace( + rid="rid-chunked", + is_chunked=1, + cp_hicache_prepared_backup=None, + ) + + with self.assertLogs( + "sglang.srt.mem_cache.hiradix_cache", level="WARNING" + ) as captured: + cache.prepare_write_backup_for_req(req) + + self.assertTrue( + any( + "[CP_HICACHE_FALLBACK][prepare_write_backup_skipped]" in message + and "reason=chunked_req" in message + for message in captured.output + ) + ) + + def test_write_backup_deterministic_eviction_avoids_reserve_all_reduce(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.page_size = 1 + cache.tp_world_size = 2 + cache.tp_group = object() + cache.cache_controller = FakeReserveWriteController( + [ + lambda device_indices: make_write_reservation( + device_indices, node_id=130, host_start=100 + ), + ] + ) + cache.cache_controller.cp_shared_kv_layout = FakeCpLayout(cp_size=1, cp_rank=0) + cache.token_to_kv_pool_host = types.SimpleNamespace(size=16) + evictions = [] + cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: ( + evictions.append((required, synchronize_across_ranks)) or required + ) + cache.evictable_host_leaves = set() + cache.eviction_strategy = FakeEvictionStrategy() + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache._record_remove_event = lambda node: None + cache._update_host_leaf_status = lambda node: None + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + cache.inc_node_lock_ref = lambda node: None + + root = TreeNode() + root.key = RadixKey(token_ids=[], extra_key=None) + root.value = [] + cache.root_node = root + evictable_node = TreeNode() + evictable_node.parent = root + evictable_node.key = RadixKey(token_ids=[2], extra_key=None) + evictable_node.value = None + evictable_node.host_len = 16 + evictable_node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=16, + owned_positions=torch.arange(16, dtype=torch.int64), + host_indices=torch.arange(90, 106, dtype=torch.int64), + page_owners=torch.zeros(16, dtype=torch.int8), + page_size=1, + ) + root.children[2] = evictable_node + cache.evictable_host_leaves.add(evictable_node) node = TreeNode() node.id = 130 node.value = torch.arange(16, dtype=torch.int64) - with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce): + with patch( + "torch.distributed.all_reduce", + side_effect=AssertionError("reserve admission must not all_reduce"), + ): backed_len = cache.write_backup(node) self.assertEqual(backed_len, 16) - # The first local success must be released before all ranks enter the - # collective host eviction/retry branch. Otherwise CP ranks can diverge: - # successful ranks submit backup while failing ranks enter eviction - # all_reduce, which matches the observed Gloo 4-vs-1 mismatch. self.assertEqual( cache.cache_controller.evicted_host_indices[0].tolist(), list(range(90, 106)), ) - self.assertEqual(evictions, [(4, True)]) + self.assertEqual(evictions, []) self.assertEqual(len(cache.cache_controller.submitted), 1) self.assertEqual( cache.cache_controller.submitted[0].metadata.host_indices.tolist(), @@ -1111,7 +1222,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(cache.cache_controller.ack_write_queue, []) self.assertEqual(evicted, [reservation.metadata]) - def test_write_backup_cp_retry_failure_leaves_node_device_only(self): + def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.cache_controller = FakeReserveWriteController( @@ -1130,9 +1241,12 @@ class TestHiRadixCacheCPBackup(CustomTestCase): node.id = 124 node.value = torch.arange(16, dtype=torch.int64) - backed_len = cache.write_backup(node) + with self.assertRaisesRegex( + RuntimeError, + "planner predicted no eviction need", + ): + cache.write_backup(node) - self.assertEqual(backed_len, 0) self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) self.assertEqual(cache.pending_host_backups, {}) @@ -1229,6 +1343,109 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertIsNotNone(node.cp_hicache) self.assertEqual(node.cp_hicache.host_indices.tolist(), [55]) + def test_cp_owner_token_counts_use_page_owners_and_page_size(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=4, cp_rank=0) + ) + + counts = cache._cp_owner_token_counts( + torch.tensor([0, 1, 1, 2], dtype=torch.int8), + page_size=64, + ) + + self.assertEqual(counts, (64, 128, 64, 0)) + + def test_cp_metadata_count_asserts_local_owner_lane(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=3, cp_rank=1), + has_draft_hicache=True, + ) + metadata = CpHiCacheNodeMetadata( + logical_len=192, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(100, 164, dtype=torch.int64), + draft_host_indices=torch.arange(200, 264, dtype=torch.int64), + page_owners=torch.tensor([0, 1, 2], dtype=torch.int8), + page_size=64, + ) + + counts = cache._cp_assert_metadata_counts(metadata, context="unit") + + self.assertEqual(counts, (64, 64, 64)) + + def test_cp_metadata_count_asserts_mismatched_local_slots(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=1), + has_draft_hicache=False, + ) + metadata = CpHiCacheNodeMetadata( + logical_len=128, + owned_positions=torch.empty((0,), dtype=torch.int64), + host_indices=torch.empty((0,), dtype=torch.int64), + page_owners=torch.tensor([0, 1], dtype=torch.int8), + page_size=64, + ) + + with self.assertRaisesRegex(RuntimeError, "owner lane mismatch"): + cache._cp_assert_metadata_counts(metadata, context="unit") + + def test_cp_host_capacity_snapshot_counts_committed_and_pending_draft(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0), + has_draft_hicache=True, + draft_mem_pool_host=types.SimpleNamespace(size=512), + ) + cache.token_to_kv_pool_host = types.SimpleNamespace(size=512) + cache.pending_host_backups = {} + root = TreeNode() + root.children = {} + cache.root_node = root + + committed = TreeNode() + committed.id = 501 + committed.parent = root + committed.key = RadixKey([1]) + committed.host_len = 128 + committed.cp_hicache = CpHiCacheNodeMetadata( + logical_len=128, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(64, dtype=torch.int64), + draft_host_indices=torch.arange(100, 164, dtype=torch.int64), + page_owners=torch.tensor([0, 1], dtype=torch.int8), + page_size=64, + ) + root.children[1] = committed + + pending = TreeNode() + pending.id = 502 + pending_metadata = CpHiCacheNodeMetadata( + logical_len=64, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(200, 264, dtype=torch.int64), + draft_host_indices=torch.arange(300, 364, dtype=torch.int64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + cache.pending_host_backups[pending.id] = PendingHiCacheBackup( + node=pending, + metadata=pending_metadata, + logical_len=64, + ) + + snapshot = cache._cp_host_capacity_snapshot() + + self.assertEqual(snapshot.target_capacity, (512, 512)) + self.assertEqual(snapshot.draft_capacity, (512, 512)) + self.assertEqual(snapshot.committed_target, (64, 64)) + self.assertEqual(snapshot.committed_draft, (64, 64)) + self.assertEqual(snapshot.pending_target, (64, 0)) + self.assertEqual(snapshot.pending_draft, (64, 0)) + class TestHiRadixCacheCPSplitEvict(CustomTestCase): def test_split_node_splits_cp_metadata_by_owned_positions(self): @@ -1558,6 +1775,92 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): self.assertEqual(node.host_len, 4) self.assertIsNotNone(node.cp_hicache) + def test_cp_host_eviction_plan_is_stable_across_leaf_insertion_order(self): + def make_cache(order): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0), + has_draft_hicache=False, + ) + cache.eviction_strategy = FakeEvictionStrategy() + cache.pending_host_backups = {} + cache.root_node = TreeNode() + cache.root_node.children = {} + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache._clear_pin = lambda node: None + nodes = [] + for node_id in (42, 41): + node = TreeNode(id=node_id) + node.parent = cache.root_node + node.key = RadixKey([node_id]) + node.value = None + node.lock_ref = 0 + node.host_ref_counter = 0 + node.host_len = 64 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=64, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(node_id * 100, node_id * 100 + 64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + cache.root_node.children[node_id] = node + nodes.append(node) + cache.evictable_host_leaves = {nodes[i] for i in order} + return cache + + first = make_cache([0, 1])._plan_cp_host_evictions((64, 0)) + second = make_cache([1, 0])._plan_cp_host_evictions((64, 0)) + + self.assertEqual([node.id for node in first.victims], [41]) + self.assertEqual([node.id for node in second.victims], [41]) + self.assertEqual(first.planned_freed, (64, 0)) + self.assertEqual(first.remaining_deficit, (0, 0)) + + def test_cp_host_eviction_plan_skips_pending_backup_nodes(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = types.SimpleNamespace( + cp_shared_kv_layout=FakeCpLayout(cp_size=1, cp_rank=0), + has_draft_hicache=False, + ) + cache.eviction_strategy = FakeEvictionStrategy() + cache.root_node = TreeNode() + cache.root_node.children = {} + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.pending_host_backups = {} + cache.evictable_host_leaves = set() + + for node_id in (10, 11): + node = TreeNode(id=node_id) + node.parent = cache.root_node + node.key = RadixKey([node_id]) + node.value = None + node.lock_ref = 0 + node.host_ref_counter = 0 + node.host_len = 64 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=64, + owned_positions=torch.arange(64, dtype=torch.int64), + host_indices=torch.arange(node_id * 100, node_id * 100 + 64), + page_owners=torch.tensor([0], dtype=torch.int8), + page_size=64, + ) + cache.root_node.children[node_id] = node + cache.evictable_host_leaves.add(node) + if node_id == 10: + cache.pending_host_backups[node.id] = PendingHiCacheBackup( + node=node, + metadata=node.cp_hicache, + logical_len=64, + ) + + plan = cache._plan_cp_host_evictions((64,)) + + self.assertEqual([node.id for node in plan.victims], [11]) + self.assertEqual(plan.planned_freed, (64,)) + class TestHiRadixCacheCPLoadBack(CustomTestCase): def test_cp_load_back_uses_host_len_not_host_value(self):