执行摘要

• 一句话：重构HiddenStatesConnector保存流程，支持chunked prefill和每请求配置
• 推荐动作：对于使用hidden states extraction的开发者，此PR移除了重要的限制（chunked prefill），建议更新配置并测试。对于核心开发者，此PR的异步保存设计和锁机制值得参考。Review中的安全和正确性讨论也很有价值。

功能与动机

当前ExampleHiddenStatesConnector在wait_for_save中保存hidden states，但在chunked prefill下该方法每个请求被多次调用。现有的解决方案是禁用chunked prefill。本PR通过将保存步骤移到仅调用一次的get_finished中移除此限制。

实现拆解

1. 数据模型重构：用PendingSave dataclass（含req_id, filename, token_ids, block_ids）替换旧的ReqMeta。元数据结构从requests列表改为pending_saves列表 + new_req_filenames字典。
2. 保存时序迁移：
   • save_kv_layer → 空操作（hidden states 已在 forward 时由CacheOnlyAttentionLayer缓存）。
   • wait_for_save → 仅为新请求预创建.lock文件。
   • request_finished → 创建PendingSave，返回(True, hidden_states_path)以延迟 block 释放。
   • build_connector_meta → 将_pending_saves打包进元数据发送给 worker。
   • get_finished → 读取 pending saves，启动异步 DtoH 拷贝和线程池写，完成后释放锁文件并返回完成的请求 ID。
3. Per-request 选项：通过extra_args → kv_transfer_params支持hidden_states_path（自定义路径）和include_output_tokens（包含输出 token hidden states）。新增allow_custom_save_path全局配置（默认 False），打开时允许客户端指定路径并输出警告。
4. 配置清理：删除vllm/config/vllm.py中禁用 chunked prefill 的强制检查，相关测试和 benchmark 移除了enable_chunked_prefill=False。
5. 测试与文档：集成测试启用 chunked prefill，新增 chunked prefill 场景和 TP=2 测试；示例脚本展示新选项；文档同步更新。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py（模块 KV连接器；类别 source；类型 core-logic；符号 ReqMeta, PendingSave, make_meta, add_request）: 核心逻辑重写：改变保存流程，新增PendingSave类，重构metadata，添加per-request选项
• tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py（模块 集成测试；类别 test；类型 test-coverage；符号 test_extract_hidden_states_tp2）: 测试覆盖：新增chunked prefill测试场景和TP>2测试，使用load_hidden_states替代直接safe_open
• examples/features/speculative_decoding/extract_hidden_states_offline.py（模块 示例；类别 source；类型 dependency-wiring）: 示例更新：展示per-request hidden_states_path和include_output_tokens用法，不再禁用chunked prefill
• vllm/config/vllm.py（模块 配置；类别 source；类型 core-logic；符号 _post_init_kv_transfer_config）: 移除ExampleHiddenStatesConnector强制禁用chunked prefill的检查代码
• benchmarks/benchmark_hidden_state_extraction.py（模块 基准测试；类别 source；类型 core-logic）: 移除enable_chunked_prefill=False配置以匹配新逻辑
• tests/v1/spec_decode/test_extract_hidden_states.py（模块 测试；类别 test；类型 test-coverage）: 移除enable_chunked_prefill=False配置以匹配新逻辑
• docs/features/speculative_decoding/extract_hidden_states.md（模块 文档；类别 docs；类型 documentation）: 文档更新：更新hidden states extraction配置说明和示例

关键符号：PendingSave, ExampleHiddenStatesConnectorMetadata, build_connector_meta, request_finished, get_finished, wait_for_save, _submit_async_write, save_kv_layer

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

核心逻辑重写：改变保存流程，新增PendingSave类，重构metadata，添加per-request选项

# PendingSave 表示一个待保存的 hidden states 提取请求，
# 保存时机从 wait_for_save 推迟到 get_finished。
@dataclass
class PendingSave:
    req_id: str # 请求 ID
    filename: str # 输出文件路径
    token_ids: torch.Tensor # 所有 token ID
    block_ids: list[int] # KV cache block ID 列表
​
@dataclass
class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):
    pending_saves: list[PendingSave] = field(default_factory=list)
    # 新请求的 req_id → filename 映射，worker 会预创建锁文件
    new_req_filenames: dict[str, str] = field(default_factory=dict)
​
def build_connector_meta(self, scheduler_output, **kwargs):
    # 打包 pending saves 和新请求文件名到 metadata 中发送给 worker
    meta = ExampleHiddenStatesConnectorMetadata()
    meta.pending_saves = list(self._pending_saves.values())
    self._pending_saves.clear()
    for req_id, req_data in scheduler_output.new_requests:
        filename = self._resolve_filename(req_id, req_data)
        meta.new_req_filenames[req_id] = filename
        self._request_filenames[req_id] = filename
    return meta
​
def request_finished(self, *, request, blocks, virtual_engine, **kwargs):
    # 创建 PendingSave 对象，延迟释放 blocks
    filename = self._request_filenames.pop(request.req_id)
    pending = PendingSave(
        req_id=request.req_id,
        filename=filename,
        token_ids=self._get_token_ids(request), # 含可选 output tokens
        block_ids=[b.block_id for b in blocks],
    )
    self._pending_saves[request.req_id] = pending
    return True, {'hidden_states_path': filename}

评论区精华

• 安全路径遍历（depthfirst-app[bot]）：客户端提供的hidden_states_path未验证，可能导致任意文件写入。作者通过添加allow_custom_save_path配置项（默认False）来缓解；启用时会输出警告。
• 排序连续性假设（mgoin）：询问block_id * block_size + offset顺序是否在所有模型（含hybrid/mamba）中安全。作者确认该公式是标准slot mapping，与attention后端一致。
• 拒绝draft token（mgoin）：担心all_token_ids包含被拒绝的draft token导致缺少缓存。作者澄清draft token与accepted token分开存储，all_token_ids只含已接受的token。

• TP>1锁文件竞争（shanjiaz）：所有TP rank同时创建相同锁文件导致冲突。作者修复为仅rank0创建锁文件并执行写入。

• 安全：hidden_states_path未验证导致路径遍历 (security): 通过新增allow_custom_save_path配置项（默认False），启用时输出警告，但若启用仍有风险。

• 正确性：block_id排序连续性和混合模型兼容性 (correctness): 作者确认该公式是标准slot mapping，与vLLM所有attention后端一致，写入和读取使用相同方式。
• 正确性：获取的all_token_ids是否包含被拒绝的draft token (correctness): 作者澄清draft tokens存储在separate list，返回的all_token_ids只包含accepted tokens，不存在缺失缓存问题。
• 测试：要求添加显式的chunked prefill测试 (testing): 作者添加了对应测试场景，包括非顺序层、chunked prefill和per-request选项。
• Bugfix: TP>1时多个rank重复创建锁文件导致冲突 (bugfix): 作者修复为仅TP rank0创建锁文件并执行写入，其余rank跳过。

风险与影响

• 风险：
  • 安全风险：当allow_custom_save_path=True时仍存在路径遍历风险，建议生产环境保持默认False，或进一步添加路径白名单。
  • 异步错误处理：_submit_async_write中线程池写操作可能静默失败，需要更完善的错误传播和重试机制。
  • 回归影响：request_finished现在返回True并延迟释放block，可能影响缓存管理器的工作流程，但已有单元覆盖。
  • 兼容性：移除了enable_chunked_prefill=False强制项，现有客户端若显式设置enable_chunked_prefill=False仍可工作，但建议移除以避免冲突。
• 影响：
  • 用户/开发者：现在可以在启用chunked prefill的情况下使用hidden states extraction，提高吞吐；支持每个请求独立配置保存路径和是否包含输出tokens。
  • 系统性能：文件写入移至get_finished（只调用一次），减少了重复的磁盘I/O和锁操作；TP>1场景下只有rank0写入，消除了重复操作。
  • 团队维护：代码复杂度增加，但流程更清晰；新配置项增加了测试面。
  • 风险标记：安全路径遍历（已缓解）, TP>1竞态（已修复）, 异步错误传播, 配置向后兼容

关联脉络

• PR #44733 [KV offload] Parallel-agnostic fs-tier cache for single full-attention group: 同属 kv-connector 模块，涉及 KV 缓存管理异步化设计
• PR #44243 [PD][Core] Fix Mamba prefix cache hit rate in PD disaggregation: 同属 KV Connector 模块，涉及 hidden states 相关的前缀缓存