chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Plan batch-aware CP shared-KV execution

Browse Source 
The existing CP shared-KV notes still reflected the earlier multi-slot-first conclusion, while current short cache-hit extends require true ForwardBatch>1 planning. This commit records the batch-aware contract, sequential phase plan, and parallel workstream split so implementation can proceed without rediscovering single-request blockers or accidentally removing guards into unsafe fallback paths.

Constraint: Current NSA in-seq CP shared-KV paths assume batch_size == 1 across metadata, owner-lane allocation, direct write, current reuse, prefetch, top-k, and EAGLE/draft collection.

Rejected: Remove batch_size guards directly | would treat multiple requests as one long sequence and corrupt request boundaries and page-owner state.

Rejected: Keep only the old multi-slot recommendation | online extend sizes around 200-2000 tokens make true batching a required throughput direction.

Confidence: high

Scope-risk: narrow

Directive: Implement bs>1 by per-request page-aligned planning before flattening; do not add planner collectives or silent legacy fallbacks.

Tested: Placeholder scan and referenced-path existence check for all four docs; git diff --cached --check.

Not-tested: Runtime behavior; documentation-only change.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-03 00:37:54 +08:00
parent 937f89ef89
commit 0158e28689
4 changed files with 1784 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md
View File

@@ -2,6 +2,13 @@
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.
@@ -175,7 +182,7 @@ Relevant code:
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- `_get_topk_in_seq_cp_pair(...)`
- TODO: `support mutil-batch`
- 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`.

476
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md Normal file
View File

@@ -0,0 +1,476 @@
# 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 is single-request
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:1764-1904`
- `_current_extend_kv_reuse_miss_reason()` returns `batch_size_not_one`.
- `current_extend_kv_rows_for_reuse()` requires one extend length.
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:309-380`
- Index partial-current compose requires one positive page-aligned prefix.
- `python/sglang/srt/layers/attention/nsa_backend.py:1940-2020`
- MLA partial-current compose has the same one-prefix fail-fast contract.
Implication: Phase 4/5 must first keep sync full materialize correct, then add batched current/partial-current reuse.
### 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.

646
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md Normal file
View File

@@ -0,0 +1,646 @@
# NSA Prefill CP Shared-KV Batch Size > 1 中文实现计划
> **给后续 agent/工程师的执行要求：** 实现本文档时必须先使用 `superpowers:subagent-driven-development` 或 `superpowers:executing-plans`，按 phase 和测试逐步推进。本文是实现合同与 phase 规划，不是已经完成的代码实现。
**目标：** 在 `--nsa-prefill-cp-mode in-seq-split` + CP shared KV + HiCache 场景下支持真正的 `ForwardBatch(batch_size > 1)`，让线上 cache-hit 后 `200-2000` token 的短 extend 可以通过 batch 填满 compute，同时不破坏 owner-lane、page-aligned cache、direct write、HiCache load/back up、EAGLE/draft 的正确性。
**核心架构：** 所有 CP 规划先以 request 为单位计算，再 flatten 到现有 tensor/kernel 接口。cache 管理仍以 page 为最小单位；每个 request 独立拥有 prefix、extend、page extent、split list、page owner、top-k 范围和 current reuse 区间。规划只能依赖本地 CPU metadata 和统一配置，不允许新增 collective 来“确认计划是否一致”。
**当前约束：** 先保证 target model 同步路径正确，再恢复 current/partial-current reuse，再处理 EAGLE/draft，最后打开 MLA/index L1 prefetch。不能通过删除 `batch_size != 1` guard 来“强行开启”。
- 并行派工版计划：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md`。
---
## 1. 背景与优先级变化
旧文档 `docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md` 当时建议先做 multi-slot overlap，因为那时认为 true batch 的改动面太大，且每个 extend 较长时单请求还能提供足够 compute。
现在观测到线上流量有明显变化：cache hit 很高时，extend 经常只有 `200-2000` tokens。这个量级下，即使 L2->L1 prefetch、L1 materialize、async backup 做得比较好，单请求本身也可能打不满 GPU compute。因此 true `ForwardBatch(batch_size > 1)` 变成必要吞吐优化。
但当前代码里 NSA in-seq CP shared-KV 的多个关键路径仍是 `batch_size == 1` 合同。如果直接把 guard 去掉，flatten 后的多个 request 会被当成一条长序列，导致：
- request boundary 丢失；
- causal/top-k 范围错误；
- page owner 与 compute rank 不一致；
- direct write 写入非 owner lane；
- partial-current reuse 把不同 request 的 current suffix 混在一起；
- EAGLE/draft last-token collect 取错 request 的最后 token；
- HiCache 会正确保存一个已经错误的 owner pattern。
所以 bs>1 不能做成“放开 guard”，必须做成“每个 request 独立规划，最后 flatten”。
---
## 2. 不可破坏的合同
### 2.1 request boundary 必须保留
每个 request 都必须独立记录：
- `extend_prefix_len`；
- `extend_len`；
- page-aligned extent；
- `split_list`；
- `zigzag_index`；
- `kv_len_prev/next`；
- `actual_seq_q_prev/next`；
- current suffix rows；
- prefix pages；
- page owner list。
Flatten 后的 tensor 可以共用一个 buffer，但 metadata 不能退化成单个全局标量。
### 2.2 page 是 cache 最小单位
CP shared KV + HiCache 已经被修正为 page-aligned cache 合同：
```text
valid tokens <= physical page-rounded tokens
padding <= page_size - 1
```
举例：
```text
extend_len = 100
page_size = 64
physical coverage = 128 tokens = 2 pages
```
不能因为 `cp_size=8` 就补到 8 pages 或 16 pages，也不能因短 tail fallback 到 token-balanced split。
### 2.3 owner-lane 计划必须与 flattened out_cache_loc page 顺序一致
每个 request 独立生成 page owners，然后按 batch request 顺序拼接：
```text
batch_page_owners = owners(req0) + owners(req1) + ... + owners(reqN)
```
`out_cache_loc` 的 page 顺序必须与这个 owner list 一致。后续 direct write、HiCache metadata、load back 都依赖这个顺序。
### 2.4 规划不允许新增 collective
所有 CP ranks 必须根据相同的 CPU metadata 和配置独立得出同一个 plan。不能通过新增 all-reduce/all-gather 来同步“这个 batch 的 planning 结果”。
允许后续 materialize/reduce 使用已有 collective，但 collective 顺序必须在所有 rank 上完全一致。
### 2.5 fallback 必须显式
支持范围内不允许 silent fallback。开发阶段建议 fail-fast：
```text
[CP_SHARED_KV_FAIL_FAST][batch_gt1_*]
```
生产阶段如果后续允许 fallback，也必须 warning 级别、限频、带明确原因和性能影响。
### 2.6 target / draft / prefetch 的启用顺序
推荐启用顺序：
```text
target sync correctness
-> target current/partial-current reuse
-> HiCache write/load validation
-> EAGLE/draft mirror
-> batched MLA/index prefetch
```
不要在 target sync 路径没有证明正确前启用 draft 或 async prefetch。
---
## 3. 已确认的代码事实
这一节是检查记录，避免后续反复搜索同一问题。
### C1. generic scheduler 已经能构造 multi-request extend
相关代码：
- `python/sglang/srt/managers/schedule_batch.py:1481-1540`
- `ScheduleBatch.prepare_for_extend()` 会构造 flattened `input_ids`；
- 同时保留 per-request `seq_lens`、`prefix_lens`、`extend_lens`、`extend_num_tokens`。
- `python/sglang/srt/model_executor/forward_batch_info.py:449-500`
- `ForwardBatch.init_new()` 设置 `batch_size=len(batch.seq_lens)`。
结论：bs>1 的 blocker 不在 generic batching，而在 NSA CP shared-KV 专用路径。
### C2. in-seq CP metadata 是单请求结构
相关代码：
- `python/sglang/srt/layers/attention/nsa/utils.py:1094-1244`
- `prepare_input_dp_with_cp_dsa()` 注释为 `# just support batch = 1`；
- `bs_per_cp_group = 1`；
- 只生成一个 `split_list`、一个 `zigzag_index`、一组 `kv_len_prev/next`。
- `python/sglang/srt/layers/attention/nsa/utils.py:203-244`
- `NSAContextParallelMetadata` 当前字段基本是 scalar / single-request。
结论：Phase 1 必须扩展 per-request metadata。
### C3. page-aligned split 只在 `len(...) == 1` 时启用
相关代码：
- `python/sglang/srt/layers/attention/nsa/utils.py:300-390`
- `_build_in_seq_split_for_forward_batch()` 只有在 `len(extend_seq_lens_cpu) == 1` 且 `len(extend_prefix_lens_cpu) == 1` 时使用 page-aligned split；
- 否则 fallback 到 token-balanced split。
结论：bs>1 下 token-balanced fallback 会破坏 CP shared KV page owner 合同，必须替换为 batched page-aligned split。
### C4. `can_cp_split()` 对 bs>1 的判断可能错误
相关代码：
- `python/sglang/srt/layers/attention/nsa/utils.py:460-483`
- `can_cp_split(seq_len, ...)` 使用 flattened total tokens；
- `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` 只支持单请求判断。
风险：多个很短 suffix 合并后 total tokens 可能足够触发 CP split，但每个 request 单独看都小于 `cp_size` pages。这会制造大量 zero segment / 不匹配路径。
### C5. owner-lane allocation 是单请求结构
相关代码：
- `python/sglang/srt/mem_cache/common.py:455-565`
- 只有 `len(prefix_lens_cpu) == 1` 时才生成 `page_compute_owners`；
- bs>1 记录 `multi_batch`，然后 fallback 到 legacy allocation。
- `python/sglang/srt/mem_cache/allocator.py:981-1048`
- `alloc_extend_compute_owner()` 对 batch size 不是 1 直接 raise。
结论：Phase 2 必须实现 batched owner-lane allocation。legacy allocation 对 CP shared KV direct write 是 correctness risk。
### C6. direct write local loc 是单请求 split
相关代码：
- `python/sglang/srt/layers/attention/nsa/utils.py:562-648`
- `get_cp_shared_kv_local_out_cache_loc()` 使用单个 `metadata.split_list` split 整个 flattened `out_cache_loc`；
- 然后检查本 rank local loc 是否属于当前 owner lane。
- 调用者：
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`；
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`。
结论：Phase 3 必须以 request 为单位 split `out_cache_loc`，再拼接本 rank local loc。
### C7. current / partial-current reuse 是单请求结构
相关代码：
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:1764-1904`
- `_current_extend_kv_reuse_miss_reason()` 返回 `batch_size_not_one`；
- `current_extend_kv_rows_for_reuse()` 只接受一个 extend length。
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:309-380`
- index partial-current compose 要求单个 positive page-aligned prefix。
- `python/sglang/srt/layers/attention/nsa_backend.py:1940-2020`
- MLA partial-current compose 同样是单 prefix 合同。
结论：Phase 5 需要做 batched current suffix slicing 和 remap。Phase 4 可先只做 sync correctness。
### C8. L1 shared-KV prefetch 是单请求结构
相关代码：
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:360-455`
- `CpSharedKVMlaPrefetcher.maybe_create()` 拒绝 `batch_size != 1`。
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:1135-1235`
- `CpSharedKVIndexPrefetcher.maybe_create()` 拒绝 `batch_size != 1`。
结论：Phase 8 再恢复 batched prefetch。不能在 sync path 未正确时先开 prefetch。
### C9. top-k in-seq CP pair 是单请求结构
相关代码：
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:1371-1410`
- `_get_topk_in_seq_cp_pair()` 使用 scalar `kv_len_prev/next` 和 scalar `actual_seq_q_prev/next`；
- 源码里记录 multi-batch 尚未支持。
结论：Phase 4 必须改成 request-local prev/next slices。
### C10. EAGLE / draft local path 对单请求 metadata 敏感
相关代码：
- `python/sglang/srt/models/deepseek_nextn.py:292-399`
- draft path 使用 `cp_split_and_rebuild_1d`、`cp_split_and_rebuild_position`、`cp_split_and_rebuild_data`、`cp_collect_last_token_hidden()`。
- `python/sglang/srt/layers/attention/nsa/utils.py:1333-1358`
- exact last-token owner/offset 只对 `bs == 1` 生效。
结论：Phase 7 必须修复 draft per-request split 和 last-token collect。target 未过前不要打开 draft bs>1 local path。
### C11. HiCache load 较接近 batch-capable，但依赖 write-time metadata
相关代码：
- `python/sglang/srt/managers/cache_controller.py:1338-1498`
- `load_cp()` 可以拼接多个 node 的 `page_owners`，然后 `alloc_pages_with_owners()`。
- `python/sglang/srt/managers/scheduler.py:2630-2655`
- `_prepare_hicache_write_backups_before_forward()` 已经按 request 循环调用 `prepare_write_backup_for_req(req)`。
结论：HiCache 不是第一 blocker。但如果 write-time batch metadata 错了，HiCache 会复现错误 owner pattern，因此 Phase 6 必须验证。
---
## 4. 目标 metadata 设计
建议扩展 `NSAContextParallelMetadata`，保留单请求兼容字段，同时新增 batch-aware 字段。
```text
NSAContextParallelMetadata
# 旧字段：batch_size == 1 时继续保持兼容
split_list
zigzag_index
kv_len_prev
kv_len_next
actual_seq_q_prev
actual_seq_q_next
# 新字段：batch-aware
batch_size: int
request_token_offsets: List[int]
request_padded_token_offsets: List[int]
request_page_offsets: List[int]
request_split_infos: List[PageAlignedInSeqSplitInfo]
request_split_lists: List[List[int]]
request_zigzag_indices: List[List[int]]
request_rank_local_tokens: List[int]
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]
flat_zigzag_index: List[int]
```
兼容规则：
```text
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]
batch_size > 1:
旧 scalar-only consumer 必须 fail-fast，直到对应 phase 改成 per-request consumer
```
关键点：不要把 batch 的 flattened tokens 当作一条 request 来生成单个 split list。
---
## 5. Phase 规划
### Phase 0：硬 guard 与 red tests
目标：在正式支持 bs>1 前，先防止 silent incorrect fallback。
修改文件：
- `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`
工作内容：
1. 添加 red tests，证明当前 CP shared-KV bs>1 不允许 silent token-balanced split。
2. 添加 red tests，证明当前 owner-lane allocation 不允许 silent legacy allocation。
3. 对不支持的 bs>1 CP shared-KV path 加明确 fail-fast。
4. 如果多个文件重复判断，可增加一个轻量 helper，例如 `cp_shared_kv_batching_supported(forward_batch)`。
退出标准：
- bs>1 CP shared-KV batch 要么进入新 phase path，要么抛出明确 `[CP_SHARED_KV_FAIL_FAST][batch_gt1_*]`。
- 不允许只打 debug log 后返回 `None` 进入慢 fallback。
### Phase 1：batch-aware page-aligned CP split metadata
目标：每个 request 独立生成 page-aligned in-seq CP split metadata。
修改文件：
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `test/registered/unit/layers/test_nsa_cp_utils.py`
工作内容：
1. 增加纯 CPU helper，例如：
```text
build_batch_page_aligned_in_seq_split_plan(
extend_lens,
prefix_lens,
page_size,
cp_size,
cp_rank,
)
```
2. 对每个 request 调用已有 `build_page_aligned_in_seq_split_list()`。
3. 生成 per-request 与 flat 两套视图：
- per-request split list；
- per-request page start/end；
- per-rank local valid token count；
- flattened segment offsets。
4. 更新 `prepare_input_dp_with_cp_dsa()`，让 CP shared-KV bs>1 使用新 metadata。
5. 保留 bs=1 的旧字段兼容。
测试重点：
```text
extend_lens=[100, 513]
prefix_lens=[40320, 8192]
page_size=64
cp_size=8
```
期望：
- req0 padded 到 2 pages / 128 tokens；
- req1 padded 到 9 pages / 576 tokens；
- 不补到 `cp_size` 或 `2 * cp_size` pages；
- zero segment 合法；
- non-page-aligned prefix 在进入 CP split 前 floor，否则 fail-fast。
退出标准：
- metadata 能表达 `batch_size=2`，且不会把两个 request 当成一条长序列。
### Phase 2：batch-aware owner-lane allocation
目标：每个 request 的新 page 都从对应 compute-owner lane 分配。
修改文件：
- `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`
工作内容：
1. 新增 batched owner builder：对每个 request 调用 `build_in_seq_page_compute_owners()`，再按 request 顺序拼接。
2. 新增 `alloc_extend_compute_owner_batched()`，或泛化 `alloc_extend_compute_owner()` 支持多请求。
3. 验证 `alloc_extend_naive()` 对 multi-request 是否按 flattened request order 填充 `out_indices`；如果不能保证，则封装 per-request offset 写入。
4. 更新 `alloc_paged_token_slots_extend()`，让 CP shared-KV bs>1 走 owner-lane allocation，不再走 `multi_batch` legacy fallback。
测试重点：
- 两个 request prefix/extend 不同，拼接 owner list 长度等于新分配 page 数。
- `out_cache_loc` 的 page owner 与拼接 owner list 一致。
- owner-lane 不足时抛 `KVCapacityWaitError`，不 fallback 到 legacy allocation。
退出标准：
- target bs>1 extend 可以正确分配 owner-lane pages。
### Phase 3：batch-aware CP split/rebuild 与 direct write
目标：当前所有用单个 `split_list` split tensor 的路径改成 per-request split，然后拼接本 rank local rows。
修改文件：
- `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`
工作内容：
1. 添加 `cp_split_and_rebuild_batched_1d()`。
2. 添加 `cp_split_and_rebuild_batched_data()`。
3. 让 `cp_split_and_rebuild_1d/data()` 在 metadata `batch_size > 1` 时 dispatch 到 batched helper。
4. 更新 `get_cp_shared_kv_local_out_cache_loc()`：
- 按 request 切 `out_cache_loc`；
- 每个 request 内按其 split/zigzag 取本 rank local loc；
- 拼接后再做 owner-lane validation。
5. 确保 `cp_local_physical_out_cache_loc` cache 不会误复用旧 batch metadata。
测试重点：
- 每个 CP rank 的 local `out_cache_loc` 等于 req0 local + req1 local。
- 任一 request 出现 owner mismatch 都拒绝 direct write。
- bs=1 行为保持一致。
退出标准：
- target MLA KV 和 index direct-write 可以使用 batched local loc。
### Phase 4：target index/top-k 同步正确性
目标：不依赖 async prefetch，先让 target index/top-k 在 bs>1 下正确。
修改文件：
- `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`
工作内容：
1. 把 `_get_topk_in_seq_cp_pair()` 从 scalar `kv_len_prev/next` 改成 per-request loop。
2. 每个 request 生成自己的 prev/next top-k pair。
3. 按 flattened local query order 重新拼接 top-k result。
4. index materialize 先走同步路径：构造 per-request prefix page range，再 flatten。
5. Phase 4 不打开 bs>1 index prefetch。
测试重点：
- bs=2 top-k 结果与两个 sequential bs=1 请求一致。
- 一个 request 无 prefix、另一个 request 有 page-aligned prefix 时不互相污染。
退出标准：
- target index/top-k 有 CPU/unit 等价证明。
### Phase 5：target MLA current / partial-current reuse
目标：恢复 bs>1 下 current-only 和 partial-current fast path。
修改文件：
- `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`
工作内容：
1. 移除 `batch_size_not_one` 的硬拒绝，改成 batched validation。
2. 校验 flattened `out_cache_loc` 覆盖每个 request 的 current suffix。
3. 为每个 request 记录 current rows 与 current loc slice。
4. 构造 batched current-loc remap。
5. partial-current compose 支持 per-request prefix pages + current slices。
6. padded tail rows 必须 mask，不可暴露给 attention。
测试重点：
- current-only bs=2 不 materialize shared prefix。
- partial-current bs=2 正确拼接 prefix materialized rows 与 valid current rows。
- 短 suffix 小于 `cp_size` pages 仍正确。
退出标准：
- target MLA 在支持 case 下使用 current/partial-current reuse。
### Phase 6：HiCache / radix batched write-load 验证
目标：证明 batched write 后的 HiCache metadata 能被正确 load back。
修改文件：
- `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`
工作内容：
1. 检查 `prepare_write_backup_for_req(req)` 在 batched allocation 下记录：
- `page_owners`；
- `owned_positions`；
- `valid_len`；
- `padded_len`。
2. 添加两个 node 同时 `load_cp()` 的测试。
3. 验证 chunked prefill tail page：
- 要么提前 floor 到 page boundary；
- 要么遇到 in-flight backup split 时 defer，不改树。
4. 确认 `load_cp()` 返回的 visible device indices 顺序符合 radix/request 顺序。
测试重点：
- 两个 host hit node 一次 load，所有 CP rank 上 owner pattern 都匹配。
- chunked prefill 不因 in-flight backup split 崩溃。
退出标准：
- HiCache target nodes 对 bs>1 安全。
### Phase 7：EAGLE / draft bs>1 支持
目标：draft KV 与 target KV 保持镜像合同，hidden state collection 按 request 独立处理。
修改文件：
- `python/sglang/srt/models/deepseek_nextn.py`
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `test/registered/unit/layers/test_nsa_cp_utils.py`
工作内容：
1. draft local input/position/spec-hidden 使用 batch-aware split helper。
2. 实现 per-request in-seq last-token owner/offset collect。
3. draft 不生成自己的 owner plan，一切跟随 target。
4. target bs>1 ETE 通过后，再打开 draft local path。
测试重点：
- `cp_collect_last_token_hidden()` 对 bs=2 返回两个 request 的 last hidden。
- draft local path 与 full path 在小 deterministic tensor 上一致。
退出标准：
- EAGLE accept length 不因 bs>1 cache-hit 流量掉到 1。
### Phase 8：batched L1 shared-KV prefetch
目标：在 sync correctness 已证明后，恢复 bs>1 的 async L1 prefix prefetch。
修改文件：
- `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`
工作内容：
1. `CpSharedKVMlaPrefetcher.maybe_create()` 支持 per-request prefix page descriptors。
2. `CpSharedKVIndexPrefetcher.maybe_create()` 支持 per-request prefix page descriptors。
3. 构造 deterministic flattened prefetch plan：request id 优先，再 page id。
4. 保持 materialize kernel 的 stream 策略与前面已有结论一致：占 SM 的 materialize 不应错误并发。
5. 不新增 planner collective；已有 reduce/materialize collective 顺序必须所有 rank 一致。
测试重点：
- bs=2 且两个 prefix 都 page-aligned、超过阈值时 prefetcher 被创建。
- consume 结果与 sync materialize 相同。
- tiny prefix 关闭 prefetch 时，不影响 current reuse / full current reuse。
退出标准：
- 支持范围内 bs>1 cache-hit path 不再出现 sync prefix materialize hot-path fallback。
### Phase 9：远端 CUDA / ETE / 性能验证
目标：确认 correctness 与是否值得生产启用。
只在远端测试：
- host：`g0034`；
- container：`sglang-glm5-dev-2`；
- container path：`/sgl-workspace/sglang-tai`；
- 不在本地测试 CUDA。
场景：
1. bs=2/4 target-only cache miss / cache hit。
2. bs=2/4 EAGLE enabled。
3. prefix hit 90%+，extend sizes：`200, 512, 1024, 2000`。
4. mixed prefix lengths。
5. 一个 request zero prefix，另一个 request host hit。
6. HiCache L1 eviction 后从 L2 host load back。
7. chunked prefill enabled / disabled。
8. FP8 KV cache enabled / disabled。
必须收集的证据：
- 支持范围内无 `[CP_SHARED_KV_FALLBACK]` / `[CP_SHARED_KV_FAIL_FAST]`。
- `accept len` 不再异常掉到 1。
- 无 output length 0 的 runtime/server path 问题。
- cache hit 非 0，radix entries 保持 page-aligned。
- 短 extend 下吞吐高于 sequential bs=1，否则不 production-enable。
---
## 6. 推荐第一段实现范围
第一段只做：
```text
Phase 0 -> Phase 2
目标只覆盖 target model
不打开 draft local path
不打开 async prefetch
不新增 collective
```
原因：如果 owner-lane allocation 错了，后续 direct write、HiCache backup/load、prefetch 都会稳定复现错误状态。先把 metadata + owner-lane allocation 锁死，是最小但最关键的正确性边界。
---
## 7. 主要风险
1. **tiny suffix zero segment：** 小于 `cp_size` pages 的 request 会产生很多 0-token segment。所有 split/rebuild/top-k path 必须接受 zero segment。
2. **prefix 对齐：** radix/HiCache 命中应 floor 到最近 page boundary。non-page-aligned prefix 进入 CP split 前必须 fail-fast。
3. **chunked prefill：** unfinished request 可能尝试 split in-flight backup node。策略仍是 defer，不动树。
4. **FP8 KV cache：** 必须保持 page-slot compose，不能退化成 compact token buffer。
5. **workspace pressure：** bs>1 materialize buffer 更大，必须测显存和临时 buffer，不只看 latency。
6. **collective 顺序：** 所有 CP rank 必须按完全相同的 layer/request 顺序执行 collective。
7. **fallback 可见性：** 开发阶段优先 fail-fast；生产 fallback 必须 warning 且限频。
---
## 8. 立即执行项
1. 写 Phase 0 red tests：bs>1 CP shared-KV 不能 silent token-balanced split。
2. 写 Phase 0 red tests：bs>1 CP shared-KV 不能 silent legacy owner allocation。
3. 写 Phase 1 纯 CPU batch split-plan tests。
4. 写 Phase 2 batched owner-plan tests，验证 owner list 顺序等于 flattened `out_cache_loc` page 顺序。
5. 通过测试后再改 `prepare_input_dp_with_cp_dsa()` 与 allocator runtime path。

654
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md Normal file
View File

@@ -0,0 +1,654 @@
# NSA Prefill CP Shared-KV Batch Size > 1 并行实施计划
> **目标读者：** 这版文档用于把 bs>1 支持拆给多人/多 agent 并行做。它不是单线程 phase 文档，而是按可并行工作流拆分：CP metadata、allocator、fast path/runtime、TAI kernel、HiCache、EAGLE、ETE 验证。
**最终目标：** true `ForwardBatch(batch_size > 1)` 在 NSA in-seq CP shared-KV + HiCache 下进入 fast path，而不是为了支持 batch 退回慢路径。底层 TAI/SGL kernel、materialize、direct write、load/back up、current/partial-current reuse 都应能利用 bs>1 带来的更大批量和更少 launch/API overhead。
---
## 1. 最终形态
最终系统应该长这样：
```text
ScheduleBatch.prepare_for_extend()
-> 生成真实 ForwardBatch(batch_size=N)
CP Batch Planner
-> 每个 request 独立 page-aligned split
-> 每个 request 独立 page owner plan
-> 生成 per-request metadata + flattened descriptors
Batched Owner-Lane Allocator
-> 按 flattened page owner plan 分配 out_cache_loc
-> 每个 request 的新 page 都落到对应 compute owner lane
Target Forward Fast Path
-> cp_split/rebuild 按 request 切，再 flatten local rows
-> MLA KV direct write 使用 batched local loc
-> index direct write 使用 batched local loc
-> top-k/index/attention 使用 batched metadata，不混 request boundary
Current / Partial-Current Reuse
-> 每个 request 独立 prefix pages + current suffix rows
-> batched descriptor 一次构造，一次或少量 kernel/collective 消费
-> 不因为 bs>1 禁用 current reuse
HiCache / L2-L1 / Backup
-> load、prefetch、backup 使用同一套 batched page descriptors
-> per-layer async backup/load 可以一次处理多个 request
-> host/device page owner metadata 与 target plan 一致
EAGLE / Draft
-> draft KV 镜像 target KV plan
-> draft input/spec hidden/last token collect 按 request 独立处理
TAI/SGL Kernels
-> 支持 variable-length descriptor arrays
-> bs=1 是 bs>N 的退化情况
-> bf16/fp8 e4m3 都走同一类 fast path
```
核心原则：**不要为 bs>1 引入一套慢兼容路径。bs=1 和 bs>1 应该尽量共享同一个 batched fast path，bs=1 只是 descriptor 数量较少。**
---
## 2. 并行工作流总览
可以拆成 8 个并行 workstreams：
| Workstream | 负责人类型 | 主要产物 | 可否立即开始 | 依赖 |
| --- | --- | --- | --- | --- |
| W0 合同与红线测试 | runtime/test | fail-fast + red tests + fast-path contract | 可以 | 无 |
| W1 CP metadata / batch planner | runtime | batched CP plan 数据结构 | 可以 | W0 |
| W2 batched owner-lane allocation | mem_cache/runtime | batch owner plan + allocator | 可以 | W1 接口定义 |
| W3 CP split/rebuild + direct write | runtime/model | batched local rows + direct write fast path | 需 W1 | W1、W2 |
| W4 current / partial-current reuse | runtime/attention | batched current reuse fast path | 需 W1 | W1、W3 |
| W5 底层 TAI/SGL kernel 支持 | kernel/perf | variable-length batched kernels + benchmarks | 可以 | descriptor 草案来自 W1/W3/W4 |
| W6 HiCache / load / backup / prefetch pipeline | runtime/cache | batched load/backup/prefetch 接入 | 需 W1/W2 | W1、W2、W5 部分 |
| W7 EAGLE/draft + ETE/perf | model/test | draft bs>1 + accept len/output len 验证 | 需 target path | W1-W6 |
其中 W1/W2/W5 可以最早并行启动；W3/W4/W6 依赖 W1/W2 的 metadata/allocator 合同；W7 最后接入。
---
## 3. W0：合同与红线测试
### 目标
先把“不允许 silent fallback”的红线锁住，避免其他 workstream 改动时把 bs>1 误接到慢路径或错误路径。
### 修改范围
- `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`
### 交付物
1. bs>1 CP shared-KV 不能 silent token-balanced split。
2. bs>1 CP shared-KV 不能 silent legacy allocation。
3. direct write、current reuse、prefetch 遇到 unsupported bs>1 必须 fail-fast 或 warning 级显式 fallback。
4. 所有 fail-fast reason 统一命名：
```text
[CP_SHARED_KV_FAIL_FAST][batch_gt1_*]
[CP_SHARED_KV_FALLBACK][batch_gt1_*]
```
### 完成定义
- 当前不支持的 bs>1 case 不会悄悄跑旧路径。
- 后续每个 workstream 打开一个支持点时，都能删除对应 fail-fast 并补 fast-path 测试。
---
## 4. W1：CP metadata / batch planner 改造
### 目标
把当前单请求 `NSAContextParallelMetadata` 改造成 batch-aware metadata。所有 request 独立规划，最后生成 flattened view 给 runtime/kernel 消费。
### 修改范围
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `test/registered/unit/layers/test_nsa_cp_utils.py`
### 新数据结构建议
新增一个逻辑结构，可以是 dataclass，也可以先作为 `NSAContextParallelMetadata` 的字段：
```text
CPSharedKVBatchPlan
batch_size
page_size
cp_size
cp_rank
request_token_offsets
request_padded_token_offsets
request_page_offsets
request_extend_lens
request_prefix_lens
request_padded_pages
request_padded_tokens
request_split_lists
request_zigzag_indices
request_segment_page_starts
request_segment_page_ends
request_kv_len_prev
request_kv_len_next
request_actual_seq_q_prev
request_actual_seq_q_next
request_rank_local_tokens
flat_split_list
flat_zigzag_index
flat_segment_request_ids
flat_segment_offsets
```
### 关键规则
1. 每个 request 调用 page-aligned split helper。
2. 每个 request 的 tail 最多补 `page_size - 1` tokens。
3. 不补到 `cp_size` pages，不补到 `2 * cp_size` pages。
4. zero-token segment 合法。
5. batch flattened plan 不能跨 request 合并 segment。
6. bs=1 继续兼容现有 scalar fields。
### 对外接口
W1 应给其他 workstream 提供稳定接口：
```text
get_cp_shared_kv_batch_plan(forward_batch) -> CPSharedKVBatchPlan
split_tensor_by_cp_batch_plan(tensor, plan, mode=1d/data/position)
build_flat_page_owner_plan(plan) -> List[int]
```
接口名可以不同，但语义必须稳定。
### 测试
- `extend_lens=[100, 513]`，`page_size=64`：分别 padded 到 2 pages 与 9 pages。
- `batch_size=2` 不把两个 request 合成一条长序列。
- `prefix_lens` 不同也能生成独立 segment page ranges。
- short suffix 产生 zero segments，但 metadata 合法。
### 完成定义
- W2/W3/W4/W6 可以只依赖 batch plan，不需要自己重新推导 prefix/extend/page owner。
---
## 5. W2：allocation 支持 batch size
### 目标
让 CP shared-KV bs>1 也走 owner-lane allocation。不能 fallback 到 legacy allocation。
### 修改范围
- `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`
### 输入
来自 W1 的 per-request plan：
```text
request_extend_lens
request_prefix_lens
request_padded_pages
request_page_offsets
```
### 输出
```text
flat_page_compute_owners: List[int]
out_cache_loc: torch.Tensor # flattened request order
```
### 实现要求
1. 对每个 request 独立生成 page owners。
2. 按 request 顺序拼接 page owners。
3. allocator 选择的 logical page id 必须满足：
```text
(page_id - 1) % cp_size == expected_owner
```
4. `alloc_extend_naive()` 如果天然支持 multi-request flattened order，就复用；否则包一层 per-request allocation 写入。
5. owner-lane 不足时触发已有 eviction/capacity wait 机制，不走 legacy fallback。
### 测试
- 两个 request 不同 prefix/extend，验证 `out_cache_loc` page owner 与 flat owner list 一致。
- owner-lane exhausted 时抛 `KVCapacityWaitError`。
- bs=1 结果与当前实现一致。
### 完成定义
- CP shared-KV bs>1 进入 allocator 时不再出现 `multi_batch` fallback。
- W3 direct write 可以信任 `out_cache_loc` 的 owner-lane 合同。
---
## 6. W3：CP split/rebuild + direct write fast path
### 目标
让 MLA KV 和 index KV direct write 在 bs>1 下仍然走 fast path。这里是把 W1/W2 的 plan 真正接入 model/runtime 的关键。
### 修改范围
- `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`
### 实现要求
1. `cp_split_and_rebuild_1d/data/position` 支持 batch plan。
2. 每个 request 内仍按 in-seq zigzag 取本 rank 的两个 segment。
3. 多个 request 的 local rows 按 request 顺序拼接。
4. `get_cp_shared_kv_local_out_cache_loc()` 使用 batch plan，不使用单个 global split list。
5. direct write 前验证 local loc owner lane。
6. bs=1 继续复用同一套 batched helper，避免两套逻辑长期分叉。
### 测试
- 对 cp_rank=0..cp_size-1，验证 local loc 精确等于 per-request local rows 拼接。
- 任意 request owner mismatch，direct write fail-fast。
- MLA KV direct write 与 index direct write 都不 fallback。
### 完成定义
- target model bs>1 的 KV/index 写入可以使用 direct write fast path。
---
## 7. W4：partial reuse / current reuse 支持 bs>1
### 目标
bs>1 不应关闭 current reuse。cache-hit 短 extend 的收益来自“prefix 来自 cache，current suffix 直接复用本轮计算结果”，如果 bs>1 触发 partial/current reuse fallback，吞吐可能不会提升。
### 修改范围
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
- `python/sglang/srt/layers/attention/nsa_backend.py`
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
### 实现要求
1. `current_extend_kv_rows_for_reuse()` 改成返回 batched current slices，而不是单个 row count。
2. 每个 request 独立计算：
```text
prefix_pages
current_valid_rows
current_locs
current_kv rows
```
3. partial-current compose 使用 batch descriptors：
```text
materialized prefix pages + current suffix rows
```
4. padded tail rows 不可暴露给 attention。
5. current-only bs>1 也应进入 fast path。
6. index current reuse 与 MLA current reuse 的 descriptor 尽量复用。
### 与 W5 的接口
W4 应把 descriptor 定义交给 W5：
```text
BatchedCurrentReuseDescriptor
request_id
prefix_page_start
prefix_pages
current_loc_start
current_valid_rows
current_padded_rows
output_dense_offset
```
### 测试
- bs=2 current-only：不 materialize prefix。
- bs=2 partial-current：prefix pages + valid current rows 正确拼接。
- 一个 request tiny suffix，一个 request 长 suffix，互不污染。
### 完成定义
- `batch_size > 1` 不再因为 current reuse guard 走 sync slow path。
---
## 8. W5：底层 TAI/SGL kernel 支持
### 目标
底层操作必须能吃到 bs>1 的收益。bs>1 的意义不只是 Python 上合并请求，而是让 H2D/D2H/materialize/direct write 等底层操作能批量处理 descriptor，降低 CPU launch/API overhead，提高 PCIe/HBM/NVLink 利用率。
### 范围
- `tai-kernel` 中已有 kvcache / IPC / direct path kernel。
- SGLang runtime 中调用 TAI kernel 的 wrapper。
- benchmark 脚本。
### 必须支持的 kernel 形态
1. **variable-length descriptor arrays**
- 不因不同 batch/request length 反复 JIT。
- descriptor count、page count、row count 都是运行时参数。
2. **batched page-slot compose**
- 输入多个 request 的 prefix/current descriptors。
- 输出一个 dense buffer 或 page-slot buffer。
- 保持 page as minimum cache unit。
3. **batched direct write / store**
- MLA KV store。
- index K/scale store。
- 支持 bf16 与 fp8 e4m3。
4. **batched D2H backup**
- per-layer backup 可一次处理多个 request 的 local owned pages。
- 不退回大量 torch indexing。
5. **batched H2D load / prefetch**
- L2 host -> L1 device 的 load/prefetch 使用 descriptor batch。
- direct backend 下避免 CPU indices H2D 传输；CPU indices 不支持时 fail-fast。
6. **layout 支持**
- 当前 `page_first_direct` 必须支持。
- `layer_page_first` 可以作为性能优化候选，但不能阻塞 bs>1 正确性。
### descriptor 设计原则
descriptor 应该由 W1/W3/W4 生成，W5 只消费，不重新推导 request 逻辑：
```text
struct PageCopyDesc {
src_base / dst_base
src_page_id / dst_page_id
row_offset
valid_rows
padded_rows
layer_id
dtype
request_id
}
```
实际 C++/CUDA 字段可不同，但必须覆盖这些语义。
### benchmark 要求
覆盖：
```text
batch_size: 1, 2, 4, 8
extend tokens: 200, 512, 1024, 2000, 4096, 8192
prefix pages: 0, 8, 64, 512, 2048
page layout: contiguous, owner-lane, random/scattered
cache dtype: bf16, fp8 e4m3
backend: page_first_direct, optional layer_page_first
```
指标：
- kernel time；
- effective GB/s；
- launch count；
- CPU submit time；
- HBM/PCIe/NVLink 带宽估算；
- 与当前 torch/indexing 路径对比。
### 完成定义
- bs>1 descriptor batch 比 N 次 bs=1 调用更快。
- 小 extend 也不因 launch overhead 抵消 batch 收益。
- kernel 支持变长，不依赖固定 shape 反复 JIT。
---
## 9. W6：HiCache / load / backup / prefetch pipeline
### 目标
把 W1-W5 的 batch descriptors 接到 HiCache 生命周期：backup、evict、load、prefetch 都不能因为 bs>1 fallback。
### 修改范围
- `python/sglang/srt/mem_cache/hiradix_cache.py`
- `python/sglang/srt/managers/cache_controller.py`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
- `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
- `test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py`
### 实现要求
1. write-time metadata 每个 request 独立保存：
```text
page_owners
owned_positions
valid_len
padded_len
host_indices
draft_host_indices
```
2. `load_cp(nodes_to_load)` 已经接近 batch-capable，但要验证多 node、多 request owner pattern。
3. L2->L1 load/prefetch 使用 batched descriptors。
4. per-layer backup 使用 batched descriptors，避免每 request 单独 launch。
5. chunked prefill 遇到 in-flight backup split 仍然 defer，不动树。
6. host full / evict 不应破坏 page owner metadata。
### 与 W5 的接口
W6 需要 W5 提供：
```text
submit_batched_h2d_load(descs, stream)
submit_batched_d2h_backup(descs, stream)
```
如果 W5 尚未完成，W6 可以先用同步正确路径，但必须保留 warning，不能把它当最终 fast path。
### 测试
- 两个 request 写入 HiCache 后，分别 load back，owner pattern 不变。
- 两个 node 一次 `load_cp()`，visible device indices 顺序正确。
- L1 eviction 后从 L2 host hit 恢复。
- chunked prefill 不再触发 pending backup split crash。
### 完成定义
- bs>1 下 HiCache load/back up 不 fallback，不丢 owner metadata。
---
## 10. W7：EAGLE / draft 与 ETE/perf 验证
### 目标
target path 正确后，再恢复 EAGLE/draft，并做远端 ETE/perf 验证。
### 修改范围
- `python/sglang/srt/models/deepseek_nextn.py`
- `python/sglang/srt/layers/attention/nsa/utils.py`
- ETE / replay / benchmark 脚本
### draft 实现要求
1. draft KV plan 镜像 target KV plan。
2. draft 不生成自己的 owner-lane allocation。
3. draft input ids / positions / spec hidden states 按 request split。
4. `cp_collect_last_token_hidden()` 返回每个 request 的 last hidden。
5. EAGLE accept length 不能因为 cache hit/bs>1 掉到 1。
### ETE 验证场景
远端：
```text
host: g0034
container: sglang-glm5-dev-2
path: /sgl-workspace/sglang-tai
```
场景：
- bs=2/4/8 target-only；
- bs=2/4/8 EAGLE enabled；
- cache hit 90%+；
- extend tokens：200、512、1024、2000；
- chunked prefill on/off；
- fp8 kv cache on/off；
- HiCache L1 eviction + L2 load back；
- mixed prefix lengths；
- zero-prefix + host-hit request 混 batch。
必须观察：
- no fallback / fail-fast in supported case；
- accept len 正常；
- output len 不为 0；
- cache hit 正常；
- throughput 高于 sequential bs=1；
- 显存无异常增长；
- host memory / evict 行为稳定。
---
## 11. 并行依赖图
```text
W0 合同测试
├─> W1 CP metadata / planner
│ ├─> W2 allocation
│ │ ├─> W3 direct write
│ │ │ ├─> W4 current/partial reuse
│ │ │ └─> W6 HiCache load/backup
│ │ └─> W6 HiCache owner metadata
│ └─> W5 kernel descriptor design
│ ├─> W4 kernel-backed reuse
│ └─> W6 kernel-backed load/backup/prefetch
└─> W7 ETE harness can start early, but pass criteria waits for W1-W6
W7 EAGLE/draft waits for target W1-W6 correctness.
```
可立即并行启动：
```text
W0 red tests / fail-fast
W1 batch planner CPU helper
W2 allocator batch-owner tests and API sketch
W5 TAI kernel descriptor + benchmark harness
W7 ETE harness preparation
```
需要等待 W1/W2 后再正式接入：
```text
W3 direct write
W4 current/partial-current reuse
W6 HiCache pipeline
W7 EAGLE/draft
```
---
## 12. fast path 收敛要求
最终不应该留下大量“bs=1 fast path、bs>1 slow path”的分裂逻辑。理想收敛方式：
```text
bs=1 -> 生成 batch_size=1 的 CPSharedKVBatchPlan
bs>1 -> 生成 batch_size=N 的 CPSharedKVBatchPlan
runtime / kernel 都消费 CPSharedKVBatchPlan descriptors
```
允许短期保留单请求兼容字段，但后续应逐步让以下路径都消费 batched descriptors：
- CP split/rebuild；
- owner-lane allocation；
- direct write local loc；
- MLA current/partial-current reuse；
- index current/partial-current reuse；
- L1 prefix prefetch；
- L2->L1 load；
- per-layer D2H backup；
- HiCache metadata write/load；
- EAGLE/draft local path。
完成后，bs>1 的收益应该体现在：
1. 更大的 compute batch，短 extend 更容易填满 GPU。
2. 更少 Python/runtime 调用次数。
3. 更少 kernel launch 次数。
4. 更大的 H2D/D2H/materialize 单次任务，提升带宽利用率。
5. current/partial-current reuse 不因 batch 被禁用。
6. HiCache load/back up 不因 batch 退回低效路径。
---
## 13. 每个 workstream 的交付格式
每个负责人提交时必须给出：
```text
改动文件
新增/修改测试
支持的 case
明确不支持的 case
是否新增 collective
是否新增 fallback
是否影响 bs=1
远端是否 CUDA/ETE 验证
剩余风险
```
其中“是否新增 collective”必须明确回答。默认原则是不新增 planner collective。
---
## 14. 建议第一批派工
第一批可以这样并行：
1. **工程师 A：W1 CP metadata / batch planner**
- 只做 CPU helper + metadata + unit tests。
- 不接 model forward。
2. **工程师 B：W2 owner-lane allocation**
- 基于 W1 暂定接口写 owner list builder 和 allocator tests。
- 可先用 mock plan，等 W1 接口稳定后替换。
3. **工程师 C：W5 TAI kernel benchmark / descriptor**
- 先不接 runtime。
- 设计 variable-length descriptor benchmark，覆盖 bf16/fp8、page_first_direct、random/owner-lane pages。
第一批完成后，再开始 W3/W4/W6/W7 draft 接入。