# 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
  request_rank_local_offsets
  request_actual_seq_q_prev_cu_tensor
  request_actual_seq_q_next_cu_tensor

  # phase1 narrow output collection
  request_last_token_owner
  request_last_token_local_offset
  output_collect_mode  # narrow_last_token | full_rerange | unsupported_fail_fast

  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。
7. phase1 narrow-output 优化不能因为 bs>1 回退：普通 prefill 必须通过 `request_last_token_owner/local_offset` 只收集每个 request 的 last hidden。
8. bs>1 的 full rerange/logprob/hidden capture 在 batch-aware full rerange 完成前必须 fail-fast，不能使用 scalar `cp_all_gather_rerange_output()`。

### 对外接口

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 合同。

### 当前实现状态与性能风险

PR #10 已把 W2 正确性路径接入：

- bs>1 会按 request 独立调用 in-seq page owner planner；
- owner list 按 request order flatten；
- `alloc_extend_compute_owner()` 删除 bs=1 hard guard；
- allocator 复用已有 `alloc_extend_naive()` 生成 flattened `out_cache_loc`；
- supported CP shared-KV case 不再走 `multi_batch` legacy fallback。

但当前实现仍有一段 Python/control-path descriptor overhead，主要包括：

1. `alloc_paged_token_slots_extend()` 每次 allocation 从 CPU tensor `.item()` 构造
   `extend_lens` / `extend_prefix_lens` Python lists。
2. `build_batch_in_seq_page_compute_owners()` 用 Python loop 生成
   `flat_page_compute_owners: List[int]`。
3. `_select_compute_owner_pages()` 再用 Python loop 构造
   `required_by_owner` / `positions_by_owner`。
4. 每个 owner lane 会临时构造 `position_tensor`，再把 lane pages scatter 回
   `selected_pages`。

这比旧的全量 free-page scan / `torch.isin` / sort-merge 路径轻，因为 allocator
仍使用 `_owner_free_pages` / `_owner_release_pages` 本地 bucket 记账；但它还不是最终
descriptor 化 fast path。在线上大量 200-2000 token 短 extend 叠 batch 时，allocation
频率高，这段 Python list/tensor descriptor 构造可能重新成为 CPU hot path。

后续优化方向：

```text
W1 batch plan 直接携带 page-owner descriptor
  -> W2 allocator 消费 tensor/descriptor 而不是重新 Python 推导
  -> selected_pages gather/scatter 用固定 buffer 或 kernel 完成
  -> 避免每 tick/request 反复创建 Python lists 和临时 position tensors
```

这个风险不影响当前 W2 correctness，但需要在 bs>1 ETE/perf 阶段用 micro benchmark
和线上 trace 单独验证。

---

## 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。

### 当前状态

- target 侧 batch metadata、batch split helper、per-request last-token collect 已有单测覆盖。
- EAGLE/draft prefill local path 复用上述 batch-aware split helper：
  - `cp_split_and_rebuild_1d()` 处理 draft input ids；
  - `cp_split_and_rebuild_position()` 处理 positions；
  - `_get_cp_local_spec_hidden_states()` 接受已 CP-local 的 draft hidden，或对 full hidden 走 batch-aware split；
  - `cp_collect_last_token_hidden()` 按 request owner/offset collect last hidden。
- bs>1 时不允许 silent fallback：如果 CP draft shared-KV 已开启但 spec hidden、embedding pad metadata、input embeds 形状不满足 batch fast path，直接
  `[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_*]` 报错。bs=1 兼容 fallback 暂时保留。
- scheduler 的 bs>1 admission gate 仍未打开；打开前必须完成下面 ETE 场景，尤其是 EAGLE accept length 与 output len。
- 2026-06-04 远端启动失败记录：
  - 症状：pd warmup 阶段 EAGLE target forward 在 `cp_split_and_rebuild_data()` 抛
    `[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch] input tokens=8 expected=4`。
  - 根因：启动/warmup 或 speculative buffer 会把 `input_ids` 做 batch 尾部静态 padding；
    CP batch plan 的 `request_extend_lens` 仍只描述真实 valid rows。`split_kind="compute"` 应该丢弃尾部静态 padding，
    再由 CP split helper 自己生成 dummy compute rows；不能把 pad token embedding 当成 dummy compute data。
  - 合同订正：`split_kind="compute"` 允许两类尾部截断：
    1) CP compute padding 上限内的 dummy rows；
    2) `prepare_mlp_sync_batch()` 产生、并由 `forward_batch.extend_num_tokens` 显式给出上限的全局 static padding rows。
    `split_kind="valid"` 仍 fail-fast，避免 direct-write/out_cache_loc 静默丢 token。

### 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 接入。

---

## 15. 2026-06-03 W5/W4 kernel 接入状态更新

### 已接入并验证

1. **NSA index top-k 的 batch cp_index 调用已从 per request/segment MQA top-k 收敛为单次调用。**
   - SGLang 路径：`python/sglang/srt/layers/attention/nsa/nsa_indexer.py::_get_topk_in_seq_cp_pair_batch`。
   - 语义：多个 request 的 prev/next CP segment 先 compact 成一个 `cp_index` descriptor，再调用一次 `_get_topk_ragged_with_cp`。
   - 已有测试：`test_indexer_in_seq_cp_pair_batch_*` 与 `test_indexer_ragged_cp_index_current_batch_does_not_materialize`。

2. **index partial/current compose 的 current-slot fill 已接入 TAI kernel。**
   - SGLang 路径：`cp_shared_kv_runtime.fill_current_index_page_slots`。
   - TAI kernel：`tai_kernel.nsa_prefill.cp_shared_kv_materialize.fill_current_index_page_slots`。
   - 输入是 flatten 后的 current rows，可一次处理 bs>1 的 current suffix；不再在 CUDA 上默认走 PyTorch advanced indexing。
   - fallback 仍为显式 warning：`[CP_SHARED_KV_FALLBACK][tai_materialize] reason=index_current_fill_*`。
   - 远端 CUDA smoke：通过 64/page slot compose 的 K/scale 字节布局检查。

3. **shared-index batch cp_index 的 K/S + range prepare 已接入 TAI batch descriptor kernel。**
   - SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_index_batch`。
   - TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_kv_and_range_batch`。
   - descriptor：`batch_indices, kv_lens, q_starts, q_lens, k_bases, q_bases, total_kv_len, total_q_count, max_kv_len, max_q_len`。
   - 关键约束：`max_kv_len/max_q_len` 由 Python planner 传入，避免在 wrapper 内对 CUDA tensor 做 `.max().item()` 同步。
   - 输出合同：最终仍是按 segment concat 的 dense K buffer / scale buffer；`ks` 为 segment 的 K base；`ke_offset` 为 segment-local causal end offset，调用端计算 `ke = ks + ke_offset`。
   - 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。

### 仍需继续处理的 kernel/runtime 缺口

1. **current-only cp_index 的 K/S compact copy 已接入 TAI batch kernel。**
   - SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_current_index_batch`。
   - TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_current_kv_and_range_batch`。
   - 语义：从 flatten current rows 按 segment `current_base/current_len` 拷到 concat K/S buffer，同时生成 `ks/ke_offset`。
   - 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。

2. **HiCache per-layer D2H backup 仍按 reservation/node 提交。**
   - `scheduler._prepare_hicache_write_backups_before_forward()` 仍 `for req in batch.reqs: prepare_write_backup_for_req(req)`。
   - `cache_controller.submit_write_cp_layer()` 每个 reservation 在每层调用 host pool backup。
   - 底层 TAI direct transfer kernel 接收 flatten indices，本身可以吃 batched descriptor；但 runtime 还没有把多个 request/reservation 合并成一个 layer-level descriptor。

3. **L2->L1 load/prefetch 需要继续确认实际 queue 合并粒度。**
   - `cache_controller.load_cp()` / `start_loading()` 已比 backup 更接近 batched queue，但仍需用 ETE/NVTX 确认 layer-level descriptor 是否足够大，是否存在 per-node launch。

4. **MLA current/partial-current 和 direct store 已有 flatten kernel 基础，但还需用 bs>1 ETE 验证没有 fallback hot path。**
   - 特别关注 FP8/BF16 两种 dtype、draft/EAGLE 路径、prefetch 关闭时 full/current reuse 是否仍启用。

---

## 16. 2026-06-03 W6 HiCache load/backup batch 化设计补充

### 当前结论

1. **L2->L1 load-back 的传输层已经接近 batch 化。**
   - `HiRadixCache.load_back()` 仍按 request 调 `cache_controller.load_cp()`，但 `HiCacheController.start_loading()` 会对 `load_queue` / `draft_load_queue` 做 `CacheOperation.merge_ops()`。
   - 因此只要 scheduler 允许同一个 prefill batch 内多个 request 都在 `ready_to_load_host_cache()` 前完成 `init_load_back()`，目标 KV 和 draft KV 的 H2D per-layer transfer 会按 layer 合并成一次 `load_to_device_per_layer()` 调用。
   - 已有低风险验证点：补一个单测，先连续调用两次 `load_cp()`，再调用一次 `start_loading()`，断言每层只有一次 target load、一次 draft load，且 host/device indices 为两个 request 的 concat。

2. **D2H backup 仍没有吃到 bs>1 的收益。**
   - `scheduler._prepare_hicache_write_backups_before_forward()` 当前仍然 `for req in batch.reqs: prepare_write_backup_for_req(req)`，这是正确的 per-request reservation/metadata 生命周期。
   - 真正的问题在 `HiCacheController.on_layer_end()`：它遍历 `pending_layer_writes`，对每个 reservation 单独调用 `submit_write_cp_layer()`。
   - `submit_write_cp_layer()` 又对每个 reservation / layer 分别调用 `mem_pool_host.backup_from_device_per_layer()` 和 draft host pool backup。
   - 底层 TAI direct PF/LF transfer kernel 接收 flatten indices，本身可以处理多个 request 的 concat descriptor；当前没有收益是 runtime 没有把多个 reservation 合并为 layer-level descriptor。

3. **scheduler 仍有 CP bs>1 admission gate。**
   - `schedule_policy.py` 里仍有 `# TODO support cp with multiple requests`，当 `nsa_prefill_cp_in_seq_split` 或 `prefill_context_parallel_enabled` 打开时，`len(self.can_run_list) >= 1` 会拒绝继续加 request。
   - 所以 W6 可以先把 load/backup runtime 改成 batch-ready 并用 unit test 锁住；真正 ETE 吃到收益还依赖后续移除/替换 admission gate。

### 推荐实现：保持 radix/metadata per request，传输按 layer 聚合

目标不是把 HiCache node 合并，而是只合并底层 transfer descriptor：

- reservation、host slot、draft host slot、`CpHiCacheNodeMetadata` 仍每个 request 独立；
- radix tree attach/rollback/ack 仍按 node_id 独立；
- layer-end hook 对同一 source/同一 layer 下所有 pending reservation 做一次 concat transfer；
- final ack 可以继续 per node append，避免改变 `writing_check()` 的队列语义；
- 不新增 collective，不改变 owner-lane 分布语义。

### W6A：为现有 load-back batch 行为补测试

文件：`test/registered/unit/managers/test_hicache_controller_cp.py`

新增测试目标：

1. `load_cp(req_a)` + `load_cp(req_b)` + 单次 `start_loading()`：
   - `host_pool.loads` 每个 layer 只有一次；
   - 传入的 host/device indices 是两个 request 的 concat；
   - `ack_load_queue[0].node_ids == [node_a, node_b]`；
   - draft host pool 同样是每层一次。
2. zero-owned request 参与 batch 时仍不触发实际 host load，但 node_id 保留在 ack 中。

### W6B：backup layer-level grouping

文件：`python/sglang/srt/managers/cache_controller.py`

计划改造：

1. 抽出 helper：`_submit_write_cp_layer_states(states, layer_id, submit_target, submit_draft)`。
2. `submit_write_cp_layer(reservation, ...)` 保持公开 API，但内部只传单个 state 给 helper，保证 bs=1 行为不变。
3. `on_layer_end(layer_id, source)` 不再对每个 reservation 单独 submit，而是收集当前 pending states，一次调用 helper。
4. helper 内部：
   - 计算每个 state 是否需要 target/draft；
   - 记录一个 layer event，并挂到所有参与 states 上保活；
   - 在 `write_stream` 上等待每个 state 的 `start_event` 和 group layer event；
   - 对 target states concat `state.host_indices` / `state.physical_device_indices`，调用一次 target `backup_from_device_per_layer()`；
   - 对 draft states concat `state.draft_host_indices` / 同一 physical indices，调用一次 draft backup；
   - 标记每个 state 的 target/draft layer completed；
   - 对完成全部层的 state 逐个 record finish_event、record_stream、append ack。

### W6B 单测要求

文件：`test/registered/unit/managers/test_hicache_controller_cp.py`

1. 两个 target-only reservation 同一 layer-end：
   - `allocator.device_pool.notify_layer_end_for_backup(0)` 后 `host_pool.layer_backups` 只新增一条；
   - indices 为两个 reservation concat；
   - 未完成所有 layer 前 ack 为空；
   - 最后一层后两个 node_id 都释放为 ack。
2. target + draft：
   - target source layer-end 只备 target；draft source layer-end 只备 draft；
   - 每个 source 每层只产生一次 backup call；
   - final ack 等 target/draft 都完成。
3. zero-owned reservation：
   - 不增加实际 backup call；
   - 仍按层完成并最终 ack。
4. 直接调用 `submit_write_cp_layer()` 的旧路径仍保持单 reservation 行为。

### W6C：后续优化，不作为第一步

1. `load_cp()` 的 owner-lane allocation 仍按 request 调用；如果 admission 放开后 CPU overhead 明显，再做 batch owner allocation planner。
2. `CacheOperation.merge_ops()` 仍会 `torch.cat` 二次 concat；如果 profile 显示 CPU cat 明显，再引入预分配 descriptor workspace。
3. backup grouping 首版不把 ack 合并成单个 `HiCacheAck`，避免影响 `writing_check()` 与 rollback unattached prepared backup 的语义。
4. 不引入新的 all-reduce/all-gather；W6 只改变 descriptor 聚合粒度。

### W6A 完成状态

---

## 25. 2026-06-06 bs>1 打开后性能回退排查记录

用户确认：`batch_size == 1` 也生成并消费 `CPSharedKVBatchPlan` 是预期收敛方向，
不是本轮性能问题本身。后续不要再把“bs=1 走 batch-plan”当成需要回退的 bug。

最新远端日志（`/mnt/beegfs/cjy/log/sglang_cp_hicache_20260605_164245.log`）
显示：

- 没有 `[HiCache-load] slow load_back planning` 或容量等待类日志；
- fallback 主要是 `[CP_SHARED_KV_FALLBACK][tai_ipc_materialize] reason=paged_start_slot_nonzero`，
  每 rank 限频后共 64 条，说明 index/page-buffer materialize 的非 0 起始 slot 会退到
  local materialize + collective；
- CP0 batch 分布仍以 bs=1 为主，少量 bs=2/3，因此不能简单用“真实 batch 数变大”
  解释全部慢点；
- debug 日志本身会污染性能测试，但不是一个足够可靠的根因解释。

当前高优先级性能嫌疑：

1. **bs>1 partial-current sync compose 的 row-major prefix slot span 可能过度覆盖。**
   `build_batch_prefix_slot_span()` 为了用一个 contiguous span 覆盖所有 request prefix，
   在 row-major page table 下会包含 request 之间的 current/空洞 slot。cache-hit 高、
   prefix 分布不均时，prefix materialize/all-reduce 的页数可能远大于真实 prefix 页数。
   这会落在每层 critical path 上，尤其在 bs>1 async prefetch 还没恢复时更明显。

2. **index/top-k batch path 仍有 per-layer Python/Torch descriptor overhead。**
   `_get_topk_in_seq_cp_pair_batch()` 每层构造 Python list、`torch.tensor(..., device)`、
   `torch.cat`、`torch.full`，再 scatter compact topk。TAI batch kernel 已减少部分
   K/S copy，但 descriptor 构造还不是 per-forward 预计算。

3. **tiny extend compute padding 会把 `< cp_size` pages 的 request 补到 `cp_size` pages。**
   例如 page=64、cp=8 时，200 token 会以 512 token compute 形态参与 CP split。
   这对 200-512 token 的线上短 extend 可能抵消 batch 的 compute 填充收益。

4. **bs>1 L1 prefix prefetch 仍未实现。**
   `CpSharedKVMlaPrefetcher` / `CpSharedKVIndexPrefetcher` 仍要求
   `forward_batch.batch_size == 1`。如果开启真实 bs>1 后 cache-hit 高，prefix 准备
   会更多落到同步 compose 路径，而不是目标的 async L2->L1/L1 materialize 流水线。

下一步验证顺序：

1. 加或复用低频 perf counter，记录每个 forward 的 prefix true pages、prefix span pages、
   current span pages、batch size、extend/prefix lens，不在每层刷屏；
2. microbench `_get_topk_in_seq_cp_pair_batch()` descriptor/cat/scatter 与
   `split_tensor_by_cp_batch_plan()` 在 200/512/1k/2k extend、bs=1/2/4/5 下的 CPU submit
   和 GPU elapsed；
3. 若证实 prefix span 过度覆盖，优先改为 compact prefix descriptor 或支持 nonzero-slot
   TAI IPC materialize，避免 row-major gap 被 all-reduce；
4. 若证实 descriptor overhead 主导，先把 index batch descriptor 缓存在 `ForwardBatch`
   上，避免 78 层重复构造。

已补 characterization tests，锁住现有 L2->L1 load-back batch 行为：

- `test_cp_start_loading_batches_multiple_load_cp_requests_with_draft`
  - 两个 `load_cp()` request 在一次 `start_loading()` 中合并；
  - target 每层一次 `load_to_device_per_layer()`；draft 每层一次；
  - host/device descriptors 为两个 request 的 concat；
  - ack 保留两个 node_id。
- `test_cp_start_loading_keeps_zero_owned_load_ack_in_batched_load`
  - zero-owned request 不产生实际 host transfer；
  - 但 zero-owned node_id 仍保留在 batched load ack 中。

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestHiCacheControllerCPLoad::test_cp_start_loading_batches_multiple_load_cp_requests_with_draft \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestHiCacheControllerCPLoad::test_cp_start_loading_keeps_zero_owned_load_ack_in_batched_load
=> 2 passed, 3 warnings
```

本地限制：本地 pytest 仍因缺少 `orjson` 无法收集该测试文件；本地只做了 `py_compile`。

### W6B 完成状态

已实现 CP HiCache per-layer backup 的 layer-level grouping：

- `HiCacheController.on_layer_end()` 不再逐 reservation 调 `submit_write_cp_layer()`；现在对当前 pending states 做一次 grouped submit。
- reservation、radix metadata、host slot、draft host slot、ack node_id 仍保持 per request。
- 每个 layer/source 的底层 transfer descriptor 会按 pending reservation concat：
  - target source：一次 target `backup_from_device_per_layer()`；
  - draft source：一次 draft `backup_from_device_per_layer()`；
  - zero-owned state 不参与实际 transfer，但继续按 layer 完成并最终 ack。
- `submit_write_cp_layer()` 公开 API 保留，内部走单 state grouped helper，bs=1 兼容路径不变。
- 没有新增 collective；没有修改 owner-lane 语义；没有修改 tai-kernel。TAI direct transfer kernel 已支持 flatten indices，runtime concat 后即可吃到 bs>1 descriptor 收益。

新增测试：

- `test_cp_layer_hook_groups_target_backups_across_pending_reservations`
- `test_cp_layer_hook_groups_target_and_draft_backups_by_source`
- `test_cp_layer_hook_keeps_zero_owned_ack_in_grouped_backup`

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_hicache_controller_cp.py
=> 75 passed, 3 warnings
```

补充验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
=> 10 passed, 3 warnings
```

## 18. 2026-06-04 scheduler admission gate 首版打开

目标：先允许 NSA in-seq CP shared-KV 在 scheduler 层组成真实 multi-request prefill batch，同时用显式参数把首版风险收窄。

### 参数

1. `--enable-cp-shared-kv-prefill-bs-gt1`
   - 默认关闭；关闭时 CP prefill 仍保持旧的 `batch_size == 1` admission gate。
   - 打开后，且 `--enable-nsa-prefill-cp-shared-kv` 同时开启时，scheduler 可以把多个 waiting requests 放进同一个 prefill batch；普通 CP 不因该参数打开 bs>1。
2. `--cp-shared-kv-prefill-max-batch-requests`
   - 可选正整数；限制同一个 CP shared-KV prefill batch 的 request 数。
   - 不设置时只受通用 `--prefill-max-requests` 约束。
3. `--cp-shared-kv-prefill-max-total-extend-tokens`
   - 可选正整数；限制同一个 CP shared-KV prefill batch 的 **page-aligned extend tokens** 总量。
   - 单个 request 自身超过该值时仍允许独占 batch，避免 scheduler deadlock；真正是否能运行仍由 KV allocator 容量判断。

### 容量语义

scheduler admission 不做 bf16/fp8 字节级估算。原因是实际 KV pool 容量已经在初始化时按 dtype、page size、model shape、mem fraction 转换为 token/page capacity。admission 只消费现有 `PrefillAdder` 的 token/page budget：

- `rem_input_tokens` / `rem_total_tokens` / `cur_rem_tokens` 继续作为通用容量预算；
- `alloc_for_extend()` 仍是最终容量仲裁点；
- CP owner-lane 不足仍通过 `KVCapacityWaitError` deferred，不在 scheduler 新增 collective 或重复推导 dtype bytes。

这样 bf16 与 fp8 都走同一个 admission 逻辑，差异由底层 allocator 的实际可用 token capacity 体现。

### 风险与后续

1. `cp_shared_kv_prefill_max_total_extend_tokens` 首版使用 page-aligned extend 累计，和 allocator 的最小 page 单位一致；日志/指标里如果要展示用户侧 token，需要另加 valid-token 统计。
2. 该 gate 只控制 scheduler 组 batch，不解决所有 bs>1 runtime correctness；打开 ETE 时仍应保留现有 fail-fast，特别是 draft/EAGLE、logprob、hidden capture、compute padding 路径。
3. 如果 ETE 中 `KVCapacityWaitError` 频繁出现，下一步应把 owner-lane free pages/evictable pages 的 batch precheck 前移到 scheduler，而不是引入 all-reduce。

## 19. 2026-06-04 L2 host write admission batch 化

### 背景

打开 scheduler bs>1 gate 后，CP HiCache write-through 的 radix node / host reservation 仍然必须保持 per request 独立；但 host 侧容量 admission 如果仍按 request 串行执行，会在同一个 prefill batch 内反复触发 L2 host free-room eviction：

- correctness 没问题，因为 `reserve_write_cp()` 最终会消费真实 host allocator slot；
- 性能和碎片化较差，因为每个 request 都可能单独触发 free-room watermark admission；正确公式应在 `available - required < trigger_room` 时触发，并 evict `max(required - available, target_room)`，其中 `target_room` 是最小 eviction chunk，不是目标 available 水位；
- 这与 L1/L2 free-room 的目标相反：应尽量一次 evict 到目标余量，降低 eviction 频率并给后续 allocate 更连续的空间。

### 实现原则

只把 **host write admission / eviction planning** 聚合到 batch level，不合并 request 生命周期：

- `CpHiCacheNodeMetadata`、host slot reservation、draft host slot、radix attach、rollback、ack 仍 per request；
- batch admission 只聚合多个 request 的 `required_by_owner`，用一次 `_cp_build_write_admission_from_required()` 计算 host free-room deficit；
- 如果 batch admission 需要 eviction，则只在 forward 前触发一次 `_evict_cp_host_for_write_admission()`；
- admission 成功后，每个 request 继续调用 `reserve_write_cp()` 拿自己的 host slots，但跳过重复的 per-request initial admission；
- 如果 batch admission 无法满足 deficit，会显式打 `[CP_HICACHE_FALLBACK][prepare_write_backups_batch_admission_failed]`，再回到 per-request reservation 路径，不 silent fallback；
- 不新增 collective，不改变 CP owner-lane 分布，不改变 TAI kernel 接口。

### 完成状态

代码路径：

- `Scheduler._prepare_hicache_write_backups_before_forward()` 优先调用 `prepare_write_backups_for_reqs(batch.reqs)`；
- `SessionAwareCache.prepare_write_backups_for_reqs()` 透传到 inner cache，缺失 batch API 时才逐 req fallback；
- `HiRadixCache.prepare_write_backups_for_reqs()`：
  - 先构造每个 request 的 `CpWriteBackupCandidate`；
  - 聚合所有 candidate 的 `_cp_required_host_tokens_by_rank(kv_indices)`；
  - 对聚合后的 `required_by_owner` 做一次 host write admission / eviction；
  - 再逐 request 做实际 reservation 和 per-layer backup registration。

新增测试：

- `test_cp_host_write_batch_admission_uses_aggregate_required`
  - 验证 batch admission 对聚合 required 使用 free-room 公式。
- `test_prepare_write_backups_for_reqs_runs_one_batch_admission`
  - 验证两个 request 只触发一次 batch admission / eviction，最终仍各自得到 prepared backup，并提交两个独立 reservation。

验证：

```text
python -m py_compile \
  python/sglang/srt/mem_cache/hiradix_cache.py \
  python/sglang/srt/mem_cache/session_aware_cache.py \
  python/sglang/srt/managers/scheduler.py \
  test/registered/unit/mem_cache/test_cp_hicache_metadata.py

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_hicache_metadata.py::TestCpHiCacheFreeRoom
=> 8 passed, 3 warnings

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_hicache_metadata.py
=> 117 passed, 5 warnings
```

## 20. 2026-06-04 direct-write out_cache_loc 尾部 static padding 修正

### 失败现象

远端 warmup/EAGLE target forward 进入 NSA indexer direct-write 后失败：

```text
[CP_SHARED_KV_FAIL_FAST][direct_write]
reason=batch_split_out_cache_len_mismatch
split_list tokens=4 out_cache_loc tokens=8
```

栈在：

```text
forward_mla.py -> nsa_indexer.py::_store_cp_shared_local_index_k_cache
-> get_cp_shared_kv_local_out_cache_loc()
```

这是前一轮 `cp_split_and_rebuild_data()` input-token mismatch 的同类问题，但发生在 cache write loc 边界：

- `CPSharedKVBatchPlan.request_extend_lens` 描述 valid rows；
- model-runner/speculative warmup 可能在 flattened batch 尾部追加一段全局 static padding rows；
- `out_cache_loc` 也可能携带这段尾部 padding loc；
- CP shared-KV direct-write 是 valid-token 写入，不能把 dummy compute rows 写进 KV/index cache。

### 合同修正

保持 `split_tensor_by_cp_batch_plan(..., split_kind="valid")` 严格：

- valid split helper 仍然拒绝输入长度超过 `sum(request_extend_lens)`；
- direct-write 边界负责先剥离已知的 **全局尾部 static padding locs**；
- 只有在 `batch_plan.compute_padding_enabled=True`，且
  `valid_tokens < out_cache_loc.numel() <= sum(request_compute_padded_tokens)` 时允许截尾；
- 其它 mismatch 继续 fail-fast，不做 silent fallback。

这样保证：

1. compute path 仍可吃到 bs>1 padding 后的堆叠计算；
2. valid cache write 不会写 dummy rows；
3. 如果未来改成 request 内部 interleaved padding layout，当前截尾规则会 fail-fast，而不是错误写 cache。

### 验证

新增回归：

- `test_local_out_cache_loc_ignores_trailing_static_padding_locs`
  - 构造 `valid_locs=4`，`out_cache_loc=8`，后 4 个为 static padding loc；
  - 验证 local direct-write 只返回前 4 个 valid loc。

远端 targeted 验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_ignores_trailing_static_padding_locs \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_valid_kind_rejects_trailing_padding_rows \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_uses_valid_rows_under_compute_padding
=> 4 passed, 5 warnings
```

远端相关文件验证：

```text
python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py
=> 73 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
=> 37 passed, 3 warnings

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 107 passed, 1 deselected, 5 warnings, 2 subtests passed
```

未完成验证：

- `test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel`
  在远端 installed-kernel self-test 处卡住；这是 TAI kernel self-test/import 路径，
  与本次 `get_cp_shared_kv_local_out_cache_loc()` 尾部 loc 截断不是同一逻辑路径，
  但后续需要单独处理，避免完整 runtime suite 长时间挂住。

## 21. 2026-06-04 warmup 卡死与 EAGLE draft compute padding 修正

### 现象

远端 warmup 初看卡在：

```text
Start of pd disaggregation warmup ...
[CacheCtrl-write] submit_write_cp_per_layer registered ...
[HiCache-write] prepared CP per-layer backup before forward ...
CUTE_DSL WARNING Unexpected error during package walk: cutlass.cute.experimental
```

`py-spy` 看到 scheduler rank0 实际阻塞在：

```text
torch/utils/file_baton.py:50 wait
torch/utils/cpp_extension.py load
tai_kernel/nsa_prefill/_extension_loader.py load_tai_kernel_extension
in_seq_all_gather_rerange_cuda
```

本次不是 HiCache ack 卡死，而是 TAI cpp extension JIT cache 目录存在 stale `lock`：

```text
/root/.cache/tai-kernel/ops/aa52c7047416/lock
```

同目录 `.so` 已存在且没有活跃 `ninja/nvcc`，移除 stale lock 后 warmup 继续。

### 继续暴露的真实错误

移除 stale lock 后，warmup 在 EAGLE draft extend 失败：

```text
RuntimeError: Trying to create tensor with negative dimension -56: [-56, 6144]

forward_batch_info.py:_pad_inputs_to_size
spec_info.hidden_states = self._pad_tensor_to_size(spec_info.hidden_states, num_tokens)
```

根因：

- target 侧 `capture_draft_hidden_states=True` 在 `DeepseekV2Model.forward()` 的 CP output collect 之前抓取 hidden；
- compute padding 后这个 side-channel 是 CP-local compute rows，例如 warmup 单请求 `hidden_states.shape[0] == 64`；
- model-runner/speculative warmup 的 draft input 仍带全局 static padding token 数，例如 `num_tokens == 8`；
- `ForwardBatch._pad_inputs_to_size()` 试图把 64 行 hidden pad 到 8 行，导致负维度；
- 更深一层问题是 `ForwardMode.DRAFT_EXTEND` 没被视为 `context_parallel_extend`，所以 DeepSeek NextN draft 模型不会进入 CP-local draft fast path，CP-local side-channel 与非 CP draft input 语义不一致。

### 合同修正

EAGLE/NextN draft 在 `SGLANG_CP_DRAFT_SHARED_KV=1` 且 `uses_cp_shared_kv=True` 时必须跟 target 保持 CP-local 语义：

1. `can_cp_split()` 对 CP shared-KV draft extend 返回 true，使 draft 模型也构建 NSA CP metadata；
2. `nsa_use_prefill_cp()` 对 CP shared-KV draft extend 返回 true，使 `DeepseekModelNextN` 进入 local draft path；
3. `_pad_inputs_to_size()` 遇到 CP shared-KV draft hidden rows 已经大于全局 static padded `num_tokens` 时，不再尝试缩短/重 pad hidden；保留原始 CP-local rows，让 draft 模型按 CP split 后的 local input 使用。

这不是打开新的 scheduler bs>1 行为；它修正的是 EAGLE draft 在 CP shared-KV + compute padding 下的已有语义。

### 回归

新增：

- `test_can_cp_split_enables_cp_draft_shared_kv_draft_extend`
- `test_nsa_use_prefill_cp_enables_cp_draft_shared_kv_draft_extend`
- `test_cp_draft_padding_keeps_local_hidden_when_static_tokens_are_shorter`

远端 RED：三个测试在修复前分别失败为 `can_cp_split=False`、`nsa_use_prefill_cp=False`、负维度 `-56`。

远端 GREEN：

```text
python -m py_compile \
  python/sglang/srt/layers/attention/nsa/utils.py \
  python/sglang/srt/model_executor/forward_batch_info.py

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py
=> 76 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 144 passed, 1 deselected, 3 warnings, 2 subtests passed
```

未验证：

- 最新代码还没重新跑完整 ETE warmup；需要用户重启服务后再看是否越过 draft extend。
- stale TAI JIT lock 是环境/loader 问题，不应由业务逻辑修复掩盖；后续应在 TAI loader 层处理 stale lock 或超时诊断。

### 2026-06-04 订正：上一版 guard 仍过窄

远端新一轮 warmup 仍在相同位置失败：

```text
RuntimeError: Trying to create tensor with negative dimension -56: [-56, 6144]

forward_batch_info.py:prepare_mlp_sync_batch
forward_batch_info.py:_pad_inputs_to_size
spec_info.hidden_states = self._pad_tensor_to_size(spec_info.hidden_states, num_tokens)
```

新证据：

- 远端实际加载的 `forward_batch_info.py` 已包含上一版 guard；
- 因此不是同步问题，也不是 HiCache ack；
- 失败仍说明 `spec_info.hidden_states.shape[0] == 64` 且
  MLP sync static `num_tokens == 8`；
- 上一版 guard 依赖 `uses_cp_shared_kv && SGLANG_CP_DRAFT_SHARED_KV`
  才保留 oversized draft hidden。这个前提过强：
  `prepare_mlp_sync_batch()` 发生在 `attn_backend.init_forward_metadata()`
  之前，draft 的 CP metadata/flag 传播不能作为此处是否允许普通 padding 的唯一依据。

上一版修正存在的问题：

- 仅通过 `spec_info.hidden_states.shape[0] > num_tokens` 推断 CP-local hidden
  过于宽泛；
- 这个条件只能证明“普通 padding 会失败”，不能证明该 hidden 的语义一定是
  CP-local draft hidden；
- 如果未来出现其它 draft path 产生 oversized hidden，长度判断会把错误 silently
  推迟到后面的模型计算，排查成本更高。

最终合同：

- CP-local draft hidden 必须由源头显式标记：
  `EagleDraftInput.cp_local_hidden_states=True`；
- 只有 target 侧实际返回 `logits_output.draft_hidden_states` 时，EAGLE worker
  才把该 marker 传给 `EagleDraftInput`；
- `_pad_inputs_to_size()` 只信这个 semantic marker，不再根据长度猜测；
- 如果没有 marker 但 `hidden_states.shape[0] > num_tokens`，直接 fail-fast：
  `[CP_SHARED_KV_FAIL_FAST][draft_hidden_static_padding_mismatch]`；
- 标记为 CP-local 的 hidden 保持原 tensor，后续由
  `DeepseekV3ForCausalLMNextN.forward()` 的 CP-local draft path 校验/消费。

新增回归：

- `test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible`
  - 构造 `cp_local_hidden_states=True` 且 CP metadata/env 尚不可见的最小 draft batch；
  - `hidden_states.shape[0]=64`，`num_tokens=8`；
  - 修复前远端 RED：`EagleDraftInput` 不接受该 marker；
  - 修复后 GREEN：保留 `(64, hidden)` hidden side-channel。
- `test_cp_draft_padding_rejects_unmarked_oversized_hidden`
  - 构造未标记 oversized hidden；
  - 修复前远端 RED：旧长度判断会静默保留；
  - 修复后 GREEN：fail-fast，不允许靠长度误判 hidden 语义。

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_rejects_unmarked_oversized_hidden
=> 2 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py
=> 78 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 144 passed, 1 deselected, 3 warnings, 2 subtests passed
```

### 2026-06-04 订正：draft hidden marker 不应依赖 forward_mode

远端最新失败栈已经从负维度变为我们新增的 fail-fast：

```text
RuntimeError: [CP_SHARED_KV_FAIL_FAST][draft_hidden_static_padding_mismatch]
hidden_tokens=64 num_tokens=8 forward_mode=1

scheduler.run_batch
model_worker.forward_batch_generation
eagle_worker.py:forward_draft_extend
draft_model_runner.forward
model_runner.py:_forward_raw
forward_batch.prepare_mlp_sync_batch
forward_batch_info.py:_pad_inputs_to_size
```

根因订正：

- `ForwardMode.EXTEND == 1`，所以 fail-fast 中的 `forward_mode=1` 不是
  `DRAFT_EXTEND`；
- 这不是说明 EAGLE draft 语义丢失，而是
  `prepare_mlp_sync_batch()` 在 `is_extend_in_batch + DP max padding` 路径下会临时把
  draft/verify/decode/idle 等 mode 改写为 `ForwardMode.EXTEND`，用于复用 extend 静态
  MLP sync padding；
- 因此 `_pad_inputs_to_size()` 不能用 `forward_mode.is_draft_extend()` 判断是否允许保留
  CP-local draft hidden；
- 正确的语义边界是：当前代码块已经由 `spec_info.is_draft_input()` 保护，是否 CP-local
  只应由 `EagleDraftInput.cp_local_hidden_states` 这个显式 marker 决定。

修正：

- `keep_cp_local_hidden = getattr(spec_info, "cp_local_hidden_states", False)`；
- 保留未标记 oversized hidden 的 fail-fast，避免再次回到长度推断；
- 新增回归覆盖实际运行形态：`forward_mode=ForwardMode.EXTEND` 且
  `EagleDraftInput.cp_local_hidden_states=True` 时，`hidden_states=(64, hidden)` 在
  `num_tokens=8` 的静态 padding 下必须保持不变。

远端 RED/GREEN 证据：

```text
# RED: 修复前新增测试复现远端 fail-fast
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite
=> FAILED with [CP_SHARED_KV_FAIL_FAST][draft_hidden_static_padding_mismatch]

# GREEN: 修复后 marker 测试 + 未标记 oversized 保护同时通过
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_rejects_unmarked_oversized_hidden
=> 2 passed, 5 warnings

# 相关套件
PYTHONPATH=python python -m pytest -q \
  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 \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 223 passed, 1 deselected, 5 warnings, 2 subtests passed
```

未验证：

- 需要用户重启远端服务并打 warmup/ETE 流量确认已经越过 EAGLE draft extend；
- 如果后续仍失败，应优先看新栈，不要回到 forward_mode 判定或长度推断。

### 2026-06-04 订正：普通 MLP sync static padding 也会进入 CP compute split

远端 warmup 已通过后，第一条真实请求在 target forward 失败：

```text
RuntimeError: [CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch]
input tokens=40392 expected=40387

model_runner.forward_extend
DeepseekV2Model.forward
cp_split_and_rebuild_data
split_tensor_by_cp_batch_plan
```

关键信号：

- 请求日志为 `40391 input + 200000 new`；
- CP plan valid extend rows 为 `40387`；
- 进入 `cp_split_and_rebuild_data()` 的 hidden rows 为 `40392`；
- `40392` 是按 `attn_tp_size/attn_cp_size` 对齐后的长度，差值 5 是
  `ForwardBatch.prepare_mlp_sync_batch()` 添加的全局尾部 static padding；
- 这次不是 tiny compute padding：`plan.compute_padding_enabled=False` 时也会发生。

上一版遗漏：

- `split_tensor_by_cp_batch_plan(split_kind="compute")` 只在
  `compute_padding_enabled=True` 时允许剥离尾部 padding；
- 但 MLP sync static padding 与 CP compute padding 是两类 padding：
  - CP compute padding：为 page/owner-lane compute 形态补 dummy rows；
  - MLP sync static padding：为 DP/TP/CP collective buffer 对齐，在 flattened batch 尾部补 rows；
- 二者不能混为一谈。即使不需要 CP compute padding，也必须允许 compute split 剥离已知的
  全局尾部 static padding。

修正合同：

- `split_kind="compute"` 可以剥离尾部 static padding，但必须由调用方显式传入
  `static_padded_tokens` 上限；
- `_cp_split_and_rebuild_batch_in_seq()` 从 `forward_batch.extend_num_tokens` 传入该上限；
  这是 `prepare_mlp_sync_batch()` 后的本地 padded token count；
- `split_kind="valid"` 仍不接受任何尾部 padding，direct-write/out_cache_loc 路径保持严格；
- 如果 input 超过 `static_padded_tokens` 或 compute padded 上限，继续 fail-fast。

新增回归：

- `test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding`
  - 构造 `extend_len=7,page_size=4,cp_size=2`，此时 `compute_padding_enabled=False`；
  - 输入 tensor 有 8 行，模拟 MLP sync static padding；
  - 修复前 RED：`input tokens=8 expected=7`；
  - 修复后 GREEN：剥离第 8 行，仅按前 7 个 valid rows 做 CP split。

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_valid_kind_rejects_trailing_padding_rows \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows
=> 3 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  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 \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 224 passed, 1 deselected, 5 warnings, 2 subtests passed
```

未验证：

- 需要重启远端服务并重放第一条真实请求，确认已经越过
  `input tokens=40392 expected=40387` 这一栈；
- 如果后续还有失败，应优先看新栈，不应回退到放宽 `valid` split。

### 2026-06-04 订正：position split 也必须接收 MLP sync static padding 上限

上一轮修复 data/1d compute split 后，远端再次失败，但栈前进到 position：

```text
RuntimeError: [CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch]
input tokens=40392 expected=40387

DeepseekV2Model.forward
positions = cp_split_and_rebuild_position(forward_batch, positions)
split_tensor_by_cp_batch_plan
```

根因：

- `cp_split_and_rebuild_data()` 和 `cp_split_and_rebuild_1d()` 通过
  `_cp_split_and_rebuild_batch_in_seq()` 传入了 `forward_batch.extend_num_tokens`；
- `cp_split_and_rebuild_position()` 是独立 wrapper，仍直接调用
  `split_tensor_by_cp_batch_plan(..., mode="position")`，没有传
  `static_padded_tokens`；
- 因此同一批 MLP sync 尾部 static padding 在 hidden data 上已被剥离，但 positions 上仍
  fail-fast。

修正：

- `cp_split_and_rebuild_position()` 也传入
  `static_padded_tokens=getattr(forward_batch, "extend_num_tokens", None)`；
- runtime 调用点复核：
  - `_cp_split_and_rebuild_batch_in_seq()`：compute data/1d，已传 static 上限；
  - `cp_split_and_rebuild_position()`：compute position，本次补齐；
  - `get_cp_shared_kv_local_out_cache_loc()`：valid cache write，仍保持严格，不传 static 上限。

新增回归：

- `test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding`
  - 构造 `extend_len=7,page_size=4,cp_size=2`，`compute_padding_enabled=False`；
  - 输入 positions 长度 8，模拟 MLP sync static padding；
  - 修复前 RED：`input tokens=8 expected=7`；
  - 修复后 GREEN：position split 与先截断到 7 行再 split 的结果一致。

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_valid_kind_rejects_trailing_padding_rows
=> 3 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  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 \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 225 passed, 1 deselected, 5 warnings, 2 subtests passed
```

未验证：

- 需要重启服务重放第一条请求，确认已经越过 `cp_split_and_rebuild_position()`。

### 2026-06-04 订正：direct-write 也会收到 MLP sync static padding，不能用 compute padding 兜底

远端再次失败，栈前进到 NSA indexer cache write：

```text
RuntimeError: [CP_SHARED_KV_FAIL_FAST][direct_write]
reason=batch_split_out_cache_len_mismatch
split_list tokens=40387 out_cache_loc tokens=40392

nsa_indexer.py::_store_cp_shared_local_index_k_cache
get_cp_shared_kv_local_out_cache_loc()
```

这次失败与 data/position split 属于同一类，但边界不同：

- `ForwardBatch.prepare_mlp_sync_batch()` 会把 `input_ids / positions / out_cache_loc`
  一起 pad 到 `extend_num_tokens`；
- `CPSharedKVBatchPlan.request_extend_lens` 仍只描述 valid token rows；
- direct-write 的 `out_cache_loc` 是 cache 写入地址，只能对应 valid rows；
- 因此 direct-write 边界必须先剥离全局尾部 static padding locs，再进入
  `split_kind="valid"`；
- 这与 CP page-tail/owner-lane compute padding 不是一回事。

修正后的三层合同：

1. **valid rows**：`sum(request_extend_lens)`，唯一允许写入 KV/index cache 的 rows；
2. **CP compute padding / page tail**：按 request/page/owner-lane 补 dummy rows，
   只由 `split_tensor_by_cp_batch_plan(split_kind="compute")` 在 split 内部产生；
3. **MLP sync static padding**：`prepare_mlp_sync_batch()` 在 flattened batch 尾部追加，
   上限由 `forward_batch.extend_num_tokens` 显式给出。

direct-write 规则：

- 如果 `out_cache_loc.numel() == valid_tokens`，直接进入 strict valid split；
- 如果 `out_cache_loc.numel() > valid_tokens`，只有
  `out_cache_loc.numel() <= forward_batch.extend_num_tokens` 时才截尾；
- 不再把 `batch_plan.compute_padding_enabled` 当成 direct-write 截尾依据；
- 如果没有 `extend_num_tokens` 证明这段尾部来自 MLP static padding，即使
  `compute_padding_enabled=True` 也 fail-fast。

同类调用点复核：

- `_cp_split_and_rebuild_batch_in_seq()`：data/1d compute split，传入
  `extend_num_tokens`，允许剥离 MLP static padding；
- `cp_split_and_rebuild_position()`：position compute split，传入
  `extend_num_tokens`，允许剥离 MLP static padding；
- `get_cp_shared_kv_local_out_cache_loc()`：valid cache write，先用
  `extend_num_tokens` 截掉静态尾部，再调用 strict valid split；
- `select_cp_local_valid_rows_for_cache_write()`：只负责从 local compute rows 中剥离
  CP compute padding，不处理 flattened MLP static padding；该 padding 必须在 CP split 前处理。

新增回归：

- `test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding`
  - `extend_len=7,page_size=4,cp_size=2`，`compute_padding_enabled=False`；
  - `out_cache_loc` 长度 8，模拟 MLP sync static padding；
  - 修复前 RED：`split_list tokens=7 out_cache_loc tokens=8`；
  - 修复后只返回 rank-local valid tail locs。
- `test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding`
  - 验证 direct-write 不再因为 `compute_padding_enabled=True` 静默截掉未知尾部 loc。

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_ignores_trailing_static_padding_locs \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding
=> 5 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  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 \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 227 passed, 1 deselected, 5 warnings, 2 subtests passed
```

仍需 ETE 验证：

- 需要重启远端服务并重放第一条真实请求，确认已经越过
  `get_cp_shared_kv_local_out_cache_loc()` 的 mismatch；
- 如果后续还有失败，应继续按新栈定位，不要再把三类 padding 混成一个长度条件。

### 2026-06-04 订正：index top-k batch path 必须按实际 q rows 选择 valid/compute metadata

重启后远端越过 direct-write mismatch，但在 NSA indexer top-k 阶段失败：

```text
RuntimeError: [CP_SHARED_KV_FAIL_FAST][index_topk]
reason=batch_gt1_index_q_length_mismatch
batch_size=1 layer_id=0 cursor=5056 q_tokens=4995 weights_tokens=4995
```

关键观察：

- 失败请求仍是 page-tail 场景；本 rank valid q rows 为 `4995`，compute/page-tail rows 为 `5056`；
- `_get_topk_in_seq_cp_pair_batch()` 旧逻辑直接选择 `request_compute_seq_q_*`
  或 `request_actual_seq_q_*`，而 `request_actual_seq_q_*` 当前也是 compute alias；
- 但 `forward_mla.py` 中 indexer 的 q/weights 来自已经 CP split 后的 `hidden_states/q_lora` 投影，
  实际输入是 valid rows `4995`，不是 compute-padded rows `5056`；
- 因此 top-k batch path 不能从 metadata 名字推断 q layout，必须用实际
  `q_fp8.shape[0] / weights.shape[0]` 选择 valid 或 compute 视图。

修正合同：

- CP batch plan 同时暴露 valid 和 compute q metadata；
- index top-k 在 runtime 入口选择：
  - 如果 `q_tokens == sum(request_valid_seq_q_prev/next)`，按 valid segments 消费；
  - 如果 `q_tokens == sum(request_compute_seq_q_prev/next)`，按 compute segments 消费，
    但仍用 valid segment length 过滤 dummy page-tail rows；
  - 如果两者都不匹配，fail-fast 并同时打印 valid/compute 期望长度；
- 这避免把 page-tail dummy rows 强行加到 indexer q/weights 上，也避免未来真的传入
  compute-padded q 时丢失支持。

新增实现：

- `BatchTopKQueryLengths`
- `_select_batch_topk_query_lengths()`

新增回归：

- `test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias`
  - 复现线上同型：`extend_len=40387,page_size=64,cp_size=8,cp_rank=0`；
  - valid local rows = `4995`，compute local rows = `5056`；
  - 修复前 RED：helper 不存在 / 原 runtime 只能按 compute cursor；
  - 修复后 GREEN：q rows 为 4995 时选择 valid metadata；q rows 为 5056 时仍选择 compute metadata。

同类路径复核：

- `nsa_backend.forward_extend()` 不直接读取 `request_compute_seq_q_*`；它只把
  `topk_indices` pad 到当前 q 行数，因此由 indexer 返回 rows 与 q rows 对齐即可；
- `select_cp_local_valid_rows_for_cache_write()` 仍只处理 cache write 的 compute rows -> valid rows，
  不参与 top-k metadata 选择。

远端验证：

```text
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias
=> 1 passed, 5 warnings

PYTHONPATH=python python -m pytest -q \
  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 \
  -k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
=> 228 passed, 1 deselected, 5 warnings, 2 subtests passed
```

仍需 ETE 验证：

- 需要重启远端服务并重放真实请求，确认已经越过
  `batch_gt1_index_q_length_mismatch`；
- 日志中的 `[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] current_index_k must be uint8`
  是另一个性能 fast-path dtype 问题，本次未修；它当前是 warning fallback，不是本次进程退出原因。

## 26. 2026-06-06 bs>1 partial-current prefix span 修正

用户确认：`bs=1` 走 batch-plan 路径是预期行为，不应作为回退或禁用目标。

本轮性能排查确认一个明确问题：bs>1 的 partial-current 同步 compose 之前使用
`build_batch_prefix_slot_span()` 把 row-major page table 中多个 request 的 prefix 区间压成一个
bounding span。这个 span 会覆盖 request 行之间的 current/gap slot，导致 index/MLA prefix
materialize 后对并非 prefix 的页一起做 all-reduce。在线上 cache-hit-heavy、prefix 长度不一致的
batch 中，这会把本应只覆盖 prefix pages 的通信量放大。

修正：
- 保留旧的 `build_batch_prefix_slot_span()` 作为单 bounding span helper。
- 新增 `build_batch_prefix_slot_spans()`，返回每个 request prefix 的精确 slot span，并只 merge 相邻区间。
- `materialize_prefix_and_reuse_current_kv_page_slots()` /
  `materialize_prefix_and_reuse_current_index_page_slots()` 支持 `prefix_slot_spans`。
- MLA 与 index 的 bs>1 partial-current sync compose 使用精确 prefix spans；bs=1 或单 span 仍保留 IPC fast path。

权衡：
- 可能把一次 prefix all-reduce 变成至多 batch_size 次 range all-reduce；但这不超过逐 request 执行的 collective 次数，且避免 row-major gap 过通信。
- 这只修正同步 compose 的过通信；bs>1 仍有 per-layer Python descriptor 开销、bs>1 L1 async prefetch 未支持等潜在性能项，后续需要用 ETE/benchmark 继续确认。

验证：
- 远端 `g0034` container `/sgl-workspace/sglang-tai` py_compile 通过。
- 远端 `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`：113 passed, 5 warnings, 2 subtests passed。
