# 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 都有可能被替换或减少。