# NSA Prefill CP：GLM-5 EAGLE/MTP Draft Shared-KV 支持计划

## 目标

本计划在当前 `cp-hicache-host` 分支上，为 GLM-5 / `GlmMoeDsaForCausalLM` 的 EAGLE/MTP draft prefill 恢复 CP shared-KV 语义：

```text
prefill target model: NSA CP + shared KV + narrow logits path
prefill draft model:  NSA CP + shared KV + CP-local token compute
prefill -> decode:    transfer target KV 与 draft KV 的 shared-KV shards
decode:               仍按当前 decode DP 路径执行，不引入 decode CP
```

直接目标是让 draft model 不再在 prefill 侧按每个 rank 保存完整 request KV / hidden，而是和 target model 一样按 CP owner shard 写 persistent KV，同时避免 EAGLE target hidden capture 破坏 target narrow output collection。

## 当前范围

### 纳入范围

当前只支持 GLM-5 / DSA 路径：

```text
GlmMoeDsaForCausalLM
  -> draft architecture remap: DeepseekV3ForCausalLMNextN
  -> python/sglang/srt/models/deepseek_nextn.py
  -> NSA / MLA / CP shared-KV path
```

代码依据：

- `python/sglang/srt/configs/model_config.py`
  - `GlmMoeDsaForCausalLM` draft 被 remap 到 `DeepseekV3ForCausalLMNextN`。
- `python/sglang/srt/models/glm4_moe.py`
  - `GlmMoeDsaForCausalLM` 继承 `DeepseekV2ForCausalLM`。
- `python/sglang/srt/models/deepseek_nextn.py`
  - `DeepseekV3ForCausalLMNextN` 当前已有 NSA CP metadata / split / all-gather 相关骨架。

### 不纳入范围

不在本阶段支持 GLM4 MHA/Radix draft：

```text
Glm4MoeForCausalLM / Glm4MoeLiteForCausalLM
  -> draft architecture remap: Glm4MoeForCausalLMNextN
  -> RadixAttention / MHA path
```

原因：GLM4 draft 走 `Glm4MoeForCausalLMNextN`，attention 是 MHA/Radix 路径，不是当前 GLM-5 DSA 的 NSA/MLA shared-KV 路径。MHA/Radix 的 KV write 当前主要是直接按 loc 写 `k_cache[indices] = k`、`v_cache[indices] = v`，需要单独补 owner filtering 与 logical->physical remap，和本阶段目标解耦。

## 当前问题

### 1. EAGLE target 强制 FULL hidden，禁用 target narrow path

文件：

- `python/sglang/srt/speculative/eagle_worker.py`
- `python/sglang/srt/speculative/multi_layer_eagle_worker.py`
- `python/sglang/srt/speculative/eagle_worker_v2.py`
- `python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py`
- `python/sglang/srt/models/deepseek_v2.py`

当前 EAGLE target extend 会执行：

```python
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
```

这会让 target model 走 full hidden capture。target CP narrow path 的 gate 里又明确要求：

```python
forward_batch.capture_hidden_mode == CaptureHiddenMode.NULL
```

结果是：EAGLE 开启时 target 无法走 `cp_collect_last_token_hidden(...)` narrow path，而是回退到 `cp_all_gather_rerange_output(...)`。

### 2. Deepseek nextn draft CP split 发生得偏晚

文件：

- `python/sglang/srt/models/deepseek_nextn.py`

当前 draft forward 近似流程：

```text
full input_ids -> embed full hidden
full target hidden -> hnorm
cat(full embed, full target hidden)
eh_proj full tokens
CP split hidden / positions
decoder local tokens
CP all-gather output hidden
logits processor
```

这有两个问题：

1. `eh_proj` 前仍然处理 full tokens，prefill 显存与计算没有按 CP 缩小。
2. draft 输出仍执行 `cp_all_gather_rerange_output(...)`，没有复用 target narrow 思路。

### 3. Draft KV transfer 当前依赖“target/draft indices 总是共享”的隐式假设

文件：

- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`
- `python/sglang/srt/managers/scheduler.py`

当前 prefill bootstrap queue 会把 draft pool 的 contiguous buffer 追加到 transfer buffer：

```text
target token_to_kv_pool contiguous buffers
+ draft token_to_kv_pool contiguous buffers
```

但在 CP shared-KV 下必须确认：

- draft pool 是否也按 physical shared-KV capacity 初始化；
- draft KV write 是否写入 local physical shard；
- prefill->decode transfer 是否覆盖所有 draft shard；
- decode 侧是否按 logical loc 正确消费 draft KV。

## 设计原则

1. **不新增 decode CP。** 当前仅 prefill 使用 CP，decode 仍按现有 DP/TP/EAGLE 逻辑执行。
2. **不破坏 target narrow logits path。** EAGLE 需要 draft hidden，但不应该因此强制 target all-gather full hidden。
3. **draft 与 target 使用同一套 CP shared-KV ownership 语义。** logical loc 仍对外暴露，persistent KV pool 内部使用 physical shard。
4. **优先复用 NSA/MLA shared-KV 代码。** 不为 GLM-5 重做一套 MHA/Radix shared-KV writer。
5. **新路径默认 gated。** 在远端验证稳定前，通过环境变量打开，避免影响非 CP / 非 EAGLE / 非 GLM-5 路径。

## 建议环境变量

新增运行时开关：

```text
SGLANG_CP_DRAFT_SHARED_KV=0/1
```

语义：

- `0`：保持现有 EAGLE/MTP draft 行为。
- `1`：在满足以下条件时启用 GLM-5 draft CP shared-KV 路径：
  - prefill CP enabled；
  - CP shared KV enabled；
  - draft architecture 是 `DeepseekV3ForCausalLMNextN`；
  - forward mode 是 extend/context-parallel extend；
  - `nsa_cp_metadata` 能够成功构造。

新增调试开关：

```text
SGLANG_CP_DRAFT_SHARED_KV_DEBUG=0/1
```


> P0 说明：上述环境变量用于 P1+ 功能路径 gated。P0 只落 probe log，不读取新增环境变量。

后续功能路径可打印调试日志，用于确认每个 rank 的：

- draft architecture；
- token_to_kv_pool class；
- logical / physical capacity；
- full token count 与 local token count；
- target hidden capture mode；
- draft hidden input shape；
- draft KV write loc 是否 local physical loc；
- prefill transfer 中是否包含 draft KV buffers。

## 实施阶段

### P0：诊断日志

目标：先确认当前实际运行路径，不改变行为。P0 只添加无环境变量控制的 probe 日志；如日志显示现有路径不兼容，后续阶段再决定是否回退或 gated。

修改文件：

- `python/sglang/srt/speculative/eagle_worker.py`
- `python/sglang/srt/speculative/multi_layer_eagle_worker.py`
- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`

工作项：

1. 在 draft runner 初始化后打印一次 architecture 与 pool 类型。
2. 在 EAGLE target extend 打印 target capture mode 与 logits output hidden shape。
3. 在 draft extend 打印 draft input hidden shape、`input_ids` shape、`seq_lens_cpu`。
4. 在 prefill/decode transfer init 打印 target/draft KV buffer 数量与 item size。
5. 在 scheduler disaggregation 初始化时打印 draft KV pool 来源。

验收：

- 不设置任何新增环境变量时行为不变。
- 日志能确认 GLM-5 draft path 是 `DeepseekV3ForCausalLMNextN`，不是 `Glm4MoeForCausalLMNextN`。
- 日志能确认 prefill/decode transfer 是否包含 draft KV buffers。


### P0 观测结论

远端 P0 日志确认当前 GLM-5 EAGLE prefill 的实际路径：

- target prefill 使用 `CaptureHiddenMode.FULL`，所有 CP rank 都得到完整 `hidden_states=(prompt_tokens, hidden_size)`，因此 target narrow path 被禁用。
- draft architecture 是 `DeepseekV3ForCausalLMNextN`，属于本计划支持范围。
- draft runner 的 KV pool 已经是 `NSATokenToKVPool` + `CPSharedPagedTokenToKVPoolAllocator`，说明 memory pool 层具备 shared-KV 基础。
- draft extend 初始化阶段仍传入完整 `input_ids` / `spec_info.hidden_states` / `out_cache_loc`，即每个 CP rank 都准备处理完整 prompt tokens。
- `has_nsa_cp_metadata=False` 出现在 `ForwardBatch.init_new(...)` 之后、draft model forward 之前；metadata 当前由 `DeepseekV3ForCausalLMNextN.forward(...)` 内部构造，所以日志点早于 metadata 创建。真正的问题不是 metadata 永远缺失，而是 draft 在 `eh_proj` 前仍处理 full tokens。

因此 P1/P2 的修复重点是：target 输出侧先保留 CP-local draft hidden；draft 模型在 `eh_proj` 前完成 CP split，并在输出侧只收集 last-token hidden/logits，而不是 full hidden。

### P1：target narrow logits 与 draft hidden 解耦

目标：EAGLE target 仍为 draft 提供 hidden，但不再通过 `CaptureHiddenMode.FULL` 禁用 target narrow path。

修改文件：

- `python/sglang/srt/speculative/eagle_worker.py`
- `python/sglang/srt/speculative/multi_layer_eagle_worker.py`
- `python/sglang/srt/model_executor/forward_batch_info.py`
- `python/sglang/srt/layers/logits_processor.py`
  - `LogitsProcessorOutput` 定义在该文件，新增 `draft_hidden_states` side-channel 字段
- `python/sglang/srt/models/deepseek_v2.py`

设计：

增加一个独立 side-channel 字段，例如：

```python
draft_hidden_states: Optional[torch.Tensor]
```

不要新增 `CaptureHiddenMode` enum 值。原因是 `CaptureHiddenMode` 被 cuda graph / piecewise runner / schedule batch 复用，新增 enum 容易改变 graph capture 分支。

目标数据流：

```text
target model local hidden after final norm
  -> for logits: cp_collect_last_token_hidden(...) narrow path
  -> for draft: keep CP-local hidden in draft_hidden_states side-channel
```

EAGLE target extend 获取：

```text
logits_output.next_token_logits      # narrow logits path
logits_output.draft_hidden_states    # CP-local hidden for draft
```

验收：

- EAGLE 开启且 `SGLANG_CP_DRAFT_SHARED_KV=1` 时，target 不再强制 `CaptureHiddenMode.FULL`。
- `DeepseekV2Model._should_use_narrow_output_path(...)` 对 target prefill 返回 true。
- draft hidden side-channel shape 是当前 CP rank local token 数，而不是 full prompt token 数。

### P2：Deepseek nextn draft 输入改为 CP-local

目标：draft model 在 `eh_proj` 前就按 CP local tokens 运行，并在 prefill 输出侧只收集 last-token hidden/logits。

修改文件：

- `python/sglang/srt/models/deepseek_nextn.py`
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `python/sglang/srt/layers/logits_processor.py`
- `python/sglang/srt/speculative/eagle_worker.py`
- `python/sglang/srt/speculative/multi_layer_eagle_worker.py`

当前顺序：

```text
embed full input
cat full target hidden
eh_proj full tokens
CP split
decoder local tokens
CP all-gather full output hidden
```

目标顺序：

```text
prepare nsa_cp_metadata from full logical token layout
split input_ids / positions with split_list + zigzag_index
use target draft_hidden_states side-channel if already CP-local; otherwise split full target hidden
embed local input
eh_proj local tokens
run draft decoder local tokens
collect only last-token hidden via cp_collect_last_token_hidden for logits/capture
```

需要明确的约束：

- `input_ids`、`positions`、`spec_info.hidden_states` 必须使用同一套 CP split；如果 P1 已经提供 CP-local hidden，则不能重复 split hidden。
- `forward_batch.out_cache_loc` 对外仍保持 full logical loc 语义；KV write 时继续通过 shared-KV helper 计算本 rank local logical loc，再 remap 到 physical loc。不要把 `forward_batch.out_cache_loc` 直接改成本地 loc，否则 radix / PD transfer 语义会被破坏。
- `LogitsProcessor` 的 compact hidden 分支需要支持 `CaptureHiddenMode.LAST`，否则 `cp_collect_last_token_hidden(...)` 后 `hidden_states` 会被丢弃，`capture_for_decode(...)` 无法拿到 draft hidden。
- 如果 `can_cp_split(...)` 失败，必须 fallback 到现有 full-token 路径并打印 rate-limited fallback log。

验收：

- debug log 中 draft local token 数约等于 full token 数 / cp_size。
- `eh_proj` 输入 shape 使用 local token 数。
- draft prefill 不再执行 full `cp_all_gather_rerange_output(...)`。
- `capture_for_decode(...)` 能拿到 `(batch, hidden)` 的 last-token hidden。
- fallback 原因可观测，例如：`no_nsa_cp_metadata`、`unsupported_arch`、`non_extend_mode`、`cp_shared_kv_disabled`。

### P3：输出 narrow/local collection 验证与清理

目标：P2 已把 draft 输出 narrow/local collection 纳入实现；P3 只做验证、清理旧 fallback 以及补充 benchmark/profile 证据。

修改文件：

- `python/sglang/srt/models/deepseek_nextn.py`
- `python/sglang/srt/layers/attention/nsa/utils.py`
- `python/sglang/srt/layers/logits_processor.py`

当前：

```python
hidden_states = cp_all_gather_rerange_output(...)
```

目标：

```text
只收集 draft logits / capture_for_decode 需要的 last-token hidden
```

可以复用 target 的思路：

```python
cp_collect_last_token_hidden(hidden_states, forward_batch, cp_size)
```

但要先确认 `EagleDraftInput.capture_hidden_mode = CaptureHiddenMode.LAST` 对 draft logits processor 的需求：

- `next_token_logits` 需要每个 request 的最后 token hidden；
- `capture_for_decode(...)` 需要 draft decode 后续使用的 hidden/topk 信息；
- 不需要完整 prompt hidden。

验收：

- draft extend 不再执行 full `cp_all_gather_rerange_output(...)`。
- EAGLE draft next token logits 与关闭新路径时保持一致或在可接受数值误差内一致。
- `capture_for_decode(...)` 能正常生成后续 speculative decode 所需状态。

### P4a：draft KV transfer / physical shard debug 日志

目标：在 `SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1` 下补齐 draft KV shared-KV 验证闭环，不改变默认行为。

新增日志点：

- scheduler disaggregation init：target/draft KV pool class、size、page_size、layer 范围；
- prefill/decode KV manager：target/draft contiguous buffer 数量、lens、item_lens、draft split point；
- CP shared KV write：logical→physical remap、MLA KV write path、index K/scale write path；
- prefill send chunk：logical page_indices、state_indices、是否存在 draft pool；
- mooncake register/sender/transfer worker：注册 buffer 数、CP filter 后 pages、logical positions、decode dst pages；
- decode prealloc/commit：dst pages、metadata index、spec metadata shape。

这些日志用于确认：draft KV 没有拼 full KV 到某个 prefill rank，而是复用 CP shared-KV owner page filter，把各 rank owned physical pages 传到 decode full KV pool。

### P4：draft KV shared-KV physical write 与 transfer 验证

目标：确认 draft persistent KV 与 target 一样是 CP owner-sharded physical KV。

修改文件：

- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
- `python/sglang/srt/layers/attention/nsa_backend.py`
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- `python/sglang/srt/disaggregation/prefill.py`
- `python/sglang/srt/disaggregation/decode.py`

验证重点：

1. draft `NSATokenToKVPool` 是否用 physical capacity 初始化。
2. draft MLA KV write 是否走 `_maybe_filter_shared_mla_kv_write(...)`。
3. draft index K/scale write 是否走 `get_cp_shared_kv_local_physical_out_cache_loc(...)`。
4. prefill transfer 是否包含 draft physical shards。
5. decode 是否能按 logical loc 使用 transferred draft KV。

验收：

- 每个 CP rank 只写 owner shard 对应的 physical loc。
- draft KV pool memory 不再按 full logical token 数乘以 cp_size 膨胀。
- prefill->decode ETE 可通过，且 EAGLE speculative decode 正常。

## 远端验证计划

所有运行验证在远端容器执行，不在本地执行。

推荐远端环境：

```text
host: g0034 / ubuntu@10.20.32.34
container: sglang-glm5-dev-2
code dir: /sgl-workspace/sglang-tai
```

验证分层：

### 1. 静态检查

```bash
python -m compileall python/sglang/srt/speculative python/sglang/srt/models python/sglang/srt/layers/attention/nsa
```

### 2. targeted pytest

优先跑和 CP split / schedule copy / speculative 相关的小测试。若没有现成测试，需要新增 narrow unit test 覆盖：

- `ScheduleBatch.copy()` 保留 `spec_info`；
- CP split 对 `input_ids`、`positions`、`hidden_states` 一致；
- draft path fallback reason 可观测。

### 3. GLM-5 prefill CP + decode DP ETE

启动条件：

```bash
SGLANG_CP_DRAFT_SHARED_KV=1
SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1
SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1
SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1
```

观察：

- target narrow path 命中；
- draft architecture 是 `DeepseekV3ForCausalLMNextN`；
- draft local token 数约等于 full token 数 / cp_size；
- draft KV write 使用 local physical loc；
- prefill transfer 包含 draft KV；
- decode 侧没有 draft KV missing / bootstrap timeout / speculative 状态异常。

### 4. profile 验证

Nsight / 日志关注点：

- target logits 前不再出现 full hidden all-gather；
- draft `eh_proj` 和 decoder token 数下降；
- draft 输出不再 full all-gather；
- draft KV materialize/write 不再每 rank 写 full request；
- prefill 显存下降，尤其 draft KV pool 和 logits/hidden buffer 部分。

## 风险与处理

### 风险 1：side-channel hidden 与 logits processor 输出结构耦合

处理：不复用 `hidden_states` 字段表达两种语义，新增明确字段 `draft_hidden_states`。现有 EAGLE 非 CP 路径继续使用原字段。

### 风险 2：draft output narrow 后 `capture_for_decode(...)` 缺状态

处理：先在 P3 前单独梳理 `capture_for_decode(...)` 依赖字段，只收集它真正需要的 last-token hidden/topk，不保留 full prompt hidden。

### 风险 3：draft KV transfer 依赖 target/draft loc 完全一致

处理：P4 只在日志与 ETE 确认一致后打开默认路径。若不一致，增加显式 draft logical loc mapping，不继续依赖注释里的隐式假设。

### 风险 4：GLM4 MHA/Radix 被误启用

处理：`SGLANG_CP_DRAFT_SHARED_KV=1` 下如果 draft architecture 不是 `DeepseekV3ForCausalLMNextN`，直接 fallback 并打印：

```text
CP draft shared KV fallback: unsupported draft architecture
```

## 完成标准

本阶段完成需要同时满足：

1. GLM-5 EAGLE/MTP prefill CP ETE 通过。
2. target EAGLE 开启时仍命中 CP narrow logits path。
3. draft model 在 `eh_proj` 前已经使用 CP-local token / hidden。
4. draft persistent KV 使用 CP shared-KV physical shard，而不是每 rank full KV。
5. prefill->decode draft KV transfer 正常。
6. fallback 路径安全、可观测，非 GLM-5 draft 不误走新路径。

## 2026-06-01：Draft layer as forward-layer in prefetch/backup pipeline plan

### 背景

当前 EAGLE / NextN draft 已经作为 target HiCache 的 shadow payload 接入了一部分生命周期：

- `CacheController.attach_draft_pool()` 会把 draft device/host pool 绑定到 target controller，并注册同一个 `LayerDoneCounter`。
- draft decoder layer end 会通过 `notify_layer_end_for_backup()` 进入 `CacheController.on_layer_end(layer_id, source="draft")`，再提交 draft D2H backup。
- `CacheController.start_loading()` 已经会把 target load queue 与 draft load queue 一起处理。

但 CP shared-KV 的 index/MLA prefetch 当前仍把 draft 排除在外：

- `nsa_backend.init_forward_metadata()` 对 `cp_shared_kv_is_draft_input(forward_batch)` 显式禁用 prefetcher 创建。
- `cp_shared_kv_should_prefetch_next_layer()` 对 draft input 返回 `False`。
- draft/NextN 只有一个 executable layer，`layer_id=0`，不能直接复用 target 的 `next_layer` prefetch 语义。

### 设计目标

把 draft 层视为 target forward 后的一个逻辑 forward layer：

```text
target layer 0
...
target layer N-1
draft layer 0
```

但第一阶段不重构全局 event/counter timeline；先让 draft layer 拥有独立的 same-layer prefetch 合同：

```text
draft metadata ready
  -> start draft current-layer index/MLA prefix prefetch
  -> draft layer 0 consumes prefix + current KV
  -> draft layer 0 end submits async backup
```

### 分阶段计划

#### P0：锁住当前 draft 合同

- 测试 draft input 不应走 target-style next-layer prefetch。
- 测试 draft layer id 只允许 `0`。
- 测试 draft backup notifier 仍通过 `source="draft"` 触发。
- 保留 fail-fast，不允许静默 fallback 到 target prefetch 路径。

#### P1：新增 draft same-layer prefetch API

新增独立语义，避免误用 target `start_next_layer_prefix()`：

```python
start_current_layer_prefix(layer_id=0, token_to_kv_pool=draft_pool)
```

约束：

- 只允许 draft / NextN path 使用。
- 只允许 `layer_id == 0`。
- metadata 必须已经构造完成。
- 不满足条件时明确 warning/fail-fast，不做静默慢 fallback。

#### P2：把 draft prefetch 启动点放到 draft metadata ready 后、attention 前

目标顺序：

```text
draft init_forward_metadata
pre-attention hook:
  start index prefix prefetch
  start MLA prefix prefetch
draft indexer / attention consume:
  wait pending event if needed
```

这和 target 当前的 attention-front prefetch 位置保持一致，但语义是 current-layer，不是 next-layer。

#### P3：先不重构 backup 主流程

当前 draft backup 已经在 draft layer end 触发。第一阶段只验证：

- 是否确实 per-layer async 提交；
- 是否还有 catch-up fallback；
- final ack 是否仍等待 target + draft payload 完成。

若发现 draft backup 延迟到 request 结束，再单独修复。

#### P4：后续再评估 target-last-layer -> draft-layer bridge

更激进的优化是让 target 最后一层提前发起 draft layer 0 prefetch。但这需要更早构造 draft metadata，而 draft forward batch 当前依赖 target 输出 hidden / verified token，风险较高。

暂不作为第一阶段实现项。

#### P5：后续再统一虚拟 layer timeline

理想模型是：

```text
event 0..N-1: target layers
event N:      draft layer 0
```

但这会影响 `LayerDoneCounter`、draft H2D wait、backup ack、target/draft strong-sync 等多个路径。只有当 P1-P3 证明收益不足或现有 event 绑定导致 correctness/perf 问题时再做。

### 风险记录

- 直接打开 draft prefetcher 可能复现之前 EAGLE hang；draft 没有 next-layer，必须单独建 current-layer API。
- draft metadata 生成时间晚于 target metadata，target last-layer 提前 prefetch 需要额外桥接。
- index/MLA prefix materialize 会占 SM，短 draft 层可能出现启动开销大于 overlap 收益，需要 benchmark/trace 验证。
- 任何 fallback 必须 warning/fail-fast，不能静默降级。
