执行摘要

• 一句话：修复 breakable CUDA graph 下 DSV4 indexer 打分被跳过
• 推荐动作：值得精读。改动虽小，但精准打击了 CUDA graph capture 语义与运行时动态条件分支的冲突，是 breakable graphs 场景下极具代表性的正确性修复。建议与 #52448 issue 的根因剖析、#52401 的 region 选择方案、#51318 的元数据回退一起阅读，能完整理解 DSV4 稀疏 MLA 在 CUDA graph 下的演进脉络。

功能与动机

PR body 明确指出：#49486 引入的 host 端短上下文捷径会跳过 learned indexer scoring，而 #51430 把 DeepseekV4Indexer.forward 移入了 CUDA graph 捕获区。breakable PIECEWISE 捕获使用短 dummy 元数据，导致捷径被烘焙进图，之后图重放于长 cached prefix 时（超过 2048 tokens，C4 层候选数 > 512）本应打分的路径却固定选择候选 0..511。这对应 #52448 的现象：finish_reason="length"、空 content、无 tool_calls，且 reasoning_content 陷入不输出 think-end token 的循环，单个请求烧满 max_tokens=16384。PR body 还说明 #52401 只扩大 MRV1 的 eager region，而 DSpark 强制 MRV2，因此该路径未被覆盖。

实现拆解

1. 定位问题点：vllm/models/deepseek_v4/attention.py 的 DeepseekV4Indexer.forward 中，短上下文捷径原条件为 indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens，满足时直接调用 _fill_short_context_topk_indices 填充 topk 索引并返回 None, None, None。
2. 修改条件：在原判断上追加 and not torch.cuda.is_current_stream_capturing()。当 CUDA stream 处于 graph capture 状态时，强制进入完整 indexer scoring 路径（wq_b_and_q_quant → fused_indexer_q_rope_quant → sparse top-k），确保捕获进图的控制流与长上下文重放语义一致。
3. 保留 eager 行为：非捕获状态下仍走捷径，短 prompt 的 host 端加速与 k cache 构建逻辑（_fill_short_context_topk_indices）不变。
4. 测试配套：本 PR 未新增测试文件。作者用现有 tests/v1/cudagraph/test_breakable_cudagraph.py 回归（14 passed），并在 2× NVIDIA B200、TP2+EP、DSpark-7、FULL_AND_PIECEWISE 上做真实负载验证：6144/6543 token 提示词的输出从错误循环恢复为稳定重复，且 #52448 workload 8 波 × 32 并发请求 0/256 触发 length-cap runaways。若团队要求更强回归保护，建议后续补一条“捕获期间禁用捷径”的显式断言。

关键文件：

• vllm/models/deepseek_v4/attention.py（模块 稀疏注意力；类别 source；类型 core-logic；符号 DeepseekV4Indexer.forward, _fill_short_context_topk_indices）: 唯一变更文件。DeepseekV4Indexer.forward 的短上下文捷径增加 not torch.cuda.is_current_stream_capturing() 条件，修复 breakable CUDA graph 重放长 cached prefix 时 indexer scoring 被错误跳过的数据正确性问题，是本次 bugfix 的核心。

关键符号：DeepseekV4Indexer.forward

关键源码片段

vllm/models/deepseek_v4/attention.py

唯一变更文件。DeepseekV4Indexer.forward 的短上下文捷径增加 not torch.cuda.is_current_stream_capturing() 条件，修复 breakable CUDA graph 重放长 cached prefix 时 indexer scoring 被错误跳过的数据正确性问题，是本次 bugfix 的核心。

def forward(
    self,
    hidden_states: torch.Tensor,
    qr: torch.Tensor,
    compressed_kv_score: torch.Tensor,
    indexer_weights: torch.Tensor,
    positions: torch.Tensor,
    rotary_emb: nn.Module,
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
    compressor = self.compressor
​
    attn_metadata = get_forward_context().attn_metadata
    if isinstance(attn_metadata, dict):
        indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
        # 短上下文捷径：候选数 <= topk 时每个候选都会被选中，无需 indexer 打分，
        # 但仍需构建 k cache。注意捕获期禁用：breakable PIECEWISE 捕获使用短
        # dummy 元数据，若捷径被烘焙进 CUDA graph，重放长 cached prefix（候选数
        # 超过 topk）时 C4 层会固定选择前 topk 个候选，造成 #52448 的
        # think-until-cap / 空 turn。eager 执行不受影响，仍走捷径。
        if (
            indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens
            and not torch.cuda.is_current_stream_capturing()
        ):
            compressor(compressed_kv_score, positions, rotary_emb)
            assert self.topk_indices_buffer is not None
            num_tokens = (
                indexer_metadata.num_decode_tokens
                + indexer_metadata.num_prefill_tokens
            )
            if num_tokens > 0:
                # 用 Triton 内核把“全选”语义写入 topk 索引缓冲区。
                _fill_short_context_topk_indices[(num_tokens,)](
                    self.topk_indices_buffer,
                    positions,
                    TOP_K=self.topk_tokens,
                    COMPRESS_RATIO=self.compress_ratio,
                    PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens),
                    num_warps=8,
                )
            return None, None, None
        # 否则进入完整 indexer scoring 路径：wq_b_and_q_quant 产出 Q / RoPE /
        # 量化融合结果，再由 sparse top-k 内核选出真正的前 topk 候选。

评论区精华

PR 没有实质性的 design review 评论：zyongye 在第二次 CI（Buildkite #84134）通过后直接批准合并；claude[bot] 仅提示本仓库配置为手动 review。aoshen02 在关联 issue 评论中提出 “Confirm when implementing Batch invariance with dpsk v4 flash base.”，希望后续实现 batch invariance（批大小无关）特性时同步确认本修复路径的行为，该确认请求在评论区未收到显式回复。

• Batch invariance 与 dpsk v4 flash base 的行为确认 (question): 评论区未见对该确认请求的显式回复；合并者 zyongye 在第二次 CI 通过后直接批准合并。

风险与影响

• 风险：
  1. 捕获语义依赖：修复正确性完全依赖 torch.cuda.is_current_stream_capturing() 的语义。若未来 indexer forward 被移出捕获区或 capture 检测机制变化，条件会失效；反之若捕获期有其他 host 动态分支被烘焙进图，同类问题可能复发。
  2. 跨平台差异：改动调用 CUDA 专属 API。DSV4 在 ROCm 上目前默认走 MRV1 的宽 eager region（见 #52401），大概率不受影响，但本 PR 未在非 NVIDIA 平台验证，属低风险盲区。
  3. 测试覆盖缺口：没有新增针对“捕获期禁用捷径”的独立单测，回归保护依赖现有 test_breakable_cudagraph.py 与人工负载验证。
  4. 性能影响：仅在 capture（warmup）阶段多执行一次完整打分，运行时重放无额外开销；eager 短上下文路径完全保留。
     - 影响：对用户：修复 DeepSeek-V4-Flash 在 DSpark + breakable CUDA graphs + 并发场景下的空 turn / think-until-cap 问题，避免请求烧满 max_tokens 后返回低质量结果，显著改善输出正确性与有效吞吐。对系统：仅在 CUDA graph 捕获阶段增加一次 indexer 打分，运行期行为与 eager 一致，无额外延迟。对团队：确立了“CUDA graph 捕获期禁用依赖运行时动态条件的 host 捷径”这一编码约束，为其他模型在 breakable graph 中做 host 分支提供警示；影响范围局限在 DSV4 的 CUDA + MRV2 + breakable graph 路径。
     - 风险标记：CUDA 捕获期语义依赖, 缺少专项单测, Eager 与 Graph 分支分叉, ROCm 兼容性待确认

关联脉络

• PR #51318 [Bugfix][DSv4] Revert adaptive C128A metadata packing: 同属 DSV4 稀疏 MLA 在 CUDA graph 下元数据/控制流错位的修复。PR body 明确说明 #51318 是 revert adaptive C128A packing，与本问题不同，且 #52448 在已包含 #51318 的 main 上仍可复现，佐证本修复的必要性。
• PR #52401 [Bugfix] Pick the DeepSeek V4 eager cudagraph region per model runner: #52401 通过扩大 MRV1 的 eager region 修复 DSV4 输出损坏，但 PR body 指出它不覆盖 DSpark 强制 MRV2 的路径；本 PR 与它互补，共同覆盖 MRV1/MRV2 两条模型运行路径。
• PR #51538 [Bugfix] Make DSV4 sparse MLA work end-to-end for plain decode, MTP, and DSpark: 让 DSV4 稀疏 MLA 在 decode/MTP/DSpark 下端到端可用，与本次 indexer scoring 修复同属稀疏 MLA 正确性演进脉络，可对照理解该特性的复杂度。