# PR #45100 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Bugfix] Avoid racy accepted counts in async spec decode
- 合并时间：2026-06-22 16:53
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/45100

---

# 执行摘要

- 一句话：修复异步推测解码 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 计数竞争

```python
# 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

```python
# 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