From 8571fe0cd93beae350537f1ce838497ad55175fb Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 27 May 2026 04:28:50 +0800 Subject: [PATCH] Share CP HiCache host budget across target and draft KV CP HiCache previously let the target host pool and the draft/MTP host pool each consume the full --hicache-size budget. With EAGLE/MTP enabled this doubled per-rank host allocation and could kill scheduler ranks during startup before Python emitted a traceback. The cache now treats target KV and draft KV as one logical host-cache object: target and draft capacities are computed from one per-rank byte budget, draft may receive more token capacity when its per-token footprint is smaller, and draft attachment remains tied to target residency. Constraint: --hicache-size is a per-rank host budget and must not be multiplied by attaching draft KV. Rejected: Give draft another independent --hicache-size allocation | repeats the observed host OOM failure mode. Rejected: Disable draft HiCache attachment under CP | avoids OOM but breaks target/draft cache-hit consistency for MTP. Confidence: medium Scope-risk: moderate Directive: Keep target and draft KV as one logical HiCache object; do not let draft host allocation consume an independent full hicache-size budget. Tested: python -m py_compile on modified scheduler/cache/test files Tested: remote g0034 container PYTHONPATH=python python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q (45 passed) Not-tested: full multi-rank GLM5 server restart after clearing existing remote router/defunct process state Co-authored-by: OmX --- ...nsa_prefill_cp_hicache_async_state_sync.md | 1176 +++++++++++++++++ python/sglang/srt/managers/scheduler.py | 12 + .../sglang/srt/mem_cache/cache_init_params.py | 3 +- python/sglang/srt/mem_cache/hiradix_cache.py | 132 +- .../sglang/srt/mem_cache/memory_pool_host.py | 21 +- .../mem_cache/test_cp_hicache_metadata.py | 33 +- 6 files changed, 1365 insertions(+), 12 deletions(-) create mode 100644 docs/advanced_features/nsa_prefill_cp_hicache_async_state_sync.md diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_async_state_sync.md b/docs/advanced_features/nsa_prefill_cp_hicache_async_state_sync.md new file mode 100644 index 000000000..2df3587ea --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_hicache_async_state_sync.md @@ -0,0 +1,1176 @@ +# NSA Prefill CP HiCache Async State Synchronization Notes + +Date: 2026-05-27 + +Branch context: `cp-hicache-refactor` around `d14c02b0dc`. + +This note records the current understanding of CP shared-KV + HiCache +backup/load/evict behavior and the synchronization model needed to make it +both correct and performant. + +## Problem Statement + +With CP shared KV enabled, the radix cache and scheduler operate in a full +logical KV address space, while every CP rank owns only a physical shard of the +device and host KV pools. HiCache backup, load, and eviction therefore have two +different state domains: + +- **Logical state** must be identical on every CP rank. + - Radix tree topology. + - Node residency state. + - `TreeNode.value` logical loc tensor. + - `host_len`, `cp_hicache`, `page_owners`. + - Lock/ref counters that affect eviction eligibility. + - Request prefix indices and scheduler admission decisions. +- **Physical state** is intentionally rank-local. + - Physical device pages owned by the local CP rank. + - Physical host slots allocated by the local CP rank. + - Zero-owned ranks for a given logical node. + - Draft/MTP host slots for locally owned pages. + +The correct model is therefore: + +> Logical operations must be submitted and committed in the same order on all +> CP ranks; physical transfers may be sparse and rank-local. + +The current code partially implements this, but the async lifecycle is still +spread across `HiRadixCache`, `HiCacheController`, allocator state, and CUDA +events. This makes performance tuning risky because reducing synchronization +can easily introduce cross-rank state divergence. + +## Target Contract + +The intended design for the current CP HiCache stage is: + +1. **Radix tree is synchronous.** + - All CP ranks apply radix match/insert/split/evict and residency metadata + changes in the same order. + - Background transfer completion must not independently mutate radix + topology or make partially transferred nodes visible. + +2. **Physical transfer is asynchronous.** + - D2H backup and H2D load may run on transfer streams and overlap with + scheduler progress or layer compute. + - Async completion only satisfies an already-submitted logical operation. + It does not decide which logical node exists, which node is visible, or + which victim is chosen. + +3. **Draft/MTP KV is strongly consistent with target KV.** + - Draft KV is a shadow payload of the target radix node. + - Target and draft residency transition together: host valid, device valid, + eviction, load, backup, split, and rollback are all atomic at the logical + node level. + - If target and draft disagree, the node must be treated as invalid for + speculative decoding. + - `--hicache-size` is the per-rank host budget for the whole logical cache + object, not a budget to be independently duplicated by target and draft. + When draft/MTP HiCache is attached, target and draft host pools must be + sized from this shared budget. Draft may have a larger token capacity when + its per-token footprint is smaller, but it must follow target residency + and allocation decisions. + +4. **Cache-hit load is layer-wise.** + - On host hit, device allocation and radix state are resolved + synchronously, then each layer's target/draft payload is loaded from host + to device asynchronously. + - The model waits per layer through the existing layer-ready mechanism. + - The request may not observe a layer before both target and draft payloads + for that layer are ready. + +5. **Backup should become layer-wise.** + - After each layer stores its target/draft KV for the computed suffix, that + layer can be backed up to host asynchronously. + - The radix node becomes host-visible only after all required layer backups + have completed. + - This keeps the logical behavior close to `write_through`, but reduces the + tail cost by overlapping D2H with remaining layer compute. + +In short: + +```text +sync radix/control plane +async transfer/data plane +target KV and draft KV are one logical cache object +per-layer host load on hit +per-layer host backup after layer KV store +``` + +## Current Code Facts + +### CP owner mapping + +`CpSharedKVLayout` defines page ownership by logical page id: + +- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py` + - `owner_for_logical_pages()` + - `owned_by_this_rank()` + - `logical_locs_to_physical()` + +Logical page 0 is padding. Real pages are distributed by: + +```text +owner = (logical_page - 1) % cp_size +physical_page = (logical_page - 1) // cp_size + 1 +``` + +This is the basis for all CP HiCache physical filtering. + +### Backup path + +Main paths: + +- `HiRadixCache.write_backup()` + - `python/sglang/srt/mem_cache/hiradix_cache.py` +- `HiCacheController._write_cp()` + - `python/sglang/srt/managers/cache_controller.py` +- `HiCacheController.start_writing()` + - `python/sglang/srt/managers/cache_controller.py` + +Current behavior: + +1. `TreeNode.value` is full logical locs. +2. `_write_cp()` filters local owned positions. +3. Local logical locs are translated to physical device locs. +4. Host pool allocation is only for local physical slots. +5. Target and draft host slots are reserved together when draft HiCache is + registered. +6. D2H is submitted on `write_stream`. +7. Zero-owned ranks enqueue a completed no-op write ack. +8. `writing_check()` uses `all_reduce(MIN)` over locally completed ack count so + all ranks pop the same logical write ack prefix. + +This gives asynchronous data movement, but logical backup state is updated in +`write_backup()` before a full CP-group commit protocol exists. + +Another important asymmetry: backup currently calls +`backup_from_device_all_layer()`, while load uses `load_to_device_per_layer()`. +So load has layer-wise overlap with compute, but backup is submitted as an +all-layer D2H operation after the node has been inserted. The current backup +path can overlap with later scheduler work, but it does not overlap each +already-finished layer's D2H with the remaining layers of the same prefill +forward. + +### Load path + +Main paths: + +- `HiRadixCache.load_back()` + - `python/sglang/srt/mem_cache/hiradix_cache.py` +- `HiCacheController.load_cp()` + - `python/sglang/srt/managers/cache_controller.py` +- `HiCacheController.start_loading()` + - `python/sglang/srt/managers/cache_controller.py` +- `LayerDoneCounter` + - `python/sglang/srt/managers/cache_controller.py` +- KV pool layer access waits + - `python/sglang/srt/mem_cache/memory_pool.py` + +Current behavior: + +1. Load-back walks evicted host-backed radix nodes. +2. CP mode reconstructs the write-time owner pattern from + `node.cp_hicache.page_owners`. +3. `alloc_pages_with_owners()` allocates fresh logical pages whose owner lanes + match the backed-up metadata. +4. Only local owned positions are transferred from host to device. +5. H2D load is submitted per layer on `load_stream`. +6. `LayerDoneCounter` records a producer id. +7. Forward receives `hicache_consumer_index`. +8. Attention reads call `wait_until(layer_id)`, allowing layer-wise load/compute + overlap instead of blocking on the full load. + +This is the strongest current async mechanism: load can overlap with prefill +compute at layer granularity. + +### Device eviction + +Main path: + +- `HiRadixCache.evict()` + - `python/sglang/srt/mem_cache/hiradix_cache.py` + +Current behavior: + +- If a node is already host-backed, device eviction frees device slots + immediately and keeps host metadata. +- If `write_policy == "write_back"` and a node is not host-backed, + `write_backup(..., write_back=True)` is called and then + `writing_check(write_back=True)` synchronizes all write completion before + freeing device memory. + +This means write-back eviction is not truly asynchronous in the hot allocation +path. It can enqueue D2H asynchronously, but then blocks before device memory is +released. + +### Host eviction + +Main path: + +- `HiRadixCache._evict_host_for_physical_slots()` + - `python/sglang/srt/mem_cache/hiradix_cache.py` + +Current behavior: + +- Victims are selected from logical host leaves. +- Each rank frees its own physical host slots through `evict_cp_host()`. +- With `synchronize_across_ranks=True`, a loop calls `all_reduce(MIN)` on a + local done flag until all ranks have freed enough local physical host slots. + +This preserves CP correctness under uneven per-rank host pressure, but the +current loop can produce frequent CPU all-reduces when host memory is near +capacity. + +## Key Invariants to Preserve + +1. **Radix logical state is replicated.** + - All ranks must insert, split, evict, load, and mark backup state for the + same logical nodes in the same order. + +2. **Physical ownership is rank-local.** + - A rank transfers only pages it owns. + - Zero-owned ranks still participate in logical op ordering with no-op + acks. + +3. **Host metadata must preserve write-time owner pattern.** + - `page_owners` is required because a later load allocates fresh logical + pages. The new allocation must reproduce the write-time owner pattern or + `owned_positions` will refer to wrong physical shards. + +4. **Target and draft KV must be atomic.** + - For EAGLE/MTP, target KV and draft KV must be backed up, loaded, and + evicted under the same logical operation. Partial target-only or + draft-only success can corrupt accept rate after cache hit. + +5. **Async transfer source/destination lifetime must be protected.** + - A device node under D2H write must not be evicted until the write ack is + committed. + - A device node under H2D load must not be used by a layer until that layer's + load event is complete. + +6. **Failure must be globally resolved.** + - If any rank cannot reserve host slots, reserve draft slots, allocate owner + lanes, or submit a needed transfer, all ranks must either commit or abort + the logical operation together. + +## Current Design Boundary + +The first stabilization pass should intentionally keep the control plane small. +Do not include TBO, decode-side behavior, L3/storage, or cross-instance cache +sharing in this stage. The scope is: + +- prefill-side CP shared KV, +- L1 device KV and L2 host HiCache, +- target KV plus EAGLE/MTP draft KV, +- synchronous radix tree mutation, +- asynchronous low-level D2H/H2D data movement. + +The control-plane rule is: + +> Radix tree operations are synchronous; only physical KV movement is async. + +The synchronization rule is: + +> Prefer deterministic-by-construction state transitions over explicit +> all-rank confirmation. + +This means the scheduler/radix layer remains the sole owner of logical cache +state. CUDA events and background transfer callbacks may report progress, but +they must not independently mutate radix topology, node residency, evictable +sets, or request-visible cache-hit state. + +Synchronous control-plane operations include: + +- radix match, insert, split, and eviction metadata updates, +- host/device residency state transitions, +- target/draft metadata commit or rollback, +- logical lock/ref updates, +- evictable leaf set updates, +- device and host slot reservation decisions, +- global commit/abort decisions after slow-path failure. + +Asynchronous data-plane operations include: + +- target KV D2H backup, +- draft KV D2H backup, +- target KV H2D load, +- draft KV H2D load, +- layer-wise load readiness events. + +This is deliberately conservative. It preserves CP rank determinism while still +allowing backup and load transfers to overlap with scheduler progress or model +compute. + +## Synchronization Philosophy + +A healthy CP HiCache design should not need all ranks to explicitly vote on +every normal cache transition. The common path should be correct because every +rank runs the same deterministic logical state machine: + +```text +same request order +same radix mutations +same logical page ids +same owner pattern +same victim selection +same target/draft atomic op boundaries +``` + +Given those inputs, each rank can independently derive: + +- which logical node is affected, +- which physical pages it owns, +- how many local host/device slots are required, +- which local slots can be freed, +- whether it is a zero-owned participant, +- which local async event corresponds to the logical op. + +All-reduce should therefore be a **slow-path safety mechanism**, not the normal +source of truth. It is appropriate only when a condition is genuinely +rank-local and cannot be derived from replicated logical state: + +- local physical host/device allocation fails after deterministic eviction, +- local transfer submission fails, +- background IO reports partial/failure status, +- a debug/invariant check wants to prove the ranks are still aligned, +- shutdown/recovery needs to drain unknown pending state. + +It should not be used merely to decide routine facts such as "which node was +evicted", "which prefix of write acks is logically complete", or "which host +range belongs to this request" if those facts can be represented by a shared +logical op id and deterministic local event progress. + +Practical target: + +- **Fast path:** no collective communication. +- **Normal async completion:** commit by deterministic op order and local event + readiness, with zero-owned ranks carrying no-op events. +- **Slow path:** one coarse global success/abort gate after local retry has + failed, not repeated per node or per scheduler tick. +- **Debug mode:** optional all-rank invariant checks, off the performance path. + +This does not mean "never all-reduce". It means all-reduce is reserved for +cases where the program cannot otherwise know whether every rank is still able +to complete the same logical operation. If ordinary progress requires frequent +all-reduce, the local state machine is carrying insufficient information. + +### Replicated capacity ledger + +To avoid using all-reduce as an admission oracle, the scheduler should maintain +a replicated logical capacity ledger. The ledger is updated only by synchronous +logical ops, not by background transfer callbacks. + +For CP shared KV, the ledger should track at least: + +```text +device_free_pages_by_owner +device_evictable_pages_by_owner +host_free_pages_by_owner +host_evictable_pages_by_owner +pending_load_pages_by_owner +pending_backup_pages_by_owner +``` + +The owner dimension is necessary because a load-back allocation is not just +"N tokens"; it requires a specific owner-lane pattern. The host dimension is +necessary because CP host pools are physically local even though radix state is +logical and replicated. + +With this ledger, common decisions become deterministic: + +- whether a host-hit load can reserve device pages, +- whether a cold extend can reserve device pages, +- how many host victims are needed before backup, +- which logical host victims should be selected, +- whether a request should be admitted or kept queued. + +The actual local allocator then becomes an executor of a decision already made +by the ledger. If the local allocator fails despite the ledger saying success, +that is either a bug, stale accounting, or an external failure. It should go to +slow-path recovery/debug validation, not become the normal admission mechanism. + +### Visibility gates instead of confirmation votes + +Async operations should become request-visible through deterministic gates: + +```text +submit logical op in same order on every rank +reserve ledger capacity synchronously +submit local transfer or no-op event +later, when a request/victim decision needs the result: + wait/query the local event for that op + expose the next deterministic op state +``` + +This avoids asking every rank "are you done yet?" on every tick. A faster rank +may finish its local D2H/H2D earlier, but it should not make a different +logical scheduling decision. It can only expose the result at a deterministic +visibility point for the same logical op. + +Examples: + +- Backup is submitted as `HOST_BACKUP_PENDING`. It becomes `HOST_VALID` only + at a backup visibility gate. Until then, host-hit and host-evict logic treats + it as pending. +- Load is submitted as `DEVICE_LOAD_PENDING`. The owning batch may consume it + through `hicache_consumer_index`; other requests cannot treat it as ordinary + `DEVICE_VALID` until the load visibility gate commits. +- Host eviction selects the same logical victims from the replicated host leaf + set and ledger. Each rank frees only its local host slots. There is no need + to vote on the victim list. + +The important restriction is that visibility gates must be tied to logical +operation order, not wall-clock completion. Local event completion is a data +dependency; logical visibility remains controlled by the radix/scheduler state +machine. + +## MTP/Draft KV Pool Model + +EAGLE/MTP draft KV should be treated as a shadow payload of the target KV +radix node, not as an independent cache. + +The draft KV pool must not maintain its own: + +- radix tree, +- hit/miss decision, +- backup/load/evict state, +- owner pattern, +- prefix length, +- protection state, +- logical commit/abort decision. + +The target radix node is the source of truth. If MTP is enabled, target and +draft residency are a single atomic logical state: + +```text +target HOST_VALID <=> draft HOST_VALID +target DEVICE_VALID <=> draft DEVICE_VALID +``` + +Partial states such as target-host-valid but draft-host-absent are invalid for +request-visible cache hits, because decode acceptance depends on target and +draft KV representing the same prefix. + +The target node metadata therefore owns both target and draft physical payloads: + +```text +CpHiCacheNodeMetadata: + logical_len + owned_positions + page_owners + host_indices # target host slots, local physical shard + draft_host_indices # draft host slots, local physical shard when MTP enabled +``` + +For zero-owned ranks, both `host_indices` and `draft_host_indices` are empty. +For non-zero-owned ranks under MTP, the target and draft host index arrays must +have the same owned-position cardinality. + +Lifecycle implications: + +- Backup reserves target and draft host slots together. +- Draft reservation or D2H submission failure rolls back target reservation. +- Load reuses the target node's owner pattern and owned positions for draft + H2D. +- Device eviction frees target and draft device payload together. +- Host eviction frees target and draft host payload together. +- Radix split must split target and draft metadata together, or drain the + pending op before splitting. + +This avoids the earlier failure mode where cache-hit prefix target KV was +restored but draft/MTP KV was stale, absent, or read from a mismatched owner +layout, causing speculative accept length to collapse after HiCache hits. + +## Eviction Scenarios In Scope + +For the first implementation stage, only these eviction cases need to be +handled: + +1. **Host-hit load-back needs device KV but device space is insufficient.** + - Evict only device cache that is host-backed, unlocked, not under D2H + backup, and not a load destination. + - Do not affect the currently running prefill request. + - Use owner-lane-aware capacity checks before committing load-back. + +2. **Cold prefill extend needs device KV but device space is insufficient.** + - Evict host-backed device cache first. + - Use the same CP owner-lane rules as load-back. + - If insufficient, reject/fallback at scheduling time rather than partially + committing a batch. + +3. **Write-through backup needs host KV but host space is insufficient.** + - Evict host cache only from nodes that are host-valid, device-absent, not a + load source, and not a backup destination. + - Free target and draft host slots atomically. + - Select the same logical victims across ranks; each rank frees only its + local physical host slots. + +4. **Low-frequency proactive host free-room maintenance.** + - Optional after correctness is stable. + - Should run in coarse epochs, not every scheduler tick. + +Out of scope for this stage: + +- decode prealloc and decode-side append pressure, +- TBO/two-batch overlap reserve accounting, +- L3/storage write-back or prefetch, +- async write-back demotion as a required emergency path, +- cross-instance cache sharing. + +## Backup/Load Overlap Cases to Support + +There are two distinct overlap opportunities. They should not be conflated +because their correctness constraints differ. + +### Case A: computed suffix backup while the same request is still forwarding + +For a new request or cache-miss suffix, each layer's KV becomes immutable once +that layer has stored its KV for the current prefill. In principle, that +layer's D2H backup can start before later layers finish computing: + +```text +layer i stores target KV + -> record layer-i store completion + -> D2H backup layer i on write_stream +layer i+1 compute continues on compute stream +``` + +This can reduce backup tail latency and make host cache visible sooner after +the request completes. It does not allow device KV to be evicted earlier for +the same logical node unless the implementation supports partially host-valid +nodes, which is not recommended for the first stage. + +Required conditions for per-layer backup: + +- A layer completion event must order D2H after the layer's KV store. +- Target and draft/MTP backup must both participate in the same logical backup + op. +- NSA indexer K/scale payload must be backed up with the corresponding layer, + not forgotten behind the normal KV path. +- The node becomes request-visible `HOST_VALID` only after every required + target and draft layer has completed and the global CP commit succeeds. +- Host eviction must not reclaim any layer slice of a pending backup + destination. + +The current host-pool API does not expose a generic +`backup_from_device_per_layer()` equivalent to `load_to_device_per_layer()`. +Adding per-layer backup therefore requires a host-pool API extension for MHA, +MLA, and NSA indexer payloads. Until then, all-layer backup should remain the +correctness baseline. + +### Case B: host/storage prefetch while the request later backs up computed KV + +If a request already has host-cache or storage-cache prefix, the system may +prefetch additional pages into host while the model recomputes or extends a +suffix. Later, the computed suffix may also be backed up to host. This creates +a conflict domain over both host memory and radix ownership: + +```text +storage/L3 prefetch writes host slots for logical pages P +compute finishes and write-through backup also writes host slots for P +``` + +The safe rule is: + +> Compute-derived backup wins over speculative prefetch for the same logical +> token range. + +Reason: compute-derived KV is exactly the data used by this request's forward, +while prefetch is opportunistic and can be partial, revoked, or stale with +respect to the current radix split state. + +Practical consequences: + +- Prefetch must be request-local and non-authoritative until insertion commit. +- Backup insertion for an overlapping logical range should either cancel the + overlapping prefetch or make prefetch insertion skip/free the already-backed + range. +- Prefetch host slots and backup host slots must both participate in host + memory pressure accounting. +- Host eviction must not choose active prefetch destinations or active backup + destinations. +- Radix insertion remains synchronous; background prefetch threads can only + deliver completed bytes and completion counts. + +Current code already has pieces of this model: prefetch protects the anchor +host node, records `ongoing_prefetch`, and `_insert_helper_host()` skips matched +prefix ranges when inserting prefetched host data. For CP shared KV + MTP, +however, the same principle needs to include target/draft atomicity and owner +metadata, not just host byte insertion. + +## Current Risk Areas + +### 1. Backup lacks an explicit global success commit gate + +`write_backup()` sets `node.host_len` and `node.cp_hicache` after local success. +If one rank succeeds and another fails after retry, logical host-backed state +can diverge unless all ranks reach the same outcome through surrounding control +flow. + +Current host eviction retry improves this, but it is not a complete transaction +protocol. A stronger model should separate: + +```text +local reserve success -> global success check -> logical metadata commit +``` + +### 2. Load failure can become rank-local + +`load_cp()` returns `None` if owner-lane allocation fails. `load_back()` retries +with lane-aware eviction, then returns `None` and falls back to recompute if it +still fails. + +The expected case is that all ranks make the same decision because the logical +allocator is replicated. Under pressure or bugs, this assumption can break. +The slow path should use a global abort/commit gate after retry. + +### 3. Host eviction synchronization is too fine-grained + +`_evict_host_for_physical_slots()` can all-reduce inside the victim loop. This +is correct but expensive when host memory is close to full and many small +evictions happen. + +A batched epoch model should reduce synchronization frequency. + +### 4. Write-back eviction is blocking + +When `write_policy == "write_back"`, current device eviction waits for D2H +completion before freeing device slots. This is safe but does not provide the +performance benefit expected from asynchronous demotion. + +Async demotion is possible, but it cannot help an immediate allocation failure +unless the system can tolerate delayed device release. Emergency eviction may +still need to drop cache instead of waiting. + +### 5. Write-through can create excessive ack all-reduce pressure + +`write_through` uses threshold 1, so every inserted node can trigger backup. +`writing_check()` then may all-reduce frequently even if only a small number of +acks have completed. This can become a scheduler hot path. + +### 6. Radix node split can race with async node state + +Radix tree maintenance is not just metadata bookkeeping under async HiCache. +The tree can split nodes during `match_prefix()` or `insert()`: + +- `HiRadixCache.match_prefix()` may call `_match_prefix_helper()`, which can + call `_split_node()` on a partial match. +- `HiRadixCache.insert()` can also call `_split_node()` on a partial match. +- `_split_node()` copies `child.lock_ref` into the newly inserted parent node + and splits `value`, `cp_hicache`, `host_len`, and hash metadata when present. + +This is correct for quiescent nodes, but it is risky for nodes with in-flight +async operations. + +Example risk: + +1. A new node is inserted. +2. `write_through` starts async backup for that node. +3. `write_backup()` records `ongoing_write_through[child.id] = child` and uses + `inc_node_lock_ref(child)`. +4. Before the D2H ack commits, another request partially matches the node and + `_split_node()` inserts a new parent. +5. The new parent inherits `lock_ref`, and CP host metadata is split between + new parent and child. +6. The write ack still only references the original child id. +7. `writing_check()` releases the child lock, but has no ack entry for the new + parent lock. + +This can leave the split parent permanently protected or otherwise make +`evictable_size_`, `protected_size_`, and leaf sets diverge from actual radix +state. A path lock is less exposed because `dec_lock_ref(child)` walks through +the inserted parent after split, but the write-through single-node lock path is +specifically vulnerable unless split is made op-aware. + +This suggests a general rule: + +> A radix node with in-flight async state must either be unsplittable until the +> operation commits, or the async operation must be explicitly split/aliased +> together with the radix node. + +### 7. GPU memory admission tracks KV tokens, not all transient memory + +Current scheduling largely reasons in token-pool units: + +- `PrefillAdder.rem_total_tokens` uses allocator `available_size()` plus radix + `evictable_size()`, then subtracts running/new request offsets. +- `_get_load_back_mem_quota()` reserves `real_input_tokens + page_size` before + deciding how much host cache can be loaded back. +- `alloc_paged_token_slots_extend()` evicts by token count before allocating + extend KV slots. +- CP owner-lane retry uses `_evict_for_compute_owner_lanes()`. + +This protects the KV allocator from overcommit, but it is not a full CUDA memory +budget. It does not directly account for: + +- model forward temporary buffers, +- logits / hidden-state capture, +- MLP / DeepGEMM / EP communication scratch, +- CUDA graph memory, +- two-batch overlap staging, +- multiple concurrent HiCache load producers, +- draft/MTP compute-side temporary memory. + +The model runner profiles KV token capacity from `mem_fraction_static`, then CP +shared KV expands logical capacity by `attn_cp_size` while keeping the physical +pool at the profiled size. That means async HiCache must leave explicit +headroom at the scheduler layer for non-KV transient memory; otherwise token +accounting can look valid while CUDA still OOMs in a later model kernel. + +## Recommended State Model + +Introduce a logical async operation journal for CP HiCache. The journal can be +implemented incrementally around current queues, but the state model should be +explicit. + +Each operation should contain: + +```text +op_id +kind: BACKUP | LOAD | DEVICE_EVICT | HOST_EVICT +node_ids +logical_len +page_owners hash / metadata +target_physical_count +draft_physical_count +zero_owned flag +state: PREPARED | SUBMITTED | LOCAL_DONE | GLOBAL_COMMITTED | ABORTED +stream_event +``` + +Rules: + +1. All ranks append the same logical ops in the same order. +2. Physical transfer can be zero length. +3. Commit only advances a contiguous op prefix. +4. Commit is logical; transfer completion is local. +5. Rollback is local physical cleanup plus global logical abort. +6. Radix split/merge must update or block pending ops that reference affected + nodes. + +## Proposed Synchronization Policy + +### Fast path: deterministic, no collective sync + +Use only local events and existing streams for: + +- D2H backup submission. +- H2D load submission. +- Per-layer load/compute overlap. +- Zero-owned no-op transfer. +- Device eviction of nodes already host-backed. + +Fast-path correctness comes from logical op ids, radix order, and owner-pattern +metadata. A rank does not ask other ranks whether a page belongs to it; it +derives that from `page_owners`. A rank does not ask which node is being +committed; it consumes the next logical op in the same queue. + +### Normal commit path: local-event prefix commit + +For async backup/load completion, the desired normal path is: + +```text +local event for op N is complete +zero-owned ranks use an already-complete event +logical op queue is consumed in order +``` + +The commit operation should advance only a contiguous local completed prefix. +If all ranks have the same logical op order and every op has a local event, +then each rank can commit the same semantic prefix without per-tick global +voting. If a rank's local event is not ready, that rank simply does not expose +the later logical state yet; request scheduling must respect its local pending +state. + +### Batched validation path: amortized all-reduce + +Use `all_reduce(MIN)` on a monotonically increasing completed-op prefix only +when useful: + +- Debug/invariant validation is enabled. +- Local completed prefix has advanced by a large threshold and the system wants + to assert cross-rank alignment. +- Ongoing write count exceeds a safety threshold. +- Scheduler is about to make a memory-sensitive admission decision. +- Memory pressure requires releasing write locks. +- Idle checks need to drain async ops. + +Avoid all-reduce on every scheduler tick when no local progress happened. + +### Slow path: global success/abort gate + +Use collective sync only after local retry paths: + +- Host reserve failed and host eviction/retry was attempted. +- Draft host reserve failed. +- Owner-lane device allocation failed and lane-aware eviction/retry was + attempted. +- Write-back demotion cannot reserve host memory. + +Protocol: + +```text +local_success = reserve/alloc/submit success +global_success = all_reduce_min(local_success) +if global_success: + commit logical op +else: + rollback local physical reservations + abort logical op identically on all ranks +``` + +## Proposed Backup Lifecycle + +```text +RESIDENT + -> BACKUP_PREPARED reserve target/draft host slots locally + -> BACKUP_SUBMITTED enqueue D2H on write_stream + -> BACKUP_LOCAL_DONE finish_event ready locally + -> HOST_VALID global contiguous commit releases device write lock +``` + +Failure handling: + +- Host reserve failure enters host eviction slow path. +- Draft reserve failure must roll back target reserve. +- Global failure leaves node device-resident but not host-backed. +- Write-through failure should skip backup rather than block request progress. +- Write-back failure should not assert; it needs a clear fallback path. + +Request-visible state rule: + +- `BACKUP_PREPARED` and `BACKUP_SUBMITTED` must not be treated as + `HOST_VALID`. +- The device source remains protected until the global contiguous backup commit + releases the write lock. +- Host eviction must not select backup destinations that have not globally + committed. + +For the first implementation, this can be represented without a large new +journal by adding an explicit "pending host backup" state around the existing +`ongoing_write_through` and `ack_write_queue` path. The important behavioral +change is to distinguish: + +```text +host slots reserved and D2H submitted +``` + +from: + +```text +host bytes are valid and request-visible +``` + +## Proposed Load Lifecycle + +```text +HOST_VALID + DEVICE_EVICTED + -> LOAD_PREPARED alloc_pages_with_owners(page_owners) + -> LOAD_SUBMITTED enqueue per-layer H2D on load_stream + -> LAYER_READY[i] LayerDoneCounter event for layer i + -> DEVICE_VALID final load ack releases load lock +``` + +Failure handling: + +- Owner-lane allocation failure enters lane-aware eviction slow path. +- After retry, global success/abort gate decides whether all ranks load or all + ranks recompute. +- If aborted, radix host state remains valid, but no device state is committed. + +Request-visible state rule: + +- `LOAD_PREPARED` and `LOAD_SUBMITTED` are valid only for the batch that owns + the returned `hicache_consumer_index`. +- Other requests should not treat the loaded node as generally + `DEVICE_VALID` until the load ack commits. +- Host metadata remains valid throughout load; aborting a load should fall back + to recompute without invalidating the host cache. + +Current code already has a per-layer data dependency through +`LayerDoneCounter`, which is the right data-plane primitive. The missing +control-plane distinction is that `node.value` is assigned before the H2D load +fully completes. That is acceptable only if the scheduler ensures every +consumer of those device locs carries the same `hicache_consumer_index`, or if +the node remains in a `DEVICE_LOAD_PENDING` logical state until the load ack +commits. + +For the first implementation, prefer the second rule: + +```text +load owner batch may consume DEVICE_LOAD_PENDING with producer id +all other requests see the node as host-valid but not device-valid +``` + +This keeps layer-wise overlap for the triggering batch while preventing later +requests from racing on partially loaded device KV. + +## Proposed Eviction Lifecycle + +### Device eviction + +Host-backed node: + +```text +DEVICE_VALID + HOST_VALID + -> DEVICE_EVICTED + HOST_VALID +``` + +This remains cheap and synchronous in metadata because no transfer is required. + +Node under backup: + +```text +BACKUP_SUBMITTED + -> not evictable until backup commit or abort +``` + +Write-back demotion: + +```text +DEVICE_VALID + HOST_INVALID + -> DEMOTE_SUBMITTED + -> HOST_VALID + DEVICE_EVICTED +``` + +This can be async for proactive demotion. For emergency allocation, the system +may need to drop cache instead of waiting for demotion to finish. + +### Host eviction + +Use batched epochs: + +```text +collect victims by logical priority +free local physical target/draft host slots +accumulate local_freed_slots +all_reduce_min(done) only after a batch/epoch +commit logical host removal for the common evicted prefix +``` + +This should replace per-node or per-iteration all-reduce in the common +host-pressure path. + +## Radix Tree Async Safety Model + +The radix tree should treat async backup/load/evict state as part of node +residency, not as external controller-only state. + +Recommended per-node fields conceptually: + +```text +device_state: ABSENT | VALID | LOADING | WRITING_SOURCE | DEMOTE_PENDING +host_state: ABSENT | WRITING | VALID | EVICTING +pending_ops: set/op refs +split_epoch: integer +``` + +The implementation does not need exactly these fields, but it needs equivalent +semantics. + +### Safe split options + +There are three viable designs. + +#### Option A: Drain-before-split + +Before `_split_node()` splits a node with pending async state, force the related +op to commit or abort. + +Pros: + +- Smallest correctness surface. +- Easy to reason about. +- Good first implementation. + +Cons: + +- Partial-prefix traffic can block on D2H/H2D completion. +- May reduce write-through/load overlap in workloads with frequent shared + prefix refinement. + +#### Option B: Op-aware split + +When `_split_node()` splits a node with pending backup/load state, split the +pending op's logical segment as well. + +Requirements: + +- The ack must release all locks created by the original op after split. +- Host metadata must be split in the same way as radix metadata. +- The journal must understand that one submitted physical transfer corresponds + to multiple post-split logical nodes. +- Commit must be idempotent if a later split happens before the ack. + +Pros: + +- Preserves maximum async overlap. + +Cons: + +- Highest complexity. +- Easy to get lock/ref accounting wrong. + +#### Option C: Pending segment alias + +Represent an async operation as an immutable logical segment object independent +of the current radix node object. A split creates two radix nodes that both +reference slices of the same pending segment. The op ack commits the segment, +then materializes per-node state through the current aliases. + +Pros: + +- Cleaner than mutating ack lists on every split. +- Good long-term model if async operations become more pervasive. + +Cons: + +- Requires a larger refactor around `ongoing_write_through`, + `ongoing_load_back`, and split metadata. + +Recommended path: + +1. Start with **Option A** for backup/write-through nodes to eliminate the + highest-risk lock leak. +2. Keep path-lock load behavior as-is initially, but add assertions/tests around + splitting loaded/in-flight nodes. +3. Move toward **Option C** only if drain-before-split shows measurable + performance loss. + +### Leaf set and accounting rules + +Every async state transition must update the same accounting surfaces: + +- `evictable_size_` +- `protected_size_` +- `evictable_leaves` +- `evictable_host_leaves` +- `lock_ref` +- `host_ref_counter` + +The recent class of bugs where `evictable_size_` claimed capacity but +`evictable_leaves` had no usable victims shows that size counters and leaf sets +must be updated atomically with node residency changes. Async journal commit +should own these updates rather than scattering them across controller callbacks +and radix helpers. + +## GPU Memory Margin Policy + +Token-pool accounting should be treated as necessary but not sufficient. + +Recommended device memory budget layers: + +1. **Hard KV allocator budget** + - Existing `available_size + evictable_size - offsets`. + - Must remain the source of truth for KV slot allocation. + +2. **Owner-lane budget** + - CP shared KV must check not only total logical tokens but also per-owner + lane free/release pages. + - Load-back and extend should both use owner-lane retry before falling back + or aborting. + +3. **Async in-flight reserve** + - Loaded host hits occupy device KV slots before forward consumes them. + - Write locks reduce evictable KV slots until D2H commits. + - Future two-batch overlap must reserve for all active load producers and + next-batch extend slots, not just the current batch. + +4. **Non-KV CUDA headroom** + - Reserve a fixed or profiled margin for logits, hidden states, MLP/EP + scratch, CUDA graph memory, and speculative decoding temporaries. + - This margin should be outside radix token accounting. + +5. **Emergency reserve** + - Keep at least one page per owner lane, or a small per-lane percentage, as + an emergency reserve for forward progress. + - Avoid consuming the last owner-lane pages for opportunistic load-back. + +Practical admission rule: + +```text +allow load-back only if: + total_available_plus_evictable + - real_extend_tokens + - async_inflight_reserved_tokens + - emergency_owner_lane_reserve + >= host_hit_tokens + +and each required owner lane has enough free/releasable pages after reserve. +``` + +For two-batch overlap, the scheduler should reserve the maximum of: + +- current batch real extend tokens, +- current batch host load-back tokens, +- next prepared batch real extend tokens, +- next prepared batch host load-back tokens, +- per-slot logits/hidden temporary budget if those tensors are retained across + overlap boundaries. + +This is intentionally more conservative than pure token accounting. It is the +price of avoiding late CUDA OOM after the scheduler has already committed a +batch. + +## Implementation Direction + +Recommended order: + +1. **Correctness gates first** + - Backup global success/rollback after reserve retry. + - Load global abort after owner-lane retry. + - Target/draft atomic reserve and rollback. + - Drain or make split-aware any radix node with in-flight backup state. + +2. **Reduce synchronization frequency** + - Batch write ack commits by completed prefix. + - Skip all-reduce when local completed prefix has not changed. + - Batch host eviction into epochs. + - Move host eviction synchronization from per-victim to per-epoch. + +3. **Then consider async write-back demotion** + - Keep write-through and write-through-selective stable first. + - Add write-back async demotion only with explicit emergency fallback. + +4. **Add explicit memory reserves** + - Owner-lane emergency reserve. + - In-flight load/write reserve. + - Non-KV CUDA headroom for overlap/speculative paths. + +## Open Questions + +1. What is the acceptable maximum delay before write locks are released under + write-through? This determines ack batching thresholds. +2. Should host eviction be driven only on allocation failure, or should it keep a + low-frequency free-room background target? +3. For write-back, should emergency eviction drop cache immediately when host + reserve fails, or block until demotion completes? +4. Should the op journal be a new abstraction, or should existing + `ack_write_queue`, `ack_load_queue`, `ongoing_write_through`, and + `ongoing_load_back` be tightened first? +5. Should `_split_node()` drain pending backup ops, or should pending ops become + split-aware? +6. How much non-KV CUDA headroom is required for GLM5 + EAGLE/MTP + EP under + two-batch overlap? + +## Summary + +The current code already has the right low-level pieces: + +- rank-local CP physical filtering, +- zero-owned no-op acks, +- async write/load streams, +- per-layer load/compute overlap, +- owner-pattern-preserving reload. + +The missing piece is a unified logical transaction model. Without it, making +backup/load/evict more asynchronous increases the risk of cross-rank divergence. +The desired direction is: + +> Fast path uses local async events only; commit path uses batched prefix +> all-reduce; failure path uses explicit global success/abort; physical pool +> operations remain rank-local. + +Additionally, radix tree mutation and GPU memory admission must be part of the +same design. Async state cannot be treated as controller-private state because +radix split, lock release, leaf eligibility, and memory quota decisions all +depend on it. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 57a6d9341..be59d995f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -613,6 +613,17 @@ class Scheduler( DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args) self.draft_worker = DraftWorkerClass(**draft_worker_kwargs) + def _get_draft_token_to_kv_pool_for_hicache(self): + if self.draft_worker is None or self.spec_algorithm.is_ngram(): + return None + if self.spec_algorithm.supports_spec_v2() and self.enable_overlap: + if self.server_args.enable_multi_layer_eagle: + draft_runner = self.draft_worker.draft_worker.draft_runner_list[0] + else: + draft_runner = self.draft_worker.draft_worker.draft_runner + return draft_runner.token_to_kv_pool + return self.draft_worker.model_runner.token_to_kv_pool + def init_model_worker(self): self.init_tp_model_worker() self.maybe_init_draft_worker() @@ -727,6 +738,7 @@ class Scheduler( pp_size=self.pp_size, chunked_prefill_size=server_args.chunked_prefill_size, sliding_window_size=self.sliding_window_size, + draft_token_to_kv_pool=self._get_draft_token_to_kv_pool_for_hicache(), ) if ( diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index 6de3f984f..08d01be9b 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -1,7 +1,7 @@ from __future__ import annotations import dataclasses -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional import torch @@ -33,6 +33,7 @@ class CacheInitParams: chunked_prefill_size: Optional[int] = None sliding_window_size: Optional[int] = None + draft_token_to_kv_pool: Optional[Any] = None # Time-to-live for cache entries in seconds. If None, TTL is disabled. cache_ttl_seconds: Optional[float] = None diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e79b60675..eaf23cbbc 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -53,6 +53,70 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _estimate_hicache_size_per_token(kv_cache) -> int: + dtype_size = kv_cache.store_dtype.itemsize + if isinstance(kv_cache, MHATokenToKVPool): + return kv_cache.head_dim * kv_cache.head_num * kv_cache.layer_num * dtype_size * 2 + if isinstance(kv_cache, NSATokenToKVPool): + indexer_size_per_token = ( + kv_cache.index_head_dim + + kv_cache.index_head_dim // kv_cache.quant_block_size * 4 + ) + return ( + kv_cache.kv_cache_dim * dtype_size * kv_cache.layer_num + + indexer_size_per_token + * kv_cache.layer_num + * NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize + ) + if isinstance(kv_cache, MLATokenToKVPool): + kv_cache_dim = kv_cache.kv_lora_rank + kv_cache.qk_rope_head_dim + return kv_cache_dim * dtype_size * kv_cache.layer_num + raise ValueError(f"Unsupported HiCache KV pool type: {type(kv_cache).__name__}") + + +def _compute_shared_hicache_token_capacities( + *, + total_host_bytes: int, + target_size_per_token: int, + draft_size_per_token: int, + page_size: int, +) -> Tuple[int, int]: + if total_host_bytes <= 0: + raise ValueError(f"total_host_bytes must be positive, got {total_host_bytes}") + if target_size_per_token <= 0: + raise ValueError( + f"target_size_per_token must be positive, got {target_size_per_token}" + ) + if draft_size_per_token <= 0: + raise ValueError( + f"draft_size_per_token must be positive, got {draft_size_per_token}" + ) + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + + target_pages = total_host_bytes // ( + (target_size_per_token + draft_size_per_token) * page_size + ) + if target_pages <= 0: + raise ValueError( + "HiCache shared target+draft budget cannot hold a single page: " + f"total_host_bytes={total_host_bytes}, " + f"target_size_per_token={target_size_per_token}, " + f"draft_size_per_token={draft_size_per_token}, page_size={page_size}" + ) + + target_tokens = int(target_pages * page_size) + remaining_bytes = total_host_bytes - target_tokens * target_size_per_token + draft_pages = remaining_bytes // (draft_size_per_token * page_size) + draft_tokens = int(draft_pages * page_size) + if draft_tokens < target_tokens: + raise ValueError( + "HiCache shared budget computation produced draft capacity below target " + f"capacity: target_tokens={target_tokens}, draft_tokens={draft_tokens}" + ) + return target_tokens, draft_tokens + + @dataclass class CpHiCacheNodeMetadata: logical_len: int @@ -170,43 +234,56 @@ class CpHiCacheNodeMetadata: class HiRadixCache(RadixCache): - def _create_token_to_kv_pool_host(self, kv_cache, server_args: ServerArgs): + def _create_token_to_kv_pool_host( + self, + kv_cache, + server_args: ServerArgs, + *, + host_size: Optional[int] = None, + host_token_capacity: Optional[int] = None, + ): from sglang.srt.mem_cache.memory_pool_host import ( MHATokenToKVPoolHost, MLATokenToKVPoolHost, NSATokenToKVPoolHost, ) + host_size = server_args.hicache_size if host_size is None else host_size if isinstance(kv_cache, MHATokenToKVPool): return MHATokenToKVPoolHost( kv_cache, server_args.hicache_ratio, - server_args.hicache_size, + host_size, self.page_size, server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, + host_token_capacity=host_token_capacity, ) if isinstance(kv_cache, NSATokenToKVPool): return NSATokenToKVPoolHost( kv_cache, server_args.hicache_ratio, - server_args.hicache_size, + host_size, self.page_size, server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, + host_token_capacity=host_token_capacity, ) if isinstance(kv_cache, MLATokenToKVPool): return MLATokenToKVPoolHost( kv_cache, server_args.hicache_ratio, - server_args.hicache_size, + host_size, self.page_size, server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, + host_token_capacity=host_token_capacity, ) raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools") - def attach_draft_kv_pool(self, draft_token_to_kv_pool) -> None: + def attach_draft_kv_pool( + self, draft_token_to_kv_pool, host_token_capacity: Optional[int] = None + ) -> None: if not self._uses_cp_hicache or draft_token_to_kv_pool is None: return if ( @@ -215,7 +292,10 @@ class HiRadixCache(RadixCache): ): return draft_token_to_kv_pool_host = self._create_token_to_kv_pool_host( - draft_token_to_kv_pool, self._server_args + draft_token_to_kv_pool, + self._server_args, + host_size=0 if host_token_capacity is not None else None, + host_token_capacity=host_token_capacity, ) self.cache_controller.attach_draft_pool( draft_token_to_kv_pool, draft_token_to_kv_pool_host @@ -252,8 +332,41 @@ class HiRadixCache(RadixCache): "CP shared KV HiCache host integration requires NSATokenToKVPool." ) self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() + draft_token_to_kv_pool = getattr(params, "draft_token_to_kv_pool", None) + target_host_token_capacity = None + draft_host_token_capacity = None + if ( + self._uses_cp_hicache + and draft_token_to_kv_pool is not None + and server_args.hicache_size > 0 + ): + target_size_per_token = _estimate_hicache_size_per_token(self.kv_cache) + draft_size_per_token = _estimate_hicache_size_per_token( + draft_token_to_kv_pool + ) + target_host_token_capacity, draft_host_token_capacity = ( + _compute_shared_hicache_token_capacities( + total_host_bytes=int(server_args.hicache_size * 1e9), + target_size_per_token=target_size_per_token, + draft_size_per_token=draft_size_per_token, + page_size=self.page_size, + ) + ) + logger.info( + "[HiCache-draft] sharing --hicache-size across target+draft: " + "budget_gb=%d target_size_per_token=%d draft_size_per_token=%d " + "target_tokens=%d draft_tokens=%d", + server_args.hicache_size, + target_size_per_token, + draft_size_per_token, + target_host_token_capacity, + draft_host_token_capacity, + ) self.token_to_kv_pool_host = self._create_token_to_kv_pool_host( - self.kv_cache, server_args + self.kv_cache, + server_args, + host_size=0 if target_host_token_capacity is not None else None, + host_token_capacity=target_host_token_capacity, ) self.tp_group = params.tp_cache_group @@ -302,6 +415,11 @@ class HiRadixCache(RadixCache): enable_storage_metrics=self.enable_storage_metrics, cp_shared_kv_layout=cp_shared_kv_layout, ) + if draft_token_to_kv_pool is not None: + self.attach_draft_kv_pool( + draft_token_to_kv_pool, + host_token_capacity=draft_host_token_capacity, + ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, prefetch_threshold=prefetch_threshold, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index aa4a72fdb..72c466c17 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -147,6 +147,7 @@ class HostKVCache(abc.ABC): pin_memory: bool, device: str, allocator_type: str = "default", + host_token_capacity: Optional[int] = None, ): self.device_pool = device_pool self.page_size = page_size @@ -157,12 +158,20 @@ class HostKVCache(abc.ABC): self.dtype = device_pool.store_dtype self.size_per_token = self.get_size_per_token() - if host_size > 0: + if host_token_capacity is not None: + self.size = int(host_token_capacity) + elif host_size > 0: self.size = int(host_size * 1e9 // self.size_per_token) else: self.size = int(device_pool.size * host_to_device_ratio) - # Align up the host memory pool size to the page size - self.page_num = self.size // self.page_size + 1 + # Align up the host memory pool size to the page size. Preserve the + # historical extra page for ratio/byte-budget paths, but treat an + # explicit token capacity as already being the logical budget chosen by + # the caller. + if host_token_capacity is not None: + self.page_num = (self.size + self.page_size - 1) // self.page_size + else: + self.page_num = self.size // self.page_size + 1 self.size = self.page_num * self.page_size self.start_layer = device_pool.start_layer self.end_layer = device_pool.end_layer @@ -286,6 +295,7 @@ class MHATokenToKVPoolHost(HostKVCache): pin_memory: bool = True, device: str = "cpu", allocator_type: str = "default", + host_token_capacity: Optional[int] = None, ): super().__init__( device_pool, @@ -296,6 +306,7 @@ class MHATokenToKVPoolHost(HostKVCache): pin_memory, device, allocator_type, + host_token_capacity=host_token_capacity, ) self.element_dim = self.device_pool.head_num * self.device_pool.head_dim self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( @@ -747,6 +758,7 @@ class MLATokenToKVPoolHost(HostKVCache): device: str = "cpu", allocator_type: str = "default", override_kv_cache_dim: Optional[int] = None, + host_token_capacity: Optional[int] = None, ): self.override_kv_cache_dim = override_kv_cache_dim super().__init__( @@ -758,6 +770,7 @@ class MLATokenToKVPoolHost(HostKVCache): pin_memory, device, allocator_type, + host_token_capacity=host_token_capacity, ) self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] self.data_ptrs = torch.tensor( @@ -1088,6 +1101,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): pin_memory: bool = True, device: str = "cpu", allocator_type: str = "default", + host_token_capacity: Optional[int] = None, ): # Initialize indexer metadata before HostKVCache.__init__ calls get_size_per_token. self.index_head_dim = device_pool.index_head_dim @@ -1107,6 +1121,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device, allocator_type, override_kv_cache_dim=device_pool.kv_cache_dim, + host_token_capacity=host_token_capacity, ) self.indexer_page_stride_size = ( self.indexer_size_per_token * self.page_size * self.indexer_dtype.itemsize 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 3dbb0ae2d..ac9a8b0a0 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -91,7 +91,11 @@ for _schema in ( raise from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams -from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache +from sglang.srt.mem_cache.hiradix_cache import ( + CpHiCacheNodeMetadata, + HiRadixCache, + _compute_shared_hicache_token_capacities, +) from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -117,6 +121,33 @@ class TestCpHiCacheImports(CustomTestCase): ) +class TestHiRadixCacheCPDraftHostPool(CustomTestCase): + def test_shared_budget_keeps_draft_at_least_target_capacity(self): + target_tokens, draft_tokens = _compute_shared_hicache_token_capacities( + total_host_bytes=1000, + target_size_per_token=6, + draft_size_per_token=2, + page_size=10, + ) + + self.assertEqual(target_tokens, 120) + self.assertEqual(draft_tokens, 140) + self.assertGreaterEqual(draft_tokens, target_tokens) + self.assertLessEqual(target_tokens * 6 + draft_tokens * 2, 1000) + + def test_shared_budget_handles_equal_target_and_draft_size(self): + target_tokens, draft_tokens = _compute_shared_hicache_token_capacities( + total_host_bytes=1000, + target_size_per_token=6, + draft_size_per_token=6, + page_size=10, + ) + + self.assertEqual(target_tokens, 80) + self.assertEqual(draft_tokens, 80) + self.assertLessEqual(target_tokens * 6 + draft_tokens * 6, 1000) + + class TestCpHiCacheNodeMetadata(CustomTestCase): def test_split_zero_len_moves_all_positions_to_child(self): metadata = CpHiCacheNodeMetadata(