diff --git a/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md b/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md new file mode 100644 index 000000000..f0fd9817d --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md @@ -0,0 +1,351 @@ +# NSA Prefill CP / CP Shared KV Batch Size > 1 Exploration + +Date: 2026-05-28 + +Scope: + +- NSA prefill context parallelism. +- `--nsa-prefill-cp-mode in-seq-split`. +- CP shared KV. +- HiCache L2/L1 load path. +- MTP / EAGLE draft KV integration. + +This note records why true `ForwardBatch(batch_size > 1)` is not currently a +safe near-term throughput lever for the CP shared-KV path, and why the safer +near-term direction is multi-slot overlap where each slot keeps +`batch_size == 1`. + +## Motivation + +In the cache-hit-heavy workload, observed prefix cache hit can be above 90%. +The remaining extend length is often only 2K-10K tokens. A single request then +does not always provide enough work to fill the GPU. It is natural to ask +whether prefill should group multiple short cache-hit extends into one batch. + +SGLang's generic prefill scheduler and `ScheduleBatch.prepare_for_extend()` can +construct a multi-request extend batch. The blocker is the specialized NSA +prefill CP + CP shared-KV path layered underneath it. + +## Terminology + +### True batch + +One `ForwardBatch` contains multiple sequences: + +```text +ForwardBatch(batch_size=N) +``` + +All NSA metadata, page tables, top-k transforms, materialize remaps, direct +writes, and MTP draft paths must be batch-aware. + +### Multi-slot overlap + +Multiple independent request slots are scheduled together, but each slot owns +its own `ScheduleBatch`, `ModelWorkerBatch`, and `ForwardBatch`: + +```text +slot0: ForwardBatch(batch_size=1) +slot1: ForwardBatch(batch_size=1) +... +``` + +This is the Phase9 direction. It keeps current single-request CP invariants +while allowing communication/materialize from one slot to overlap useful work +from another slot. + +## Current evidence: true batch > 1 is not supported on the CP shared-KV path + +### 1. In-seq CP metadata is single-request by construction + +`prepare_input_dp_with_cp_dsa()` builds the zigzag split for one full-length +request. + +Relevant code: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `prepare_input_dp_with_cp_dsa(...)` + - comment: `# just support batch = 1` + - `bs_per_cp_group = 1` + - computes one `split_list`, one `zigzag_index`, one pair of + `kv_len_prev/next`, and one pair of `actual_seq_q_prev/next`. + +Implication: + +For true batch > 1, concatenating multiple requests and passing their total +token count into this function would treat them as one long sequence. That does +not preserve per-request causal boundaries, per-request prefix/suffix splits, +or per-request owner-lane page semantics. + +### 2. Page-aligned split only activates for batch size 1 + +The shared-KV path depends on page-aligned in-seq split so direct writes land in +owner lanes and later materialize/load paths can reproduce the same logical +page ownership. + +Relevant code: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `_build_in_seq_split_for_forward_batch(...)` + - page-aligned split requires: + +```python +len(forward_batch.extend_seq_lens_cpu) == 1 +len(forward_batch.extend_prefix_lens_cpu) == 1 +``` + +If those conditions fail, the code falls back to token-balanced split metadata. + +Implication: + +True batch > 1 loses the page-aligned split contract required by the current +CP shared-KV direct-write path. + +### 3. Compute-owner allocation supports only batch size 1 + +CP shared KV relies on owner-lane-aware allocation: newly allocated logical +pages should come from the lane owned by the CP rank that computes the page. + +Relevant code: + +- `python/sglang/srt/mem_cache/allocator.py` + - `alloc_extend_compute_owner(...)` + - raises if `len(prefix_lens_cpu) != 1` or `len(seq_lens_cpu) != 1`. +- `python/sglang/srt/mem_cache/common.py` + - `alloc_paged_token_slots_extend(...)` + - only builds `page_compute_owners` when `len(prefix_lens_cpu) == 1`. + - batch > 1 records reason `multi_batch` and falls back to legacy allocation. + +Implication: + +True batch > 1 either fails the owner-lane allocator or falls back to legacy +allocation. That fallback weakens the invariant that direct write, HiCache +backup/load, and later CP materialize assume: logical page owner must match the +physical owner lane. + +### 4. Direct-write path depends on single-request page-aligned metadata + +The MLA KV and NSA index direct-write paths get this rank's local output locs +through: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `get_cp_shared_kv_local_out_cache_loc(...)` + - checks `metadata.page_aligned`. + - checks `sum(metadata.split_list) == out_cache_loc.numel()`. + - uses `cp_split_and_rebuild_1d(...)` with the single-request CP metadata. + - validates local loc ownership by `layout.owned_by_this_rank(...)`. + +Callers: + +- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` + - `_maybe_write_cp_shared_local_mla_kv(...)` +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + - `_store_cp_shared_local_index_k_cache(...)` + +Implication: + +If true batch > 1 does not produce per-request page-aligned metadata and +owner-lane-aligned `out_cache_loc`, direct write should fall back. This removes +a major reason CP shared KV is fast. + +### 5. Phase8 prefix prefetch is explicitly batch size 1 + +The current L1 shared-KV prefix prefetch path is conservative and only supports +one request. + +Relevant code: + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` + - `CpSharedKVMlaPrefetcher.maybe_create(...)` + - rejects `forward_batch.batch_size != 1`. + - requires `len(extend_prefix_lens_cpu) == 1`. +- same file: + - `CpSharedKVIndexPrefetcher.maybe_create(...)` + - same batch-size and prefix-lens restrictions. + +Implication: + +If true batch > 1 is used in a cache-hit-heavy workload, the exact path we need +to hide -- prefix materialize/reduce from L1 shared KV -- is disabled. The code +falls back to synchronous materialize in the attention path. + +### 6. NSA top-k in-seq CP pair path has single-batch assumptions + +Relevant code: + +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + - `_get_topk_in_seq_cp_pair(...)` + - TODO: `support mutil-batch` + - uses one `kv_len_prev/next` pair and one `actual_seq_q_prev/next` pair from + `forward_batch.nsa_cp_metadata`. + +Implication: + +True batch > 1 would require per-request CP pair metadata and top-k outputs +that preserve each sequence's independent query/key ranges. + +### 7. MTP / EAGLE draft path depends on the same single-request CP metadata + +Relevant code: + +- `python/sglang/srt/models/deepseek_nextn.py` + - draft model local path uses `cp_split_and_rebuild_1d(...)`, + `cp_split_and_rebuild_data(...)`, and `cp_collect_last_token_hidden(...)`. +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `_in_seq_collect_last_token(...)` only uses the exact in-seq owner/offset + calculation when `bs == 1`. + +Implication: + +Even if target-model CP shared KV were made batch-aware, MTP/EAGLE would still +need per-request draft hidden state slicing and per-request last-token +collection semantics before true batch > 1 can be declared safe. + +### 8. HiCache load is closer to batch-capable, but it inherits owner semantics + +The CP HiCache load path can merge load ops across nodes: + +- `python/sglang/srt/managers/cache_controller.py` + - `load_cp(nodes_to_load, ...)` + - collects `page_owners` across nodes. + - allocates fresh device pages through `alloc_pages_with_owners(...)`. + - appends target and draft load ops. + - `start_loading()` merges load queues and performs per-layer load. + +This part is not the primary batch-size blocker. + +However, it assumes every backed-up node carries valid CP `page_owners` and +owned-position metadata. If the original true-batch allocation/split path does +not preserve owner-lane semantics, HiCache load can faithfully reproduce an +already-wrong owner pattern. + +## Why true batch may not improve this workload immediately + +For cache-hit-heavy short extends, batching increases query work, but it can +also increase or re-expose prefix materialize/reduce work: + +```text +cache-hit prefix pages + -> L2 host to L1 device load + -> L1 shared-KV materialize/reduce + -> top-k / attention +``` + +If true batch disables Phase8 prefetch and direct write, the critical path may +become worse even though more requests are grouped. + +The throughput bottleneck is likely not only "too few query tokens". It is also +"attention-visible shared-KV prefix preparation is not fully hidden". + +## Recommended near-term direction + +Use multi-slot overlap before true batch > 1. + +Rationale: + +1. Keeps current per-slot `batch_size == 1` invariants. +2. Preserves page-aligned split, owner-lane allocation, direct write, and + Phase8 prefetch eligibility. +3. Allows slot A's attention/MLP to hide slot B's CP materialize/prefetch/load + work. +4. Avoids redesigning `NSAContextParallelMetadata`, top-k, MTP hidden slicing, + and direct-write remaps in one large step. + +This matches the Phase9 plan: + +- `docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md` + - distinguishes multi-slot from true `ForwardBatch(batch_size=N)`. + - explicitly keeps each initial slot at `batch_size == 1`. + - defers true batch support until batch-aware NSA CP metadata exists. + +## Requirements for future true batch > 1 support + +True batch support should be treated as a separate phase with explicit +fail-fast gates until all requirements below are implemented. + +### P0: hard guard / loud fallback + +When `uses_cp_shared_kv` and `batch_size > 1`, the system should not silently +fall back to a slow or semantically weaker path. + +Acceptable behavior during development: + +- fail fast for unsupported combinations; or +- emit a loud warning with exact fallback reason and expected performance loss. + +### P1: batch-aware NSA CP metadata + +`NSAContextParallelMetadata` needs per-request fields: + +- per-request `split_list`; +- per-request `zigzag_index`; +- per-request `kv_len_prev/next`; +- per-request `actual_seq_q_prev/next`; +- per-request page-aligned segment ranges; +- a flattened representation usable by kernels without losing request + boundaries. + +### P2: batch-aware page-aligned split and owner-lane allocation + +For each request in the batch: + +1. compute page-aligned split independently; +2. build page owners for that request's suffix; +3. concatenate page-owner plans in the same order as `out_cache_loc`; +4. allocate pages with matching owner lanes; +5. preserve per-request offsets for direct write and materialize remap. + +### P3: batch-aware direct write + +`get_cp_shared_kv_local_out_cache_loc(...)` cannot treat `out_cache_loc` as one +single in-seq split. It needs to produce local locs by request, then concatenate +them in local compute order while preserving owner-lane validation. + +### P4: batch-aware Phase8 L1 shared-KV prefetch + +The MLA and index prefetchers need: + +- per-request prefix page ranges; +- per-request prefix page alignment checks; +- a batched slot-remap layout that preserves the current dense-page semantics; +- deterministic collective ordering across ranks. + +### P5: batch-aware top-k CP pair path + +`_get_topk_in_seq_cp_pair(...)` needs per-request CP pair metadata rather than +one global `kv_len_prev/next` pair. + +### P6: batch-aware MTP / EAGLE local path + +Draft model CP local input, spec hidden states, and last-token hidden +collection need per-request CP split and gather semantics. + +### P7: tests + +Minimum tests before enabling true batch > 1: + +- two cache-hit requests with different prefix lengths; +- page-aligned and non-page-aligned suffixes; +- short radix-hit suffix with replicated compute path; +- HiCache host hit + load-back; +- draft/MTP enabled with target/draft KV mirrored; +- compare batch > 1 output against sequential batch-size-1 execution; +- verify no direct-write fallback under supported cases. + +## Practical conclusion + +For the current high-cache-hit workload, true `ForwardBatch(batch_size > 1)` is +not the right first optimization. The current CP shared-KV path has hard and +soft batch-size-1 assumptions in split metadata, owner-lane allocation, +direct-write, prefix prefetch, top-k, and draft MTP. + +The lower-risk path is: + +```text +multi-slot CP overlap first + -> each slot keeps batch_size == 1 + -> preserve shared-KV correctness invariants + -> overlap CP materialize/load/prefix-prefetch with peer slot compute + -> only later add true batch-aware CP metadata +``` + diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_load_prefetch_overlap_notes.md b/docs/advanced_features/nsa_prefill_cp_hicache_load_prefetch_overlap_notes.md new file mode 100644 index 000000000..c8435cc32 --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_hicache_load_prefetch_overlap_notes.md @@ -0,0 +1,1162 @@ +# NSA Prefill CP HiCache Load and Shared-KV Prefetch Overlap Notes + +Date: 2026-05-28 + +Branch context: `cjy-cp-refactor` after the direct/page_first_direct HiCache +load/backup work. + +This note records the current understanding of cache-hit prefill when all of the +following are enabled: + +```text +--enable-nsa-prefill-context-parallel +--nsa-prefill-cp-mode in-seq-split +--enable-nsa-prefill-cp-shared-kv +--enable-hierarchical-cache +--hicache-io-backend direct +--hicache-mem-layout page_first_direct +SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1 +``` + +The important conclusion is that **HiCache L2->L1 load overlap is necessary but +not sufficient**. In CP shared-KV mode, attention cannot consume the restored L1 +KV shard directly. It must still run the shared-KV compatibility path that +materializes dense prefix buffers and performs CP collectives. Therefore the +cache-hit critical path has two distinct prefetch/load stages: + +1. **HiCache host L2 -> device L1 shard load**. +2. **CP shared-KV L1 shard -> attention-visible dense prefix materialization**. + +Optimizing only stage 1 can leave the real stall in stage 2. + +--- + +## 1. Current cache-hit pipeline + +### 1.1 L2 -> L1 HiCache load + +On a host-cache hit, radix/scheduler logic first restores host-backed nodes into +fresh device KV slots: + +```text +schedule_policy.py + AddReq -> tree_cache.init_load_back(...) + +hiradix_cache.py + HiRadixCache.load_back(...) + -> cache_controller.load_cp(...) + +cache_controller.py + HiCacheController.load_cp(...) + -> alloc_pages_with_owners(page_owners) + -> enqueue target/draft load ops + +scheduler.py + ready_to_load_host_cache() + -> cache_controller.start_loading() +``` + +For CP shared KV, `load_cp()` preserves the write-time owner pattern via +`node.cp_hicache.page_owners`. This is required because a later cache hit +allocates fresh logical pages; if the owner sequence changes, each rank would load +its host bytes into slots interpreted as another rank's shard. + +With `direct/page_first_direct`, each layer uses the TAI direct H2D path: + +```text +memory_pool_host.py + load_to_device_per_layer(..., io_backend="direct") + -> tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_pf_lf(...) +``` + +The TAI op uses `cudaMemcpyBatchAsync`, so this path does not spend SM cycles on +an H2D copy kernel when the runtime capability is available. It is submitted on +SGLang's `load_stream`. + +### 1.2 Layer-wise wait point + +`start_loading()` submits per-layer H2D work and records one layer event per +layer. The forward path receives `hicache_consumer_index`, and the KV pool waits +when a layer's buffer is actually accessed: + +```text +memory_pool.py + get_key_buffer(layer_id) / get_value_buffer(layer_id) + -> layer_transfer_counter.wait_until(layer_id) +``` + +So the current L2->L1 behavior is: + +```text +load_stream: load L0 -> event0 -> load L1 -> event1 -> load L2 -> ... +forward_stream: wait L0 -> compute L0 -> wait L1 -> compute L1 -> ... +``` + +This provides layer-wise overlap, but it is not a complete cache-hit solution for +shared KV because restored L1 KV remains sharded. + +### 1.3 CP shared-KV dense prefix materialization + +In CP shared-KV mode, persistent KV/index cache is sharded by owner lane. Before +MLA/NSA attention can use a historical prefix, each layer still needs to build an +attention-visible dense view: + +```text +shared physical L1 shards + -> local materialize/remap into dense prefix buffer + -> CP all-reduce / range reduce + -> attention consumes dense rows/pages +``` + +Phase 8 adds one-layer-ahead prefix prefetch for two payloads: + +- MLA KV prefix. +- NSA index K/scale prefix. + +The intended Phase 8 flow is: + +```text +Layer L prepare: + start Layer L+1 prefix index + MLA materialize prefetch + +Layer L attention: + consume Layer L prefetched prefix if available + synchronously materialize/reduce current suffix only + run attention + +Layer L+1 attention: + wait for Layer L+1 prefix prefetch only if it is not done yet + synchronously materialize/reduce current suffix only +``` + +This means a cache-hit layer may have both: + +```text +HiCache L2->L1 H2D readiness wait +CP shared-KV prefix materialization/reduce readiness wait +``` + +A useful cache-hit overlap design must reason about both. + +--- + +## 2. Why the current overlap is weaker than expected + +The current direct/page_first_direct HiCache load path overlaps only the physical +host-to-device restore. It does not by itself hide the shared-KV compatibility +work that follows. + +### 2.1 Front-loaded L2->L1 submit + +`start_loading()` is called before `prepare_for_extend()`/forward execution and +submits per-layer load ops in a Python loop. The copy work is asynchronous, but +submission and index preparation are front-loaded before the model starts doing +layer compute. + +This is acceptable for correctness, but it means L2->L1 overlap is not the same +as "submit exactly at layer boundary". The first layer still waits, and CPU +submit overhead is paid before useful model compute can hide it. + +### 2.2 Shared-KV prefetch depends on restored L1 KV buffers + +The Phase 8 prefetcher eventually reads `token_to_kv_pool.get_key_buffer(...)` or +`get_index_k_with_scale_buffer(...)` for the next layer. In a HiCache hit, those +buffer getters are also guarded by the HiCache layer-transfer counter. Therefore +the CP shared-KV prefetch launch can be implicitly delayed by the L2->L1 layer +load event. + +Effective dependency: + +```text +L2 host page -> L1 physical shard loaded + -> shared-KV prefix prefetch can read that layer's shard + -> dense prefix materialize + CP reduce + -> attention can consume prefix +``` + +If L2->L1 load for layer `L+1` is late, starting the shared-KV prefetch for +layer `L+1` from layer `L` may still block at the buffer getter and lose the +intended overlap window. + +### 2.3 Prefix prefetch and HiCache load use different readiness abstractions + +HiCache load readiness is represented by `LayerDoneCounter` events attached to +KV pool access. Shared-KV prefetch readiness is represented by prefetch handles +and CUDA events inside `CpSharedKVMlaPrefetcher` and +`CpSharedKVIndexPrefetcher`. + +These mechanisms are both correct locally, but they are not currently planned as +a single cache-hit pipeline. As a result, each stage can appear overlapped in +isolation while the composed pipeline still serializes at consumption points. + +### 2.4 Bandwidth and stream contention remain possible + +The direct/page_first_direct L2->L1 load uses copy-engine H2D, while shared-KV +materialization/reduce uses GPU work and NCCL/PyNccl collectives. They are not +the same operation, but in production they can still contend for: + +- host pinned-memory bandwidth, +- PCIe/NVLink copy resources, +- NVLink/NCCL collective bandwidth, +- stream scheduling and event dependencies, +- CPU submission time. + +Seeing no end-to-end cache-hit speedup after L2->L1 load overlap does not prove +that L2->L1 overlap is broken; it can mean the post-load shared-KV prefix path is +now the limiting stage. + +--- + +## 3. Improved overlap model + +For CP shared KV + HiCache, the useful mental model is a two-stage prefetch DAG, +not a single H2D load: + +```text +For layer L+1 on a cache-hit request: + +Stage A: HiCache restore + host L2 page_first_direct target/draft/index payload + -> direct H2D into owner-lane L1 physical slots + -> layer-load event A[L+1] + +Stage B: shared-KV prefix prefetch + wait/read L1 physical shard for L+1 prefix + -> local dense prefix materialize/remap + -> async CP reduce / prefix event B[L+1] + +Stage C: layer L+1 attention consume + wait A[L+1] only when raw L1 buffer is needed + wait B[L+1] only when dense prefix is needed + materialize current/suffix rows synchronously + run attention +``` + +The target is to launch Stage A and Stage B early enough that Stage C rarely +waits on either event. + +### Desired staggered timeline + +The expected cache-hit pipeline is **two-layer-ahead L2->L1 prefetch** plus the +existing **one-layer-ahead shared-KV L1 prefix prefetch**: + +```text +Layer L-2 / batch setup: + start L2->L1 HiCache load for Layer L + target event: raw owner-lane L1 shard for Layer L becomes ready before Layer L-1 + +Layer L-1 prepare: + start Layer L shared-KV prefix prefetch from the restored L1 shard + prefetch_stream waits on Layer L L2->L1 load event if needed + forward_stream should not block just to launch this prefetch + +Layer L-1 attention / MLP: + hide Layer L shared-KV prefix materialize and CP reduce + +Layer L attention: + consume prefetched prefix + only do current/suffix materialize on the critical path +``` + +Equivalently, for a layer `X` that will consume cache-hit prefix data: + +```text +Stage A: L2->L1 load for X starts at X-2 or earlier +Stage B: L1 shared-KV prefetch X starts at X-1 +Stage C: attention consume X runs at X +``` + +The two-layer lead for Stage A is not an arbitrary latency knob. It follows from +the dependency chain: Stage B reads L1 physical shards, so Stage A must already be +far enough ahead that Stage B can launch without stealing the current layer's +compute window. + +### Practical implication + +`SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1` should be evaluated together with +HiCache load state. A prefetcher that starts "one layer ahead" but immediately +waits for the L2->L1 load of that next layer has less overlap than expected. + +The current all-layer `start_loading()` submission is not by itself proof that +Stage A has the required two-layer lead. The useful property is readiness at the +Stage-B launch point: by Layer `L-1` prepare, Layer `L`'s L2->L1 load should be +complete or represented as a local CUDA event that the prefetch stream can wait +on without blocking the forward stream. + +--- + +## 4. Concrete improvement directions + +### 4.1 Make cache-hit prefetch dependency explicit + +Do not treat HiCache load and shared-KV prefetch as independent optimizations. +Represent their dependency explicitly in the design: + +```text +shared-KV prefix prefetch for layer X requires L2->L1 load event for layer X +``` + +This does not necessarily require a new collective. It can be expressed through +local CUDA events because the dependency is local to each rank's physical shard. + +### 4.2 Start L2->L1 load early enough for next-layer shared-KV prefetch + +The current load stream submits all layers before forward, which is already early +for the copy itself. The remaining question is whether CPU submission and copy +engine progress are far enough ahead of Phase 8's next-layer prefetch hook. + +The useful measurement is not just "load submitted". It is: + +```text +At Layer L-1 prepare, is L2->L1 load for Layer L already complete or close enough +that shared-KV prefix prefetch can launch without stalling? +``` + +### 4.3 Chain shared-KV prefix prefetch after load events without blocking the +current layer + +If the next-layer L1 shard is not ready yet, the prefetch path should record a +stream wait on the load event and enqueue the prefetch work on the prefetch +stream, rather than blocking the current forward stream while calling a KV buffer +getter. + +Conceptually, when Layer `L-1` wants to prefetch Layer `L`: + +```text +prefetch_stream.wait_event(load_event[L]) +prefetch_stream: materialize prefix for L + async reduce +forward_stream: continue Layer L-1 compute +``` + +This preserves the dependency while keeping the current layer's compute window +available for overlap. + +A direct call path that goes through `get_key_buffer(L)` or +`get_index_k_with_scale_buffer(L)` can accidentally turn this into a forward +stream wait, because those accessors enforce `LayerDoneCounter.wait_until(L)`. +The improved design therefore needs a non-blocking way for the shared-KV +prefetcher to depend on the HiCache load event, or a raw-buffer access path used +only after the prefetch stream has been ordered behind that event. + +### 4.4 Keep suffix/current materialization separate + +Prefix is safe to prefetch because it already exists in the persistent pool after +cache hit/load. Current/suffix pages for layer `L+1` are not ready until layer +`L+1` prepare writes them. The existing Phase 8 split remains correct: + +```text +prefix: prefetchable +current/suffix: materialize synchronously at consume time +``` + +Do not try to prefetch current/suffix before the layer has written its local KV. + +### 4.5 Target and draft remain one logical cache object + +The L2->L1 load dependency applies to target and draft/MTP payloads. Shared-KV +prefix materialization currently concerns target attention buffers, but cache-hit +correctness and speculative accept rate still require draft KV residency to move +with target KV. Any improvement must preserve: + +```text +target load/backup/evict visibility == draft load/backup/evict visibility +``` + +--- + +## 5. What to measure next + +A benchmark or trace should separate the following timings per layer: + +1. L2->L1 H2D load submit time. +2. L2->L1 H2D load completion event time. +3. Shared-KV prefix prefetch launch time. +4. Shared-KV prefix reduce completion time. +5. Layer attention consume wait time. +6. Current/suffix materialize/reduce time. + +The key diagnostic questions are: + +- At `Layer L-1` prepare, is `load_event[L]` already complete, or can the + prefetch stream wait on it without blocking the forward stream? +- Did the L2->L1 load for `Layer L` start at `Layer L-2` or earlier in the + effective timeline, not merely get enqueued somewhere before forward? +- If the load is not complete, does the current forward stream block while trying + to start `Layer L` shared-KV prefetch? +- At `Layer L` attention, is the prefix prefetch event complete? +- Is the remaining stall mostly L2->L1 H2D, prefix materialize/reduce, or suffix + materialize/reduce? +- Does enabling `direct/page_first_direct` reduce H2D time but expose the + shared-KV prefix reduce as the new bottleneck? + +Without these separated timings, an end-to-end "cache hit has no speedup" result +is ambiguous. + +--- + +## 6. Current confidence and limits + +High-confidence facts from code/docs: + +- HiCache CP load is per-layer and uses `LayerDoneCounter` for layer-wise waits. +- `direct/page_first_direct` load routes to TAI `cudaMemcpyBatchAsync` H2D. +- Phase 8 shared-KV prefetch is one-layer-ahead and handles both MLA KV prefix + and NSA index prefix. +- Shared-KV prefetch consumes L1 KV/index buffers and then performs dense + materialization plus CP reduce. + +Medium-confidence inference: + +- In cache-hit workloads, Stage B can dominate after Stage A is optimized. This + explains why improving L2->L1 load alone may not improve end-to-end latency. + +Unknowns requiring measurement: + +- How often next-layer shared-KV prefetch blocks on the HiCache layer load event. +- Whether the main stall is CPU submit, H2D bandwidth, NCCL reduce, materialize + local copy, or suffix/current path. +- Whether current production traces show copy/collective contention or simply a + missing/late prefetch launch. + + +--- + +## 7. Proposed implementation design: make Stage-B prefetch wait on Stage-A locally + +### 7.1 Current concrete conflict points + +The conflict is not that L2->L1 and L1 shared-KV prefetch are conceptually +incompatible. The conflict is that the current prefetch code enters the same +synchronizing accessors used by the attention consume path. + +Current blocking points: + +```text +CpSharedKVMlaPrefetcher.maybe_create(...) + token_to_kv_pool.get_key_buffer(first_layer_id) + +CpSharedKVMlaPrefetcher.start_next_layer_prefix(...) + token_to_kv_pool.get_key_buffer(next_layer_id) + +CpSharedKVIndexPrefetcher.maybe_create(...) + token_to_kv_pool.get_index_k_with_scale_buffer(first_layer_id) + +CpSharedKVIndexPrefetcher.start_next_layer_prefix(...) + token_to_kv_pool.get_index_k_with_scale_buffer(next_layer_id) +``` + +Those getters intentionally wait on HiCache layer-load readiness: + +```text +get_key_buffer(layer_id) + -> layer_transfer_counter.wait_until(layer_id - start_layer) + +get_index_k_with_scale_buffer(layer_id) + -> layer_transfer_counter.wait_until(layer_id - start_layer) +``` + +That is correct for the final consumer. It is wrong for prefetch launch when the +call is made from the forward path, because the wait is attached to the current +stream instead of the prefetch stream. + +A second issue is that `start_next_layer_prefix()` currently runs local prefix +materialization before creating the pending handle. Because it is called from +layer-forward hooks, this materialization also tends to execute on the current +forward stream instead of the shared-KV prefetch stream. + +Net effect: + +```text +Layer L forward hook wants to prefetch Layer L+1 + -> calls synchronizing getter for L+1 + -> forward stream waits for L2->L1 load_event[L+1] + -> local prefix materialization runs on forward stream + -> async reduce is only launched after that +``` + +The desired effect is: + +```text +Layer L forward hook wants to prefetch Layer L+1 + -> prefetch stream waits for L2->L1 load_event[L+1] + -> local prefix materialization runs on prefetch stream + -> prefix reduce runs on the same FIFO prefetch stream + -> forward stream continues Layer L work +``` + +### 7.2 Design invariant + +For cache-hit CP shared KV, the Stage-B shared-KV prefetch for layer `X` must have +an explicit local dependency on the Stage-A HiCache L2->L1 load event for layer +`X`. + +Stage-A must also finish **capacity reservation in CP owner-lane vector form** +before any Stage-B prefetch is considered valid. In other words, the prefetch +repair is not only a stream/event repair: + +```text +old capacity model: + total_available_tokens >= total_required_tokens + +new CP shared-KV capacity model: + for each owner lane o: + available_pages_by_owner[o] >= required_pages_by_owner[o] +``` + +The shared-KV prefix prefetcher must never discover L1 capacity shortage. If it +does, the bug is in Stage-A reservation/eviction, not in Stage-B prefetch. + +This must be local CUDA stream/event ordering only: + +- no new CP all-reduce, +- no cross-rank polling, +- no CPU synchronization, +- no silent fallback to sync full materialization. + +The only ordering needed is local because each rank only reads its own restored +owner-lane shard before the existing shared-KV collective materializes the dense +prefix. + +### 7.3 Narrow API addition + +Add a narrow prefetch-only KV-pool API instead of weakening the normal getters. +The normal getters should continue to block because attention consume must remain +safe. + +Suggested shape: + +```python +def wait_layer_transfer_on_stream(self, layer_id: int, stream) -> None: + """If a HiCache layer transfer is active, make `stream` wait for this layer. + + This must not wait on the caller/current stream unless `stream` is the current + stream by construction. + """ + + +def get_key_buffer_for_prefetch(self, layer_id: int, stream) -> torch.Tensor: + self.wait_layer_transfer_on_stream(layer_id, stream) + return self._get_key_buffer_no_transfer_wait(layer_id) + + +def get_index_k_with_scale_buffer_for_prefetch(self, layer_id: int, stream) -> torch.Tensor: + self.wait_layer_transfer_on_stream(layer_id, stream) + return self._get_index_k_with_scale_buffer_no_transfer_wait(layer_id) +``` + +Implementation notes: + +- `LayerDoneCounter` already owns per-layer events through + `events[consumer_index].load_events[layer_index]`. +- Add a method that records `stream.wait_event(...)` for every active consumer + index, mirroring `wait_until()` but using the supplied stream. +- For layers outside pool range, preserve existing `layer_in_pool` checks. +- If no `layer_transfer_counter` exists, return raw buffers directly. +- Do not expose raw getters broadly. Keep them private or name them + `*_for_prefetch` to prevent accidental unsafe attention use. + +For MLA pools where `_get_key_buffer()` does not already exist, add a private raw +helper and keep public `get_key_buffer()` behavior unchanged. + +For NSA state/index pools, add a raw helper for +`index_k_with_scale_buffer[layer_id - start_layer]` and route only prefetch +through it. + +### 7.4 Prefetcher changes + +For both MLA and index prefetchers: + +1. Stop using synchronizing getters in `maybe_create()` and + `start_next_layer_prefix()`. +2. Use the prefetch-safe getter with the prefetcher's FIFO stream. +3. Move local prefix materialization into `with torch.cuda.stream(self.stream)`. +4. Keep the existing single FIFO stream shared by index and MLA prefix prefetch. + This avoids rank skew across PyNccl collectives. +5. Continue to create one handle per layer and consume it with the existing + `handle.event` wait at attention time. + +The new `start_next_layer_prefix()` shape should be: + +```text +start_next_layer_prefix(next_layer_id): + if disabled/already_started/out_of_pool: return + + kv_or_index_buffer = pool.get_*_for_prefetch(next_layer_id, self.stream) + + with prefetch_stream: + # prefetch_stream is already ordered after load_event[next_layer_id] + dense_buffer = allocate dense prefix buffer + materialize local prefix pages/rows into dense_buffer + event = async CP reduce on dense prefix range + + handles[next_layer_id] = handle(event=event, dense_buffer=...) +``` + +This removes the current two-step behavior where local materialization happens +before `launch_pending_reduce()` and only the reduce is moved to the prefetch +stream. If we keep `launch_pending_reduce()` as a helper, it should assume the +local materialization is already enqueued on the same prefetch stream, not on the +forward stream. + +### 7.5 First-layer behavior + +Layer 0 / first layer is special: + +- There is usually little or no preceding layer compute to hide its prefetch. +- It may still need to wait for L2->L1 load before attention can consume cache-hit + prefix. + +The design should still route first-layer initialization through the same +prefetch-safe API so the behavior is uniform and does not accidentally attach +load waits to the wrong stream. We should not promise speedup for layer 0; the +main objective is avoiding unnecessary serialization for layer 1+. + +### 7.6 Fallback and failure policy + +Expected no-op cases may stay quiet or debug-level: + +- no prefetcher because CP shared KV is disabled, +- prefix too small for prefetch threshold, +- layer is outside this pool. + +Unexpected cases must be warning-level and explicit: + +- prefetch-safe getter unavailable while HiCache layer transfer is active, +- async reduce unavailable, +- any exception in local prefix materialize/reduce launch, +- falling back from prefetched prefix to synchronous full materialize for layer > + start_layer. + +The warning should include enough fields to distinguish L2->L1 conflict from +normal prefix miss: + +```text +cp_rank, layer_id, start_layer, has_layer_transfer_counter, +consumer_indices, prefix_pages, total_slots, path=mla|index +``` + +Do not add environment gates for correctness warnings. Temporary high-frequency +profiling logs can be gated or removed after use. + +### 7.7 Tests before CUDA ETE + +Local/unit tests should not require CUDA kernels. The important behavior can be +locked with fakes: + +1. **No synchronizing getter in prefetch path** + - fake pool where `get_key_buffer()` raises, + - `get_key_buffer_for_prefetch()` records the stream argument and returns a + fake tensor, + - verify `CpSharedKVMlaPrefetcher.start_next_layer_prefix()` does not call the + blocking getter. + +2. **Layer event is attached to prefetch stream** + - fake `LayerDoneCounter.wait_until_on_stream(layer, stream)` records calls, + - verify `next_layer_id - start_layer` is used. + +3. **Index path parity** + - same test for `get_index_k_with_scale_buffer_for_prefetch()`. + +4. **Consume path remains blocking-safe** + - normal `consume()` still accepts the already-prefetched handle, + - sync fallback still uses normal getter/materialize and emits warning when it + happens after start layer. + +Remote CUDA/ETE validation should then check: + +- no warnings about prefetch-safe getter fallback, +- reduced or eliminated `consume_miss` for layer > start layer on cache-hit runs, +- no new collective beyond the existing prefix reduce collectives, +- no worse MTP accept rate after cache hit, +- first-layer wait remains expected, later-layer waits shrink. + +### 7.8 Non-goals for this change + +Do not solve these in the first patch: + +- true `batch_size > 1` CP shared-KV prefill, +- changing radix tree or host-cache residency policy, +- changing owner-lane allocation semantics, +- making current/suffix KV prefetchable, +- adding new CP/global synchronization to prove readiness. + +This patch should only fix the composition between existing HiCache layer-load +readiness and existing shared-KV prefix prefetch. + +--- + +## 8. L1 capacity reservation plan for L2->L1 load prefetch + +### 8.1 Current code path + +Current cache-hit load-back path is: + +```text +schedule_policy.py + AddReq.add_one_req(...) + -> tree_cache.init_load_back( + InitLoadBackParams(last_host_node, host_hit_length, mem_quota) + ) + +hiradix_cache.py + HiRadixCache.init_load_back(...) + -> load_back(last_host_node, mem_quota) + + HiRadixCache.load_back(...) + -> collect evicted host-backed nodes_to_load + -> inc_lock_ref(ancestor_node) for quota/protection during planning + -> cache_controller.load_cp(nodes_to_load) + -> on owner-lane allocation failure, retry after _evict_for_compute_owner_lanes(...) + -> set loaded_node.value to reserved device indices + -> ongoing_load_back[last_hit_node.id] = last_hit_node + -> inc_lock_ref(last_hit_node) until load ACK + +scheduler.py + new_batch.hicache_consumer_index = tree_cache.ready_to_load_host_cache() + +hiradix_cache.py + ready_to_load_host_cache() + -> cache_controller.start_loading() + +cache_controller.py + start_loading() + -> submit per-layer host->device copies on load_stream + -> record LayerDoneCounter event for each target layer + -> append ack_load_queue + +hiradix_cache.py + loading_check() + -> query finish_event + -> dec_lock_ref(loaded node) +``` + +The important existing properties are good: + +- Device pages are allocated before async H2D starts. +- `loaded_node.value` is assigned before forward, so radix state knows the node is + now device-resident. +- `last_hit_node` is locked while H2D is in flight, so normal device eviction + should not evict the just-loaded prefix before the request consumes it. +- CP load uses `alloc_pages_with_owners(page_owners)` to reproduce the write-time + owner-lane pattern. + +### 8.2 Current gap + +The current planning is still mostly total-token based: + +```text +schedule_policy._get_load_back_mem_quota(...) + available_device_tokens + evictable_tokens - request_reserve_tokens +``` + +For CP shared KV this is insufficient. Load-back capacity is owner-lane +constrained: + +```text +required_by_owner[o] <= available_by_owner[o] +``` + +not just: + +```text +sum(required_by_owner) <= sum(available_by_owner) +``` + +Current `load_cp()` handles this reactively: + +1. try `alloc_pages_with_owners(page_owners)`, +2. if it returns `None`, call `_evict_for_compute_owner_lanes(...)`, +3. retry allocation. + +That retry path is better than plain total-token eviction, but it is still not a +complete capacity plan: + +- the first allocation attempt is expected to fail under lane skew; +- `_evict_for_compute_owner_lanes()` computes lane deficits, but currently calls + normal `tree_cache.evict(EvictParams(num_tokens=...))`, so eviction candidates + are still selected by the cache eviction policy rather than by contribution to + the deficient owner lanes; +- it may over-evict healthy lanes and still fail to free enough pages in the + deficient lane; +- failure diagnostics do not consistently show `required/available/deficit` per + owner lane. + +The explicit design decision is to replace this scalar admission condition: + +```text +available_tokens + evictable_tokens >= required_tokens +``` + +with a CP owner-lane vector condition: + +```text +required_pages_by_owner: int[cp_size] +available_pages_by_owner: int[cp_size] +deficit_pages_by_owner: int[cp_size] + +all(deficit_pages_by_owner[o] == 0 for o in range(cp_size)) +``` + +This is a per-owner-lane model, not a new cross-rank synchronization model. All +CP ranks should derive the same `page_owners` and the same logical allocator page +ownership. If ranks disagree, that is allocator/radix state divergence and must +be surfaced as a correctness bug rather than hidden behind another collective. + +### 8.3 Desired invariant + +L2->L1 prefetch should never make capacity decisions per layer. + +Before submitting any H2D copy, the scheduler/cache layer must have completed: + +```text +host-hit nodes_to_load selected + -> page_owners collected + -> required_pages_by_owner / available_pages_by_owner / deficit_pages_by_owner computed + -> owner-lane vector capacity checked + -> owner-lane deficits evicted, if possible + -> device pages reserved with matching owner pattern + -> radix nodes marked device-resident and locked + -> only then submit async per-layer load +``` + +Per-layer load and later shared-KV prefix prefetch should only consume already +reserved device slots and local CUDA readiness events. + +### 8.4 Proposed implementation units + +#### Unit A: owner-lane load-back planning helper + +Add a small helper around CP load-back planning, probably in +`hiradix_cache.py` or `common.py`: + +```python +@dataclass +class CpLoadBackPlan: + page_owners: list[int] + required_by_owner: list[int] + available_by_owner: list[int] + deficit_by_owner: list[int] + host_hit_len: int +``` + +Responsibilities: + +- collect `page_owners` from `node.cp_hicache.page_owners`, +- compute `required_by_owner`, `available_by_owner`, `deficit_by_owner` using + allocator APIs, +- validate that every node has CP HiCache metadata, +- produce warning fields for all fail-closed cases. + +This helper is read-only and should not allocate or evict. + +#### Unit B: owner-lane-aware device eviction + +Replace the retry-only `_evict_for_compute_owner_lanes()` behavior for CP +HiCache load-back with a real owner-lane-aware eviction path. + +The owner-lane-aware eviction should: + +1. inspect evictable radix leaves whose `lock_ref == 0`, +2. for each candidate with `node.value`, compute the owner-lane contribution of + its pages: + +```text +page_id = logical_token_loc // page_size +owner = (page_id - 1) % cp_size +``` + +3. rank candidates primarily by contribution to deficient lanes and secondarily + by existing eviction policy priority, +4. evict candidates until all lane deficits are covered or no useful candidates + remain, +5. preserve existing write-through/write-back semantics: + - backed node: `_evict_backuped(node)`, + - unbacked node under write-through: `_evict_regular(node)`, + - write-back policy: use existing `write_backup(..., write_back=True)` path + only with explicit warnings if it becomes synchronous or expensive. + +The first patch can be conservative: + +- only specialize CP shared-KV + HiCache load-back; +- keep generic eviction unchanged; +- if owner-aware eviction cannot prove progress, fail closed and log a warning. + +No collective is needed. All owner-lane accounting is local because allocator +page ids and CP size are local and deterministic. + +#### Unit C: reserve/lock API boundary + +Keep `cache_controller.load_cp()` focused on allocation/enqueue bookkeeping, but +make its contract explicit: + +```text +load_cp(nodes_to_load) either returns fully reserved device_indices or None. +It must not return partial reservation. +``` + +`HiRadixCache.load_back()` remains the owner of radix-state transitions: + +- temporary planning lock on `ancestor_node`, +- assign `loaded_node.value` after reservation succeeds, +- add `ongoing_load_back[last_hit_node.id]`, +- hold `inc_lock_ref(last_hit_node)` until `loading_check()` sees the H2D finish + event. + +If any exception occurs after reservation but before `ongoing_load_back` is +registered, free the reserved device indices and roll back node state. This is +important because the current path crosses Python code and async copy submission; +reservation leaks here become GPU KV leaks. + +#### Unit D: diagnostics and fail-fast policy + +Capacity failure should be warning-level and structured. Required fields: + +```text +node_id +host_hit_len +nodes_to_load +required_by_owner +available_by_owner +deficit_by_owner +evictable_size +protected_size +ongoing_load_back_count +allocator_available_size +allocator_free_pages +allocator_release_pages +``` + +Expected skip cases may remain info-level: + +- host hit below load-back threshold, +- total host hit exceeds configured quota. + +Unexpected cases must be warning-level: + +- owner-lane reservation failed, +- owner-aware eviction made no progress, +- retry allocation still failed, +- CP metadata missing from a node selected for CP load-back. + +### 8.5 Interaction with shared-KV L1 prefix prefetch + +This L1 capacity plan is a prerequisite for the Stage-B shared-KV prefix +prefetch work in Section 7. + +Once `load_cp()` returns successfully, Stage-B prefetch must assume: + +```text +persistent L1 pages exist and are reserved +the reservation satisfied the per-CP-owner-lane vector check +radix nodes are locked against device eviction +only per-layer H2D readiness remains pending +``` + +Therefore Stage-B must not trigger device eviction. If Stage-B cannot access a +layer's persistent shard, that is a bug in reservation/event ordering, not a +signal to evict more memory. + +### 8.6 Implementation plan + +#### P0: Convert CP HiCache load-back admission from scalar to owner-lane vector + +Files: + +- `python/sglang/srt/managers/schedule_policy.py` +- `python/sglang/srt/mem_cache/hiradix_cache.py` +- `python/sglang/srt/mem_cache/allocator.py` + +Change: + +- keep generic scheduler budget as a coarse scalar admission check; +- for CP HiCache load-back, treat scalar `mem_quota` only as an upper bound; +- make final load-back admission depend on: + +```text +required_pages_by_owner[o] +available_pages_by_owner[o] +deficit_pages_by_owner[o] +``` + +Acceptance: + +- a request is not admitted for CP load-back unless every owner lane can be + satisfied after targeted eviction; +- scalar total-token availability is never used as the final CP load-back proof; +- no new all-reduce or cross-rank capacity polling is introduced. + +#### P1: Document and test current reservation lifecycle + +Files: + +- `python/sglang/srt/mem_cache/hiradix_cache.py` +- `python/sglang/srt/managers/cache_controller.py` +- unit test under `test/registered/unit/mem_cache/` + +Tests: + +- CP load-back returns all-or-nothing reservation. +- `loaded_node.value` is assigned only after reservation succeeds. +- `ongoing_load_back` lock is released only through `loading_check()` ack. + +#### P2: Add owner-lane load-back stats helper + +Files: + +- likely `python/sglang/srt/mem_cache/hiradix_cache.py` +- maybe shared helper in `python/sglang/srt/mem_cache/common.py` + +Tests: + +- page owner list `[0, 1, 1, 3]` produces required counts correctly. +- allocator free/release page distribution produces available counts correctly. +- missing `cp_hicache` metadata fails closed. + +#### P3: Add owner-lane-aware eviction for CP HiCache load-back + +Files: + +- `python/sglang/srt/mem_cache/hiradix_cache.py` +- possibly `python/sglang/srt/mem_cache/common.py` if reusing from extend alloc + +Algorithm: + +```text +while deficits remain: + choose evictable leaf with max contribution to positive deficits + evict it using existing node eviction primitive + recompute available/deficit + stop on no progress +``` + +Tests: + +- skewed allocator state where total pages are enough but owner lane 0 is short. +- normal total-token eviction would not guarantee success; owner-aware eviction + selects a node containing lane-0 pages. +- locked nodes are skipped. + +#### P4: Replace reactive retry with preflight plan + targeted eviction + +Current shape: + +```text +try alloc_pages_with_owners +if None: evict_for_compute_owner_lanes and retry +``` + +Target shape: + +```text +build load-back plan +if deficit: owner-aware evict before allocation +alloc_pages_with_owners once +if None: warning + fail closed +``` + +A final retry can remain temporarily for safety, but if it triggers it must emit a +warning that the preflight plan was wrong. + +#### P5: Integrate with Stage-B prefetch API work + +After owner-lane vector reservation is deterministic, implement Section 7: + +- prefetch-safe pool getters, +- stream-specific layer-load wait, +- shared-KV prefix materialization on prefetch stream, +- no forward-stream wait just to launch prefetch. + +The Stage-B prefetcher should rely on the P0-P4 guarantee: + +```text +for every prefetched layer: + persistent L1 page slots have already been owner-lane reserved + radix nodes are locked + only the local per-layer load event may still be pending +``` + +#### P6: Remote validation + +Remote ETE checks should include: + +- cache-hit request under owner-lane skew does not OOM if enough evictable pages + exist in the right lanes; +- warning logs show per-lane fields on failure; +- no new CP/global collective is introduced; +- L2->L1 load ack releases lock; +- shared-KV prefetch no longer calls blocking getters from forward stream; +- MTP accept rate does not drop after cache hit. + +### 8.7 Non-goals + +- Do not make true batch-size > 1 prefill part of this patch. +- Do not change host-cache eviction policy. +- Do not introduce all-reduce to agree on capacity. +- Do not allow partial load-back of a host-hit node chain. +- Do not let Stage-B shared-KV prefetch perform device eviction. + +--- + +## 9. Implementation status after P5-P8 + +Implemented on `cjy-cp-refactor`: + +- CP HiCache load-back now builds a per-owner-lane plan from persisted + `node.cp_hicache.page_owners` before reserving L1 pages. +- Load-back owner-lane deficits are handled by targeted device eviction rather + than by scalar total-token eviction. If the lane vector still cannot be + satisfied, the path fails closed with `[CP_HICACHE_FALLBACK]` warning fields. +- Cold/extend allocation now tries compute-owner allocation before scalar + eviction, passes `EvictParams.owner_lane_deficits` into tree eviction, and + raises recoverable `KVCapacityWaitError` instead of silently falling back to + legacy allocation when owner lanes remain exhausted. +- Scheduler prefill allocation catches this recoverable wait, releases the + `PrefillAdder` radix locks, restores the waiting queue/chunked request state, + and retries later instead of crashing or leaking locks. +- `SGLANG_DISABLE_TAI_OWNER_SELECT` is intentionally not supported. There is no + env-gated legacy owner-selection fallback in this path; unexpected owner-lane + exhaustion is surfaced as capacity wait or warning-level fail-closed behavior. +- HiCache layer-load readiness can now be attached to an explicit CUDA stream by + `LayerDoneCounter.wait_until_on_stream(...)`. +- KV pools expose prefetch-only getters such as + `get_key_buffer_for_prefetch(layer_id, stream)` and + `get_index_k_with_scale_buffer_for_prefetch(layer_id, stream)`. Normal getters + still block on the current stream and remain the attention-consumer-safe API. +- MLA and NSA-index CP shared-KV prefetch now use the prefetch-safe getters, + enqueue local prefix materialization on the prefetch stream, and enqueue the + async CP reduce on the same FIFO prefetch stream. This avoids attaching + HiCache L2->L1 load waits or prefix materialization to the forward stream just + to launch next-layer prefetch. + +Validation performed: + +```text +PYTHONPATH=python python -m py_compile \ + python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py \ + python/sglang/srt/managers/cache_controller.py \ + python/sglang/srt/managers/scheduler.py \ + python/sglang/srt/mem_cache/base_prefix_cache.py \ + python/sglang/srt/mem_cache/common.py \ + python/sglang/srt/mem_cache/hiradix_cache.py \ + python/sglang/srt/mem_cache/memory_pool.py + +Remote g0034 container: +PYTHONPATH=python python -m pytest \ + test/registered/unit/mem_cache/test_alloc_pages_with_owners.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \ + test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q +# 95 passed +``` + +Remaining risks: + +- The new prefetch stream ordering is unit-tested with fake streams and CPU + tensors. End-to-end CUDA traces are still required to quantify whether the + forward-stream stall moved from L2->L1 load to prefix materialize/reduce or was + actually hidden. +- True `ForwardBatch(batch_size > 1)` remains unsupported for CP shared-KV. Use + multi-slot overlap as the safer next throughput direction. +- Generic non-owner paths still exist for unsupported CP-shared-KV shapes such as + multi-batch or non-page-aligned short paths. Those are explicit fallback paths, + not the intended fast path. 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 7f79b33d5..0474b9888 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 @@ -45,6 +45,50 @@ def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None: ) +def _prefetch_pool_get_key_buffer( + *, + token_to_kv_pool: Any, + layer_id: int, + stream: torch.cuda.Stream, + path: str, +) -> torch.Tensor: + getter = getattr(token_to_kv_pool, "get_key_buffer_for_prefetch", None) + if getter is not None: + return getter(layer_id, stream) + if getattr(token_to_kv_pool, "layer_transfer_counter", None) is not None: + logger.warning( + "[CP_SHARED_KV_FALLBACK][%s_prefetch] " + "reason=prefetch_safe_getter_unavailable layer_id=%s pool=%s " + "has_layer_transfer_counter=True", + path, + layer_id, + type(token_to_kv_pool).__name__, + ) + return token_to_kv_pool.get_key_buffer(layer_id) + + +def _prefetch_pool_get_index_buffer( + *, + token_to_kv_pool: Any, + layer_id: int, + stream: torch.cuda.Stream, +) -> torch.Tensor: + getter = getattr( + token_to_kv_pool, "get_index_k_with_scale_buffer_for_prefetch", None + ) + if getter is not None: + return getter(layer_id, stream) + if getattr(token_to_kv_pool, "layer_transfer_counter", None) is not None: + _index_prefetch_fallback_log( + "prefetch_safe_getter_unavailable", + "pool has active layer transfer counter but no prefetch-safe index " + "getter. layer_id=%s pool=%s", + layer_id, + type(token_to_kv_pool).__name__, + ) + return token_to_kv_pool.get_index_k_with_scale_buffer(layer_id=layer_id) + + def _is_cuda_stream_capturing() -> bool: if not torch.cuda.is_available(): return False @@ -235,9 +279,15 @@ class CpSharedKVMlaPrefetcher: ) return None + prefetch_stream = stream if stream is not None else torch.cuda.Stream() try: first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) - kv_cache = token_to_kv_pool.get_key_buffer(first_layer_id) + kv_cache = _prefetch_pool_get_key_buffer( + token_to_kv_pool=token_to_kv_pool, + layer_id=first_layer_id, + stream=prefetch_stream, + path="mla", + ) remap = get_or_build_shared_token_kv_slot_remap( forward_batch, kv_cache=kv_cache, @@ -276,7 +326,7 @@ class CpSharedKVMlaPrefetcher: dense_num_pages=remap.dense_num_pages, owned_prefix_pages=owned_prefix_pages, owned_total_pages=owned_total_pages, - stream=stream, + stream=prefetch_stream, ) def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool: @@ -447,7 +497,12 @@ class CpSharedKVMlaPrefetcher: return try: - kv_cache = token_to_kv_pool.get_key_buffer(next_layer_id) + kv_cache = _prefetch_pool_get_key_buffer( + token_to_kv_pool=token_to_kv_pool, + layer_id=next_layer_id, + stream=self.stream, + path="mla", + ) except Exception: logger.exception( "Failed to get next-layer KV cache for CP shared KV MLA prefetch." @@ -461,31 +516,61 @@ class CpSharedKVMlaPrefetcher: return try: - dense_kv_cache = kv_cache.new_zeros( - (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) - ) - self._log_next_layer( - next_layer_id, - "start_prefix_begin next_layer=%s start_slot=0 end_slot=%s " - "dense_rows=%s", - next_layer_id, - self.prefix_pages, - int(dense_kv_cache.shape[0]), - ) - materialize_local_token_kv_page_slots_into( - kv_cache=kv_cache, - dense_kv_cache=dense_kv_cache, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - page_size=self.page_size, - start_slot=0, - end_slot=self.prefix_pages, - ) + current_stream = torch.cuda.current_stream() + self.stream.wait_stream(current_stream) prefix_rows = slot_range_to_token_slice( self.page_size, 0, self.prefix_pages, ) + with torch.cuda.stream(self.stream): + dense_kv_cache = kv_cache.new_zeros( + (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) + ) + self._log_next_layer( + next_layer_id, + "start_prefix_begin next_layer=%s start_slot=0 end_slot=%s " + "dense_rows=%s", + next_layer_id, + self.prefix_pages, + int(dense_kv_cache.shape[0]), + ) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + start_slot=0, + end_slot=self.prefix_pages, + ) + event = _all_reduce_materialized_buffer_async( + dense_kv_cache[prefix_rows], + cp_size=self.layout.cp_size, + stream=self.stream, + nvtx_source="mla.prefetch_prefix", + nvtx_layer_id=next_layer_id, + nvtx_cp_rank=self.layout.cp_rank, + nvtx_rows=(prefix_rows.start, prefix_rows.stop), + ) + if event is None: + self.disabled = True + logger.warning( + "[CP_SHARED_KV_FALLBACK][mla_prefetch] " + "reason=async_reduce_unavailable layer_id=%s cp_rank=%s " + "cp_size=%s prefix_pages=%s total_slots=%s", + next_layer_id, + self.layout.cp_rank, + self.layout.cp_size, + self.prefix_pages, + self.total_slots, + ) + self._log_next_layer( + next_layer_id, + "start_disable reason=async_reduce_unavailable next_layer=%s", + next_layer_id, + ) + return except Exception: logger.exception("Failed to start CP shared KV MLA prefix prefetch.") self.disabled = True @@ -500,12 +585,14 @@ class CpSharedKVMlaPrefetcher: layer_id=next_layer_id, dense_kv_cache=dense_kv_cache, prefix_rows=prefix_rows, + event=event, ) self.handles[next_layer_id] = handle self.pending_attention_handle = handle self._log_next_layer( next_layer_id, - "start next_layer=%s prefix_pages=%s prefix_rows=%s dense_rows=%s", + "start next_layer=%s prefix_pages=%s prefix_rows=%s dense_rows=%s " + "reduce_enqueued=True", next_layer_id, self.prefix_pages, prefix_rows.stop - prefix_rows.start, @@ -776,10 +863,13 @@ class CpSharedKVIndexPrefetcher: ) return None + prefetch_stream = stream if stream is not None else torch.cuda.Stream() try: first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) - page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer( - layer_id=first_layer_id + page_buffer = _prefetch_pool_get_index_buffer( + token_to_kv_pool=token_to_kv_pool, + layer_id=first_layer_id, + stream=prefetch_stream, ) remap = get_or_build_shared_paged_buffer_slot_remap( forward_batch, @@ -822,7 +912,7 @@ class CpSharedKVIndexPrefetcher: dense_num_pages=remap.dense_num_pages, owned_prefix_pages=owned_prefix_pages, owned_total_pages=owned_total_pages, - stream=stream, + stream=prefetch_stream, ) def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool: @@ -990,8 +1080,10 @@ class CpSharedKVIndexPrefetcher: return try: - page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer( - layer_id=next_layer_id + page_buffer = _prefetch_pool_get_index_buffer( + token_to_kv_pool=token_to_kv_pool, + layer_id=next_layer_id, + stream=self.stream, ) except Exception: logger.exception( @@ -1006,26 +1098,56 @@ class CpSharedKVIndexPrefetcher: return try: - dense_page_buffer = page_buffer.new_zeros( - (self.dense_num_pages, *page_buffer.shape[1:]) - ) - self._log_next_layer( - next_layer_id, - "index_start_prefix_begin next_layer=%s start_slot=0 " - "end_slot=%s dense_pages=%s", - next_layer_id, - self.prefix_pages, - int(dense_page_buffer.shape[0]), - ) - materialize_local_paged_buffer_page_slots_into( - page_buffer=page_buffer, - dense_page_buffer=dense_page_buffer, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - start_slot=0, - end_slot=self.prefix_pages, - ) + current_stream = torch.cuda.current_stream() + self.stream.wait_stream(current_stream) prefix_rows = slot_range_to_page_slice(0, self.prefix_pages) + with torch.cuda.stream(self.stream): + dense_page_buffer = page_buffer.new_zeros( + (self.dense_num_pages, *page_buffer.shape[1:]) + ) + self._log_next_layer( + next_layer_id, + "index_start_prefix_begin next_layer=%s start_slot=0 " + "end_slot=%s dense_pages=%s", + next_layer_id, + self.prefix_pages, + int(dense_page_buffer.shape[0]), + ) + materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + start_slot=0, + end_slot=self.prefix_pages, + ) + event = _all_reduce_materialized_buffer_async( + dense_page_buffer[prefix_rows], + cp_size=self.layout.cp_size, + stream=self.stream, + nvtx_source="index.prefetch_prefix", + nvtx_layer_id=next_layer_id, + nvtx_cp_rank=self.layout.cp_rank, + nvtx_rows=(prefix_rows.start, prefix_rows.stop), + ) + if event is None: + self.disabled = True + _index_prefetch_fallback_log( + "async_reduce_unavailable", + "async reduce unavailable. layer_id=%s cp_rank=%s cp_size=%s " + "prefix_pages=%s total_slots=%s", + next_layer_id, + self.layout.cp_rank, + self.layout.cp_size, + self.prefix_pages, + self.total_slots, + ) + self._log_next_layer( + next_layer_id, + "index_start_disable reason=async_reduce_unavailable next_layer=%s", + next_layer_id, + ) + return except Exception: logger.exception("Failed to start CP shared KV index prefix prefetch.") self.disabled = True @@ -1040,12 +1162,14 @@ class CpSharedKVIndexPrefetcher: layer_id=next_layer_id, dense_page_buffer=dense_page_buffer, prefix_rows=prefix_rows, + event=event, ) self.handles[next_layer_id] = handle self.pending_attention_handle = handle self._log_next_layer( next_layer_id, - "index_start next_layer=%s prefix_pages=%s dense_pages=%s", + "index_start next_layer=%s prefix_pages=%s dense_pages=%s " + "reduce_enqueued=True", next_layer_id, self.prefix_pages, int(dense_page_buffer.shape[0]), diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index c39864d5a..58d4337ef 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -61,6 +61,10 @@ class LayerLoadingEvent: def wait(self, layer_index: int): device_module.current_stream().wait_event(self.load_events[layer_index]) + def wait_on_stream(self, layer_index: int, stream): + assert 0 <= layer_index < self._num_layers + stream.wait_event(self.load_events[layer_index]) + @property def finish_event(self): return self.load_events[-1] @@ -101,6 +105,12 @@ class LayerDoneCounter: for consumer_index in self.consumer_indices: self.events[consumer_index].wait(threshold) + def wait_until_on_stream(self, threshold: int, stream) -> None: + if not self.consumer_indices: + return + for consumer_index in self.consumer_indices: + self.events[consumer_index].wait_on_stream(threshold, stream) + def reset(self): self.producer_index = -1 self.consumer_index = -1 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 847ce4f89..5358a8f77 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -179,8 +179,9 @@ from sglang.srt.managers.scheduler_update_weights_mixin import ( ) from sglang.srt.managers.session_controller import SessionController from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length +from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import release_kv_cache +from sglang.srt.mem_cache.common import KVCapacityWaitError, release_kv_cache from sglang.srt.mem_cache.radix_cache import RadixCache from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors @@ -2417,6 +2418,8 @@ class Scheduler( if len(can_run_list) == 0: return None + waiting_queue_before_prepare = list(self.waiting_queue) + chunked_req_before_prepare = self.chunked_req can_run_set = set(can_run_list) self.waiting_queue = [x for x in self.waiting_queue if x not in can_run_set] if adder.preempt_list: @@ -2456,7 +2459,20 @@ class Scheduler( self.tree_cache.ready_to_load_host_cache() ) - new_batch.prepare_for_extend() + try: + new_batch.prepare_for_extend() + except KVCapacityWaitError as exc: + self._release_prefill_adder_locks( + can_run_list, skip_req=chunked_req_before_prepare + ) + self.waiting_queue = waiting_queue_before_prepare + self.chunked_req = chunked_req_before_prepare + self.running_batch.batch_is_full = True + logger.warning( + "[CP_SHARED_KV_CAPACITY_WAIT] prefill allocation deferred: %s", + exc, + ) + return None # Record prefill stats for logging after forward new_batch.prefill_stats = PrefillStats.from_adder( @@ -2485,6 +2501,38 @@ class Scheduler( return new_batch + def _release_prefill_adder_locks( + self, reqs: List[Any], *, skip_req: Optional[Any] = None + ) -> None: + """Release lock refs acquired by PrefillAdder for a failed prefill batch. + + `PrefillAdder.add_one_req` persists one `inc_lock_ref` per staged + request after its short-lived scheduling lock context exits. If KV + allocation later reports a recoverable capacity wait, the batch never + runs, so those persistent refs must be dropped before the requests are + returned to the waiting queue. + """ + + if getattr(self.tree_cache, "disable", False): + return + + use_swa_params = self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache() + for req in reqs: + if req is skip_req: + continue + last_node = getattr(req, "last_node", None) + if last_node is None: + continue + if use_swa_params: + params = DecLockRefParams( + swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None) + ) + self.tree_cache.dec_lock_ref(last_node, params) + if hasattr(req, "swa_uuid_for_lock"): + req.swa_uuid_for_lock = None + else: + self.tree_cache.dec_lock_ref(last_node) + def update_running_batch(self, batch: ScheduleBatch) -> Optional[ScheduleBatch]: """Update the current running decoding batch.""" initial_bs = batch.batch_size() diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index dea7cec91..4462e18c3 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -78,6 +78,7 @@ class EvictParams: num_tokens: int swa_num_tokens: int = 0 mamba_num: int = 0 + owner_lane_deficits: Optional[list[int]] = None @dataclasses.dataclass diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index c4a42df79..cdfda526b 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -32,6 +32,27 @@ MAMBA_STATE_PER_REQ_NO_CACHE = 1 logger = logging.getLogger(__name__) +class KVCapacityWaitError(RuntimeError): + """Recoverable KV-capacity pressure with CP owner-lane diagnostics.""" + + def __init__( + self, + message: str, + *, + required_by_owner: list[int] | None = None, + available_by_owner: list[int] | None = None, + deficit_by_owner: list[int] | None = None, + ): + self.required_by_owner = required_by_owner + self.available_by_owner = available_by_owner + self.deficit_by_owner = deficit_by_owner + super().__init__( + f"{message}; required_by_owner={required_by_owner} " + f"available_by_owner={available_by_owner} " + f"deficit_by_owner={deficit_by_owner}" + ) + + def _log_cp_shared_kv_alloc_fallback( reason: str, message: str, @@ -341,7 +362,12 @@ def _evict_for_compute_owner_lanes( before_available, evictable_size, ) - evict_result = tree_cache.evict(EvictParams(num_tokens=evict_tokens)) + evict_result = tree_cache.evict( + EvictParams( + num_tokens=evict_tokens, + owner_lane_deficits=[int(v) for v in deficits], + ) + ) after_available = allocator.available_size() evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0) logger.info( @@ -367,7 +393,7 @@ def alloc_paged_token_slots_extend( # Over estimate the number of tokens: assume each request needs a new page. allocator = tree_cache.token_to_kv_pool_allocator num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size - evict_result = evict_from_tree_cache(tree_cache, num_tokens) + evict_result = None # logger.info( # "[MemCache-alloc] alloc_paged_token_slots_extend: extend_num_tokens=%d batch_size=%d num_tokens=%d page_size=%d " @@ -420,11 +446,11 @@ def alloc_paged_token_slots_extend( "multi_batch" if len(prefix_lens_cpu) != 1 else "unknown" ) - state = None - if backup_state: - state = allocator.backup_state() - if page_compute_owners is not None: + state = None + if backup_state: + state = allocator.backup_state() + out_cache_loc = alloc_extend_compute_owner( prefix_lens, prefix_lens_cpu, @@ -460,10 +486,10 @@ def alloc_paged_token_slots_extend( required, available, deficits = compute_owner_lane_stats( page_compute_owners ) - _log_cp_shared_kv_alloc_fallback( - "owner_lane_exhausted", - "failed to allocate pages from compute-owner lanes; " - "falling back to legacy page allocation. extend_num_tokens=%s page_size=%s " + logger.warning( + "[CP_SHARED_KV_FAIL_FAST][owner_lane_exhausted] " + "failed to allocate pages from compute-owner lanes after eviction; " + "raising recoverable capacity wait. extend_num_tokens=%s page_size=%s " "required_by_owner=%s available_by_owner=%s deficit_by_owner=%s", extend_num_tokens, allocator.page_size, @@ -471,15 +497,19 @@ def alloc_paged_token_slots_extend( available, deficits, ) - out_cache_loc = allocator.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, + raise KVCapacityWaitError( + "owner_lane_exhausted: failed to allocate pages from CP " + "compute-owner lanes after owner-lane eviction", + required_by_owner=required, + available_by_owner=available, + deficit_by_owner=deficits, ) else: + evict_result = evict_from_tree_cache(tree_cache, num_tokens) + state = None + if backup_state: + state = allocator.backup_state() + if alloc_extend_compute_owner is not None: _log_cp_shared_kv_alloc_fallback( compute_owner_unavailable_reason or "compute_owner_not_available", @@ -568,6 +598,7 @@ def alloc_for_extend( extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True) # Allocate req slots + newly_allocated_reqs = [req for req in batch.reqs if req.req_pool_idx is None] req_pool_indices = alloc_req_slots( batch.req_to_token_pool, batch.reqs, batch.tree_cache ) @@ -575,23 +606,29 @@ def alloc_for_extend( req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True) # Allocate KV cache (throws exception on failure) - if batch.tree_cache.page_size == 1: - out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens) - else: - # Paged allocation - build last_loc - last_loc = [ - (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) - for t in prefix_tensors - ] - out_cache_loc = alloc_paged_token_slots_extend( - tree_cache=batch.tree_cache, - prefix_lens=prefix_lens_device, - prefix_lens_cpu=prefix_lens_cpu, - seq_lens=batch.seq_lens, - seq_lens_cpu=batch.seq_lens_cpu, - last_loc=torch.cat(last_loc), - extend_num_tokens=batch.extend_num_tokens, - ) + try: + if batch.tree_cache.page_size == 1: + out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens) + else: + # Paged allocation - build last_loc + last_loc = [ + (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) + for t in prefix_tensors + ] + out_cache_loc = alloc_paged_token_slots_extend( + tree_cache=batch.tree_cache, + prefix_lens=prefix_lens_device, + prefix_lens_cpu=prefix_lens_cpu, + seq_lens=batch.seq_lens, + seq_lens_cpu=batch.seq_lens_cpu, + last_loc=torch.cat(last_loc), + extend_num_tokens=batch.extend_num_tokens, + ) + except KVCapacityWaitError: + for req in newly_allocated_reqs: + if req.req_pool_idx is not None: + batch.req_to_token_pool.free(req) + raise # Write to req_to_token_pool write_cache_indices( diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 0ff09acac..418301926 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -33,7 +33,6 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, @@ -288,6 +287,22 @@ class CpHiCacheEvictionPlan: remaining_deficit: Tuple[int, ...] +@dataclass(frozen=True) +class CpLoadBackPlan: + page_owners: List[int] + required_by_owner: List[int] + available_by_owner: List[int] + deficit_by_owner: List[int] + host_hit_len: int + + +@dataclass(frozen=True) +class CpLoadBackEvictionPlan: + victims: Tuple[TreeNode, ...] + planned_freed_by_owner: Tuple[int, ...] + remaining_deficit_by_owner: Tuple[int, ...] + + class HiCachePendingBackupSplit(Exception): def __init__(self, node: TreeNode): self.node = node @@ -869,6 +884,249 @@ class HiRadixCache(RadixCache): remaining_deficit=tuple(deficits), ) + def _build_cp_load_back_plan( + self, nodes_to_load: List[TreeNode], *, node_id: int + ) -> CpLoadBackPlan: + allocator = self.token_to_kv_pool_allocator + compute_owner_lane_stats = getattr(allocator, "compute_owner_lane_stats", None) + if compute_owner_lane_stats is None: + raise RuntimeError( + "CP HiCache load-back requires an allocator with " + f"compute_owner_lane_stats: node_id={node_id}" + ) + + page_owners: List[int] = [] + host_hit_len = 0 + for node in nodes_to_load: + metadata = getattr(node, "cp_hicache", None) + if metadata is None: + raise RuntimeError( + "CP HiCache load-back missing cp_hicache metadata: " + f"load_node_id={getattr(node, 'id', '?')} root_node_id={node_id}" + ) + node_host_len = int(self._node_host_len(node)) + if node_host_len % self.page_size != 0: + raise RuntimeError( + "CP HiCache load-back requires page-aligned host lengths: " + f"load_node_id={getattr(node, 'id', '?')} " + f"host_len={node_host_len} page_size={self.page_size}" + ) + if int(getattr(metadata, "logical_len", node_host_len)) != node_host_len: + raise RuntimeError( + "CP HiCache load-back metadata length mismatch: " + f"load_node_id={getattr(node, 'id', '?')} " + f"host_len={node_host_len} metadata_len={metadata.logical_len}" + ) + page_owners.extend(int(owner) for owner in metadata.page_owners.tolist()) + host_hit_len += node_host_len + + expected_pages = host_hit_len // self.page_size + if len(page_owners) != expected_pages: + raise RuntimeError( + "CP HiCache load-back page owner count mismatch: " + f"node_id={node_id} host_hit_len={host_hit_len} " + f"page_size={self.page_size} page_owners={len(page_owners)}" + ) + + required, available, deficits = compute_owner_lane_stats(page_owners) + return CpLoadBackPlan( + page_owners=list(page_owners), + required_by_owner=[int(v) for v in required], + available_by_owner=[int(v) for v in available], + deficit_by_owner=[int(v) for v in deficits], + host_hit_len=host_hit_len, + ) + + def _refresh_cp_load_back_plan(self, plan: CpLoadBackPlan) -> CpLoadBackPlan: + required, available, deficits = ( + self.token_to_kv_pool_allocator.compute_owner_lane_stats(plan.page_owners) + ) + return CpLoadBackPlan( + page_owners=plan.page_owners, + required_by_owner=[int(v) for v in required], + available_by_owner=[int(v) for v in available], + deficit_by_owner=[int(v) for v in deficits], + host_hit_len=plan.host_hit_len, + ) + + def _cp_device_leaf_is_load_back_victim(self, node: TreeNode) -> bool: + if node == getattr(self, "root_node", None): + return False + if getattr(node, "lock_ref", 0) > 0: + return False + if getattr(node, "value", None) is None: + 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 + if self._is_pinned(node): + return False + return True + + def _cp_load_back_node_owner_page_counts( + self, node: TreeNode, cp_size: int + ) -> Tuple[int, ...]: + value = getattr(node, "value", None) + if value is None or int(value.numel()) == 0: + return tuple(0 for _ in range(cp_size)) + if int(value.numel()) % self.page_size != 0: + raise RuntimeError( + "CP HiCache load-back eviction requires page-aligned device values: " + f"node_id={getattr(node, 'id', '?')} value_len={value.numel()} " + f"page_size={self.page_size}" + ) + + first_locs = value[:: self.page_size] + logical_pages = torch.div(first_locs, self.page_size, rounding_mode="floor") + owners = torch.remainder(logical_pages - 1, cp_size) + return tuple(int((owners == owner).sum().item()) for owner in range(cp_size)) + + def _plan_cp_load_back_owner_lane_evictions( + self, plan: CpLoadBackPlan + ) -> CpLoadBackEvictionPlan: + deficits = [max(0, int(v)) for v in plan.deficit_by_owner] + cp_size = len(deficits) + planned_freed = [0 for _ in range(cp_size)] + victims: List[TreeNode] = [] + used_node_ids = set() + + while any(v > 0 for v in deficits): + best_node = None + best_counts = None + best_score = None + for node in list(getattr(self, "evictable_leaves", set())): + node_id = getattr(node, "id", None) + if node_id in used_node_ids: + continue + if not self._cp_device_leaf_is_load_back_victim(node): + continue + counts = self._cp_load_back_node_owner_page_counts(node, cp_size) + contribution = sum( + min(int(count), int(deficit)) + for count, deficit in zip(counts, deficits) + ) + if contribution <= 0: + continue + score = ( + -int(contribution), + self.eviction_strategy.get_priority(node), + int(node_id or 0), + ) + if best_score is None or score < best_score: + best_score = score + best_node = node + best_counts = counts + + if best_node is None or best_counts is None: + break + + victims.append(best_node) + used_node_ids.add(getattr(best_node, "id", None)) + for owner, count in enumerate(best_counts): + planned_freed[owner] += int(count) + deficits[owner] = max(0, deficits[owner] - int(count)) + + return CpLoadBackEvictionPlan( + victims=tuple(victims), + planned_freed_by_owner=tuple(planned_freed), + remaining_deficit_by_owner=tuple(deficits), + ) + + def _evict_cp_load_back_owner_lanes( + self, plan: CpLoadBackPlan, *, node_id: int + ) -> CpLoadBackPlan: + eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan) + if any(v > 0 for v in eviction_plan.remaining_deficit_by_owner): + logger.warning( + "[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_eviction_insufficient] " + "node_id=%d required_by_owner=%s available_by_owner=%s " + "deficit_by_owner=%s planned_freed_by_owner=%s " + "remaining_deficit_by_owner=%s evictable_size=%d protected_size=%d " + "victims=%s allocator_state=%s", + node_id, + plan.required_by_owner, + plan.available_by_owner, + plan.deficit_by_owner, + eviction_plan.planned_freed_by_owner, + eviction_plan.remaining_deficit_by_owner, + int(getattr(self, "evictable_size_", 0)), + int(getattr(self, "protected_size_", 0)), + [getattr(node, "id", None) for node in eviction_plan.victims], + self.token_to_kv_pool_allocator.allocator_state_str(), + ) + return self._refresh_cp_load_back_plan(plan) + + if len(eviction_plan.victims) == 0: + return plan + + num_evicted = 0 + write_back_nodes: List[TreeNode] = [] + for victim in eviction_plan.victims: + victim_id = getattr(victim, "id", None) + if not self._cp_device_leaf_is_load_back_victim(victim): + raise RuntimeError( + "CP HiCache load-back owner-lane eviction plan diverged " + f"before reservation: node_id={node_id} victim_id={victim_id}" + ) + + if victim.pin_expiry > 0 and time.monotonic() > victim.pin_expiry: + self._clear_pin(victim) + + if not self._node_backuped(victim): + if self.cache_controller.write_policy == "write_back": + logger.warning( + "[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_write_back] " + "node_id=%d victim_id=%s reason=unbacked_write_back_victim", + node_id, + victim_id, + ) + written = self.write_backup(victim, write_back=True) + if written > 0: + write_back_nodes.append(victim) + continue + num_evicted += self._evict_regular(victim) + else: + num_evicted += self._evict_backuped(victim) + + for child in victim.parent.children.values(): + if child in write_back_nodes: + continue + if not child.evicted: + break + else: + self._update_leaf_status(victim.parent) + + if write_back_nodes: + self.writing_check(write_back=True) + for victim in write_back_nodes: + if self._node_backuped(victim): + num_evicted += self._evict_backuped(victim) + + refreshed = self._refresh_cp_load_back_plan(plan) + logger.info( + "[HiCache-load] owner-lane device eviction before CP load-back: " + "node_id=%d victims=%s num_evicted=%d required_by_owner=%s " + "before_available_by_owner=%s before_deficit_by_owner=%s " + "after_available_by_owner=%s after_deficit_by_owner=%s " + "planned_freed_by_owner=%s", + node_id, + [getattr(node, "id", None) for node in eviction_plan.victims], + num_evicted, + plan.required_by_owner, + plan.available_by_owner, + plan.deficit_by_owner, + refreshed.available_by_owner, + refreshed.deficit_by_owner, + eviction_plan.planned_freed_by_owner, + ) + return refreshed + def _cp_capacity_debug_enabled(self) -> bool: return bool( getattr(self, "_cp_hicache_capacity_debug", False) @@ -2470,14 +2728,28 @@ class HiRadixCache(RadixCache): result = self.inc_lock_ref(ancester_node) delta = result.delta - # load it all or not at all - host_hit_len = sum(self._node_host_len(n) for n in nodes_to_load) + # load it all or not at all. The scalar length remains only a + # coarse upper bound; final CP admission is the owner-lane vector + # check below. + try: + load_back_plan = self._build_cp_load_back_plan( + nodes_to_load, node_id=last_hit_node.id + ) + except Exception: + self.dec_lock_ref(ancester_node) + raise + host_hit_len = load_back_plan.host_hit_len logger.info( - "[HiCache-load] load_back CP: node_id=%d nodes_to_load=%d host_hit_len=%d threshold=%d", + "[HiCache-load] load_back CP: node_id=%d nodes_to_load=%d " + "host_hit_len=%d threshold=%d required_by_owner=%s " + "available_by_owner=%s deficit_by_owner=%s", last_hit_node.id, len(nodes_to_load), host_hit_len, self.load_back_threshold, + load_back_plan.required_by_owner, + load_back_plan.available_by_owner, + load_back_plan.deficit_by_owner, ) if host_hit_len < self.load_back_threshold or ( host_hit_len > mem_quota + delta if mem_quota is not None else False @@ -2491,32 +2763,48 @@ class HiRadixCache(RadixCache): ) return None + if any(v > 0 for v in load_back_plan.deficit_by_owner): + load_back_plan = self._evict_cp_load_back_owner_lanes( + load_back_plan, node_id=last_hit_node.id + ) + + if any(v > 0 for v in load_back_plan.deficit_by_owner): + self.dec_lock_ref(ancester_node) + logger.warning( + "[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_capacity_failed] " + "node_id=%d host_hit_len=%d nodes_to_load=%d " + "required_by_owner=%s available_by_owner=%s " + "deficit_by_owner=%s evictable_size=%d protected_size=%d " + "ongoing_load_back_count=%d allocator_state=%s", + last_hit_node.id, + host_hit_len, + len(nodes_to_load), + load_back_plan.required_by_owner, + load_back_plan.available_by_owner, + load_back_plan.deficit_by_owner, + int(getattr(self, "evictable_size_", 0)), + int(getattr(self, "protected_size_", 0)), + len(getattr(self, "ongoing_load_back", {})), + self.token_to_kv_pool_allocator.allocator_state_str(), + ) + return None + device_indices = self.cache_controller.load_cp( nodes_to_load, node_id=last_hit_node.id ) if device_indices is None: - logger.info( - "[HiCache-load] load_back CP retry with lane-aware eviction: " - "node_id=%d tokens_needed=%d", + failed_plan = self._refresh_cp_load_back_plan(load_back_plan) + logger.warning( + "[CP_HICACHE_FALLBACK][cp_load_back_preflight_mismatch] " + "node_id=%d host_hit_len=%d required_by_owner=%s " + "available_by_owner=%s deficit_by_owner=%s " + "allocator_state=%s", last_hit_node.id, host_hit_len, - ) - # Lane-aware eviction: alloc_pages_with_owners failed because - # at least one owner lane is short of free pages. Targeted - # eviction frees pages from THOSE specific lanes, leaving - # other lanes untouched. Mirrors cold prefill's retry path - # in common.py:alloc_paged_token_slots_extend. - _retry_page_owners: List[int] = [] - for _node in nodes_to_load: - # page_owners is CPU int8; tolist() returns list[int]. - _retry_page_owners.extend(_node.cp_hicache.page_owners.tolist()) - _evict_for_compute_owner_lanes( - tree_cache=self, - allocator=self.token_to_kv_pool_allocator, - page_compute_owners=_retry_page_owners, - ) - device_indices = self.cache_controller.load_cp( - nodes_to_load, node_id=last_hit_node.id + failed_plan.required_by_owner, + failed_plan.available_by_owner, + failed_plan.deficit_by_owner, + self.token_to_kv_pool_allocator.allocator_state_str(), ) self.dec_lock_ref(ancester_node) if device_indices is None: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 65688fcd3..3026bc720 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -746,6 +746,32 @@ class KVCache(abc.ABC): def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter): self.layer_transfer_counter = layer_transfer_counter + def wait_layer_transfer_on_stream(self, layer_id: int, stream) -> None: + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until_on_stream( + layer_id - self.start_layer, stream + ) + + def get_key_buffer_for_prefetch(self, layer_id: int, stream) -> torch.Tensor: + self.wait_layer_transfer_on_stream(layer_id, stream) + raw_getter = getattr(self, "_get_key_buffer", None) + if raw_getter is None: + raise NotImplementedError( + f"{type(self).__name__} does not expose a non-blocking key-buffer getter" + ) + return raw_getter(layer_id) + + def get_index_k_with_scale_buffer_for_prefetch( + self, layer_id: int, stream + ) -> torch.Tensor: + self.wait_layer_transfer_on_stream(layer_id, stream) + raw_getter = getattr(self, "_get_index_k_with_scale_buffer", None) + if raw_getter is None: + raise NotImplementedError( + f"{type(self).__name__} does not expose a non-blocking index-buffer getter" + ) + return raw_getter(layer_id) + def register_layer_backup_notifier(self, notifier): self.layer_backup_notifiers.append(notifier) @@ -1379,6 +1405,10 @@ class HybridLinearKVPool(KVCache): layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_key_buffer(layer_id) + def get_key_buffer_for_prefetch(self, layer_id: int, stream): + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_key_buffer_for_prefetch(layer_id, stream) + def get_value_buffer(self, layer_id: int): layer_id = self._transfer_full_attention_id(layer_id) return self.full_kv_pool.get_value_buffer(layer_id) @@ -1552,25 +1582,29 @@ class MLATokenToKVPool(KVCache): return _copy_buf_infos(self._contiguous_buf_infos) - def get_key_buffer(self, layer_id: int): - if self.layer_transfer_counter is not None: - self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - + def _get_key_buffer(self, layer_id: int): if self.store_dtype != self.dtype: return self.kv_buffer[layer_id - self.start_layer].view(self.dtype) return self.kv_buffer[layer_id - self.start_layer] - def get_value_buffer(self, layer_id: int): + def get_key_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + return self._get_key_buffer(layer_id) + def _get_value_buffer(self, layer_id: int): if self.store_dtype != self.dtype: return self.kv_buffer[layer_id - self.start_layer][ ..., : self.kv_lora_rank ].view(self.dtype) return self.kv_buffer[layer_id - self.start_layer][..., : self.kv_lora_rank] + def get_value_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + return self._get_value_buffer(layer_id) + def get_kv_buffer(self, layer_id: int): return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id) @@ -1724,10 +1758,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): del self.kv_scale_buffer self._contiguous_buf_infos = None - def get_key_buffer(self, layer_id: int): - if self.layer_transfer_counter is not None: - self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - + def _get_key_buffer(self, layer_id: int): if self.store_dtype != self.dtype: cache_k_nope_fp4 = self.kv_buffer[layer_id - self.start_layer].view( torch.uint8 @@ -1743,6 +1774,11 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): return self.kv_buffer[layer_id - self.start_layer] + def get_key_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + return self._get_key_buffer(layer_id) + def set_kv_buffer( self, layer: RadixAttention, @@ -1893,10 +1929,13 @@ class NSATokenToKVPool(MLATokenToKVPool): ] self._finalize_allocation_log(size) + def _get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: + return self.index_k_with_scale_buffer[layer_id - self.start_layer] + def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - return self.index_k_with_scale_buffer[layer_id - self.start_layer] + return self._get_index_k_with_scale_buffer(layer_id) def get_index_k_continuous( self, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py new file mode 100644 index 000000000..b3aa49ca8 --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py @@ -0,0 +1,318 @@ +import functools +import json +import sys +import types +import unittest + +import torch + +if "pybase64" not in sys.modules: + pybase64_stub = types.ModuleType("pybase64") + pybase64_stub.b64encode = lambda *args, **kwargs: b"" + pybase64_stub.b64decode = lambda *args, **kwargs: b"" + sys.modules["pybase64"] = pybase64_stub +if "orjson" not in sys.modules: + orjson_stub = types.ModuleType("orjson") + orjson_stub.loads = lambda data, *args, **kwargs: json.loads( + data.decode() if isinstance(data, (bytes, bytearray)) else data + ) + orjson_stub.dumps = lambda obj, *args, **kwargs: json.dumps(obj).encode() + sys.modules["orjson"] = orjson_stub +sgl_kernel_stub = sys.modules.setdefault("sgl_kernel", types.ModuleType("sgl_kernel")) +sgl_kernel_stub.__file__ = getattr(sgl_kernel_stub, "__file__", "sgl_kernel_stub.py") +sgl_kernel_stub.__path__ = getattr(sgl_kernel_stub, "__path__", []) +if not hasattr(sgl_kernel_stub, "__getattr__"): + + def _sgl_kernel_getattr(name): + if name.startswith("__"): + raise AttributeError(name) + fn = lambda *args, **kwargs: None + setattr(sgl_kernel_stub, name, fn) + return fn + + sgl_kernel_stub.__getattr__ = _sgl_kernel_getattr +for _name in ( + "sgl_per_token_group_quant_8bit", + "sgl_per_token_group_quant_fp8", + "sgl_per_token_quant_fp8", + "fp8_blockwise_scaled_mm", + "fp8_scaled_mm", + "silu_and_mul", +): + if not hasattr(sgl_kernel_stub, _name): + setattr(sgl_kernel_stub, _name, lambda *args, **kwargs: None) +quantization_stub = sys.modules.setdefault( + "sgl_kernel.quantization", types.ModuleType("sgl_kernel.quantization") +) +for _name in ( + "ggml_dequantize", + "ggml_moe_a8", + "ggml_moe_a8_vec", + "ggml_moe_get_block_size", + "ggml_mul_mat_a8", + "ggml_mul_mat_vec_a8", +): + if not hasattr(quantization_stub, _name): + setattr(quantization_stub, _name, lambda *args, **kwargs: None) +kvcacheio_stub = sys.modules.setdefault( + "sgl_kernel.kvcacheio", types.ModuleType("sgl_kernel.kvcacheio") +) +for _name in ( + "transfer_kv_all_layer", + "transfer_kv_all_layer_direct_lf_pf", + "transfer_kv_all_layer_lf_pf", + "transfer_kv_all_layer_lf_ph", + "transfer_kv_all_layer_mla", + "transfer_kv_all_layer_mla_lf_pf", + "transfer_kv_direct", + "transfer_kv_per_layer", + "transfer_kv_per_layer_direct_pf_lf", + "transfer_kv_per_layer_mla", + "transfer_kv_per_layer_mla_pf_lf", + "transfer_kv_per_layer_pf_lf", + "transfer_kv_per_layer_ph_lf", +): + if not hasattr(kvcacheio_stub, _name): + setattr(kvcacheio_stub, _name, lambda *args, **kwargs: None) +_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") +for _schema in ( + "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", + "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()", + "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()", + "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor", + "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor", +): + try: + _sgl_kernel_lib.define(_schema) + except RuntimeError as exc: + if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): + raise + +from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator +from sglang.srt.mem_cache.hiradix_cache import ( + CpHiCacheNodeMetadata, + HiRadixCache, +) +from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, get_child_key +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + + +class _FakeLayout: + cp_size = 4 + cp_rank = 0 + + +class _FakeEvictionStrategy: + def get_priority(self, node): + return getattr(node, "priority", 0) + + +class _FakeController: + write_policy = "write_through" + has_draft_hicache = False + cp_shared_kv_layout = _FakeLayout() + + def __init__(self, allocator): + self.allocator = allocator + self.load_calls = 0 + self.evicted_device_indices = [] + self.ack_load_queue = [] + self.force_load_none = False + + def load_cp(self, nodes_to_load, node_id=-1): + page_owners = [] + for node in nodes_to_load: + page_owners.extend(node.cp_hicache.page_owners.tolist()) + _, _, deficits = self.allocator.compute_owner_lane_stats(page_owners) + if any(deficits): + raise AssertionError( + f"load_cp called before owner-lane deficits were evicted: {deficits}" + ) + self.load_calls += 1 + if self.force_load_none: + return None + return self.allocator.alloc_pages_with_owners(page_owners) + + def evict_device(self, indices): + self.evicted_device_indices.append(indices.clone()) + self.allocator.free(indices) + return int(indices.numel()) + + +def _make_allocator(page_size=4, cp_size=4): + allocator = CPSharedPagedTokenToKVPoolAllocator( + logical_size=page_size * 16, + physical_size=page_size * 4, + page_size=page_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + cp_size=cp_size, + cp_rank=0, + ) + allocator.release_pages = torch.empty((0,), dtype=torch.int64) + return allocator + + +def _metadata(page_owners, page_size=4): + logical_len = len(page_owners) * page_size + owned_positions = [] + for page_idx, owner in enumerate(page_owners): + if owner == 0: + owned_positions.extend( + range(page_idx * page_size, (page_idx + 1) * page_size) + ) + return CpHiCacheNodeMetadata( + logical_len=logical_len, + owned_positions=torch.tensor(owned_positions, dtype=torch.int64), + host_indices=torch.arange(len(owned_positions), dtype=torch.int64), + page_owners=torch.tensor(page_owners, dtype=torch.int8), + page_size=page_size, + ) + + +def _make_node(node_id, key_start, page_owners, *, value=None, priority=0): + page_size = 4 + node = TreeNode(id=node_id, priority=priority) + node.key = RadixKey(token_ids=list(range(key_start, key_start + len(page_owners) * page_size))) + node.value = value + node.host_len = len(page_owners) * page_size + node.cp_hicache = _metadata(page_owners, page_size=page_size) + return node + + +def _attach_child(cache, parent, child): + child.parent = parent + parent.children[cache.get_child_key_fn(child.key)] = child + + +def _make_cache(allocator): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.token_to_kv_pool_allocator = allocator + cache.page_size = allocator.page_size + cache.cache_controller = _FakeController(allocator) + cache.root_node = TreeNode(id=0, priority=-999) + cache.root_node.key = RadixKey(token_ids=[]) + cache.root_node.value = [] + cache.root_node.lock_ref = 1 + cache.evictable_leaves = set() + cache.evictable_host_leaves = set() + cache.evictable_size_ = 0 + cache.protected_size_ = 0 + cache.ongoing_load_back = {} + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + cache.load_back_threshold = 0 + cache.metrics_collector = None + cache.eviction_strategy = _FakeEvictionStrategy() + cache.get_child_key_fn = functools.partial(get_child_key, page_size=allocator.page_size) + cache.enable_kv_cache_events = False + return cache + + +class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase): + def test_load_back_plan_reports_owner_lane_vectors(self): + allocator = _make_allocator() + allocator.free_pages = torch.tensor([1, 2], dtype=torch.int64) + allocator.release_pages = torch.tensor([4], dtype=torch.int64) + cache = _make_cache(allocator) + node = _make_node(10, 100, [0, 1, 1, 3]) + + plan = cache._build_cp_load_back_plan([node], node_id=node.id) + + self.assertEqual(plan.page_owners, [0, 1, 1, 3]) + self.assertEqual(plan.required_by_owner, [1, 2, 0, 1]) + self.assertEqual(plan.available_by_owner, [1, 1, 0, 1]) + self.assertEqual(plan.deficit_by_owner, [0, 1, 0, 0]) + self.assertEqual(plan.host_hit_len, 16) + + def test_load_back_plan_fails_closed_without_cp_metadata(self): + allocator = _make_allocator() + cache = _make_cache(allocator) + node = TreeNode(id=11) + node.host_len = 4 + node.cp_hicache = None + + with self.assertRaisesRegex(RuntimeError, "missing cp_hicache metadata"): + cache._build_cp_load_back_plan([node], node_id=node.id) + + def test_load_back_evicts_owner_lane_deficit_before_allocating(self): + allocator = _make_allocator() + # Only owner lane 1 is initially available. The load-back target needs + # lanes [0, 1], so calling load_cp before targeted eviction is a bug. + allocator.free_pages = torch.tensor([2], dtype=torch.int64) + cache = _make_cache(allocator) + + victim = _make_node( + 20, + 200, + [0], + value=torch.tensor([4, 5, 6, 7], dtype=torch.int64), + priority=0, + ) + _attach_child(cache, cache.root_node, victim) + cache.evictable_leaves.add(victim) + cache.evictable_size_ = len(victim.key) + + target = _make_node(21, 300, [0, 1], value=None, priority=10) + _attach_child(cache, cache.root_node, target) + + loaded = cache.load_back(target, mem_quota=100) + + self.assertIsNotNone(loaded) + self.assertEqual(cache.cache_controller.load_calls, 1) + self.assertEqual(len(cache.cache_controller.evicted_device_indices), 1) + self.assertEqual(cache.cache_controller.evicted_device_indices[0].tolist(), [4, 5, 6, 7]) + self.assertEqual(target.value.tolist(), loaded.tolist()) + self.assertIn(target.id, cache.ongoing_load_back) + + def test_load_back_failure_leaves_node_unassigned_and_unlocked(self): + allocator = _make_allocator() + allocator.free_pages = torch.tensor([1, 2], dtype=torch.int64) + cache = _make_cache(allocator) + cache.cache_controller.force_load_none = True + + target = _make_node(30, 400, [0, 1], value=None) + _attach_child(cache, cache.root_node, target) + + loaded = cache.load_back(target, mem_quota=100) + + self.assertIsNone(loaded) + self.assertIsNone(target.value) + self.assertEqual(target.lock_ref, 0) + self.assertNotIn(target.id, cache.ongoing_load_back) + + def test_load_back_lock_is_released_only_by_loading_ack(self): + class ReadyEvent: + def query(self): + return True + + allocator = _make_allocator() + allocator.free_pages = torch.tensor([1, 2], dtype=torch.int64) + cache = _make_cache(allocator) + + target = _make_node(31, 500, [0, 1], value=None) + _attach_child(cache, cache.root_node, target) + + loaded = cache.load_back(target, mem_quota=100) + + self.assertIsNotNone(loaded) + self.assertEqual(target.lock_ref, 1) + self.assertIn(target.id, cache.ongoing_load_back) + + cache.cache_controller.ack_load_queue.append((None, ReadyEvent(), [target.id])) + cache.loading_check() + + self.assertEqual(target.lock_ref, 0) + self.assertNotIn(target.id, cache.ongoing_load_back) + self.assertEqual(cache.cache_controller.ack_load_queue, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index 0dad64de5..f084e8578 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py @@ -4,6 +4,18 @@ from unittest.mock import patch import numpy as np import torch +_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") +try: + _sgl_kernel_lib.define( + "moe_fused_gate(Tensor input_tensor, Tensor? bias, int num_expert_group, " + "int topk_group, int topk, int num_fused_shared_experts, " + "float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) " + "-> (Tensor, Tensor)" + ) +except RuntimeError as exc: + if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): + raise + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -435,6 +447,75 @@ class TestCPSharedPagedAllocator(CustomTestCase): self.assertEqual(allocator.calls, 1) self.assertEqual(out.numel(), page_size * 8) + def test_compute_owner_alloc_skips_aggregate_evict_before_owner_attempt(self): + from types import SimpleNamespace + + from sglang.srt.mem_cache import common + + page_size = 64 + + class FakeAllocator: + def __init__(self): + self.page_size = page_size + self.cp_size = 4 + self.calls = 0 + + def available_size(self): + # Aggregate capacity looks insufficient, but the owner-aware + # allocator can satisfy the request from the right lanes. + # Aggregate eviction before this attempt would evict unrelated + # lanes and destroy cache locality. + return 0 + + def alloc_extend_compute_owner( + self, + _prefix_lens, + _prefix_lens_cpu, + _seq_lens, + _seq_lens_cpu, + _last_loc, + extend_num_tokens, + _page_compute_owners, + ): + self.calls += 1 + return torch.arange(extend_num_tokens, dtype=torch.int64) + + def alloc_extend(self, *_args, **_kwargs): + raise AssertionError("legacy allocation should not be used") + + class FakeTreeCache: + def __init__(self, allocator): + self.token_to_kv_pool_allocator = allocator + + def is_chunk_cache(self): + return False + + def evict(self, _params): + raise AssertionError( + "aggregate eviction should not run before owner-aware allocation" + ) + + allocator = FakeAllocator() + server_args = SimpleNamespace( + enable_nsa_prefill_cp_shared_kv=True, + enable_nsa_prefill_context_parallel=True, + nsa_prefill_cp_mode="in-seq-split", + ) + + with patch.object(common, "get_global_server_args", return_value=server_args): + out = common.alloc_paged_token_slots_extend( + tree_cache=FakeTreeCache(allocator), + prefix_lens=torch.tensor([0], dtype=torch.int64), + prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + seq_lens=torch.tensor([page_size * 8], dtype=torch.int64), + seq_lens_cpu=torch.tensor([page_size * 8], dtype=torch.int64), + last_loc=torch.tensor([-1], dtype=torch.int64), + extend_num_tokens=page_size * 8, + ) + + self.assertEqual(allocator.calls, 1) + self.assertEqual(out.numel(), page_size * 8) + def test_compute_owner_lane_eviction_recovers_exhausted_owner_lane(self): from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import EvictResult @@ -508,6 +589,253 @@ class TestCPSharedPagedAllocator(CustomTestCase): ) self.assertIsNotNone(locs) + def test_compute_owner_lane_eviction_passes_deficits_to_tree_cache(self): + from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator + from sglang.srt.mem_cache.base_prefix_cache import EvictResult + from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes + + page_size = 64 + cp_size = 4 + allocator = CPSharedPagedTokenToKVPoolAllocator( + logical_size=page_size * 16, + physical_size=page_size * 4, + page_size=page_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + cp_size=cp_size, + cp_rank=0, + ) + allocator.free_pages = torch.empty((0,), dtype=torch.int64) + allocator.release_pages = torch.empty((0,), dtype=torch.int64) + + class FakeTreeCache: + def __init__(self): + self.owner_deficits = [] + + def is_chunk_cache(self): + return False + + def evictable_size(self): + return page_size * cp_size + + def evict(self, params): + deficits = getattr(params, "owner_lane_deficits", None) + self.owner_deficits.append(deficits) + if deficits == [1, 0, 0, 0]: + allocator.free(torch.tensor([page_size], dtype=torch.int64)) + else: + allocator.free(torch.tensor([page_size * 2], dtype=torch.int64)) + return EvictResult(num_tokens_evicted=page_size) + + tree_cache = FakeTreeCache() + _evict_for_compute_owner_lanes( + tree_cache=tree_cache, + allocator=allocator, + page_compute_owners=[0], + ) + + self.assertEqual(tree_cache.owner_deficits[0], [1, 0, 0, 0]) + locs = allocator.alloc_extend_compute_owner( + prefix_lens=torch.tensor([0], dtype=torch.int64), + prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + seq_lens=torch.tensor([page_size], dtype=torch.int64), + seq_lens_cpu=torch.tensor([page_size], dtype=torch.int64), + last_loc=torch.tensor([-1], dtype=torch.int64), + extend_num_tokens=page_size, + page_compute_owners=[0], + ) + self.assertIsNotNone(locs) + + def test_compute_owner_capacity_wait_reports_owner_lane_deficits(self): + from types import SimpleNamespace + + from sglang.srt.mem_cache import common + from sglang.srt.mem_cache.base_prefix_cache import EvictResult + + page_size = 64 + + class FakeAllocator: + def __init__(self): + self.page_size = page_size + self.cp_size = 4 + + def available_size(self): + return self.page_size * 4 + + def allocator_state_str(self): + return "allocator_state_for_test" + + def alloc_extend_compute_owner(self, *_args, **_kwargs): + return None + + def compute_owner_lane_stats(self, _page_compute_owners): + return [2, 2, 2, 2], [2, 2, 2, 0], [0, 0, 0, 2] + + def alloc_extend(self, *_args, **_kwargs): + raise AssertionError("legacy allocation should not be used") + + class FakeTreeCache: + def __init__(self): + self.token_to_kv_pool_allocator = FakeAllocator() + self.evict_params = [] + + def is_chunk_cache(self): + return False + + def evictable_size(self): + return page_size * 8 + + def evict(self, params): + self.evict_params.append(params) + return EvictResult(num_tokens_evicted=0) + + def pretty_print(self): + raise AssertionError("recoverable capacity wait should not dump tree") + + server_args = SimpleNamespace( + enable_nsa_prefill_cp_shared_kv=True, + enable_nsa_prefill_context_parallel=True, + nsa_prefill_cp_mode="in-seq-split", + ) + tree_cache = FakeTreeCache() + + with patch.object(common, "get_global_server_args", return_value=server_args): + with self.assertRaises(common.KVCapacityWaitError) as cm: + common.alloc_paged_token_slots_extend( + tree_cache=tree_cache, + prefix_lens=torch.tensor([0], dtype=torch.int64), + prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + seq_lens=torch.tensor([page_size * 8], dtype=torch.int64), + seq_lens_cpu=torch.tensor([page_size * 8], dtype=torch.int64), + last_loc=torch.tensor([-1], dtype=torch.int64), + extend_num_tokens=page_size * 8, + ) + + err = cm.exception + self.assertEqual(err.required_by_owner, [2, 2, 2, 2]) + self.assertEqual(err.available_by_owner, [2, 2, 2, 0]) + self.assertEqual(err.deficit_by_owner, [0, 0, 0, 2]) + self.assertIn("owner_lane_exhausted", str(err)) + self.assertIn("deficit_by_owner=[0, 0, 0, 2]", str(err)) + self.assertIn( + [0, 0, 0, 2], + [params.owner_lane_deficits for params in tree_cache.evict_params], + ) + + def test_alloc_for_extend_releases_req_slots_on_recoverable_capacity_wait(self): + from types import SimpleNamespace + + from sglang.srt.mem_cache import common + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + + req_to_token_pool = ReqToTokenPool( + size=1, + max_context_len=128, + device="cpu", + enable_memory_saver=False, + ) + req = SimpleNamespace( + req_pool_idx=None, + is_chunked=0, + kv_committed_len=0, + prefix_indices=torch.empty((0,), dtype=torch.int64), + ) + batch = SimpleNamespace( + maybe_evict_swa=lambda: None, + reqs=[req], + prefix_lens=[0], + extend_lens=[64], + device="cpu", + req_to_token_pool=req_to_token_pool, + tree_cache=SimpleNamespace(page_size=64), + seq_lens=torch.tensor([64], dtype=torch.int64), + seq_lens_cpu=torch.tensor([64], dtype=torch.int64), + extend_num_tokens=64, + ) + wait_error = common.KVCapacityWaitError( + "owner_lane_exhausted for test", + required_by_owner=[1, 0], + available_by_owner=[0, 0], + deficit_by_owner=[1, 0], + ) + + with patch.object( + common, + "alloc_paged_token_slots_extend", + side_effect=wait_error, + ): + with self.assertRaises(common.KVCapacityWaitError): + common.alloc_for_extend(batch) + + self.assertIsNone(req.req_pool_idx) + self.assertEqual(req_to_token_pool.available_size(), 1) + + def test_scheduler_capacity_wait_rollback_releases_adder_locks(self): + from types import SimpleNamespace + + from sglang.srt.managers.scheduler import Scheduler + + class FakeTreeCache: + def __init__(self): + self.dec_calls = [] + + def supports_swa(self): + return False + + def is_tree_cache(self): + return True + + def dec_lock_ref(self, node, params=None): + self.dec_calls.append((node, params)) + + scheduler = object.__new__(Scheduler) + scheduler.tree_cache = FakeTreeCache() + req = SimpleNamespace(last_node="scheduled-node", swa_uuid_for_lock=None) + chunked_req = SimpleNamespace(last_node="chunked-node", swa_uuid_for_lock=None) + + scheduler._release_prefill_adder_locks( + [req, chunked_req], + skip_req=chunked_req, + ) + + self.assertEqual( + [(node, params) for node, params in scheduler.tree_cache.dec_calls], + [("scheduled-node", None)], + ) + + def test_scheduler_capacity_wait_rollback_releases_swa_lock_uuid(self): + from types import SimpleNamespace + + from sglang.srt.managers.scheduler import Scheduler + + class FakeTreeCache: + def __init__(self): + self.dec_calls = [] + + def supports_swa(self): + return True + + def is_tree_cache(self): + return True + + def dec_lock_ref(self, node, params=None): + self.dec_calls.append((node, params)) + + scheduler = object.__new__(Scheduler) + scheduler.tree_cache = FakeTreeCache() + req = SimpleNamespace(last_node="scheduled-node", swa_uuid_for_lock=17) + + scheduler._release_prefill_adder_locks([req]) + + self.assertEqual(scheduler.tree_cache.dec_calls[0][0], "scheduled-node") + self.assertEqual( + scheduler.tree_cache.dec_calls[0][1].swa_uuid_for_lock, + 17, + ) + self.assertIsNone(req.swa_uuid_for_lock) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 4c7fb66ec..2fe31e1c6 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -1,9 +1,30 @@ import unittest +import sys +import types from types import SimpleNamespace from unittest.mock import patch import torch +_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") +try: + _sgl_kernel_lib.define( + "moe_fused_gate(Tensor input_tensor, Tensor? bias, int num_expert_group, " + "int topk_group, int topk, int num_fused_shared_experts, " + "float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) " + "-> (Tensor, Tensor)" + ) +except RuntimeError as exc: + if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): + raise + +flash_attn_stub = sys.modules.setdefault( + "sgl_kernel.flash_attn", types.ModuleType("sgl_kernel.flash_attn") +) +for _name in ("flash_attn_varlen_func", "flash_attn_with_kvcache"): + if not hasattr(flash_attn_stub, _name): + setattr(flash_attn_stub, _name, lambda *args, **kwargs: None) + from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=1, suite="stage-a-test-cpu") @@ -14,7 +35,7 @@ def _identity_all_reduce(buffer, *args, **kwargs): class TestCpSharedKVRuntimeHelpers(unittest.TestCase): - def test_mla_prefetch_materializes_on_current_stream_and_defers_reduce_to_launch( + def test_mla_prefetch_materializes_and_reduces_on_prefetch_stream( self, ): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -48,8 +69,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def __init__(self): self.kv_cache = torch.zeros((32, 1, 2), dtype=torch.float32) + self.prefetch_getter_streams = [] def get_key_buffer(self, layer_id): + raise AssertionError("prefetch must not call blocking get_key_buffer") + + def get_key_buffer_for_prefetch(self, layer_id, stream): + self.prefetch_getter_streams.append((layer_id, stream.name)) return self.kv_cache active_stream = ["current"] @@ -89,23 +115,28 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): "_all_reduce_materialized_buffer_async", side_effect=record_reduce, ): + pool = FakePool() prefetcher.start_next_layer_prefix( next_layer_id=1, - token_to_kv_pool=FakePool(), + token_to_kv_pool=pool, ) - self.assertEqual(calls, [("materialize", "current")]) - self.assertEqual(prefetch_stream.waited, []) + self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) + self.assertEqual( + calls, + [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], + ) + self.assertEqual(prefetch_stream.waited, ["current"]) prefetcher.launch_pending_reduce() self.assertEqual( calls, - [("materialize", "current"), ("reduce", "prefetch", "prefetch")], + [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], ) self.assertEqual(prefetch_stream.waited, ["current"]) - def test_index_prefetch_materializes_on_current_stream_and_defers_reduce_to_launch( + def test_index_prefetch_materializes_and_reduces_on_prefetch_stream( self, ): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -139,8 +170,15 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def __init__(self): self.page_buffer = torch.zeros((16, 3), dtype=torch.uint8) + self.prefetch_getter_streams = [] def get_index_k_with_scale_buffer(self, layer_id): + raise AssertionError( + "prefetch must not call blocking get_index_k_with_scale_buffer" + ) + + def get_index_k_with_scale_buffer_for_prefetch(self, layer_id, stream): + self.prefetch_getter_streams.append((layer_id, stream.name)) return self.page_buffer active_stream = ["current"] @@ -179,22 +217,55 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): "_all_reduce_materialized_buffer_async", side_effect=record_reduce, ): + pool = FakePool() prefetcher.start_next_layer_prefix( next_layer_id=1, - token_to_kv_pool=FakePool(), + token_to_kv_pool=pool, ) - self.assertEqual(calls, [("materialize", "current")]) - self.assertEqual(prefetch_stream.waited, []) + self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) + self.assertEqual( + calls, + [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], + ) + self.assertEqual(prefetch_stream.waited, ["current"]) prefetcher.launch_pending_reduce() self.assertEqual( calls, - [("materialize", "current"), ("reduce", "prefetch", "prefetch")], + [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], ) self.assertEqual(prefetch_stream.waited, ["current"]) + def test_mla_pool_prefetch_getter_orders_layer_transfer_on_prefetch_stream(self): + from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool + + class FakeCounter: + def __init__(self): + self.calls = [] + + def wait_until(self, threshold): + raise AssertionError("prefetch getter must not wait on current stream") + + def wait_until_on_stream(self, threshold, stream): + self.calls.append((threshold, stream.name)) + + class FakeStream: + name = "prefetch" + + pool = MLATokenToKVPool.__new__(MLATokenToKVPool) + pool.start_layer = 2 + pool.layer_transfer_counter = FakeCounter() + pool.store_dtype = torch.float32 + pool.dtype = torch.float32 + pool.kv_buffer = [torch.ones((4, 1), dtype=torch.float32) for _ in range(3)] + + out = pool.get_key_buffer_for_prefetch(3, FakeStream()) + + self.assertIs(out, pool.kv_buffer[1]) + self.assertEqual(pool.layer_transfer_counter.calls, [(1, "prefetch")]) + def test_all_reduce_uses_group_fast_path_for_float_buffers(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -1782,12 +1853,12 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): self.assertIs(dense_buffer, fallback_buffer) self.assertIs(dense_pages, fallback_pages) - logger.info.assert_called_once() + logger.warning.assert_called_once() self.assertIn( - "CP shared KV index prefetch fallback", - logger.info.call_args.args[0], + "[CP_SHARED_KV_FALLBACK][index_prefetch]", + logger.warning.call_args.args[0], ) - self.assertIn("consume_miss", logger.info.call_args.args[1]) + self.assertIn("consume_miss", logger.warning.call_args.args[1]) def test_index_prefetch_first_layer_miss_does_not_log_fallback(self): from sglang.srt.layers.attention.nsa import nsa_indexer @@ -1835,7 +1906,7 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): logical_page_table=logical_pages, ) - logger.info.assert_not_called() + logger.warning.assert_not_called() def test_index_prefetch_create_skip_logs_fallback_when_enabled(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -1847,7 +1918,7 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): ), patch.object( prefetch.torch.cuda, "is_available", return_value=False ), patch.object( - prefetch.logger, "info" + prefetch.logger, "warning" ) as logger: result = prefetch.CpSharedKVIndexPrefetcher.maybe_create( forward_batch=SimpleNamespace(), @@ -1858,7 +1929,7 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): self.assertIsNone(result) logger.assert_called_once() self.assertIn( - "CP shared KV index prefetch fallback", + "[CP_SHARED_KV_FALLBACK][index_prefetch]", logger.call_args.args[0], ) self.assertIn("cuda_unavailable_or_stream_capturing", logger.call_args.args[1])