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

docs(cp): add NSA prefill CP dual all-gather explanation and Phase 1 plan

Browse Source 
- nsa_prefill_cp_all_gather.md: explains why there are two all-gathers
  in current NSA prefill CP (per-layer KV/index gather vs tail output
  hidden gather), with data flow diagrams and involved file lists
- nsa_prefill_cp_phase1_narrow_output_collection.md: detailed Phase 1
  design plan covering batch-level eligibility, fallback conditions,
  affected modules, KV transfer impact, risks, and test matrix
This commit is contained in:
laoyao0822
2026-04-22 04:55:15 +08:00
parent bc9a9f128b
commit e7a8c4b52f
2 changed files with 970 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
docs/advanced_features/nsa_prefill_cp_all_gather.md Normal file
View File

@@ -0,0 +1,290 @@
# NSA Prefill CP 中两次 All-Gather 的作用说明
本文档描述 **当前 SGLang 实现** 中，GLM 路径在 **NSA Prefill Context Parallel (CP)** 下为什么会出现两类 `all-gather`，它们分别服务于什么语义边界，以及涉及哪些关键文件。
> 本文档只描述 **当前行为**，不讨论重构后的目标设计。
## 适用范围
- 只讨论 **prefill CP**
- 不讨论 decode 开启 CP 的情况
- 这里的 GLM 实际走的是：
- `python/sglang/srt/models/glm4_moe.py`
- `GlmMoeDsaForCausalLM`
- 继承 `DeepseekV2ForCausalLM`
- 因此，相关逻辑实际落在 **DeepSeek-NSA** 这套实现上，而不是独立的 `glm5.py`
## 一句话总结
当前实现里的两次 gather 不是在重复做同一件事：
1. **attention 内的 gather**：为当前层的 **K 侧 / index 视图** 服务
2. **prefill 尾部的 gather**：为最终 **hidden_states -> logits/logprobs** 服务
更简洁地说：
> 第一次是给 attention/index 用，第二次是给输出/logprob 用。
---
## 总体流程图
```text
[Request enters prefill]
|
v
ScheduleBatch.prepare_for_extend
|
+--> alloc_for_extend()
| - 分配 out_cache_loc
| - write_cache_indices()
| - req_to_token[req, token_pos] = loc
|
v
DeepseekV2ForCausalLM.forward
|
+--> prepare_input_dp_with_cp_dsa()
| - 生成 nsa_cp_metadata
|
v
DeepseekV2Model.forward
|
+--> cp_split_and_rebuild_data(hidden_states)
+--> cp_split_and_rebuild_position(positions)
| => 输入进入 CP-local token 布局
|
v
[每一层 NSA/MLA attention]
|
+--> 本地算 q / latent_cache / k_nope / k_pe
|
+--> Gather #1
| - 重建当前层需要的 K / latent KV / index key 全局顺序视图
|
+--> NSA index / topk / page-table
|
+--> 写入本地 KV pool
|
v
[所有层结束]
|
+--> Gather #2
| - 恢复最终 hidden_states 的原 token 顺序
|
+--> logits_processor
| - next token logits
| - input logprobs
| - top_logprobs / token_ids_logprobs
|
v
[Prefill ends]
```
---
## 第一次 All-Gather：attention 内的 gather
### 是为了什么
在 CP 模式下，prefill 一开始就会把 token 重新切分并重排，以便均衡负载。
但是当前 NSA/MLA 的实现仍然需要：
- 正确顺序的 `K`
- 正确顺序的 `latent KV`
- 正确顺序的 `index key`
只有这样，后续的：
- MLA attention
- NSA index / top-k
- page-table 生成
才能继续工作。
所以 attention 内部需要先 gather 一次，把当前层所需的 K 侧信息恢复成全局逻辑顺序。
### gather 的对象
主要是两类：
- `latent_cache / k_nope / k_pe`
- `indexer key`
### 关键调用点
#### 1. MLA KV 重建
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py:302-306`
- `python/sglang/srt/models/deepseek_v2.py:1475-1487`
这里 `rebuild_cp_kv_cache()` 会调用：
- `cp_all_gather_rerange_output(...)`
把局部的 MLA KV 重新拼成全局逻辑顺序。
#### 2. NSA indexer key 重建
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:336-344`
这里同样会：
- `cp_all_gather_rerange_output(key, ...)`
把 indexer 用的 key 恢复成全局顺序，再去做 top-k/index。
### 关键实现 helper
- `python/sglang/srt/layers/attention/nsa/utils.py:341-404`
这里的 `cp_all_gather_rerange_output(...)` 注释已经明确说明，它会把被：
- `round-robin-split`
- `in-seq-split`
打散过的 token / block 重新恢复成原始逻辑顺序。
### 本质语义
这次 gather 的本质是：
> **KV / index 语义 gather**
它不是为了最终输出，也不是为了 logits。
---
## 第二次 All-Gather：prefill 尾部的 gather
### 是为了什么
即使 attention 内部已经把当前层需要的 K/index 问题解决了，**模型最终输出的 `hidden_states` 仍然还处在 CP 重排后的布局里**。
但下游的：
- `logits_processor`
- input logprobs
- top_logprobs
- token_ids_logprobs
都默认假设：
> `hidden_states` 已经按原始 flatten token 顺序排好
因此，在模型尾部还要再 gather 一次，把最终 hidden 恢复回原 token 顺序。
### gather 的对象
- 最终层输出的 `hidden_states`
### 关键调用点
#### 1. 模型尾部 gather
- `python/sglang/srt/models/deepseek_v2.py:2046-2053`
这里直接做：
- `cp_all_gather_rerange_output(hidden_states, ...)`
#### 2. 下游消费者
- `python/sglang/srt/layers/logits_processor.py:422-527`
这里可以直接看到：
- 无 input logprob 时，会按每个请求最后 token 的原始位置取 `hidden_states`
- 有 input logprob 时，会按原始 flatten token 顺序切 `hidden_states`
例如：
- `422-448`：按每请求最后 token 位置取 hidden
- `449-527`：按原始 token 顺序构造 `sample_indices` / `input_logprob_indices` / `token_to_seq_idx`
### 本质语义
这次 gather 的本质是：
> **output / logprob 语义 gather**
它不是为了 attention 本身，而是为了让输出层继续消费“原 token 顺序 hidden”。
---
## 为什么会有两次 gather
因为它们处理的是两个不同的边界：
### 第一次 gather
- 发生在每层 attention/indexer 路径中
- 解决的是：
- 当前层要怎么“看见”完整且按正确顺序排列的 K / key / index 视图
### 第二次 gather
- 发生在 prefill 尾部
- 解决的是：
- 最终 `hidden_states` 怎么重新和原始 token 顺序对齐
- 这样 `logits/logprobs` 才能对得上输入
所以当前实现下，这两次 gather 不是冗余，而是分别服务于：
1. **attention/index 边界**
2. **output/logprob 边界**
---
## 涉及的关键文件
| 文件 | 作用 |
|---|---|
| `python/sglang/srt/models/glm4_moe.py` | GLM 入口；`GlmMoeDsaForCausalLM` 继承 `DeepseekV2ForCausalLM` |
| `python/sglang/srt/models/deepseek_v2.py` | CP 主模型逻辑、MLA KV 重建、最终 hidden gather |
| `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` | MLA prepare 阶段触发 KV 重建 |
| `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | index key gather、top-k/index 相关逻辑 |
| `python/sglang/srt/layers/attention/nsa/utils.py` | `cp_all_gather_rerange_output(...)` 与 split/rerange 逻辑 |
| `python/sglang/srt/layers/attention/nsa_backend.py` | NSA metadata、topk/page-table、KV 写入 |
| `python/sglang/srt/layers/logits_processor.py` | logits/logprob 消费原始 token 顺序 hidden |
| `python/sglang/srt/managers/schedule_batch.py` | 构造 prefill/logprob 元数据 |
| `python/sglang/srt/model_executor/forward_batch_info.py` | batch 侧输出张量和 metadata |
| `python/sglang/srt/mem_cache/common.py` | `alloc_for_extend()` / `write_cache_indices()` |
| `python/sglang/srt/mem_cache/memory_pool.py` | `set_mla_kv_buffer()` / `get_key_buffer()` |
---
## 最容易混淆的一点
`cp_all_gather_rerange_output(...)` 这个 helper 会被复用。
所以不能只看“是不是调用了同一个函数”，而要看它 gather 的到底是什么 tensor：
- 如果 gather 的是：
- `latent_cache`
- `key`
- `k_nope / k_pe`
- 那它属于 **attention/KV/index gather**
- 如果 gather 的是：
- `hidden_states`
- 那它属于 **输出/logprob gather**
---
## 一个更准确的口头总结
如果要用一句简单但准确的话向别人解释当前实现：
> **每层 attention 中间的 all-gather，是为了重建当前层需要的完整 K/index 视图；请求结尾的 all-gather，是为了把最终 hidden 恢复成原 token 顺序，好让 logits/logprobs 对得上输入。**
---
## 备注
- 本文档只描述 **当前实现**，不代表所有 context parallel 设计都必须这样。
- 如果未来改成：
- ring attention
- sharded KV
- shard-aware logits/logprob 输出
那么这两类 gather 都有可能被替换或减少。

680
docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md Normal file
View File

@@ -0,0 +1,680 @@
# NSA Prefill CP Phase 1 计划：收窄 Prefill 尾部输出收集，移除全量 Hidden All-Gather
本文档给出一个 **Phase 1 级别** 的详细计划：
> 先不改每层 attention 内部的 KV / index gather，也不改 KV cache 的底层 shard 语义；
> 第一阶段只尝试把 **prefill 尾部那次“全量 hidden all-gather + restore 原 token 顺序”** 改掉，
> 换成 **更窄的输出收集路径**。
本文档面向当前实际模型路径：
- `python/sglang/srt/models/glm4_moe.py`
- `GlmMoeDsaForCausalLM`
- 继承 `DeepseekV2ForCausalLM`
- 实际落在 **DeepSeek-NSA prefill CP** 路径
---
## 1. 背景与问题定义
### 1.1 当前观察到的热点
当前在 NSA prefill CP 模式下，每个请求在 **prefill 最后一个 forward** 结束前，会出现一次明显耗时的 `all-gather`。
这次通信当前发生在：
- `python/sglang/srt/models/deepseek_v2.py:2046-2053`
核心逻辑是：
```python
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
```
### 1.2 问题的准确表述
这个热点经常会被口头描述成：
> “请求结束后，需要 all-gather 所有 logits”
但从代码语义上更准确的说法是：
> **当前实现会在 prefill 尾部 all-gather 全量 `hidden_states`，恢复原 token 顺序后，再交给 `logits_processor` 生成 logits/logprobs。**
所以 Phase 1 不是改 “logits all-gather”，而是改：
> **prefill 尾部的 hidden/output 收集方式**
---
## 2. 范围与前提
### 2.1 Phase 1 只做什么
Phase 1 只处理：
- **prefill 尾部的全量 hidden gather**
- `hidden_states -> logits_processor` 这条输出消费链
- 与该链条直接相关的请求元数据、fallback 逻辑、PD/spec hidden side-channel
### 2.2 Phase 1 不做什么
以下内容明确 **不属于本阶段**：
- attention 内部 `KV / key / index` 的 gather
- `rebuild_cp_kv_cache()` 的语义重构
- `nsa_indexer` 的 distributed top-k / local-topk merge
- `KV cache` 改成真正 shard-aware 存储
- `req_to_token_pool / out_cache_loc / page_table` 语义重构
- decode 侧 CP-aware / shard-aware runtime
- ring attention / Helix / DCP
### 2.3 当前默认假设
Phase 1 默认采用以下判断：
1. **普通 prefill 生成 batch**（batch 内所有请求都不返回 input logprobs、也不要求 full hidden）
- 只需要 **每个请求最后一个 token 的 hidden**
2. **只要 batch 内有任意一个请求** 需要以下任一能力：
- `input logprob / top_logprobs / token_ids_logprobs`
- `return_hidden_states`
- multi-item scoring
- EAGLE target prefill / hidden side-channel
- chunked / mixed / padded-static 等未明确验证场景
- 则整个 batch 继续回退到当前 **全量 hidden gather** 路径
也就是说，Phase 1 的核心思想是：
> **只优化满足条件的普通 prefill batch；任何需要完整顺序语义的 batch 都先保留旧路径。**
---
## 3. 当前实现数据流
建议先配合阅读：
- `docs/advanced_features/nsa_prefill_cp_all_gather.md`
### 3.1 当前 prefill 尾部数据流
```text
局部 hidden_states (CP 打散布局)
|
v
deepseek_v2.py
cp_all_gather_rerange_output(hidden_states)
|
v
全量、按原 token 顺序排列的 hidden_states
|
v
logits_processor._get_pruned_states()
|
+--> 无 input logprob：只取每请求最后 token
|
+--> 有 input logprob：按原始 flatten token 顺序切多个 token 子集
|
v
next_token_logits / input_token_logprobs / hidden capture
```
### 3.2 当前下游为什么依赖“原 token 顺序 hidden”
关键逻辑在：
- `python/sglang/srt/layers/logits_processor.py:399-538`
其中：
#### 普通 prefill（无 input logprob）
- `422-448`
- 实际只使用：
- 每个请求最后一个 token 的 hidden
#### prefill + input logprob
- `449-527`
- 明确依赖：
- 原始 flatten token 顺序
- `sample_indices`
- `input_logprob_indices`
- `token_to_seq_idx`
所以：
> **当前全量 gather 并不是 attention 理论要求，而是输出消费接口要求。**
---
## 4. Phase 1 目标
### 4.1 核心目标
对以下 **batch 级条件同时成立** 的情况：
- `forward_mode.is_extend()`
- `extend_return_logprob == False`
- `capture_hidden_mode == NULL`
- 非 multi-item scoring
- 非 EAGLE target prefill
- 非 PD + EAGLE hidden side-channel
- 非 chunked / mixed / padded-static 等未明确验证场景
不再执行：
- `cp_all_gather_rerange_output(hidden_states)` 全量 restore
改成：
- **只收集每个请求最后一个 token 的 hidden**
- 直接供 `logits_processor` 生成 next-token logits
### 4.2 目标收益
预期收益包括：
- 降低 prefill 尾部一次性通信量
- 降低临时全量 hidden buffer 的显存峰值
- 降低 `hidden_states` 全量 restore 带来的 stream/同步开销
---
## 5. 推荐方案：Narrow Output Collection
### 5.1 方案概念
把当前的：
> `full hidden gather -> prune`
改成：
> `local hidden -> collect only needed hidden -> logits`
### 5.2 具体策略
#### 路径 A：普通生成 batch（推荐优化路径）
条件：
- `forward_mode.is_extend()`
- `not extend_return_logprob`
- `capture_hidden_mode == CaptureHiddenMode.NULL`
- 非 multi-item scoring
- 非 EAGLE target prefill
- 非 PD + EAGLE hidden side-channel
- 非 chunked / mixed / padded-static 等未明确验证场景
行为：
- 不再 full gather 所有 hidden
- 只计算并收集：
- **每个请求最后 token 对应的 hidden**
#### 路径 B：需要精确 token 顺序语义的 batch（保留旧路径）
包括但不限于：
- `extend_return_logprob=True`
- `return_hidden_states=True` → `capture_hidden_mode = FULL`
- 见 `schedule_batch.py:2309-2319`
- `token_ids_logprobs`
- `top_logprobs`
- prefill-only multi-item scoring
- speculative/EAGLE target prefill 依赖 hidden 的路径
- PD + EAGLE hidden side-channel 需要 `output_hidden_states`
- chunked / mixed / padded-static 等暂不纳入 Phase 1 的场景
行为：
- 保持当前：
- `cp_all_gather_rerange_output(hidden_states)`
### 5.3 为什么这是一条低风险路径
因为当前代码已经证明：
- `logits_processor.py:422-448`
- 在“无 input logprob”的 prefill 分支里，本来就只需要最后 token 的 hidden
所以第一阶段只是：
> **把“先 gather 全量 hidden 再截最后 token”改成“直接收最后 token hidden”。**
---
## 6. 受影响文件与模块
下面按“必须改 / 必须验证 / 暂不改”分组。
### 6.1 必须改
#### 1) `python/sglang/srt/models/deepseek_v2.py`
职责：
- 当前在模型尾部执行全量 hidden gather
Phase 1 改动：
- 把尾部逻辑改成双路径：
- full restore path
- narrow output collection path
- 引入对请求能力的 gating
可能新增：
- `collect_last_token_hidden_only(...)`
- 或调用 utils 中的新 helper
---
#### 2) `python/sglang/srt/layers/attention/nsa/utils.py`
职责：
- 当前提供 `cp_all_gather_rerange_output(...)`
Phase 1 改动：
- 新增更窄的收集 helper，例如：
- `cp_collect_last_token_hidden(...)`
- `cp_collect_selected_hidden(...)`
要求：
- 利用现有 CP metadata 推导“每请求最后 token 所属 rank 和 local offset”
- 不再默认按全量 token 恢复顺序
---
#### 3) `python/sglang/srt/layers/logits_processor.py`
职责：
- 消费 `hidden_states`
- 生成 logits/logprobs/hidden capture
Phase 1 改动：
- 支持两种输入模式：
1. 原始全量 hidden（兼容路径）
2. compact per-request hidden（优化路径）
至少需要明确：
- `_get_pruned_states()` 的输入约定
- 是否新增 compact 模式分支
---
#### 4) `python/sglang/srt/model_executor/forward_batch_info.py`
职责：
- 承载 forward batch 元数据
Phase 1 改动：
- 可能需要新增 output-collection 相关 metadata：
- 每请求最后 token 在 local shard 中的位置
- compact 输出模式开关
---
#### 5) `python/sglang/srt/managers/schedule_batch.py`
职责：
- 生成 `ModelWorkerBatch`
- 设置：
- `capture_hidden_mode`
- `return_logprob`
- `is_prefill_only`
Phase 1 改动：
- 明确 narrow-path 的进入条件
- 明确 fallback 条件
---
### 6.2 必须验证（可能不用改，至少要确认）
#### 6) `python/sglang/srt/managers/tp_worker.py`
职责：
- `forward_batch_generation()`
- 处理 `GenerationBatchResult`
要验证：
- `logits_output.hidden_states` 变成 compact 形状后，是否仍满足 prefill-only 路径
---
#### 7) `python/sglang/srt/managers/scheduler_output_processor_mixin.py`
职责：
- 把 `LogitsProcessorOutput` 回填到请求对象
关键代码：
- `205-219`
- 当前 `return_hidden_states` 会按 `len(req.origin_input_ids)` 切 full hidden
Phase 1 要验证：
- fallback 时保持不变
- narrow path 不能误走到这里的 full hidden slicing 逻辑
说明：
- 当前普通路径真正的主要风险面仍在 `logits_processor.py`
- `scheduler_output_processor_mixin.py` 更像是 fallback 路径的契约校验点
---
#### 8) `python/sglang/srt/disaggregation/utils.py`
职责：
- `MetadataBuffers`
- `output_hidden_states`
关键代码：
- `252-265`
- 当前只有在 `req.hidden_states_tensor is not None` 时，才把 hidden 写到 metadata buffer
Phase 1 要验证：
- 普通 prefill 路径改小 gather后，PD + EAGLE hidden side-channel 所需 `output_hidden_states` 契约是否保持
- 明确该契约当前是“每请求一个 hidden vector”，而不是全量 token hidden
---
#### 9) `python/sglang/srt/disaggregation/prefill.py` / `decode.py`
职责：
- KV transfer / metadata transfer
Phase 1 要验证：
- 普通 KV transfer 不依赖 full hidden gather
- 但 PD + EAGLE 路径的 metadata buffer 里确实含 `output_hidden_states`
---
#### 10) `python/sglang/srt/speculative/eagle_worker.py` / `eagle_worker_v2.py`
职责：
- EAGLE target prefill / draft 相关 hidden capture
Phase 1 要验证：
- 当前这些路径是否显式要求 `CaptureHiddenMode.FULL`
- 如果要求，则必须整体回退到 full gather 路径
---
#### 11) `python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py`
职责：
- PD decode 与 prebuilt batch 交接
Phase 1 要验证：
- compact output path 不应破坏现有 prebuilt / metadata buffer 契约
---
### 6.3 暂不改（但需明确不在本期范围）
以下模块当前只做背景参考，**不应在 Phase 1 里动**：
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- `python/sglang/srt/layers/attention/nsa_backend.py`
- `python/sglang/srt/mem_cache/common.py`
- `python/sglang/srt/mem_cache/memory_pool.py`
- `python/sglang/srt/mem_cache/allocator.py`
- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`
这些属于：
- attention 内部 KV/index gather
- KV allocator / page_table / out_cache_loc
- decode 读取本地 KV pool 的语义
都应留给后续 Phase 2/3。
---
## 7. KV Cache 传输与 Disaggregation 影响
### 7.1 Phase 1 对 KV 传输的原则
Phase 1 **不改变** 当前 KV cache 传输格式与 allocator 语义。
理由：
- 当前目标只是收窄 prefill 尾部输出 gather
- 不是改：
- `out_cache_loc`
- `req_to_token`
- `page_table`
- `token_to_kv_pool`
也就是说，当前：
- `mem_cache/common.py:328-391`
- `alloc_for_extend()`
的输出仍然保留：
> `req_to_token[req, token_pos] = 本地可直接读的 loc`
### 7.2 对 PD disaggregation 的具体影响
#### 标准 KV transfer
- 仍由 `PrefillBootstrapQueue` / `CommonKVManager` 传递 KV pool buffer
- Phase 1 不修改该传输格式
#### metadata buffer
- `disaggregation/utils.py` 里 `MetadataBuffers` 包含：
- `output_ids`
- `output_token_logprobs_*`
- `output_top_logprobs_*`
- `output_hidden_states`
其中：
- `output_hidden_states` 只在 `req.hidden_states_tensor is not None` 时使用
- 当前主要服务于 **PD + EAGLE 的 hidden side-channel**
### 7.3 Phase 1 的明确结论
> **Phase 1 不应该改 KV transfer 协议本身。**
但必须明确：
- PD + EAGLE 路径如果依赖 `output_hidden_states`
- 则该路径必须继续使用 full hidden gather fallback
---
## 8. 可能触发的问题 / 风险清单
### 8.1 输出正确性风险
1. **最后 token 定位错误**
- round-robin / in-seq-split 下，最后 token 所属 rank / local offset 计算错
2. **padding / static len 场景错位**
- `logits_processor.py:427-441`
3. **chunked prefill 语义错位**
4. **mixed batch / padded-static 场景错位**
### 8.2 logprob 风险
1. `extend_return_logprob=True` 却误走 narrow path
2. `token_ids_logprobs` / `top_logprobs` 仍按 full hidden 顺序切片
3. `sample_indices / input_logprob_indices` 和 compact hidden 不匹配
### 8.3 hidden capture 风险
1. `return_hidden_states=True` 被错误走窄路径
2. `capture_hidden_mode=FULL` 语义被破坏
3. `scheduler_output_processor_mixin.py` 仍按 full token 数切 hidden
### 8.4 speculative / EAGLE 风险
1. `spec_info.hidden_states` 依赖旧 hidden capture 语义
2. `output_topk_p / output_topk_index / output_hidden_states` 的 side-channel 失配
3. `eagle_worker.py / eagle_worker_v2.py` 强制 FULL capture 的路径被误走优化分支
### 8.5 disaggregation 风险
1. PD + EAGLE hidden side-channel 的 metadata buffer 契约被破坏
2. prefill-only 请求与 decode-ready 请求混合批时的输出形状不一致
### 8.6 性能风险
1. 实现 narrow collect 但仍然隐式产生全量临时 buffer
2. 为了定位最后 token，引入额外 CPU/GPU 同步
3. compact gather 触发更多小通信，吞掉收益
---
## 9. 推荐实施顺序
### 步骤 1：把问题定义收紧
先在代码和文档里统一口径：
- 当前优化对象是：
- **prefill 尾部 full hidden gather**
- 不是：
- attention 内部 KV gather
- KV transfer 协议
### 步骤 2：新增 output collection helper
在 `nsa/utils.py` 增加：
- `cp_collect_last_token_hidden(...)`
或
- `cp_collect_selected_hidden(...)`
要求：
- 输入仍基于现有 CP metadata
- 输出是 compact hidden
### 步骤 3：在 `deepseek_v2.py` 尾部引入双路径
- 普通生成 batch：走 compact collect
- 任何触发 fallback 条件的 batch：走 full restore fallback
### 步骤 4：让 `logits_processor.py` 接受 compact hidden
最优先支持：
- prefill without input logprobs
暂不支持：
- 全量 logprob 场景（保留 fallback）
### 步骤 5：补齐 scheduler / worker / PD 边界保护
确保：
- `return_hidden_states=True` 必回退
- `extend_return_logprob=True` 必回退
- multi-item scoring 必回退
- EAGLE target prefill / PD + EAGLE hidden side-channel 必回退
- chunked / mixed / padded-static 等未覆盖场景先回退
---
## 10. 验证与测试矩阵
### 10.1 必测功能用例
#### A. 普通 prefill 生成（目标路径）
- `return_logprob=False`
- `return_hidden_states=False`
- 非 multi-item scoring
- 非 EAGLE target prefill
- CP 开启
- 验证：
- next token logits 正确
- 输出 token 正确
- 通信量下降
#### B. prefill + input logprob（fallback）
- `return_logprob=True`
- 验证 fallback 仍走原 full gather 路径
#### C. `return_hidden_states=True`（fallback）
- 验证 hidden 输出完整正确
#### D. prefill-only multi-item scoring（fallback）
- 验证仍走 full restore 路径
#### E. PD + EAGLE hidden side-channel（fallback）
- 验证 `output_hidden_states` metadata buffer 契约不变
#### F. EAGLE target prefill（fallback）
- 验证 `CaptureHiddenMode.FULL` 语义不变
#### G. round-robin / in-seq-split 两种模式
- 验证最后 token owner / local offset 计算无误
#### H. mixed / chunked / padded-static
- 若未纳入第一阶段实现，验证是否正确回退
### 10.2 建议复用/补充的测试位置
- `test/registered/8-gpu-models/test_deepseek_v32_cp_single_node.py`
- `test/registered/core/test_hidden_states.py`
- `test/registered/openai_server/features/test_openai_server_hidden_states.py`
- `test/registered/openai_server/basic/test_protocol.py`
- `test/registered/unit/managers/test_scheduler_output_processor_mixin.py`
### 10.3 性能验收指标
至少应采集：
- prefill TTFT
- 尾部 gather 时间占比
- hidden gather bytes
- 峰值显存
---
## 11. Phase 1 完成标准
当以下条件同时满足，可认为 Phase 1 完成：
1. 普通 prefill 请求不再执行全量 hidden gather
2. next-token logits 与 baseline 一致
3. `return_logprob / return_hidden_states / PD+spec` 仍可正确回退
4. KV transfer 协议与 decode 读取语义未被破坏
5. TTFT 或尾部 gather 时间有可测下降
---
## 12. 后续阶段（本期不做）
### Phase 2
- prefill shard-KV 保存
- 不再在每层 gather 完整 KV 再写本地 pool
### Phase 3
- decode shard-aware KV/page_table/runtime
- 去掉 prefill finalize 时为 decode 做的 KV materialize
### Phase 4
- logits/logprob 全面 shard-aware
- full hidden restore fallback 范围进一步缩小甚至消失
---
## 13. 最终建议
这次 Phase 1 的最佳工程策略是：
> **把它当作“输出接口解耦”问题来做，而不是当作“KV shard 重构”问题来做。**
也就是说，先解决：
- 为什么当前必须全量 gather hidden
- 如何只收集最后 token hidden
- 哪些请求必须回退
而不是第一阶段就去碰：
- KV allocator
- page_table
- req_to_token
- decode shard-aware runtime
这是当前代码库里风险最低、最容易验证收益的一步。