执行摘要

• 一句话：修复Eagle3/DFLASH推测解码在CUDA图捕获时辅助隐藏状态未启用的问题。
• 推荐动作：该PR值得精读，特别是对于涉及CUDA图捕获和推测解码的开发者。关注点包括：初始化顺序的重要性、配置集中化的设计决策、以及如何避免重复调用导致的参数不一致。

功能与动机

根据PR body描述，之前set_eagle3_layers_to_capture()在initialize()中CUDA图捕获之后被调用，导致捕获的图运行时未启用辅助隐藏状态捕获，对于Eagle3会导致运行时接受长度为零。CudaGraphRunner中的变通方案调用了无参数的set_eagle3_layers_to_capture()，使用了默认层ID而非配置指定的ID，破坏了具有自定义eagle_aux_hidden_state_layer_ids的模型。

实现拆解

1. 重构配置逻辑：在model_runner.py中创建新方法init_aux_hidden_state_capture()，将Eagle3和DFLASH的辅助隐藏状态捕获配置（set_eagle3_layers_to_capture和set_dflash_layers_to_capture）集中于此，并添加文档说明必须在CUDA图捕获前调用。
2. 调整初始化顺序：在ModelRunner.initialize()中，在init_routed_experts_capturer()之后、init_device_graphs()之前插入self.init_aux_hidden_state_capture()调用，确保CUDA图捕获时辅助隐藏状态路径已启用。
3. 移除冗余代码：从ModelRunner.initialize()中删除原有的条件配置代码块，并从CudaGraphRunner.__init__()中完全移除相关配置调用，避免重复执行和参数不一致。
4. 清理遗留调用：在ModelRunner._dummy_run()中移除对set_eagle3_layers_to_capture()和set_dflash_layers_to_capture()的调用，添加注释说明配置已在initialize()中完成。

关键文件：

• python/sglang/srt/model_executor/model_runner.py（模块 模型执行器；类别 source；类型 core-logic；符号 init_aux_hidden_state_capture）: 核心变更文件，重构了辅助隐藏状态捕获的配置逻辑和初始化顺序。
• python/sglang/srt/model_executor/cuda_graph_runner.py（模块 CUDA图运行器；类别 source；类型 cleanup）: 移除了冗余的辅助隐藏状态配置调用，避免参数不一致和重复执行。

关键符号：init_aux_hidden_state_capture

关键源码片段

python/sglang/srt/model_executor/model_runner.py

核心变更文件，重构了辅助隐藏状态捕获的配置逻辑和初始化顺序。

def init_aux_hidden_state_capture(self):
    """Configure auxiliary hidden state capture for speculative decoding.
​
    Must be called before CUDA graph capture so the captured graphs
    include aux hidden state output paths.
    """
    if self.eagle_use_aux_hidden_state:
        # 设置Eagle3模型层以捕获辅助隐藏状态，使用配置的层ID列表
        self.model.set_eagle3_layers_to_capture(
            self.eagle_aux_hidden_state_layer_ids
        )
    if self.dflash_use_aux_hidden_state:
        # 检查模型是否支持DFLASH辅助隐藏状态捕获
        if not hasattr(self.model, "set_dflash_layers_to_capture"):
            raise ValueError(
                f"Model {self.model.__class__.__name__} does not implement "
                "set_dflash_layers_to_capture, which is required for DFLASH."
            )
        # 设置DFLASH模型目标层以捕获辅助隐藏状态
        self.model.set_dflash_layers_to_capture(self.dflash_target_layer_ids)

评论区精华

Review评论中，merrymercy指出旧代码在cuda_graph_runner.py和model_runner.py中两次调用类似代码且参数不同，这是错误的。结论是应该在model_runner.py::initialize中统一调用，而非在CudaGraphRunner中。这直接导致了PR的实现方案：将配置逻辑集中到init_aux_hidden_state_capture()并在正确时机调用。

• 配置调用时机与冗余问题 (correctness): 应在model_runner.py::initialize中统一调用，移除CudaGraphRunner中的冗余代码。

风险与影响

• 风险：1. 回归风险：修改了初始化顺序，若其他模块依赖init_aux_hidden_state_capture()的执行时机，可能引入意外行为。但变更范围小，且通过测试验证。
  2. 性能风险：无，仅是配置调用时机调整，不涉及运行时逻辑。
  3. 兼容性风险：对于使用自定义eagle_aux_hidden_state_layer_ids或dflash_target_layer_ids的模型，修复后能正确应用配置，但需确保配置参数传递正确。
  4. 测试覆盖：PR body提到已测试Eagle3推测解码与CUDA图，确认非零接受长度，并测试了自定义层ID配置，但未提供自动化测试变更，可能依赖现有测试套件。
• 影响：1. 用户影响：使用Eagle3或DFLASH推测解码且启用CUDA图的用户将获得正确的辅助隐藏状态捕获，避免接受长度为零或配置失效问题，提升推测解码效果。
  2. 系统影响：修复了CUDA图捕获与辅助隐藏状态配置的时序问题，确保推测解码算法在图形化运行时正常工作。
  3. 团队影响：简化了配置逻辑，移除冗余代码，提高了代码可维护性，并为未来类似功能提供了清晰模式。
• 风险标记：初始化顺序变更, 缺少测试覆盖

关联脉络

• PR #22897 streaming session: trim spec v2 overshoot in cache_finished_req: 同属推测解码（spec）相关修复，涉及缓存一致性，但本PR聚焦CUDA图捕获时序。
• PR #22862 Streaming session: fix retract tail leak via _free_tail: 同属流式会话和缓存相关修复，但本PR针对辅助隐藏状态捕获。