# PR #28343 完整报告 - 仓库：`sgl-project/sglang` - 标题：[Kimi K2.5] Fix eagle3 aux capture for tp>1 when AR fusion is enabled - 合并时间：2026-06-18 03:59 - 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/28343 --- # 执行摘要 - 一句话：修复 EAGLE3 aux 捕获在 AR 融合下的 per-rank 碎片问题 - 推荐动作：值得精读，特别是 `_post_attn_residual_is_read_only` 的设计模式可作为类似 fusion-pass 中捕获正确状态的范例。建议合并前确认是否有针对 TP>1 的集成测试覆盖，后续可考虑添加。 # 功能与动机 在 TP>1 且启用 AR 融合时，MLP all-reduce 被推迟到下一层的 RMSNorm 中融合执行，原先的 EAGLE3 aux 捕获在 `hidden_states + residual` 时尚未完成全归约，导致捕获了 per-rank 部分和，静默降低 draft 质量。PR body 的流程图清晰对比了 AR 融合开启 / 关闭的路径差异。 # 实现拆解 1. **communicator.py**- 在 `prepare_attn_and_capture_last_layer_outputs` 中添加 `quant_format` 参数并透传给 `prepare_attn`（用于 DeepSeek 的 gfx95 格式，其他调用方默认 `""`）。修改捕获逻辑：当 `gathered_last_layer_output` 是 `residual` 且 `_post_attn_residual_is_read_only` 返回 `False` 时才执行 clone，避免在 AR 融合已生成新 tensor 时不必要 clone。新增 `_post_attn_residual_is_read_only` 方法，检测当前 norm 函数是否使用 FlashInfer all-reduce-fusion 路径且未经过 input-scattered 分支，若是则 `residual` 已由融合内核生成新 tensor，是只读的，可跳过 clone。 2. **deepseek_v2.py**- `DeepseekV2DecoderLayer.forward` 新增 `captured_last_layer_outputs` 参数，将以往直接调用 `prepare_attn` 改为调用 `prepare_attn_and_capture_last_layer_outputs`，同时传入 `captured_last_layer_outputs` 和 `quant_format`。`DeepseekV2Model.forward` 移除原先使用 `all_gather` 层间手动计算 `hidden_states + residual` 的代码，改为仅在 `layers_to_capture` 的目标层向 `layer.forward` 传递 `captured_last_layer_outputs` 参数（即 `aux_hidden_states`），使捕获发生在正确的后融合节点。 关键文件： - `python/sglang/srt/layers/communicator.py`（模块 通信层；类别 source；类型 core-logic；符号 _post_attn_residual_is_read_only）: 核心修复：新增 `_post_attn_residual_is_read_only` 方法判断 residual 是否只读；修改 `prepare_attn_and_capture_last_layer_outputs` 以条件性跳过 clone，并透传 `quant_format`。 - `python/sglang/srt/models/deepseek_v2.py`（模块 模型定义；类别 source；类型 data-contract）: 模型层调用改为使用 `prepare_attn_and_capture_last_layer_outputs`，移除冗余的 `all_gather` 手动计算，并将 `captured_last_layer_outputs` 传入 `DeepseekV2DecoderLayer.forward`。 关键符号：_post_attn_residual_is_read_only, prepare_attn_and_capture_last_layer_outputs ## 关键源码片段 ### `python/sglang/srt/layers/communicator.py` 核心修复：新增 `_post_attn_residual_is_read_only` 方法判断 residual 是否只读；修改 `prepare_attn_and_capture_last_layer_outputs` 以条件性跳过 clone，并透传 `quant_format`。 ```python # python/sglang/srt/layers/communicator.py (head 版本 ) def prepare_attn_and_capture_last_layer_outputs( self, hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, post_residual_addition: Optional[torch.Tensor] = None, quant_format: str = "", # 新增参数，用于 DeepSeek gfx95 量化格式 ): hidden_states, residual = self.prepare_attn( hidden_states, residual, forward_batch, quant_format=quant_format, # 透传给 prepare_attn post_residual_addition=post_residual_addition, ) if captured_last_layer_outputs is not None: gathered_last_layer_output = self._communicate_simple_fn( hidden_states=residual, forward_batch=forward_batch, context=self._context, ) # 仅在 residual 不是只读（即可能被后续 in-place 修改）时才 clone if ( gathered_last_layer_output is residual and not self._post_attn_residual_is_read_only(residual) ): gathered_last_layer_output = residual.clone() captured_last_layer_outputs.append(gathered_last_layer_output) return hidden_states, residual def _post_attn_residual_is_read_only(self, residual: torch.Tensor) -> bool: """判断 prepare_mlp 的 post-attention RMSNorm 是否会修改 residual。 只有当使用 FlashInfer all-reduce-fusion 内核时，residual_out 是新分配的张量， 原始 residual 不会被修改，因此可以安全引用而无需 clone。 """ norm_fn = getattr( self._communicate_with_all_reduce_and_layer_norm_fn, "func", self._communicate_with_all_reduce_and_layer_norm_fn, ) # 仅当使用 gather 模式的 norm 函数时才可能走 fusion 路径 uses_gather_norm = norm_fn in ( CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual_moe, ) return ( uses_gather_norm and not get_attn_tp_context().input_scattered and apply_flashinfer_allreduce_fusion(residual.shape[0]) ) ``` ### `python/sglang/srt/models/deepseek_v2.py` 模型层调用改为使用 `prepare_attn_and_capture_last_layer_outputs`，移除冗余的 `all_gather` 手动计算，并将 `captured_last_layer_outputs` 传入 `DeepseekV2DecoderLayer.forward`。 ```python # python/sglang/srt/models/deepseek_v2.py (head 版本 , DeepseekV2DecoderLayer.forward) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], zero_allocator: BumpAllocator, gemm_output_zero_allocator: BumpAllocator = None, llama_4_scaling: Optional[torch.Tensor] = None, prev_topk_indices: Optional[torch.Tensor] = None, captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, # 新增参数 ) -> torch.Tensor: hidden_states_orig = hidden_states # 改用 prepare_attn_and_capture_last_layer_outputs， # 它会在 prepare_attn 之后捕获 post-AR 的 residual hidden_states, residual = ( self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( hidden_states, residual, forward_batch, captured_last_layer_outputs=captured_last_layer_outputs, quant_format=getattr(self, "_gfx95_quant_format", ""), ) ) # ... 后续的 self_attn, prepare_mlp 等不变 ``````python # DeepseekV2Model.forward 中循环调用 layer.forward 的部分 for i in range(self.start_layer, self.end_layer): # ... 上下文设置 ... # 原先手动的 all_gather + hidden_states + residual 代码被移除， # 改为通过 captured_last_layer_outputs 参数传递 layer = self.layers[i] hidden_states, residual, topk_indices = layer( positions, hidden_states, forward_batch, residual, zero_allocator, gemm_output_zero_allocator, llama_4_scaling, prev_topk_indices=topk_indices, captured_last_layer_outputs=( aux_hidden_states if i in self.layers_to_capture else None ), ) ``` # DeepseekV2Model.forward 中循环调用 layer.forward 的部分 for i in range(self.start_layer, self.end_layer): # ... 上下文设置 ... # 原先手动的 all_gather + hidden_states + residual 代码被移除， # 改为通过 captured_last_layer_outputs 参数传递 layer = self.layers[i] hidden_states, residual, topk_indices = layer( positions, hidden_states, forward_batch, residual, zero_allocator, gemm_output_zero_allocator, llama_4_scaling, prev_topk_indices=topk_indices, captured_last_layer_outputs=( aux_hidden_states if i in self.layers_to_capture else None ), ) ``` # 评论区精华 PR 无 review comments，仅获得 ch-wan 的 APPROVE。commit history 显示 4 次迭代：最初是简单的修复，然后改为复用融合内核的 residual_out，接着移除不必要的 clone，最后 `refactor` 整合到统一方法中。 - 暂无高价值评论线程 # 风险与影响 - 风险：核心风险在于 `_post_attn_residual_is_read_only` 的判定条件：误判为 `True` 会导致跳过必要的 clone，使后续的 `captured_last_layer_outputs` 引用可能被 in-place 操作篡改，产生静默错误。目前条件只限定了 FlashInfer all-reduce-fusion + gather_norm + 非 input_scattered 的特定组合，但若有其他后端（如 aiter）新增类似输出新 tensor 的路径，需同步更新此方法。另外，仅在 TP>1 场景暴露，单卡不会触发。 - 影响：直接影响所有使用 Kimi K2.5/K2.6 且 TP>1 并开启 AR 融合的 EAGLE3 推理场景。根据 PR 提供的基准测试，draft 接受长度提升 0.3%-4.2%（MTBench +2.5%，HumanEval +3.1%），速度无额外开销（TPOT 从 3.43 降至 3.40）。不影响单卡或 AR 融合关闭的配置。由于未携带正式单元测试，回归风险较低但依赖人工验证。 - 风险标记：核心路径变更 , 缺少测试覆盖 # 关联脉络 - PR #28465 Batch EAGLE draft/draft-extend replay memcpys via grouped foreach copy: 同样是 EAGLE3 性能优化 PR，改动了相关 eagle draft 文件，可一起关注 speculative-decoding 的整体演进。 - PR #28500 [Perf] Make spec-decode penalty H2D non-blocking and share decode cumulate path: 涉及 speculative-decoding 的另一个性能优化，与当前 PR 同属推理优化系列。