执行摘要

• 一句话：MRv2 新增 extract_hidden_states 投机解码支持
• 推荐动作：值得精读，尤其是新投机器的 propose 签名与 set_forward_context 用法，以及"draft token 恒等于目标采样 token、验证必接受"的设计思路。对后续向 MRv2 增加新的 speculative method 有直接参考价值；CUDA Graph 缺口可作为后续优化专题跟进。

功能与动机

关联 issue #49562 指出 MRv2 的 init_speculator() 未注册 extract_hidden_states，会抛出 NotImplementedError(f"{speculative_config.method} is not supported yet.")，而 MRv1 已通过 ExtractHiddenStatesProposer 支持该方法，用于 draft 模型数据生成与分析。mgoin 在评论中表示："i didn't realize this wasn't ported to mrv2 yet. We should definitely prioritize this as we increasingly have more mrv2-only models"。issue 同时指出若不做移植，下游只能依赖 out-of-tree fork 或另建独立导出 API，维护成本与范围风险都更高。

实现拆解

1. 配置放行：vllm/config/vllm.py::_get_v2_model_runner_unsupported_features 从 MRv2 不支持特性清单中移除 extract_hidden_states；tests/test_config.py 新增 test_v2_model_runner_supports_extract_hidden_states 断言 _get_v2_model_runner_unsupported_features() 返回空列表。
2. 投机器分发：vllm/v1/worker/gpu/spec_decode/__init__.py::init_speculator 在 dflash 之前新增 method == "extract_hidden_states" 分支，延迟导入并实例化 ExtractHiddenStatesSpeculator，与既有 DFlash、DSpark、MTP、EAGLE 等保持一致。
3. 目标模型输出契约：vllm/v1/worker/gpu/model_runner.py 将 extract_hidden_states 加入 use_aux_hidden_state_outputs 方法集合，目标模型在 forward 时输出指定层的 aux_hidden_states；PP 组合仍走既有 ValueError 拒绝路径。
4. 新投机器实现：新建 vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py：初始化校验 draft_sample_method 必须为 greedy、num_speculative_tokens 必须为 1、不支持 disable_padded_drafter_batch，且 draft 模型 HF 配置必须提供 eagle_aux_hidden_state_layer_ids；按层数与 hidden_size 预分配 hidden_states 缓冲；load_draft_model 以 set_model_tag("extract_hidden_states") 复用 ExtractHiddenStatesModel，load_model 校验恰好一个注意力层；init_cudagraph_manager 与 capture 为空实现，明确跳过 CUDA Graph；propose 直接取 last_sampled 作为 draft token，校验并堆叠 aux_hidden_states 后拷入缓冲，再在 set_forward_context 提供的缓存层注意力元数据下驱动模型 forward，返回 draft token（验证必接受）。
5. 测试配套：新增 tests/v1/worker/test_gpu_extract_hidden_states_speculator.py 覆盖 greedy 限制、init_speculator 分发、propose 缓存与返回逻辑、缺失 aux_hidden_states 的失败路径；PR body 报告 predictable dummy model 的 MRv2 GPU E2E 通过。

关键文件：

• vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py（模块 投机解码；类别 source；类型 core-logic；符号 ExtractHiddenStatesSpeculator, load_draft_model, load_model, init_cudagraph_manager）: 新增 MRv2 ExtractHiddenStatesSpeculator 核心实现：配置校验、加载缓存模型、在执行步把目标模型 aux hidden states 拷入预分配缓冲并通过 forward 写入缓存，同时返回必然被接受的 draft token。
• tests/v1/worker/test_gpu_extract_hidden_states_speculator.py（模块 单元测试；类别 test；类型 test-coverage；符号 _RecordingModel, test_init_requires_greedy_draft_sampling, test_init_speculator_dispatches_extract_hidden_states, test_propose_caches_hidden_states_and_returns_sampled_tokens）: 新增聚焦单元测试，覆盖 greedy 限制、init_speculator 分发、propose 缓存行为与失败路径，为后续 CG 支持和回归提供基线。
• vllm/v1/worker/gpu/model_runner.py（模块 运行器；类别 source；类型 data-contract）: 将 extract_hidden_states 加入 use_aux_hidden_state_outputs 方法集合，使目标模型在 forward 时输出 aux hidden states 供投机器消费；继续拒绝与 PP 组合。
• vllm/v1/worker/gpu/spec_decode/__init__.py（模块 投机解码；类别 source；类型 dependency-wiring；符号 init_speculator）: init_speculator 新增 extract_hidden_states 分发分支，是 MRv2 接入该投机方法的入口。
• vllm/config/vllm.py（模块 配置校验；类别 source；类型 configuration；符号 _get_v2_model_runner_unsupported_features）: 配置放行点：从 MRv2 不支持特性白名单移除 extract_hidden_states，否则引擎会直接拒绝配置。
• tests/test_config.py（模块 配置校验；类别 test；类型 test-coverage；符号 test_v2_model_runner_supports_extract_hidden_states）: 配置层回归测试，确保 MRv2 不再把 extract_hidden_states 判为不支持特性。

关键符号：ExtractHiddenStatesSpeculator.init, ExtractHiddenStatesSpeculator.propose, ExtractHiddenStatesSpeculator.load_draft_model, ExtractHiddenStatesSpeculator.load_model, ExtractHiddenStatesSpeculator.init_cudagraph_manager, ExtractHiddenStatesSpeculator.capture, init_speculator, _get_v2_model_runner_unsupported_features

评论区精华

审阅者 mgoin 在 extract_hidden_states.py:147 留下唯一 inline 评论："We should be able to add CG support in the future"，确认当前跳过 CUDA Graph 是刻意取舍而非遗漏；随后两次 APPROVED，第二次附言 "LGTM, @shanjiaz has tested e2e with speculators"。提交历史中第二笔提交 "Restrict hidden-state extraction to greedy sampling"（mgoin 与 OpenAI Codex 共同署名）把实现收敛到 greedy 采样，原因是本方法 draft token 恒等于目标采样 token，验证必然接受，probabilistic 采样无意义。CI 执行期间 zupengwang 定位到失败源于 #52131 迁移测试文件后路径失效（对应 #52939 修复），与本 PR 无关，重建后全绿。

• CUDA Graph 支持暂缺 (design): 合入时接受该取舍，CG 支持列为未来工作；propose 采用 inference_mode 并复用缓存模型 forward。
• 限制为 greedy 采样 (design): 已合入限制，并在 init 抛 ValueError，配套测试 test_init_requires_greedy_draft_sampling 覆盖。
• CI 失败与 #52131 测试文件迁移 (other): Buildkite CI #84817 全绿后由 mgoin 合入。

风险与影响

• 风险：
  1. CUDA Graph 缺失：init_cudagraph_manager 与 capture 均为空实现，extract_hidden_states 路径无法享受 CG 提速，可能影响高频采集场景的吞吐；mgoin 已标注为未来优化项。
  2. 数据契约耦合：propose 强依赖目标模型输出的 aux_hidden_states 数量与顺序，model_runner.py 的 use_aux_hidden_state_outputs 开关一旦与投机器不一致，会在运行时抛 ValueError。
  3. 配置面变严：非 greedy 采样、num_speculative_tokens != 1、启用 disable_padded_drafter_batch 的既有配置在 MRv2 下会直接失败，属于预期行为但需在文档中强调。
  4. 验证深度有限：PR body 明确未做精度或性能 benchmark；GPU E2E 使用 predictable dummy model，真实模型（如 Qwen3.5）smoke 虽过但未在最终提交上 rerun。
     - 影响：用户侧：MRv2 用户可启用 extract_hidden_states 做 draft 数据生成/分析，并与 hidden-states KV connector 兼容，填补 MRv1/MRv2 功能缺口。系统侧：MRv2 spec_decode 家族新增"纯缓存"型投机方法，init_speculator 分发模式保持一致；目标模型输出 aux_hidden_states 的开关集合被扩展。团队侧：PR 仅 6 文件、+314/-2，范围小，maintainer 明确认可并合入；后续维护点集中在 CG 支持与 PP 支持。
     - 风险标记：跳过 CUDA Graph 捕获, 依赖 aux_hidden_states 数据契约, 管线并行组合仍不支持, 缺少真实模型精度/性能基准, 限制 greedy 与单步投机

关联脉络

• 暂无明显关联 PR