执行摘要

• 一句话：支持新一代 EAGLE-3 Draft 模型
• 推荐动作：值得精读：展示了如何在不破坏向后兼容的前提下扩展 speculative decoding 架构。特别关注 nn.ModuleList 替代单层、动态 num_aux、以及归一化位置的设计权衡。与 PR#24826 配合理解可窥见 EAGLE 系列的整体演进方向。

功能与动机

即将发布论文和新模型检查点（gpt-oss-20b/120b），需要 SGLang 提供 Day-0 支持。当前 EAGLE-3 实现硬编码为单层和 3 个辅助隐藏状态，无法兼容 vLLM 等其他框架导出的多层模型。改进后可运行更多模型并支持论文中的新设计。

实现拆解

1. Draft 模型结构调整 （python/sglang/srt/models/llama_eagle3.py）：

   • 将 self.midlayer 改为 self.layers（nn.ModuleList），支持多个 Draft 层。
   • 增加 is_input_layer 标志，使输入层的 QKV 投影维度正确（2x hidden_size），其余层保持原尺寸。
   • 动态计算 num_aux_hidden_states：优先读取配置中的显式值，其次从 eagle_config.eagle_aux_hidden_state_layer_ids 推导，最后默认 3。
   • 用 fc_norm 或 use_aux_norm 标志控制是否在 FC 之前对各辅助隐藏状态施加独立的 RMSNorm。
   • 权重加载时兼容旧检查点：将 midlayer 键名映射为 layers.0。

2. 推理数据结构适配 （python/sglang/srt/speculative/eagle_info.py）：

   • EagleDraftExtendInput.hidden_size_for 方法根据 num_aux_hidden_states 和 target_hidden_size 计算辅助隐藏状态的展平宽度，替代原来的固定 3 * target_hidden。
   • 逻辑改为：非 EAGLE-3 或无 aux 模式时返回 spec_hidden_size；否则根据动态 num_aux 计算。

3. 配套重构：

   • 删除了原来的 use_aux_norm 分支，统一为 fc_norm，保持向后兼容。
   • 移除了原有硬编码的三个 aux norm 层，改用 nn.ModuleList 动态创建。

关键文件：

• python/sglang/srt/models/llama_eagle3.py（模块 Draft 模型；类别 source；类型 core-logic；符号 LlamaDecoderLayer, LlamaModel, LlamaModel.load_weights）: 核心 Draft 模型定义，改动涉及多层支持、可变 aux 数量和归一化位置，影响所有 EAGLE-3 检查点加载和推理。
• python/sglang/srt/speculative/eagle_info.py（模块 推理数据；类别 source；类型 core-logic；符号 EagleDraftExtendInput.hidden_size_for）: 定义了投机解码推理阶段的数据结构，hidden_size_for 方法根据 Draft 模型实际 aux 数量计算特征宽度，与 Draft 模型改动联动。

关键符号：LlamaDecoderLayer.init, LlamaDecoderLayer.forward, LlamaModel.init, LlamaModel.forward, LlamaModel.load_weights, EagleDraftExtendInput.hidden_size_for

评论区精华

• Qiaolin-Yu 质疑将 midlayer 替换为 layers.0 的方式“有点 hacky”。Dogacel 解释这是为了兼容不同检查点命名（如 Redhat 使用 layers.0，Lmsys 使用 midlayer），需保持向后兼容。
• Qiaolin-Yu 建议使用 hidden_states_to_aux 辅助方法简化 forward 中的分支。Dogacel 表示认同。

• kpham-sgl 指出 eagle_worker.py 和 multi_layer_eagle_worker.py 中的隐藏大小计算可以复用 PR#24826 新引入的 _get_eagle_aux_layer_count 函数。同时质疑为什么从 target config 读取 num_aux 而不是 Draft config。

• midlayer 兼容性处理 (design): 保留替换方式，因为这是最简洁的兼容方案，且只在权重加载时触发，不影响运行时。

• forward 中隐藏状态处理 (design): 未在最终补丁中看到实际修改，可能作为后续改进。讨论已接受建议意图。
• 复用 _get_eagle_aux_layer_count (refactor): 作为待办事项，未在此 PR 中集成。合并后可能在其他 PR 中重构。
• num_aux 读取来源 (question): 未解决，可能当前实现符合预期（因为 draft 模型配置继承自 target），但需确认。

风险与影响

• 风险：
  • 向后兼容风险：midlayer -> layers.0 的映射只在权重加载时处理，若其他代码路径使用 midlayer 命名可能出现 KeyError。已通过替换逻辑缓解，但需确保所有引用同步。
  • 配置解析逻辑：num_aux_hidden_states 的解析链（显式 > eagle_config > 默认）可能存在歧义，特别是当 eagle_config 缺失时。需要验证是否所有场景都能正确 fallback。
  • 性能影响：is_input_layer 条件分支在 forward 中引入额外判断，但对推理性能影响极微。
  • 测试覆盖：PR 未包含直接的新增测试，依赖已有集成测试覆盖。风险较低但需注意回归。
• 影响：
  • 用户影响：支持加载更多 EAGLE-3 变体模型（多层、可变 aux），显著提升投机解码的吞吐量和接受长度（详见 benchmark）。
  • 系统影响：Draft 模型结构改变，但保持 API 兼容。已有模型检查点（仅单层）仍可正常工作。
  • 团队影响：此 PR 为即将发布的论文和模型提供 Day-0 支持，降低后续集成成本。
  • 风险标记：向后兼容性处理, 配置解析逻辑, 缺少测试覆盖, 跨 PR 代码重复

关联脉络

• PR #24826 [Spec] Add _get_eagle_aux_layer_count helper: 该 PR 引入了 _get_eagle_aux_layer_count 辅助函数，与本 PR 中的 num_aux_hidden_states 解析逻辑重复，reviewer 建议复用。
• PR #25013 spec: route idle hidden_size via EagleDraft{,Extend}Input classmethods: 同属 speculative-decoding 模块，重构了隐藏大小路由，与本 PR 的 hidden_size_for 方法改动相关。
• PR #24262 (3/n - prefill optimize)[LoRA][MoE] Optimize virtual experts: 同属性能优化相关，但与本 PR 无直接代码依赖。