Expand prefill CP KV capacity by sharding persistent NSA KV

Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer.

Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary.

Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization
Constraint: Decode side remains non-CP and receives full KV through Mooncake
Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts
Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode
Confidence: medium
Scope-risk: broad
Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode
Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics
Tested: Remote py_compile for shared-KV touched Python files in g0034 container
Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected
Not-tested: Full GLM5 multi-node throughput/regression run after final doc update
Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends

docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md

@@ -6,6 +6,45 @@ Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它
 
 ---
 
+## 0. 当前实现状态
+
+当前代码已经落地 Phase 2 的首版兼容实现，启用开关为：
+
+```bash
+--enable-nsa-prefill-cp-shared-kv
+```
+
+首版实现的关键语义：
+
+1. **persistent KV at rest 是 sharded/shared 的**：
+   - `req_to_token`、radix cache、scheduler admission 使用 CP group 统一的 **logical loc/page**。
+   - 每个 prefill CP rank 的 `NSATokenToKVPool` 只按本 rank **physical capacity** 分配。
+   - 每个 rank 只写入自己 owner 的 logical page，对应本地 physical page。
+2. **MLA latent KV 和 NSA index K/scale 同时 shard**：
+   - latent KV 写入前会把 `out_cache_loc` 从 logical loc 过滤为本 rank owner loc，并转换为 physical loc。
+   - NSA indexer K/scale 写入使用同一 owner mapping，避免 index cache 与 latent KV 不一致。
+3. **attention runtime 仍走 full-view compatibility**：
+   - 现有 NSA attention/indexer kernel 仍假设可读 full KV/index view。
+   - shared KV 下运行时会 materialize dense full-view buffer：每 rank 拷贝本地 owner pages，其余填 0，然后 CP group all-reduce 得到完整 runtime view。
+   - 这是 Phase 2 的主要性能成本；Phase 3 才会改成 shard-aware attention/topk。
+4. **Mooncake PD transfer 已支持 prefill CP shards -> decode full KV**：
+   - prefill CP rank 发送本 rank owner physical pages。
+   - transfer chunk 携带 `logical_page_positions`，decode 侧用这些位置选择 full-layout dst pages。
+   - 因 decode 不开启 CP，prefill PD 必须设置 `SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1`。
+5. **debug/assert 默认不影响性能**：
+   - `SGLANG_DEBUG_CP_SHARED_KV=0` 时不执行 expensive checksum、CUDA tensor predicate、topk validation。
+   - `SGLANG_DEBUG_CP_SHARED_KV=1` 仅用于定位 correctness 问题，不应参与性能压测。
+
+当前仍未完成、属于 Phase 3 或后续工作：
+
+- runtime 不再 materialize full/maxlen KV。
+- NSA index/topk/attention 的 shard-aware 计算与 global topk merge。
+- `nsa_prefill_cp_mode=round-robin` 的 runtime wiring。
+- decode 侧 CP/shared KV。
+- Mooncake 之外的 PD transfer backend 完整支持。
+
+---
+
 ## 1. 适用范围
 
 Phase 2 首版只覆盖以下组合：
@@ -608,7 +647,7 @@ PD disaggregation 下，如果 prompt KV 长度超过 decode side physical KV po
 
 ## 8. 配置与保护
 
-建议新增显式开关：
+新增显式开关：
 
 ```text
 --enable-nsa-prefill-cp-shared-kv
@@ -629,7 +668,7 @@ PD disaggregation 下，如果 prompt KV 长度超过 decode side physical KV po
 
 如果任一条件不满足，启动时报明确错误。
 
-建议日志：
+建议/当前日志：
 
 ```text
 CP shared KV enabled: physical_tokens_per_rank=237312, logical_tokens=1898496, cp_size=8, shard_policy=page_interleaved
@@ -637,6 +676,37 @@ CP shared KV PD transfer: all CP ranks transfer enabled, backend=mooncake
 CP shared KV attention compatibility: full-view materialization enabled; Phase3 required for shard-aware runtime
 ```
 
+### 8.1 调试环境变量
+
+```text
+SGLANG_DEBUG_CP_SHARED_KV=1
+```
+
+用途：
+
+1. 打印 CP shared KV 相关 debug 日志，包括 persistent KV shard write、NSA index materialize、Mooncake sender filtering、runtime full-view materialize 前后 checksum。
+2. 开启 shared KV 源头定位 assert，用于区分合法 `-1` sentinel 和真正异常的负 logical page/loc：
+   - attention token materialize path 允许 `logical_locs == -1`，这是 NSA topk/page table 的 invalid sentinel。
+   - attention token materialize path 拒绝 `logical_locs < -1`。
+   - token materialize 的 `remap_logical_locs` 拒绝任何 `< 0`，因为它应来自 `metadata.page_table_1` / `req_to_token`。
+   - paged buffer materialize 拒绝 `logical_pages < 0`，因为 NSA index page table / real page table 不应携带 topk sentinel。
+   - persistent MLA KV / NSA index write filtering 拒绝 `out_cache_loc < 0`。
+
+性能影响：
+
+- 默认值为 `0`，正常线上路径只多一个 Python debug 分支，不执行 GPU tensor predicate，不引入 CUDA 同步。
+- 设置为 `1` 后，部分 assert 会对 CUDA tensor 执行 `torch.any()` / `min()` / `max()`，Python 控制流会触发同步；只能用于问题定位，不应参与性能压测。
+- Mooncake PD transfer 负页检查是 always-on CPU numpy 检查，用于防止异常 page list 被发送；成本相对 RDMA 传输可以忽略。
+
+建议使用方式：
+
+```bash
+SGLANG_DEBUG_CP_SHARED_KV=1 \
+python -m sglang.launch_server ...
+```
+
+当定位到负值源头后，应关闭该变量重新做吞吐和延迟验证。
+
 ---
 
 ## 9. 验证计划