From e7a8c4b52fce3dda0420cf875c4d65590fd75265 Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Wed, 22 Apr 2026 04:55:15 +0800
Subject: [PATCH] docs(cp): add NSA prefill CP dual all-gather explanation and
 Phase 1 plan

- 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
---
 .../nsa_prefill_cp_all_gather.md              | 290 ++++++++
 ...fill_cp_phase1_narrow_output_collection.md | 680 ++++++++++++++++++
 2 files changed, 970 insertions(+)
 create mode 100644 docs/advanced_features/nsa_prefill_cp_all_gather.md
 create mode 100644 docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md

diff --git a/docs/advanced_features/nsa_prefill_cp_all_gather.md b/docs/advanced_features/nsa_prefill_cp_all_gather.md
new file mode 100644
index 000000000..04f429e95
--- /dev/null
+++ b/docs/advanced_features/nsa_prefill_cp_all_gather.md
@@ -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 都有可能被替换或减少。
diff --git a/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md b/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md
new file mode 100644
index 000000000..f0c16b52e
--- /dev/null
+++ b/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md
@@ -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
+
+这是当前代码库里风险最低、最容易验证收益的一步。