执行摘要

本PR修复了Speculative Decoding V2中导致Eagle3模型崩溃的两个潜在bug：1）在eagle_worker_v2.py中补充设置forward_batch.return_logprob = False以保持与V1行为一致；2）在logits_processor.py中修复aux_pruned_states未填充的问题。这些修复提升了推测解码V2路径的稳定性和正确性，影响范围有限但关键，建议合并后关注相关测试结果。

功能与动机

根据PR body，当Eagle3模型启用SGLANG_ENABLE_SPEC_V2时发生崩溃。作者识别出两个问题：

• forward_batch.return_logprob = False在eagle_worker.py（V1）中已设置，但在V2版本中遗漏，可能导致参数不一致。
• aux_pruned_states在logits_processor.py的某些条件下未填充，引用作者原话："aux_pruned_states is not populated, it's populated in previous two conditions, but not for the condition we fixed."

实现拆解

1. eagle_worker_v2.py修改

在_draft_extend_for_prefill函数中添加一行代码：

forward_batch.return_logprob = False

这确保了与V1行为一致，避免因logprob返回设置导致的潜在问题。

2. logits_processor.py修改

在_get_pruned_states函数中新增逻辑：

• 初始化aux_pruned_states_lists：如果aux_hidden_states不为None，则为每个隐藏状态创建空列表。
• 在循环中填充：如果aux_pruned_states_lists不为None，遍历aux_hidden_states并切片添加到对应列表。
• 最终拼接：如果aux_pruned_states_lists不为None，将列表拼接为aux_pruned_states。
  关键代码片段：

aux_pruned_states_lists = (
    [[] for _ in aux_hidden_states]
    if aux_hidden_states is not None
    else None
)
...
if aux_pruned_states_lists is not None:
    for j, hidden in enumerate(aux_hidden_states):
        aux_pruned_states_lists[j].append(
            hidden[pt + start_len : pt + extend_len]
        )
...
if aux_pruned_states_lists is not None:
    aux_pruned_states = [torch.cat(lst) for lst in aux_pruned_states_lists]

评论区精华

Review讨论较少，只有Qiaolin-Yu的批准，没有具体技术交锋。从提交历史看，有一个合并主分支的提交（SHA: bad6521），可能涉及冲突解决，但未提供详细讨论。

风险与影响

风险

• 核心路径变更：logits_processor.py的修改涉及推测解码的核心逻辑，虽然修复了bug，但需确保不影响其他使用场景。
• 性能开销：aux_pruned_states的填充增加了列表操作和拼接，但通过条件检查（if aux_hidden_states is not None）避免了不必要的计算。
• 测试覆盖：缺少针对这两个修复的专项测试，依赖现有CI测试（已通过run-ci标签触发）。

影响

• 用户影响：直接修复Eagle3模型在SpecV2下的崩溃，提升用户体验。
• 系统影响：确保推测解码V2路径的状态一致性，避免运行时错误。
• 团队影响：维护了V1/V2行为一致性，减少后续开发混淆。

关联脉络

与近期PR的关联：

• PR#22146（隔离Spec V1路径）：同属推测解码相关修改，本PR修复V2 bug，可能补充了V2路径的完整性。
• PR#22104（重新启用SpecV2测试）：本PR的bug修复可能影响此类测试的稳定性，建议合并后验证测试结果。
  从历史PR看，推测解码（speculative-decoding）是近期活跃模块，本PR是其中重要的bugfix，有助于提升该功能的鲁棒性。