# NSA Prefill CP / CP Shared KV Batch Size > 1 Exploration

Date: 2026-05-28

## Update 2026-06-03

- This exploration is still valid as a record of the original `batch_size == 1` blockers.
- The near-term priority has changed because observed online cache-hit extends are often only `200-2000` tokens and do not fill compute.
- The active implementation phase plan is now `docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md`; Chinese version: `docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md`; parallel handoff plan: `docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md`.
- Do not implement true batch by removing guards only; use the new per-request page-aligned planning contract.

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(...)`
  - source note records multi-batch as unsupported
  - 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 zero-owner-page segments;
- 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
```

