# 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 1 当时的初始拆分。后续 shared KV 工作已经细化为：

```text
Phase 2: shared/sharded persistent KV pool
Phase 3: current chunk reuse
Phase 4: page-aligned CP token split
Phase 5: compute-owner shared KV layout
```

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

这是当前代码库里风险最低、最容易验证收益的一步。
