sglang/docs/advanced_features/nsa_prefill_cp_all_gather.md

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 用。

---

总体流程图

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