执行摘要

• 一句话：清理推测解码中的空闲存根和形状检查模式
• 推荐动作：值得精读以了解推测解码中空闲状态处理的正确模式。特别是 create_idle_input 工厂方法与裸构造函数的权衡，以及 input_ids.shape[0] 优于 numel 的理由。但需留意 review 中提到的 MultiLayerEagleWorker 潜在隐藏大小不匹配问题，建议跟进修复。

功能与动机

PR #24859 将 EagleDraftInput 拆分为单独的 EagleDraftExtendInput 后，原有的空闲存根模式（裸构造 EagleDraftInput(capture_hidden_mode=...)）会产生 None 字段，导致后续 merge_batch/filter_batch 操作崩溃。此外，input_ids.numel() 与 input_ids.shape[0] 的混用存在潜在语义不一致。

实现拆解

1. 替换空闲存根创建逻辑：在 eagle_worker.py 和 multi_layer_eagle_worker.py 的 forward_batch_generation 方法中，将裸构造 EagleDraftInput(capture_hidden_mode=CaptureHiddenMode.LAST) 替换为调用统一方法 self._draft_preprocess_idle(batch)，该方法内部使用 create_idle_input 工厂方法创建包含零长度张量的合法实例。
2. 统一形状检查方式：在三个 worker 的 forward_draft_extend_after_decode 方法和 frozen_kv_mtp_worker.py 的相同方法中，将 draft_extend_input.input_ids.numel() == 0 全部改为 draft_extend_input.input_ids.shape[0] == 0，确保语义清晰且避免因 numel 与 shape[0] 在非连续张量下可能的不一致。
3. 添加数据类文档：在 eagle_info.py 的 EagleDraftInput 类定义前插入注释，明确声明 topk_p、topk_index、hidden_states、bonus_tokens 是 batch-uniform 字段，调用者应使用工厂方法而非裸构造函数。

关键文件：

• python/sglang/srt/speculative/eagle_worker.py（模块 推测解码；类别 source；类型 core-logic；符号 forward_batch_generation, forward_draft_extend_after_decode）: V1 EAGLE worker，替换空闲存根创建并统一形状检查方式，是本次变更的核心文件之一。
• python/sglang/srt/speculative/multi_layer_eagle_worker.py（模块 推测解码；类别 source；类型 core-logic；符号 forward_batch_generation, forward_draft_extend_after_decode）: 多层 EAGLE worker，与 eagle_worker.py 同步变更，但 review 指出继承的 _draft_preprocess_idle 可能存在 hidden_size 不匹配风险。
• python/sglang/srt/speculative/frozen_kv_mtp_worker.py（模块 推测解码；类别 source；类型 core-logic；符号 forward_batch_generation, forward_draft_extend_after_decode）: FrozenKV MTP worker，变更最小（仅两处 numel 改为 shape[0]），但覆盖了第三种推测解码方案。
• python/sglang/srt/speculative/eagle_info.py（模块 推测解码；类别 source；类型 documentation；符号 EagleDraftInput）: 在 EagleDraftInput 数据类前添加文档注释，说明 batch-uniform 字段约束，帮助开发者避免误用。

关键符号：forward_batch_generation, forward_draft_extend_after_decode, _draft_preprocess_idle

关键源码片段

python/sglang/srt/speculative/eagle_worker.py

V1 EAGLE worker，替换空闲存根创建并统一形状检查方式，是本次变更的核心文件之一。

# python/sglang/srt/speculative/eagle_worker.py ( 简化片段 )
def forward_batch_generation(self, batch):
    # ... 前面省略
    with self.draft_tp_context(...):
        draft_extend_input = verify_output.draft_extend_input
        if self.server_args.enable_dp_attention or draft_extend_input.input_ids.shape[0] > 0:
            # decode 未结束：暂存 extend_input，然后获取下一轮的 draft_input
            batch.spec_info = draft_extend_input
            next_draft_input = self.forward_draft_extend_after_decode(batch)
            batch.spec_info = next_draft_input
        else:
            # 所有请求已完成且 DP attention 不强制 extend：
            # 安装一个空闲的 EagleDraftInput，使得下一轮的调度操作
            # （merge_batch / filter_batch）能看到类型正确的空张量而不是 None。
            self._draft_preprocess_idle(batch)
    # ... 后续不变
​
def forward_draft_extend_after_decode(self, batch):
    draft_extend_input = batch.spec_info
    input_is_idle = batch.forward_mode.is_idle()
    if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
        # 所有请求已验证完成：切换到空闲 ExtendInput
        batch = batch.copy()
        batch.prepare_for_idle()
        hidden_size = (
            self.model_config.hidden_size * 3
            if self.speculative_algorithm.is_eagle3()...
            else self.model_config.hidden_size
        )
        draft_extend_input = EagleDraftExtendInput.create_idle_input(
            device=self.device, hidden_size=hidden_size, dtype=self.model_config.dtype
        )
    # ... 后续逻辑

python/sglang/srt/speculative/eagle_info.py

在 EagleDraftInput 数据类前添加文档注释，说明 batch-uniform 字段约束，帮助开发者避免误用。

# python/sglang/srt/speculative/eagle_info.py
@dataclass
class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
    # 对于空闲存根请使用 create_idle_input，而不是裸构造函数：
    # filter_batch / merge_batch 会无条件地对 topk_p / topk_index /
    # hidden_states / bonus_tokens 进行切片 / 拼接。
​
    # shape: (b, topk)
    topk_p: torch.Tensor = None
    topk_index: torch.Tensor = None
    # shape: (b, hidden_size) - 每个请求一个 hidden，供 draft 前向使用。
    hidden_states: torch.Tensor = None
    capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
    # 每个请求的 bonus token，由 EagleDraftExtendInput.prepare_extend_after_decode 写入。
    bonus_tokens: torch.Tensor = None
    # ... 其余字段

评论区精华

review 中 gemini-code-assist[bot] 指出 MultiLayerEagleWorker 中 _draft_preprocess_idle 方法继承自父类，但父类使用 self.model_config.spec_hidden_size，而 MultiLayerEagleWorker 的 idle 分支使用 self.model_config.hidden_size，可能导致 AttributeError 或形状不匹配。这一潜在问题未在 PR 中修复，作者已合并 PR 但未回应。

• MultiLayerEagleWorker 中 _draft_preprocess_idle 的 hidden_size 潜在不匹配 (correctness): 未解决；PR 已合并但作者未回应此评论。

风险与影响

• 风险：MultiLayerEagleWorker 调用继承的 _draft_preprocess_idle 可能因 spec_hidden_size 缺失而崩溃（低概率，取决于 MTP 模型配置是否提供该属性）。若 spec_hidden_size 与 hidden_size 不同，会导致 merge_batch 时张量拼接形状不匹配（中等风险）。建议确认 MTP 模型配置是否始终包含 spec_hidden_size，或直接在 MultiLayerEagleWorker 中覆写该方法。
• 影响：影响范围限于 V1 EAGLE、多层 EAGLE 和 FrozenKV MTP 三种推测解码 worker。对正常推理路径无影响，仅在所有请求完成且 DP attention 未强制 extend 的边界情况下触发。改进后空闲存根的行为更健壮，可防止后续 merge_batch/filter_batch 因 None 字段崩溃。
• 风险标记：缺少测试覆盖, 潜在隐藏大小不匹配, 审核意见未解决

关联脉络

• PR #24859 [Spec V1] Split draft-extend phase from EagleDraftInput into new EagleDraftExtendInput: 本 PR 是对 #24859 的后续清理：引入 create_idle_input 工厂方法并统一形状检查，解决 #24859 中因拆分引入的空闲存根 None 字段问题。