执行摘要

• 一句话：修复 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/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/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 同属推理优化系列。