# PR #52436 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Bugfix][Spec Decode][Structured Output] DSpark: fix the grammar bitmask mapping when the draft budget is zero
- 合并时间：2026-08-16 19:09
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/52436

---

# 执行摘要

- 一句话：修复 DSpark 零验证预算下结构化输出 bitmask 行数错配崩溃
- 推荐动作：值得精读，尤其推荐关注 PR body 中“三种预算机制对比表”的根因拆解方式——它清晰展示了同一个字段在不同分支下语义漂移如何导致 host/device 契约不一致。同时建议关注 LucasWilkinson 在 #52477 中的替代方案，理解“不依赖分支字段”的架构化修复思路。

# 功能与动机

PR body 明确指出：当 DSpark 自适应验证（#47808 引入）与结构化输出组合使用时，只要草稿置信度下降、自适应验证决定本步验证 0 个草稿，就会触发 `assert num_masks == len(mapping)` 崩溃。作者强调 `draft_budget == 0` 并非边缘情况，而是 argmax 在草稿无法覆盖验证成本时的常规输出（`adaptive_verification.py:329`），在高熵 token、warm-up、grammar 约束步中会频繁出现，普通负载即可稳定复现。

# 实现拆解

本 PR 的修复分五步落地，核心是让 host 端映射行数推导不再依赖会被压缩语义污染的字段：

1. **抽取映射构造函数**：在 `vllm/v1/worker/gpu/structured_outputs.py` 新增模块级函数 `_build_grammar_mapping(req_ids, grammar_req_ids, cu_num_logits_np, num_draft_tokens_per_req, num_bonus_tokens, mask_stride)`，把原先内联在 `apply_grammar_bitmask` 中的映射构造逻辑独立出来，便于测试直接引用。

2. **行数推导策略分流**：当 `num_draft_tokens_per_req` 为 `None` 时（非自适应验证路径，`cu_num_logits_np` 未被压缩），沿用 `cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx]` 差值推导；否则改用 `int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens`，与调度器按“每请求已调度 draft 数 + 1”分配 bitmask 行的口径一致。`num_draft_tokens_per_req` 来自调度器同一份 `scheduled_spec_decode_tokens`（`model_runner.py:1132`），不会被 `compact_batch` 改写。

3. **构造器与调用点改造**：`StructuredOutputsWorker.__init__` 新增 `num_bonus_tokens` 参数并保存为实例属性；`apply_grammar_bitmask` 改为调用 `_build_grammar_mapping`，传入 `input_batch.req_ids`、`cu_num_logits_np`、`num_draft_tokens_per_req` 与 `self.num_bonus_tokens`。

4. **数据契约配套**：`vllm/v1/worker/gpu/model_runner.py` 构造 `StructuredOutputsWorker` 时新增 `num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step`，保证 bonus 数与采样器实际产出一致。

5. **测试配套**：在 `tests/v1/spec_decode/test_adaptive_verification.py` 新增 `test_zero_budget_keeps_one_grammar_row_per_scheduled_draft`，构造 3 请求（2 个验证请求各 2 个草稿 + 1 个 prefill）的零预算场景，断言 `len(mapping) == 调度器 bitmask 行数` 且精确值 `mapping == [0, 1, 2, 3, 4, 5, 6]`；作者验证无修复时该测试以 `assert 3 == 7` 失败。

关键文件：
- `vllm/v1/worker/gpu/structured_outputs.py`（模块 结构化输出；类别 source；类型 core-logic；符号 _build_grammar_mapping）: 修复核心：新增 `_build_grammar_mapping` 函数，将 bitmask 到 logits 的映射行数推导从被压缩的 `cu_num_logits_np` 改为基于 `num_draft_tokens_per_req + num_bonus_tokens`，消除三种预算机制下的行数错配。
- `tests/v1/spec_decode/test_adaptive_verification.py`（模块 自适应验证；类别 test；类型 test-coverage；符号 test_zero_budget_keeps_one_grammar_row_per_scheduled_draft）: 新增零预算场景回归测试，直接调用 `_build_grammar_mapping` 并断言 mapping 行数与调度器 bitmask 行数一致，且作者验证了无修复时测试以 `assert 3 == 7` 失败。
- `vllm/v1/worker/gpu/model_runner.py`（模块 模型运行器；类别 source；类型 data-contract）: 作为数据契约配套，构造 `StructuredOutputsWorker` 时新增 `num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step`，把 bonus 数从采样器状态注入结构化输出 worker。

关键符号：_build_grammar_mapping, apply_grammar_bitmask, test_zero_budget_keeps_one_grammar_row_per_scheduled_draft

## 关键源码片段

### `vllm/v1/worker/gpu/structured_outputs.py`

修复核心：新增 `_build_grammar_mapping` 函数，将 bitmask 到 logits 的映射行数推导从被压缩的 `cu_num_logits_np` 改为基于 `num_draft_tokens_per_req + num_bonus_tokens`，消除三种预算机制下的行数错配。

```python
# 修复核心：新增模块级函数 _build_grammar_mapping，把 bitmask -> logits 的映射
# 构造从 apply_grammar_bitmask 中独立出来，便于直接单测与后续复用。
def _build_grammar_mapping(
    req_ids: list[str],
    grammar_req_ids: list[str],
    cu_num_logits_np: np.ndarray,
    num_draft_tokens_per_req: np.ndarray | None,
    num_bonus_tokens: int,
    mask_stride: int,
) -> list[int]:
    # 背景：调度器按“每请求已调度 draft 数 + 1（bonus）”分配 bitmask 行，
    # 而 compact_batch 的零预算分支会把 cu_num_logits_np 重写为仅 bonus 布局，
    # 旧代码从重写后的 offsets 推导行数，导致 mapping 长度只有 bitmask 的几分之一。
    mapping: list[int] = []
    req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)}
    for grammar_req_id in grammar_req_ids:
        req_idx = req_id_to_idx[grammar_req_id]
        if num_draft_tokens_per_req is None:
            # 非自适应验证路径：cu_num_logits_np 未被压缩，直接取相邻 offsets 差值。
            num_positions = int(
                cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx]
            )
        else:
            # 自适应验证路径：从调度器同一来源（num_draft_tokens_per_req）推导行数，
            # 再加上 bonus token，保证三种预算机制（全量 / 部分 / 零）下都与 bitmask 对齐。
            num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens
        # 键为 (request, position) 而非绝对 logit 下标：设备端 kernel 依据 GPU 侧
        # cu_num_logits 解析真实偏移，压缩布局中没有物理行的位置由 position_is_active 掩蔽。
        mapping.extend(
            req_idx * mask_stride + position for position in range(num_positions)
        )
    return mapping

```

# 评论区精华

核心 review 交锋来自维护者 LucasWilkinson，他认可问题本身但质疑实现方式是否必须依赖 `num_draft_tokens_per_req` 分支：

> "im wondering if theres a way to fix this without having to branch on `num_draft_tokens_per_req` experimenting with that here: #52477 but overall I think this approach is reasonable (we can land this in the interim)"

作者 oops-oom 在 issue 评论区回应会先读完 #52477 再回答 `num_draft_tokens_per_req` 的问题，最终 LucasWilkinson 给予 APPROVED，结论是当前方案先落地、#52477 作为后续重构探索。

- 是否可不依赖 num_draft_tokens_per_req 分支来修复 (design): 当前基于 `num_draft_tokens_per_req + num_bonus_tokens` 的方案被认可并先落地，后续在 #52477 中探索去除分支的替代方案。
- num_draft_tokens_per_req 方案与 #52477 的对比 (question): PR 以当前方案合并，作者与维护者均知晓 #52477 的后续探索方向。

# 风险与影响

- 风险：主要风险集中在数据契约与覆盖范围：

1. **对 `num_draft_tokens_per_req` 字段的依赖**：修复正确性依赖 `input_batch.num_draft_tokens_per_req` 始终与调度器 `scheduled_spec_decode_tokens` 同源且不被压缩改写；若 MRV2 后续改造改变该字段语义或移除，需同步调整 `_build_grammar_mapping`。

2. **回退分支的脆弱性**：`num_draft_tokens_per_req is None` 分支保留了旧逻辑，若未来非自适应路径的 `cu_num_logits_np` 也被压缩，同样的错配会再次出现，当前没有防御性断言覆盖该分支。

3. **测试覆盖单一场景**：新测试只覆盖 `budget == 0` 一种机制；`0 < budget < num_drafts` 与 `budget == num_drafts` 两个分支在 `structured_outputs` 层没有独立回归测试，依赖 PR body 中的手推表格佐证。

4. **与 #52477 的并存风险**：LucasWilkinson 正在 #52477 中实验去掉分支的替代方案，两 PR 合并顺序可能导致重构冲突或语义漂移。
- 影响：影响范围集中在 DSpark 自适应验证 + 结构化输出（JSON schema / grammar）的用户路径：修复前在草稿置信度下降时稳定崩溃，修复后可正常运行且不改变采样行为。由于改动只发生在 host 端映射构造，不涉及 kernel 与 GPU 布局，对吞吐和延迟无实质影响；对团队而言，这次修复明确了“调度器按已调度数分配行、worker 端压缩布局”这一隐式契约，为 MRV2 后续迭代提供了可参考的边界案例。
- 风险标记：核心路径断言变更 , 依赖 num_draft_tokens_per_req 字段契约 , 新测试仅覆盖零预算单一场景 , 与 #52477 替代方案潜在冲突

# 关联脉络

- PR #47808 [Spec Decode] DSpark confidence-scheduled verification: 本 PR 修复的崩溃正是 #47808 引入的自适应验证行为（零预算分支重写 `cu_num_logits_np`）与结构化输出叠加产生的边界问题。
- PR #52477 Unnamed alternative fix by LucasWilkinson: LucasWilkinson 在 review 中明确提到正在 #52477 实验另一种不依赖 `num_draft_tokens_per_req` 分支的修复方案，与本 PR 同功能线。
- PR #52288 [Bugfix][Spec Decode] DSpark: inherit the target's attention backend when the speculative config names none: 同属 DSpark 调度 /worker 布局契约修复，修复注意力后端继承导致的缓存形状崩溃，与本 PR 共享相同的 spec-decode 契约风险主题。
- PR #51538 [Bugfix] Make DSV4 sparse MLA work end-to-end for plain decode, MTP, and DSpark: DSV4/DSpark 端到端修复链路上的后续修补，与本 PR 一样涉及 `model_runner` 与 spec-decode 各阶段的一致性。