Make CP HiCache residency owner-lane deterministic
CP shared KV cannot treat capacity as a scalar token count: cache-hit load-back and fresh extend allocation both have to preserve the logical page owner pattern or later direct writes, HiCache reload, and prefix materialization can read the wrong lane. This change moves the critical paths to owner-lane plans, makes owner-lane exhaustion recoverable during prefill scheduling, and routes shared-KV prefix prefetch through prefetch-stream-safe KV getters so HiCache layer-load waits do not attach to the forward stream. Constraint: CP shared KV correctness depends on page owner lane preservation across allocation, backup, load, eviction, and prefix materialization. Constraint: Avoid adding CP/global collectives for capacity agreement; derive capacity from deterministic local owner-lane state. Rejected: Keep SGLANG_DISABLE_TAI_OWNER_SELECT fallback | legacy allocation can silently break owner-lane invariants. Rejected: Scalar total-token eviction for CP HiCache load-back | total capacity can be sufficient while the required owner lane is exhausted. Confidence: medium Scope-risk: broad Directive: Do not reintroduce silent legacy fallback in owner-lane paths; unexpected owner-lane failure must be warning-level fail-closed or recoverable capacity wait. Tested: 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. Tested: Local py_compile for modified runtime/cache/scheduler modules. Not-tested: Full CUDA ETE performance trace for cache-hit overlap and MTP accept-rate impact. Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
@@ -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
|
||||
```
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user