执行摘要

• 一句话：修复异步推测解码 accepted 计数竞态
• 推荐动作：推荐精读。该 PR 展示了异步推测解码中一个微妙的竞态条件及其修复过程，涉及 GPU D2H 拷贝、input batch 行重排、cudagraph 元数据生命周期等深层机制。对于从事 speculative decoding、cudagraph 或 async scheduling 开发的工程师有重要参考价值。

功能与动机

Fix an async speculative decoding race for hybrid non-align Mamba/GDN models (observed with Qwen3.5 MTP) where _prepare_inputs consumes a stale CPU copy of num_accepted_tokens. In the failing path, _update_states_after_model_execute writes accepted-token counts from GPU to input_batch.num_accepted_tokens_cpu_tensor with a non-blocking D2H copy. Under async scheduling, the next _prepare_inputs may also remap request rows after input_batch.swap_states() / condense(). At a request's prefill-to-first-spec-decode transition, this can make GDN see another row's accepted-token count, e.g. 4 instead of 1. GDN then restores the recurrent state from the wrong speculative state slot; for Qwen3.5 this loses prompt memory and the request often ends quickly with garbled text plus EOS.

实现拆解

本 PR 通过以下步骤解决竞态问题：

1. 调整 accepted 同步条件：在 vllm/v1/worker/gpu_model_runner.py 的 _prepare_inputs 中新增条件 needs_cpu_accepted_counts，当启用异步调度且 mamba 缓存模式非 "align" 时，跳过从 CPU 同步 num_accepted_tokens，避免读取到过时值。改为默认填充 1（所有 draft 被接受），并依赖后续 update_num_computed_tokens_for_batch_change GPU 核函数根据 valid_sampled_token_count 修正正确值。
2. 保留 align 模式原有路径：当 mamba_cache_mode == "align" 时，继续原有同步逻辑，因为该模式的预处理 preprocess_mamba_all_specdec 会消费 CPU 端的 accepted 计数。
3. 修复 FULL cudagraph 的 batch size：在 vllm/v1/attention/backends/gdn_attn.py 的 build 方法中，将用于 cudagraph 的 batch_size 从 m.num_actual_tokens（token 填充后的数量）改为 m.num_reqs（请求数），因为 spec_state_indices_tensor、spec_sequence_masks、spec_query_start_loc、num_accepted_tokens 等元数据是按请求索引的，使用 token 数会导致创建多余的请求行。
4. 添加单元测试验证：在 tests/v1/attention/test_gdn_metadata_builder.py 中新增 test_full_cudagraph_spec_metadata_uses_request_count，验证 FULL cudagraph 下上述 request-indexed 张量的 shape 与 batch_size 一致。
5. 根据 review 重构：根据 tdoublep 的建议，将条件和分支反转，复用已有的 fill(1) 默认路径，消除代码重复（commit 066f691）。

关键文件：

• vllm/v1/worker/gpu_model_runner.py（模块 模型运行器；类别 source；类型 core-logic）: 主要修复点，通过条件控制避免异步下 CPU accepted 计数竞争
• vllm/v1/attention/backends/gdn_attn.py（模块 注意力后端；类别 source；类型 core-logic）: 修复 FULL cudagraph 下 batch size 错误，用 num_reqs 替代 num_actual_tokens
• tests/v1/attention/test_gdn_metadata_builder.py（模块 GDN 测试；类别 test；类型 test-coverage；符号 test_full_cudagraph_spec_metadata_uses_request_count）: 新增 FULL cudagraph 下 request-indexed 张量验证测试

关键符号：_prepare_inputs, build

关键源码片段

vllm/v1/worker/gpu_model_runner.py

主要修复点，通过条件控制避免异步下 CPU accepted 计数竞争

# Determine whether to synchronize num_accepted_tokens from CPU.
# In async non-align mode, the CPU copy may be stale, so we skip sync.
needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and not (
    self.use_async_scheduling and self.cache_config.mamba_cache_mode != "align"
)
if needs_cpu_accepted_counts:
    self.num_accepted_tokens_event.synchronize()
    # ... remap using prev_positions, etc. ...
    self.num_accepted_tokens.copy_to_gpu()
else:
    # Default: assume all draft tokens were accepted.
    # The GPU kernel (update_num_computed_tokens_for_batch_change) will
    # correct rows that had actual drafts using valid_sampled_token_count.
    self.num_accepted_tokens.np.fill(1)
    self.num_accepted_tokens.gpu.fill_(1)

该片段展示了核心逻辑：在异步非 align 模式下完全跳过 CPU 同步，采用 device 默认值并依赖后续 GPU 校正。

vllm/v1/attention/backends/gdn_attn.py

修复 FULL cudagraph 下 batch size 错误，用 num_reqs 替代 num_actual_tokens

# Prepare per-request tensors for cudagraph. m.num_actual_tokens is
# token-padded for FULL graph replay, but the GDN state/query/accepted
# metadata below is indexed by request.
batch_size = m.num_reqs
​
if (
    self.use_full_cuda_graph
    and num_prefills == 0
    and num_decodes == 0
    and num_spec_decodes <= self.decode_cudagraph_max_bs
    and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
):
    assert spec_sequence_masks is not None
    self.spec_state_indices_tensor[:num_spec_decodes].copy_(
        spec_state_indices_tensor, non_blocking=True
    )
    spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
    spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID)
    # ... similar for spec_sequence_masks, spec_query_start_loc, etc.

该片段展示了使用 m.num_reqs 作为 cudagraph 元数据的 batch size，避免 token 填充导致的行数错误。

评论区精华

ZJY0516 询问 batch_size 变更："I don't quite get this. Could you explain this a little more?"
sunnweiwei 解释："The tensors padded in this block are request-indexed, while m.num_actual_tokens is token-indexed and can be padded for FULL graph replay. For example, with mtp3, 72 reqs produce 288 decode tokens. In that case m.num_reqs == 72, while m.num_actual_tokens == 288. But tensors such as spec_state_indices_tensor, spec_sequence_masks, spec_query_start_loc, and num_accepted_tokens have one row per request, not per token. So they should be padded/sliced to m.num_reqs; using m.num_actual_tokens creates fake request rows."
结论：解释被接受，变更确认正确。

tdoublep 询问 condition 对 all 模式的影响："Is this change also needed for all mode or specifically for self.cache_config.mamba_cache_mode == None?"
sunnweiwei 回答："Both. Only 'align' is excluded, as its preprocess_mamba genuinely consumes the CPU-side counts."
结论：确认只有 align 模式需要保留 CPU 同步。

tdoublep 建议重构分支："Is there perhaps a way to rework the logic so that we can re-use this code instead of having it in two branches?"
sunnweiwei 响应："Good point! done in 066f691. Inverted the condition so it reuses the existing fallback instead of duplicating it."
结论：已完成重构，消除代码重复。

• batch_size 使用 num_reqs 而非 num_actual_tokens (design): 解释被接受，确认使用 m.num_reqs 正确。
• CPU accepted 同步条件对 mamba_cache_mode == all 等的影响 (correctness): 确认只有 align 模式需要保留 CPU 同步，其他模式均走 device-authoritative 路径。
• 重构避免两个分支的代码重复 (design): 通过反转条件复用已有 fallback 分支，消除重复。

风险与影响

• 风险：
  • 核心路径变更：修改了 _prepare_inputs，该函数在每个 step 被调用，改动可能影响异步调度下的其他模块，但通过条件隔离，仅影响非 align mamba 模式。
  • 依赖 GPU correction kernel：跳过 CPU 同步后，正确性依赖 update_num_computed_tokens_for_batch_change 核函数准确修正，若该核函数有 bug 可能引发新的错误，已有 A/B 验证表现正常。
  • align 模式保持旧路径：该模式未受影响，但若未来有类似竞态，可能未被覆盖。
  • cudagraph batch size 变更：仅影响 FULL cudagraph 且 spec decode 场景，若非 spec decode 或 PIECEWISE 模式不受影响。单元测试验证了 shape，但未测试实际数值正确性。
  • 测试覆盖有限：仅添加一个单元测试，但提供了充分的 runtime A/B 对比（20k+ 生成），多种配置下均无退化。
• 影响：
  • 用户影响：修复了 Qwen3.5-27B 等混合 Mamba/GDN 模型在异步推测解码下的生成乱码和提前结束问题，提高了生成质量。
  • 系统影响：无性能退化（A/B 验证显示吞吐微升），代码改动量小（+46/-5），兼容现有配置，对齐模式行为完全不变。
  • 团队影响：明确了 async spec decode 中 accepted 计数的 device-authoritative 原则，为后续异步路径开发提供了参考。
  • 风险标记：核心调度路径变更, 异步竞态修复, cudagraph 元数据修复

关联脉络

• 暂无明显关联 PR