执行摘要

• 一句话：DSv4 eager 路径新增模型级缓冲池复用工作区，小幅优化 TTFT
• 推荐动作：值得精读。重点学习两个设计决策：一是 _packed_size 用 max 而非 sum 合并互斥 scratch 的思想——通过在注释中明确三类用途的 C4/C128 对齐约束来证明互斥性；二是 C++ 算子 out 变体模式（..._insert_out），把 kernel 内部分配改为调用方预分配，是削减 eager 路径分配开销的通用手法。对从事 CUDA Graph eager break、MLA 类模型优化或 kernel workspace 管理的工程师有直接参考价值。结合 review 中关于接口隐式约束的讨论，也能看到 vLLM 对函数接口友好性的坚持。

功能与动机

PR body 明确这是 issue #45861（DeepSeek V4 性能优化总任务）的子任务，标题为「Optimize workspace reuse for eager break」。动机是：在 eager break（CUDA Graph 可打断路径）下，注意力模块的 q 量化、FP4 indexer、global topk 映射和 compressor 每层每个 forward 都会临时分配中间张量，分配开销叠加在 TTFT/TPOT 关键路径上。作者通过预分配模型级工作区并在各层间复用，减少 eager 路径的动态分配，同时保持精度不变（gsm8k exact_match 0.9507→0.9515）。

实现拆解

本 PR 按以下 5 步完成 workspace 复用的改造：

1. 新增模型级缓冲池 DeepseekV4EagerScratchPool（vllm/models/deepseek_v4/eager_scratch.py，全新文件 +137 行）：在构造时一次性分配两类存储——_q（bf16，(max_num_tokens, padded_q_heads, q_head_dim)）和一块 uint8 aux storage。aux storage 通过 _packed_size + _views 按 256 字节对齐切出三组视图：FP4 indexer 输出（values/scales/weights）、global topk 输出（indices/lens）、compressor scratch（fp32）。关键设计是 aux_bytes = max(三组打包大小) 而非 sum，注释说明依据是「FP4 indexer 仅 C4、global mapping 在 FP4 indexer 之后、compressor scratch 仅 C128」，三者使用时机互斥，因此取最大即可覆盖。对外提供 q_out、compressor_scratch、indexer_q_outputs、global_topk_outputs 四个按 num_tokens 缓存视图的访问器，dict 按 token 数惰性建视图，避免重复切片。

2. 在模型构造处创建并逐层传递（vllm/models/deepseek_v4/nvidia/model.py）：DeepseekV4ForCausalLM 构造时用 _select_dsv4_attn_cls(vllm_config).get_padded_num_q_heads(...) 计算 padded heads，再以 max_num_batched_tokens、config.index_n_heads、config.index_head_dim、config.index_topk 创建 self.eager_scratch_pool；仅当 not vllm_config.parallel_config.use_ubatching 时创建，并留 TODO 说明 dbo/ubatch 需要带 ubatch 维的 buffer。随后经 DeepseekV4DecoderLayer 把 pool 传入每个 attention 层。

3. 改造 attention eager break 主路径（vllm/models/deepseek_v4/attention.py + csrc/libtorch_stable/ops.h + torch_bindings.cpp）：新增 C++ 算子 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out，与既有 ..._insert 不同点在于多收一个 q_out 输出张量参数，kernel 不再内部分配 padded q，而是写入调用方传入的池视图；池不可用时回退原路径。另新增 _global_topk_output_buffers(topk_indices) 辅助方法：仅当 compress_ratio == 4 且池存在时返回 global topk 输出视图，否则返回 None，供 flashinfer_sparse / flashmla 后端把 output_buffers 透传给 kernel。

4. indexer 与 compressor 复用同一 storage（vllm/models/deepseek_v4/common/ops/fused_indexer_q.py、compressor.py）：fused_indexer_q_rope_quant 增加 output_buffers 参数，原实现里 assert use_fp4 限制在 review 中被指为接口缺陷，作者最终选择扩展支持而非仅加文档；DeepseekCompressor.forward 在 not self.overlap 且池存在时通过 extra_kwargs["compress_scratch"] 传入 compressor_scratch(num_actual)，替换原来按 _use_two_stage_fused_compressor 分支分配的临时 scratch。

5. 测试与验证配套：tests/kernels/test_fused_indexer_q_rope_quant.py 增加 output_buffers 复用分支的对照测试；tests/kernels/test_compressor_kv_cache.py 新增 test_compute_global_topk_reuses_output_buffers 验证 topk 输出缓冲确实被复用；tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py 适配新的 out 变体。PR body 给出 vllm serve deepseek-ai/DeepSeek-V4-Flash --tensor-parallel-size 4 --enable-expert-parallel --attention-backend FLASHMLA_SPARSE_DSV4 的端到端精度与性能数据。

关键文件：

• vllm/models/deepseek_v4/eager_scratch.py（模块 缓冲池；类别 source；类型 data-contract；符号 DeepseekV4EagerScratchPool, init, _packed_size, _views）: 本 PR 核心：新增 DeepseekV4EagerScratchPool，统一预分配并复用 q、FP4 indexer、global topk、compressor 四类工作区，是全部优化的承载体。
• vllm/models/deepseek_v4/attention.py（模块 注意力层；类别 source；类型 data-contract；符号 _global_topk_output_buffers, _fused_qnorm_rope_kv_insert）: 注意力层接入缓冲池：qnorm+RoPE+KV insert 路径切换为 out 变体算子，新增 _global_topk_output_buffers 供各后端复用 topk 输出。
• vllm/models/deepseek_v4/nvidia/model.py（模块 模型组装；类别 source；类型 data-contract；符号 DeepseekV4DecoderLayer, DeepseekV4ForCausalLM）: 模型装配入口：按 config 计算 padded heads 并创建 DeepseekV4EagerScratchPool，经 DecoderLayer 逐层传入 attention；非 ubatching 才启用。
• vllm/models/deepseek_v4/compressor.py（模块 压缩器；类别 source；类型 data-contract；符号 DeepseekCompressor）: compressor 在非 overlap 路径改用池内 compressor_scratch，去掉每层临时分配。
• csrc/libtorch_stable/ops.h（模块 算子绑定；类别 source；类型 core-logic；符号 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out）: 新增 fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out 声明，是算法下沉到 C++ 的关键接口变更。
• vllm/models/deepseek_v4/common/ops/fused_indexer_q.py（模块 索引器；类别 infra；类型 infrastructure；符号 fused_indexer_q_rope_quant）: fused_indexer_q_rope_quant 增加 output_buffers 参数，review 讨论后由 assert use_fp4 扩展为兼容非 FP4。
• tests/kernels/test_compressor_kv_cache.py（模块 测试；类别 test；类型 test-coverage；符号 test_compute_global_topk_reuses_output_buffers）: 新增 test_compute_global_topk_reuses_output_buffers，验证 global topk 输出缓冲被复用的核心行为。
• tests/kernels/test_fused_indexer_q_rope_quant.py（模块 测试；类别 test；类型 test-coverage）: 为 output_buffers 分支补充对照测试，覆盖 FP4 indexer 输出写入预分配视图的路径。

关键符号：DeepseekV4EagerScratchPool.init, DeepseekV4EagerScratchPool._packed_size, DeepseekV4EagerScratchPool._views, DeepseekV4EagerScratchPool.q_out, DeepseekV4EagerScratchPool.compressor_scratch, DeepseekV4EagerScratchPool.indexer_q_outputs, DeepseekV4EagerScratchPool.global_topk_outputs, DeepseekV4Attention._global_topk_output_buffers, DeepseekV4Attention._fused_qnorm_rope_kv_insert, fused_indexer_q_rope_quant, fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out

关键源码片段

vllm/models/deepseek_v4/attention.py

注意力层接入缓冲池：qnorm+RoPE+KV insert 路径切换为 out 变体算子，新增 _global_topk_output_buffers 供各后端复用 topk 输出。

    def _global_topk_output_buffers(
        self, topk_indices: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor] | None:
        # 只有 compress_ratio == 4 的路径会写 global topk 输出，
        # 且必须在池存在时才返回预分配视图，否则交给 kernel 自行分配。
        if self.compress_ratio != 4 or self.eager_scratch_pool is None:
            return None
        return self.eager_scratch_pool.global_topk_outputs(topk_indices)
​
    # 在 fp8_ds_mla UE8M0 paged 路径中：
    # Q 侧做 per-head RMSNorm（无权重）+ GPT-J RoPE，并填充 padding head 槽位；
    # KV 侧做 GPT-J RoPE + UE8M0 FP8 量化 + paged cache 写入。
    # 有池时优先把 padded q 写入池内预分配视图，避免 kernel 每次 eager break
    # 都临时分配一块 padded q 张量；池不可用时回退原始 insert 算子。
    if self.eager_scratch_pool is not None:
        q_out = self.eager_scratch_pool.q_out(q.shape[0])
        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
            q,
            kv,
            q_out,
            swa_kv_cache_2d,
            swa_metadata.slot_mapping,
            positions,
            cos_sin_cache,
            self.padded_heads,
            self.eps,
            swa_metadata.block_size,
        )
        return q_out
    return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
        q,
        kv,
        swa_kv_cache_2d,
        swa_metadata.slot_mapping,
        positions,
        cos_sin_cache,
        self.padded_heads,
        self.eps,
        swa_metadata.block_size,
    )

评论区精华

review 由 tlrmchlsmth 主导，共 3 条实质意见，全部已解决并最终 APPROVED：

• 测试中的魔术数字 == 7：tlrmchlsmth 在 tests/kernels/test_fused_indexer_q_rope_quant.py 追问为什么 num_tokens == 7 and cache_dtype == torch.float32 and use_fp4 才走 output_buffers 分支。作者回复 Fixed，未在评论中展开具体原因，但可以看出该分支用于构造一个覆盖多种 dtype/token 数组合的对比测试点。
• output_buffers 隐含 use_fp4 的接口缺陷：tlrmchlsmth 在 fused_indexer_q.py 指出「output_buffers implies use_fp4」会让函数接口「bumpy」，建议要么在 docstring 说明，要么扩展支持。作者选择后者：「Fixed, we extend the support」，去掉了 assert use_fp4 的隐式约束。
• dbo/ubatching TODO 过于简略：tlrmchlsmth 在 nvidia/model.py 建议扩写 TODO，作者同样回复 Fixed。

• CI 失败归因：作者在 review 中说明某次失败的单元测试与本 PR 无关（此前 commit 在 buildkite 已通过），属于既有 flaky。

  • 测试分支中魔术数字 == 7 的疑问 (testing): 作者回复 Fixed，调整了测试分支的触发条件或补充说明，未再引发争议。
  • output_buffers 隐含 use_fp4 的接口缺陷 (design): 作者选择扩展支持并回复「Fixed, we extend the support」，去掉 assert use_fp4 的隐式限制，接口更干净。
  • dbo/ubatching 的 TODO 说明过于简略 (documentation): 作者回复 Fixed，补全了 TODO 上下文。
  • CI 单元测试失败与本次变更无关 (testing): 被接受，未影响最终合并。

风险与影响

• 风险：
  1. 新增 C++ 算子注册面：fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out 在 ops.h 声明、torch_bindings.cpp 绑定，但 CUDA 实现是否在 ROCm 等平台可用未被本 PR 材料覆盖；DeepseekV4 路径主要面向 NVIDIA，非 NVIDIA 平台若走到该分支需确认 fallback 仍走原 ..._insert 路径。
  2. 显存一次性占用：池按 max_num_batched_tokens 预分配 q（bf16）与 aux storage，大 max_num_batched_tokens 下显存占用上升；不过 aux 取 max 而非 sum 已把三类 scratch 的重复分配消除，且原来每层临时分配总和通常更大，净效果应下降或持平。
  3. 按 num_tokens 缓存的 dict 增长：_q_outputs 等四个 dict 以 num_tokens 为键，推理过程中 token 数种类有限，但极端动态 batch 下条目可能累积；所有视图共享同一块 storage，显存不会翻倍，风险可控。
  4. dbo/ubatching 路径被排除：池创建被 not use_ubatching 条件挡住，未来启用 dbo 时该优化不生效，且 TODO 未闭环，需防止后续有人误以为所有路径都已覆盖。
  5. 硬断言：global_topk_outputs 中 assert topk == self.index_topk，若配置变更而池参数未同步，会直接崩溃；当前构造参数来自同一 config，风险低。
     - 影响：影响范围严格限定在 DeepSeek V4 模型的 NVIDIA eager 路径：注意力层、indexer、compressor 三大子模块共享一个模型级工作区，减少 eager break 下的临时分配次数。用户侧收益为 TTFT 均值约 3.8% 改善（46.41→44.65ms）、TPOT 约 0.5% 改善（5.74→5.71ms），吞吐基本持平；精度经 gsm8k 验证无回退。对团队而言，该 PR 建立了 DeepSeek V4 系列优化的可复用缓冲池抽象，后续 dbo 支持、更多 kernel 复用可在其上扩展；对其它模型无影响，通用路径完全不触碰。
     - 风险标记：新增 C++ out 算子需跨平台注册验证, 缓冲池一次性占用显存, dbo/ubatching 路径未覆盖, compress_ratio 分支含硬断言

关联脉络

• PR #52948 [Model] Support bidirectional (encoder-only) attention for DeepSeek e…: 同属 vllm/models/deepseek_v4 模块，改动 model.py / attention.py 等共享装配与注意力层接口，体现该模块持续演进。
• PR #45061 issue #45861 优化链中的首个 DSv4 性能子任务（标题在材料中未提供）: 与本次 PR 同属 issue #45861「Performance Optimization for Deepseek V4」任务列表，作者同为 yewentao256，workspace 复用是该系列的一环。