# NSA Prefill CP Shared-KV Batch Size > 1 Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: use `superpowers:subagent-driven-development` or `superpowers:executing-plans` before implementing this plan task-by-task. This document is the phase contract; implementation must still use TDD and checkpoint after each phase.

**Goal:** make `--nsa-prefill-cp-mode in-seq-split` + CP shared KV + HiCache support true `ForwardBatch(batch_size > 1)` for cache-hit-heavy short extends without falling back to non-owner-lane or synchronous slow paths.

**Architecture:** preserve the current page-as-minimum-cache-unit contract, but make every CP planning object per-request before flattening into the existing tensor/kernel interfaces. Correctness comes from deterministic local planning from `extend_prefix_lens_cpu`, `extend_seq_lens_cpu`, `out_cache_loc`, page size, and CP rank; no new collective is allowed for planning.

**Tech Stack:** SGLang Python scheduler/model runtime, NSA CP metadata, CP shared-KV runtime/prefetch, HiCache radix/cache controller, PyTorch CUDA tensors, existing SGL/Tai kernels only after CPU/unit contracts are locked.

---

## 1. Current conclusion

The old exploration note recommended multi-slot overlap first because online extends were thought to be larger and the CP shared-KV code had hard `batch_size == 1` assumptions. Current traffic changes the priority: many cache-hit extends are only about `200-2000` tokens, so a single request often does not fill compute. True batch support is now a throughput feature, not just a cleanup.

The implementation cannot simply remove `batch_size != 1` guards. The current in-seq CP shared-KV path uses one request's split metadata as if the flattened input is one sequence. With `batch_size > 1`, that would mix request boundaries and corrupt page ownership, causal/top-k ranges, partial-current reuse, and draft/EAGLE last-token collection.

## 2. Non-negotiable invariants

1. **Request boundary is never lost.** Each request owns its own prefix length, extend length, page extent, CP split list, and `kv_len_prev/next` pair.
2. **Page is the minimum cache unit.** Short tails may waste up to `page_size - 1` physical rows; they must not trigger token-balanced CP fallback.
3. **Owner-lane plan order equals flattened `out_cache_loc` page order.** For a batch, concatenate per-request owner plans in request order.
4. **All CP ranks derive the same plan locally.** Planning uses CPU length metadata and server config only; no all-reduce/all-gather for planner agreement.
5. **Fallbacks are fail-fast or explicit warnings.** Supported bs>1 paths must not silently fall back to legacy allocation, token-balanced split, or sync slow path.
6. **Target first, draft second, prefetch last.** Enable target-model sync correctness before EAGLE/draft and before async L1 prefetch.
7. **HiCache metadata mirrors write-time owner state.** Load-back can only reuse nodes whose `page_owners`, `owned_positions`, `valid_len`, and `padded_len` were written under the batch-aware contract.

## 3. Code facts already checked

This section is the ledger to avoid rediscovering the same points.

### C1. Generic scheduler can create multi-request extends

- `python/sglang/srt/managers/schedule_batch.py:1481-1540`
  - `ScheduleBatch.prepare_for_extend()` already builds flattened `input_ids`, per-request `seq_lens`, `prefix_lens`, `extend_lens`, and `extend_num_tokens`.
- `python/sglang/srt/model_executor/forward_batch_info.py:449-500`
  - `ForwardBatch.init_new()` sets `batch_size=len(batch.seq_lens)` and passes flattened tensors onward.

Implication: the blocker is not generic batching. The blocker is specialized NSA CP shared-KV metadata and cache ownership.

### C2. In-seq CP metadata is single-request

- `python/sglang/srt/layers/attention/nsa/utils.py:1094-1244`
  - `prepare_input_dp_with_cp_dsa()` has `# just support batch = 1`, sets `bs_per_cp_group = 1`, and builds one `split_list`, one `zigzag_index`, one `kv_len_prev/next` pair.
- `python/sglang/srt/layers/attention/nsa/utils.py:203-244`
  - `NSAContextParallelMetadata` stores scalar/single-request fields.

Implication: Phase 1 must introduce per-request metadata and flattened derived views.

### C3. Page-aligned split only handles `len(...) == 1`

- `python/sglang/srt/layers/attention/nsa/utils.py:300-390`
  - `_build_in_seq_split_for_forward_batch()` chooses page-aligned split only when `len(extend_seq_lens_cpu) == 1` and `len(extend_prefix_lens_cpu) == 1`.
  - Otherwise it falls back to token-balanced split.

Implication: bs>1 support must make page-aligned split batched. Token-balanced fallback is not acceptable for CP shared KV.

### C4. `can_cp_split()` can make the wrong decision for bs>1

- `python/sglang/srt/layers/attention/nsa/utils.py:460-483`
  - `can_cp_split(seq_len, ...)` uses the flattened total token count.
  - `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` returns `False` unless batch size is one.

Implication: a batch of many tiny suffixes can enter in-seq CP split based on total tokens even when each request has fewer than `cp_size` physical pages. Phase 1 must make split eligibility per request or fail-fast before entering CP split.

### C5. Owner-lane allocation is single-request

- `python/sglang/srt/mem_cache/common.py:455-565`
  - `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.
- `python/sglang/srt/mem_cache/allocator.py:981-1048`
  - `alloc_extend_compute_owner()` raises on `len(prefix_lens_cpu) != 1 or len(seq_lens_cpu) != 1`.

Implication: Phase 2 must add batch-aware owner plans and allocation. Legacy allocation is correctness-risky under CP direct write.

### C6. Direct-write local loc computation is single-request

- `python/sglang/srt/layers/attention/nsa/utils.py:562-648`
  - `get_cp_shared_kv_local_out_cache_loc()` requires one `metadata.split_list`, splits the entire flattened `out_cache_loc`, then validates owner-lane ownership.
- Callers:
  - `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
  - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`

Implication: Phase 3 must split `out_cache_loc` per request, then concatenate this rank's local locs in the same local compute order.

### C7. Current/partial-current reuse has been extended to target bs>1

- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:1764-1904`
  - `_current_extend_kv_reuse_miss_reason()` no longer rejects `batch_size > 1`; it validates batch size, prefix/extend/seq lengths, and `out_cache_loc`.
  - `current_extend_kv_rows_for_reuse()` returns flattened valid current rows from `sum(extend_seq_lens_cpu)`, so padded tails are not exposed to attention.
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:309-380`
  - Index partial-current compose uses batch prefix slot spans for bs>1.
  - Index top-k current-only bs>1 avoids materialize; partial-current bs>1 materializes once and reuses the dense buffer for all request segments.
- `python/sglang/srt/layers/attention/nsa_backend.py:1940-2020`
  - MLA partial-current compose uses batch prefix slot spans for bs>1.

Implication: target MLA/index sync-correct current-only and partial-current reuse are covered. Remaining work is kernel-backed variable-length descriptors, L1 prefetch, HiCache load/backup, and draft/EAGLE.

### C8. L1 prefix prefetch is single-request

- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:360-455`
  - `CpSharedKVMlaPrefetcher.maybe_create()` rejects `batch_size != 1`.
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:1135-1235`
  - `CpSharedKVIndexPrefetcher.maybe_create()` rejects `batch_size != 1`.

Implication: Phase 8 must add batched prefetch only after sync correctness. Otherwise bs>1 can become slower than bs=1 because prefix preparation lands on the critical path.

### C9. Top-k in-seq CP pair is single-request

- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:1371-1410`
  - `_get_topk_in_seq_cp_pair()` contains an upstream multi-batch unsupported note and uses scalar `kv_len_prev/next` plus scalar `actual_seq_q_prev/next`.

Implication: Phase 4 must make top-k operate over per-request local prev/next slices.

### C10. EAGLE/draft local path is single-request-sensitive

- `python/sglang/srt/models/deepseek_nextn.py:292-399`
  - Draft path calls `cp_split_and_rebuild_1d`, `cp_split_and_rebuild_position`, `cp_split_and_rebuild_data`, and `cp_collect_last_token_hidden()`.
- `python/sglang/srt/layers/attention/nsa/utils.py:1333-1358`
  - Exact last-token owner/offset handling only works for `bs == 1`.

Implication: Phase 7 must fix draft after target correctness. Do not use draft bs>1 local path before per-request last-token collection is proven.

### C11. HiCache load is closer, but depends on write-time metadata

- `python/sglang/srt/managers/cache_controller.py:1338-1498`
  - `load_cp()` can concatenate `page_owners` across multiple nodes and allocate with `alloc_pages_with_owners()`.
- `python/sglang/srt/managers/scheduler.py:2630-2655`
  - `_prepare_hicache_write_backups_before_forward()` already loops per request.

Implication: HiCache is not the first blocker. It should work after write-time batch metadata is correct, but needs targeted tests for multi-node load-back.

## 4. Target data model

Add batch-aware metadata without breaking existing single-request users.

```text
NSAContextParallelMetadata
  existing scalar fields                      # preserved for batch_size == 1
  batch_size: int
  request_token_offsets: List[int]            # flattened valid-token offsets
  request_padded_token_offsets: List[int]     # flattened physical page-token offsets
  request_page_offsets: List[int]             # flattened page offsets
  request_split_infos: List[PageAlignedInSeqSplitInfo]
  request_split_lists: List[List[int]]        # one cp_size*2 split per request
  request_zigzag_indices: List[List[int]]     # usually same formula, stored for clarity
  request_rank_local_tokens: List[int]        # local valid rows for this rank per request
  request_kv_len_prev: List[int]
  request_kv_len_next: List[int]
  request_actual_seq_q_prev: List[int]
  request_actual_seq_q_next: List[int]
  flat_split_list: List[int]                  # request0 segments, request1 segments, ...
  flat_zigzag_index: List[int]                # indices into flat segments
```

Single-request compatibility rule:

```text
if batch_size == 1:
  split_list == request_split_lists[0]
  zigzag_index == request_zigzag_indices[0]
  kv_len_prev == request_kv_len_prev[0]
  actual_seq_q_prev == request_actual_seq_q_prev[0]
```

For `batch_size > 1`, existing scalar fields are either derived from request 0 only for legacy readers that are gated off, or kept invalid and fail-fast when used. Prefer explicit new per-request fields over overloading scalar fields.

## 5. Phase plan

### Phase 0: hard guards and red tests

Goal: prevent silent incorrect bs>1 behavior while tests document the current breakpoints.

Files:

- `test/registered/unit/layers/test_nsa_cp_utils.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `python/sglang/srt/mem_cache/common.py`

Work:

1. Add tests proving CP shared-KV bs>1 does not take token-balanced split or legacy owner allocation silently.
2. Add a single helper such as `cp_shared_kv_batching_supported(forward_batch)` only if it reduces duplicated guards.
3. Fail-fast in unsupported bs>1 CP shared-KV paths until the corresponding phase turns that path on.

Exit criteria:

- A bs>1 CP shared-KV batch either uses the new phase path or raises an explicit `[CP_SHARED_KV_FAIL_FAST][batch_gt1_*]` reason.
- No production fast path silently returns `None` with only debug logging for bs>1.

### Phase 1: batch-aware page-aligned CP split metadata

Goal: build per-request in-seq CP split metadata while preserving request boundaries.

Files:

- `python/sglang/srt/layers/attention/nsa/utils.py`
- `test/registered/unit/layers/test_nsa_cp_utils.py`

Work:

1. Add a pure CPU helper, e.g. `build_batch_page_aligned_in_seq_split_plan(extend_lens, prefix_lens, page_size, cp_size, cp_rank)`.
2. For each request, call existing `build_page_aligned_in_seq_split_list()` with that request's valid extend length and prefix length.
3. Produce request-local and flattened views:
   - per-request split list;
   - per-request page starts/ends;
   - per-rank local valid token counts;
   - flattened segment offsets for splitting tensors.
4. Update `prepare_input_dp_with_cp_dsa()` to create the batch-aware metadata when CP shared KV is enabled.
5. Keep old scalar fields valid only for `batch_size == 1`; for `batch_size > 1`, any old scalar-only consumer should fail-fast until updated.

Tests:

- `extend_lens=[100, 513]`, `prefix_lens=[40320, 8192]`, `page_size=64`, `cp_size=8`:
  - first request pads to 2 pages / 128 physical rows;
  - second request pads to 9 pages / 576 physical rows;
  - no request pads to `cp_size` or `2*cp_size` pages.
- Mixed short suffixes produce zero-length segments without token-balanced fallback.
- Non-page-aligned prefix either is floored before this helper or raises the existing CP shared-KV fail-fast.

Exit criteria:

- Metadata can describe `batch_size=2` without treating the two requests as one long sequence.

### Phase 2: batch-aware owner-lane allocation

Goal: allocate newly written pages from the CP rank that computes each page, for every request in the batch.

Files:

- `python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py`
- `python/sglang/srt/mem_cache/common.py`
- `python/sglang/srt/mem_cache/allocator.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
- `test/registered/unit/mem_cache/test_alloc_pages_with_owners.py`

Work:

1. Add a batched owner builder that calls `build_in_seq_page_compute_owners()` per request and concatenates owner lists in flattened request order.
2. Add `alloc_extend_compute_owner_batched()` or generalize `alloc_extend_compute_owner()` to accept multi-request `prefix_lens_cpu` and `seq_lens_cpu`.
3. Verify `alloc_extend_naive()` fills `out_indices` in flattened request order for multiple sequences; if not, wrap it with explicit per-request offsets.
4. Update `alloc_paged_token_slots_extend()` to use owner-lane allocation for CP shared-KV bs>1 and remove the `multi_batch` legacy fallback for supported cases.

Tests:

- Two requests with different prefix lengths produce a concatenated owner list whose length equals total newly allocated physical pages.
- `out_cache_loc` page ids have owner modulo equal to the concatenated plan.
- Owner-lane exhaustion raises `KVCapacityWaitError`, not legacy fallback.

Exit criteria:

- A target-model bs>1 extend allocates owner-lane pages for all new request pages.

### Phase 3: batch-aware CP split/rebuild and direct write

Goal: every tensor that is currently CP-split with one `split_list` can be split per request and concatenated in rank-local compute order.

Files:

- `python/sglang/srt/layers/attention/nsa/utils.py`
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- `test/registered/unit/layers/test_nsa_cp_utils.py`

Work:

1. Add `cp_split_and_rebuild_batched_1d()` and `cp_split_and_rebuild_batched_data()` using `request_token_offsets` plus per-request split lists.
2. Make `cp_split_and_rebuild_1d/data()` dispatch to batched helpers when metadata has `batch_size > 1`.
3. Update `get_cp_shared_kv_local_out_cache_loc()` to compute local locs per request and validate owner lanes after concatenation.
4. Ensure `cp_local_physical_out_cache_loc` cache key includes a batch metadata/version identifier if needed.

Tests:

- For each CP rank, local `out_cache_loc` equals the concatenation of that rank's two zigzag segments per request.
- Direct write refuses owner mismatch for either request.
- Single-request behavior remains byte-for-byte identical.

Exit criteria:

- MLA KV and index direct-write can use batch-aware local locs for target model.

### Phase 4: target index/top-k sync correctness

Goal: make target model top-k and index buffer materialization correct for bs>1 without async prefetch.

Files:

- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
- `test/registered/unit/layers/test_nsa_topk_transform.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`

Work:

1. Replace `_get_topk_in_seq_cp_pair()` scalar `kv_len_prev/next` logic with a loop over request-local prev/next slices.
2. Reassemble top-k results in flattened local query order.
3. For index materialize, build per-request prefix page ranges and one flattened logical page table that preserves request page order.
4. Keep async index prefetch disabled for bs>1 in this phase; sync materialize must be correct first.

Tests:

- Compare batched top-k result against sequential bs=1 execution for two requests.
- One request with no prefix and one request with page-aligned prefix does not corrupt the other's top-k range.

Exit criteria:

- Target index/top-k has a CPU/unit equality proof against sequential execution.

### Phase 5: target MLA current/partial-current reuse

Goal: restore current-only and partial-current fast paths for bs>1 target MLA after sync correctness.

Files:

- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
- `python/sglang/srt/layers/attention/nsa_backend.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`

Work:

1. Replace `batch_size_not_one` in current reuse with batched validation:
   - `out_cache_loc` covers the flattened current suffix;
   - per-request current rows are known;
   - current loc slices are page-aware and mask padded tail rows.
2. Add batched current-loc remap builder that emits one remap over all request current pages.
3. Extend partial-current compose to accept per-request prefix pages and current slices.
4. Keep fail-fast for mixed unsupported shapes rather than falling back silently.

Tests:

- Current-only bs=2 reuses current KV without materializing shared prefix.
- Partial-current bs=2 materializes prefix pages and appends only valid current rows, masking padded tails.
- One short suffix under `cp_size` pages remains correct.

Exit criteria:

- Target MLA uses current/partial-current reuse in supported bs>1 cases.

### Phase 6: HiCache/radix validation under batched writes and loads

Goal: prove per-request backed nodes retain correct owner metadata and can be loaded back in batched contexts.

Files:

- `python/sglang/srt/mem_cache/hiradix_cache.py`
- `python/sglang/srt/managers/cache_controller.py`
- `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
- `test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py`

Work:

1. Confirm `prepare_write_backup_for_req(req)` records `page_owners`, `owned_positions`, `valid_len`, and `padded_len` per request under the batched allocation contract.
2. Add tests for two nodes loaded in one `load_cp()` call where each node has different owner patterns.
3. Verify chunked prefill tail nodes either floor to page boundary before split or defer split without mutating in-flight backup nodes.
4. Ensure batched load returns visible device indices in request/node order expected by radix.

Tests:

- Two-node host hit load-back produces owner-matched logical pages on every CP rank.
- In-flight backup split remains deferred and does not crash under chunked prefill.

Exit criteria:

- HiCache backed target nodes are safe for bs>1 target path.

### Phase 7: EAGLE/draft bs>1 support

Goal: make draft KV and hidden-state collection mirror target request boundaries.

Files:

- `python/sglang/srt/models/deepseek_nextn.py`
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `test/registered/unit/layers/test_nsa_cp_utils.py`

Work:

1. Make draft local input/position/spec-hidden splitting use the batch-aware CP split helpers.
2. Implement per-request in-seq last-token owner and offset collection.
3. Keep target/draft KV mirror semantics: draft does not invent its own owner plan.
4. Only after target bs>1 passes ETE, enable draft local path for bs>1.

Tests:

- `cp_collect_last_token_hidden()` returns one last hidden row per request for bs=2 across CP ranks.
- Draft local path output matches full-path output on small deterministic tensors.

Exit criteria:

- EAGLE accept length does not collapse under bs>1 cache-hit traffic.

### Phase 8: batched L1 shared-KV prefetch

Goal: reintroduce async L1 prefix prefetch for bs>1 after sync paths are correct.

Files:

- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`

Work:

1. Replace single-prefix checks in MLA and index prefetcher `maybe_create()` with per-request prefix page descriptors.
2. Build one deterministic flattened prefetch/materialize plan ordered by request id and page id.
3. Keep materialize kernels off the wrong stream when they consume SM; preserve the earlier decision that prefix materialize should not race current-layer compute unless measured.
4. Use existing collectives in the same order on every CP rank. Do not add new collective kinds for bs>1 planning.

Tests:

- Prefetcher is created for bs=2 when both prefixes are page-aligned and above threshold.
- Consume returns the same dense buffer as sync materialize.
- Prefetch disabled for tiny prefixes does not disable current reuse.

Exit criteria:

- bs>1 cache-hit path has no sync prefix materialize hot-path fallback under supported sizes.

### Phase 9: remote CUDA/ETE validation and performance gates

Goal: prove correctness and decide production enablement thresholds.

Remote-only checks:

- Run on `g0034` container path `/sgl-workspace/sglang-tai`; do not run CUDA locally.
- Use current production flags: CP shared KV, EAGLE, HiCache, `direct + page_first_direct` unless a phase explicitly tests another layout.

Scenarios:

1. bs=2/4 target-only smoke with cache miss and cache hit.
2. bs=2/4 with EAGLE enabled.
3. Prefix hit 90%+, extend sizes `200, 512, 1024, 2000`.
4. Mixed prefix lengths and one zero-prefix request.
5. HiCache host hit after L1 eviction.
6. Chunked prefill enabled and disabled.
7. FP8 KV cache enabled and disabled.

Required evidence:

- No `[CP_SHARED_KV_FALLBACK]` or `[CP_SHARED_KV_FAIL_FAST]` in supported cases.
- `accept len` does not collapse to 1 compared with baseline.
- No output length 0 caused by server/runtime path.
- Cache hit remains non-zero and radix entries remain page-aligned.
- Throughput improves versus sequential bs=1 for short extends, or the phase is not production-enabled.

## 6. Recommended first implementation slice

Start with **Phase 0 -> Phase 2 only**, target model only:

```text
batch-aware metadata + owner-lane allocation
no draft local path
no async prefetch
no new collectives
```

Reason: if owner-lane allocation is wrong, every later optimization can faithfully preserve corrupted cache state. Phase 0-2 gives the smallest correctness checkpoint: a bs>1 batch can be planned and allocated page-correctly before attention, top-k, draft, or prefetch are touched.

## 7. Risks to keep visible

1. **Tiny suffixes:** a request with fewer than `cp_size` pages creates zero-token segments. This is legal only if every downstream split helper accepts zero segments.
2. **Mixed prefix alignment:** radix/HiCache should floor hits to a page boundary before CP split. Non-page-aligned prefixes should fail fast until proven safe.
3. **Chunked prefill:** an unfinished request can try to split an in-flight backup node. Existing policy is defer/suspend split, not mutate the tree.
4. **FP8 KV cache:** page-slot compose must preserve page shape; compact token buffers caused shape failures previously.
5. **Workspace pressure:** bs>1 materialize can allocate larger dense buffers. Phase 9 must measure memory, not just latency.
6. **Collective order:** every CP rank must execute materialize/reduce collectives in identical layer/request order.
7. **Fallback noise versus visibility:** unsupported bs>1 should be fail-fast during development; production fallback, if allowed later, must be warning-level and rate-limited.

## 8. Immediate next actions

1. Add Phase 0 red tests for silent token-balanced split and legacy owner allocation under bs>1 CP shared KV.
2. Add pure CPU batch split-plan helper tests before changing runtime flow.
3. Add batched owner-plan tests that verify page-owner order equals flattened `out_cache_loc` page order.
4. Only then update `prepare_input_dp_with_cp_dsa()` and allocator runtime paths.
