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

Make CP shared-KV direct writes enforce batch-local ownership

Browse Source 
W3 needs batch-size>1 extends to use packed valid tokens while preserving per-request page boundaries. The local out_cache_loc planner now validates against the batch plan's request lengths instead of the first request's scalar split metadata, then reuses the existing batch split helper to produce this rank's logical/physical cache locs.

Direct-write failures inside the CP shared-KV contract now fail fast instead of silently falling back to legacy index/MLA stores. This exposes allocator, owner-lane, page-alignment, and shape-contract bugs early for both bs=1 and bs>1.

Constraint: bs>1 batching must not pad short requests to the longest request; only per-request page-boundary padding is allowed.

Rejected: Keep bs=1 compatibility fallback | it hides CP shared-KV contract violations and caused repeated slow-path ambiguity.

Rejected: Pad batch to max request length | wastes compute and complicates cache validity for short extends.

Confidence: high

Scope-risk: moderate

Directive: CP shared-KV contract errors should stay fail-fast; do not reintroduce silent direct-write fallback without ETE evidence and explicit warning semantics.

Tested: Remote g0034 py_compile for utils.py nsa_indexer.py forward_mla.py

Tested: Remote g0034 PYTHONPATH=python pytest test/registered/unit/layers/test_nsa_cp_utils.py -> 43 passed

Tested: Remote g0034 PYTHONPATH=python pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -> 170 passed, 2 subtests passed

Not-tested: Full ETE bs>1 serving run with live traffic
This commit is contained in:
laoyao0822
2026-06-03 02:13:57 +08:00
parent e4cf8d18b4
commit f8b4f1915e
6 changed files with 837 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md
View File

@@ -9,6 +9,7 @@
**当前约束：** 先保证 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`。
- W3/W4 target sync 细化文档：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md`。
---

637
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md Normal file
View File

@@ -0,0 +1,637 @@
# NSA Prefill CP Shared-KV bs>1 W3/W4 实现文档
> 日期：2026-06-03
> 分支：`cjy-cp-refactor`
> 当前基线：`e4cf8d18b`
> 范围：W3 `local out_cache_loc + direct write`，W4 `target index/top-k sync correctness`。
> **命名说明：** 本文的 W4 指顺序实现计划中的 **Phase 4 target index/top-k sync correctness**。并行派工文档里的 W4 是 current/partial-current reuse；那部分在本文中仍视为后续阶段，不在本轮实现范围内。
## 0. 目标和非目标
目标是在 **不等待 W2 allocator 最终完成** 的前提下，先把 target model 的 W3/W4 runtime consumer 做成 batch-aware：
1. bs>1 的 CP shared-KV direct write 使用每个 request 独立的 page-aligned split，再按 request order 拼接本 rank local rows。
2. MLA KV direct write 和 index KV direct write 共享同一个 ForwardBatch 级 page/local-loc plan。
3. target index/top-k sync path 使用 per-request metadata，不能把 batch flatten 成一条长序列。
4. W4 第一版只做 sync correctness，不启用 bs>1 current reuse、partial-current reuse、L1 prefetch、draft/EAGLE。
非目标：
- 不实现 W2 owner-lane allocator；W3 单测使用 synthetic owner-valid `out_cache_loc`。
- 不实现 bs>1 current/partial-current reuse。
- 不实现 batched MLA/index prefetch。
- 不实现 draft/EAGLE bs>1。
- 不新增 collective 来同步 batch plan。
## 1. 必须保持的核心语义
CP shared-KV 下 page 相关规划是 **request/batch 级别**，不是 layer 级别：
```text
同一个 ForwardBatch 内：
logical page id / owner / page table / local loc / physical remap 不随 layer 变化
每个 layer 内：
只有 KV bytes、index K/scale、materialized dense buffer 内容、topk/logits 变化
```
因此 W3/W4 的实现不能每层重新推导 page plan。应该在 `ForwardBatch` 上缓存：
- local logical `out_cache_loc`；
- local physical `out_cache_loc`；
- paged slot remap；
- request-local prev/next segment offsets；
- request-local page/table descriptor。
当前已有缓存字段：
- `ForwardBatch.cp_local_out_cache_loc` 和 `cp_local_physical_out_cache_loc` 已在 `forward_batch_info.py:427-428`。
- `ForwardBatch.cp_shared_kv_paged_slot_remap_key/remap` 已在 `forward_batch_info.py:432-433`。
- `get_or_build_shared_paged_buffer_slot_remap()` 已在 `cp_shared_kv_runtime.py:2486-2520` 按 key 缓存 paged remap。
W3/W4 的实现原则：
```text
prepare/batch plan 阶段确定 page 和 segment。
direct write/top-k/materialize 只消费 plan。
layer_id 只选择 layer buffer，不参与 page planning。
```
bs>1 的 padding 合同：
```text
不把短 request pad 到 batch 内最长 request 的长度。
每个 request 只独立 pad 到 page_size 边界：
valid_tokens=100, page_size=64 -> padded_tokens=128
ForwardBatch 内计算输入仍是 sum(request_extend_lens) 的 packed/ragged tokens。
request 边界由 request_extend_lens / request_split_lists / cu_seqlens 保留。
如果某条 fast path 需要等长，第一版应改为 segmented/per-request descriptor，
不能通过 max-length padding 掩盖接口缺陷。
```
## 2. 已确认的代码事实
### C1. W1 batch metadata 已经存在
`CPSharedKVBatchPlan` 在 `nsa/utils.py:280-312`，包含：
- `request_extend_lens`
- `request_prefix_lens`
- `request_padded_pages/tokens`
- `request_split_lists`
- `request_zigzag_indices`
- `request_kv_len_prev/next`
- `request_actual_seq_q_prev/next`
- `request_rank_local_offsets`
- `flat_*` views
`build_batch_page_aligned_in_seq_split_plan()` 在 `nsa/utils.py:335-490`，每个 request 独立 page-align split，再 flatten。
`prepare_input_dp_with_cp_dsa()` 在 `nsa/utils.py:1579-1646`，当 CP shared-KV 且 `len(extend_seq_lens_cpu) > 1` 时会构造 batch plan，并通过 `_build_batch_metadata_from_plan()` 写入 `NSAContextParallelMetadata`。
### C2. batch split helper 已经存在，但 W3 local loc 尚未完全使用
`split_tensor_by_cp_batch_plan()` 在 `nsa/utils.py:507-560`，会按 `request_extend_lens` 切 request，再按每个 request 的 `request_split_lists/request_zigzag_indices` 拼本 rank local rows。
`cp_split_and_rebuild_data()` / `cp_split_and_rebuild_1d()` 已在 `nsa/utils.py:951` 附近 dispatch 到 `_cp_split_and_rebuild_batch_in_seq()`。
### C3. 当前 W3 最大疏漏：`get_cp_shared_kv_local_out_cache_loc()` 仍用 scalar split 做长度检查
`get_cp_shared_kv_local_out_cache_loc()` 在 `nsa/utils.py:1039-1124`。
当前代码在 `nsa/utils.py:1093-1102`：
```python
split_tokens = sum(int(x) for x in metadata.split_list)
out_cache_tokens = int(out_cache_loc.numel())
if split_tokens != out_cache_tokens:
...
return None
```
bs>1 下 `metadata.split_list` 是 `_build_batch_metadata_from_plan()` 写入的 **first request scalar compatibility field**，见 `nsa/utils.py:728-741`。因此这个检查会错误地把 bs>1 batch 判定为 mismatch，导致 direct write fallback。
修复原则：
```text
if batch_plan.batch_size > 1:
expected = sum(batch_plan.request_extend_lens)
else:
expected = sum(metadata.split_list)
```
然后仍然使用 `cp_split_and_rebuild_1d()` / `split_tensor_by_cp_batch_plan()` 得到 local logical loc。
### C4. 当前 physical remap 已经是 ForwardBatch 级缓存
`get_cp_shared_kv_local_physical_out_cache_loc()` 在 `nsa/utils.py:1127-1173`，注释已经写明 logical loc 是 batch-scoped, not layer-scoped，并缓存到 `forward_batch.cp_local_physical_out_cache_loc`。
W3 需要保留这个设计。不要把 physical remap 做成 per-layer。
### C5. MLA direct write 主要依赖 local loc helper
`_maybe_write_cp_shared_local_mla_kv()` 在 `forward_mla.py:551-623`。
流程：
1. 通过 `get_cp_shared_kv_local_out_cache_loc()` 取得 local logical loc。
2. 校验 `k_nope/k_pe.shape[0] == local_out_cache_loc.numel()`。
3. 优先 `try_tai_fused_mla_store(...)`。
4. 否则用 `get_cp_shared_kv_local_physical_out_cache_loc()` 后 `set_mla_kv_buffer()`。
因此 W3 对 MLA 的主要实现面是 local loc helper 和测试。
### C6. index direct write 失败后目前会 silent 落回 legacy store
`_store_cp_shared_local_index_k_cache()` 在 `nsa_indexer.py:1606-1652`。
如果 local loc 不可用或 shape mismatch，它返回 `False`。调用点在 `nsa_indexer.py:1720-1731`、`1745-1756`、`1760-1771` 会继续调用 `_store_index_k_cache()`。
`_store_index_k_cache()` 如果未传 `out_loc_override`，会走 `_filter_shared_index_write()`，后者在 `nsa_indexer.py:535-579` 直接基于完整 `forward_batch.out_cache_loc` filter owned logical loc。
这对历史 bs=1 是兼容 fallback，但对当前 CP shared-KV/page-aligned/direct-write 合同不合适：W3 支持后，bs=1/bs>1 都不应 silent fallback 到 legacy filter path。否则即使 allocator/page plan 有问题，也可能被 fallback 掩盖。
W3 要求：
```text
CP shared-KV 合同内 + local loc unavailable / owner mismatch / shape mismatch:
index/MLA direct write 必须 fail-fast
非 CP shared-KV 路由:
可以返回 False 走原路径；这不是 CP shared-KV fallback
```
### C7. 当前 W4 `_get_topk_in_seq_cp_pair()` 是单请求
`_get_topk_in_seq_cp_pair()` 在 `nsa_indexer.py:1371-1455`。
它读取 scalar：
- `metadata.kv_len_prev`
- `metadata.kv_len_next`
- `metadata.actual_seq_q_prev`
- `metadata.actual_seq_q_next`
并用 `split_in_seq_cp_local_pair()` 把整个 local tensor 切成 prev/next 两段。
bs>1 下 local tensor 实际应是：
```text
req0_prev, req0_next, req1_prev, req1_next, ...
```
不能用 first request 的 scalar lengths 切整批。
### C8. `_get_topk_ragged_with_cp()` 的非 cp_index 分支硬编码 batch 0
`_get_topk_ragged_with_cp()` 在 `nsa_indexer.py:1087-1369`。
非 `cp_index` 分支在 `nsa_indexer.py:1224-1232` 使用：
```python
forward_batch.seq_lens_cpu[0]
forward_batch.extend_seq_lens_cpu[0]
```
并在 `nsa_indexer.py:1254-1287` 使用：
```python
block_tables[0]
```
W4 如果想复用这个函数处理 bs>1 的每个 request，就必须新增 `batch_idx: int = 0` 参数，并用 `block_tables[batch_idx]`、`seq_lens_cpu[batch_idx]`、`extend_seq_lens_cpu[batch_idx]`。
### C9. 不应使用现有 `cp_index` 分支作为 W4 第一版
`_get_topk_ragged_with_cp()` 的 `cp_index` 分支在 `nsa_indexer.py:1155-1222`，源码注释明确写着 `TODO Multi-batch support has accuracy issues`。
W4 第一版为了 correctness 应该避开该分支，采用 per-request/per-segment 同步调用：
```text
for req_id:
call prev segment
call next segment
```
后续性能优化可以再把多个 segment 合并成 batched top-k descriptor。
### C10. W4 current/partial-current reuse 仍是单请求合同
`_maybe_materialize_shared_index_buffer()` 在 `nsa_indexer.py:309-517`。
当 `current_index_kv is not None` 时，它要求：
```python
len(prefix_lens_cpu) == 1
positive page-aligned prefix
```
见 `nsa_indexer.py:341-355`。
因此 W4 第一版遇到 bs>1 + `current_index_kv is not None` 必须 fail-fast：
```text
[CP_SHARED_KV_FAIL_FAST][batch_gt1_index_current_reuse_unsupported]
```
不要在 W4 里临时拼 current reuse，否则会和 W5 的 partial/current reuse 工作混在一起。
### C11. `_build_batch_metadata_from_plan()` 的 batch cu tensor 不能直接传给单 segment top-k
`_build_batch_metadata_from_plan()` 在 `nsa/utils.py:832-840` 构造：
```python
request_actual_seq_q_prev_cu_tensor = [0] + cumsum(request_actual_seq_q_prev)
request_actual_seq_q_next_cu_tensor = [0] + cumsum(request_actual_seq_q_next)
```
这是 batch-level cumulative tensor，不是单个 request segment 的 `[0, segment_len]`。
W4 的 per-request/per-segment top-k 调用应该传：
```text
actual_seq_q_tensor = tensor([segment_len])
actual_seq_q_cu_tensor = tensor([0, segment_len])
```
第一版可以每次构造小 tensor；后续再优化为 metadata 预构造或 TAI descriptor。
## 3. W3 设计：batch-aware local out_cache_loc + direct write
### 3.1 输入合同
W3 依赖以下输入：
- `forward_batch.uses_cp_shared_kv == True`
- `forward_batch.cp_shared_kv_layout is not None`
- `forward_batch.nsa_cp_metadata` 包含 batch plan 或 batch fields
- `forward_batch.out_cache_loc` 是按 request order flatten 的 **valid-token loc view**
关于 W2 allocator 的交接：
```text
W3 第一版假设 out_cache_loc.numel() == sum(request_extend_lens)
```
如果 W2 allocator 以后返回 padded-token loc，则 W2 必须同时提供 valid view，或者 W3 增加 `request_valid_out_cache_loc` view；不能让 direct write 写 padded tail rows。
### 3.2 输出合同
`get_cp_shared_kv_local_out_cache_loc(forward_batch)` 返回：
```text
local logical loc = concat(local(req0), local(req1), ...)
```
其中每个 `local(req)` 按该 request 的 `request_zigzag_indices[req]` 取两个 in-seq local segment。
`get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)` 返回：
```text
layout.logical_locs_to_physical(local logical loc)
```
并在 ForwardBatch 上缓存。
### 3.3 实现步骤
#### W3-S1：新增 batch-local-loc helper
建议在 `nsa/utils.py` 增加轻量 helper：
```python
def _get_cp_shared_kv_expected_out_cache_tokens(metadata) -> int:
plan = get_cp_shared_kv_batch_plan(...)
if plan is not None and plan.batch_size > 1:
return sum(plan.request_extend_lens)
return sum(metadata.split_list)
```
也可以直接写在 `get_cp_shared_kv_local_out_cache_loc()` 里，但 helper 更容易测。
#### W3-S2：修改 length check
将 scalar-only 检查替换成：
```text
batch plan exists:
expected_tokens = sum(request_extend_lens)
log reason = batch_out_cache_len_mismatch
else:
expected_tokens = sum(split_list)
```
#### W3-S3：使用 batch split helper 生成 local logical loc
继续调用：
```python
local_out_cache_loc = cp_split_and_rebuild_1d(forward_batch, out_cache_loc.contiguous())
```
因为 `cp_split_and_rebuild_1d()` 已经会在 `metadata.batch_size > 1` 时进入 `split_tensor_by_cp_batch_plan()`。
#### W3-S4：保留 owner validation
保留：
```python
valid_locs = local_out_cache_loc[local_out_cache_loc > 0]
layout.owned_by_this_rank(valid_locs)
```
这个校验是 W2 allocator 和 W3 consumer 的关键集成检查。
#### W3-S5：index direct-write 不允许 silent fallback
在 `_store_cp_shared_local_index_k_cache()` 中，当：
```text
forward_batch.uses_cp_shared_kv
local_out_loc is None or shape mismatch
```
必须 fail-fast，而不是返回 `False` 进入 `_store_index_k_cache()` legacy path。bs=1 和 bs>1 一致；非 CP shared-KV 可以继续返回 `False` 走原路径。
#### W3-S6：MLA direct-write 不允许 silent fallback
`_maybe_write_cp_shared_local_mla_kv()` 当前返回 `False` 后，前向会触发 legacy rebuild/materialize。CP shared-KV 合同内如果 local loc 缺失或 shape mismatch，应按同样规则 fail-fast。
建议第一版：
```text
bs>1 + local loc unavailable -> RuntimeError [CP_SHARED_KV_FAIL_FAST][batch_gt1_mla_direct_write_unavailable]
bs>1 + shape mismatch -> RuntimeError [CP_SHARED_KV_FAIL_FAST][batch_gt1_mla_direct_write_shape_mismatch]
```
### 3.4 W3 测试
#### T1：batch local loc 保持 request boundary
构造：
```text
page_size=4, cp_size=2, cp_rank=1
plan:
req0 split [4,0,0,0], zigzag [1,2] -> local []
req1 split [4,4,1,0], zigzag [1,2] -> local req1 segment1 + segment2
```
或选择 cp_rank=0，构造 req0/req1 都有 local rows，便于断言：
```text
local_out_cache_loc == local(req0) + local(req1)
```
#### T2：batch owner mismatch 拒绝 direct write
synthetic `out_cache_loc` 中让某个 local page 的 `(page_id - 1) % cp_size` 不等于 `cp_rank`。
期望：
```text
get_cp_shared_kv_local_out_cache_loc() is None
warning reason local_loc_owner_mismatch
```
如果实现采用 fail-fast，则断言对应 fail-fast reason。
#### T3：batch physical loc 缓存
调用两次 `get_cp_shared_kv_local_physical_out_cache_loc()`，断言返回同一对象，并且 loc 映射符合 `CpSharedKVLayout.logical_locs_to_physical()`。
#### T4：index direct write 不 silent fallback
mock `_store_index_k_cache()`，在 bs>1 local loc 不可用时确保不会调用 legacy fallback。
#### T5：MLA direct write shape mismatch fail-fast
构造 bs>1 local loc tokens=N，但 `k_nope/k_pe` tokens != N，断言 fail-fast。
## 4. W4 设计：target index/top-k sync correctness
### 4.1 输入合同
W4 依赖：
- W1 batch metadata；
- W3 local q/weights 已按 request boundary split；
- `metadata.get_page_table_64()` 可返回 batch page table；
- 不启用 bs>1 current reuse / partial-current reuse。
### 4.2 输出合同
`_get_topk_in_seq_cp_pair()` 在 bs>1 下返回：
```text
concat(
topk(req0_prev),
topk(req0_next),
topk(req1_prev),
topk(req1_next),
...
)
```
这个顺序必须匹配 local q/weights 顺序，也就是 `split_tensor_by_cp_batch_plan()` 的输出顺序。
### 4.3 实现步骤
#### W4-S1：给 `_get_topk_in_seq_cp_pair()` 增加 batch dispatch
```python
if getattr(metadata, "batch_size", 1) > 1:
return self._get_topk_in_seq_cp_pair_batch(...)
```
保留原 scalar path。
#### W4-S2：新增 `_get_topk_in_seq_cp_pair_batch()`
伪代码：
```python
plan = get_cp_shared_kv_batch_plan(forward_batch)
if current_index_kv is not None:
fail-fast batch_gt1_index_current_reuse_unsupported
shared_block_tables = metadata.get_page_table_64()
shared_index_buffer, shared_block_tables = self._maybe_materialize_shared_index_buffer(
forward_batch, layer_id, shared_block_tables
)
cursor = 0
outputs = []
for req_id in range(plan.batch_size):
prev_len = metadata.request_actual_seq_q_prev[req_id]
next_len = metadata.request_actual_seq_q_next[req_id]
q_prev = q_fp8[cursor: cursor + prev_len]
w_prev = weights[cursor: cursor + prev_len]
cursor += prev_len
q_next = q_fp8[cursor: cursor + next_len]
w_next = weights[cursor: cursor + next_len]
cursor += next_len
outputs.append(call_segment(req_id, prev, prev_len, kv_len_prev))
outputs.append(call_segment(req_id, next, next_len, kv_len_next))
assert cursor == q_fp8.shape[0]
return cat(outputs)
```
zero segment：
```text
segment_len == 0 -> 返回 shape (0, index_topk) 的 empty tensor，不调用底层 MQA/topk
```
#### W4-S3：泛化 `_get_topk_ragged_with_cp(..., batch_idx=0)`
新增参数：
```python
batch_idx: int = 0
```
并把非 `cp_index` 分支里的 hard-coded `[0]` 改成 `[batch_idx]`：
- `forward_batch.seq_lens_cpu[batch_idx]`
- `forward_batch.extend_seq_lens_cpu[batch_idx]`
- `block_tables[batch_idx]`
scalar path 默认 `batch_idx=0`，行为不变。
#### W4-S4：单 segment cu tensor
每个 segment 调用 `_get_topk_ragged_with_cp()` 时传：
```python
actual_seq_q_tensor = torch.tensor([segment_len], device=q_fp8.device, dtype=torch.int32)
actual_seq_q_cu_tensor = torch.tensor([0, segment_len], device=q_fp8.device, dtype=torch.int32)
```
不要传 `request_actual_seq_q_prev_cu_tensor` 的 batch cumulative view。
#### W4-S5：不要使用 `cp_index` 分支
现有 `cp_index` 分支虽然看起来能表达 multi-batch，但源码已标注 accuracy issue。W4 第一版只追 correctness，明确不使用。
### 4.4 W4 测试
#### T1：bs>1 per-request prev/next 调用顺序
mock `_get_topk_ragged_with_cp()`，记录：
- `batch_idx`
- `kv_len`
- `actual_seq_q`
- q slice 内容
- weights slice 内容
- `actual_seq_q_cu_tensor`
断言调用顺序：
```text
req0 prev
req0 next
req1 prev
req1 next
```
#### T2：materialize 只做一次
mock `_maybe_materialize_shared_index_buffer()`，断言 bs>1 top-k 只调用一次 materialize，并且所有 segment 复用同一个 `shared_index_buffer/shared_block_tables`。
#### T3：zero segment 不调用 top-k kernel
构造某个 request 的 `actual_seq_q_next=0`，断言 `_get_topk_ragged_with_cp()` 不被调用，输出中该 segment 是 empty。
#### T4：bs>1 current_index_kv fail-fast
传入 `current_index_kv`，断言：
```text
[CP_SHARED_KV_FAIL_FAST][batch_gt1_index_current_reuse_unsupported]
```
#### T5：scalar tests 保持
现有：
- `test_indexer_in_seq_cp_pair_materializes_index_once_for_prev_next`
- `test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused`
必须继续通过。
## 5. 执行顺序
建议按以下顺序实施，避免把 W3/W4 的 bug 混在一起：
1. W3 tests：batch local loc / owner mismatch / physical loc cache。
2. W3 implementation：修 `get_cp_shared_kv_local_out_cache_loc()` 的 bs>1 expected length。
3. W3 direct-write guard：index/MLA bs>1 local loc unavailable 不 silent fallback。
4. W4 tests：batch top-k per-request order / materialize once / zero segment / current reuse fail-fast。
5. W4 implementation：batch dispatch + `_get_topk_ragged_with_cp(batch_idx=...)`。
6. 只跑 unit，不跑 ETE；W2 allocator 未合入前，ETE 仍可能被 allocator blocker 阻塞。
## 6. 验证命令
本地：
```bash
python -m py_compile \
python/sglang/srt/layers/attention/nsa/utils.py \
python/sglang/srt/layers/attention/nsa/nsa_indexer.py \
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py \
test/registered/unit/layers/test_nsa_cp_utils.py
```
远端容器：
```bash
cd /sgl-workspace/sglang-tai
PYTHONPATH=python python -m pytest -q \
test/registered/unit/layers/test_nsa_cp_utils.py
```
必要时追加：
```bash
PYTHONPATH=python python -m pytest -q \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
```
## 7. 风险和后续
### R1. W2 allocator 输出可能是 padded-token loc
W3 当前应按 valid-token loc 设计。如果 W2 返回 padded-token loc，direct write 不能直接消费，否则会把 padding tail 写入真实 KV/index。
需要在 W2/W3 集成时确认：
```text
out_cache_loc.numel() == sum(request_extend_lens)
```
或者 W2 提供 valid loc view。
### R2. per-segment top-k 会增加调用次数
W4 第一版 bs=N 会最多调用 `2N` 次 `_get_topk_ragged_with_cp()`。这是 correctness-first 设计。后续性能优化可以引入 batched segment descriptor 或修复 `cp_index` 分支，但不能在第一版混入。
### R3. current reuse 不支持会影响 cache-hit 性能
W4 只保证 sync correctness。cache-hit 的性能收益需要 W5 current/partial-current reuse 恢复后再评估。
### R4. fallback 策略需要保持醒目
W3/W4 支持范围内不能 silent fallback。尤其 index direct-write 当前 fallback 很隐蔽，必须在 bs>1 CP shared-KV 下收窄。
### R5. page descriptor 必须 layer-invariant
如果实现中发现某个 page/table/remap descriptor 每层重建，应记录为性能风险并优先改成 ForwardBatch 缓存。