[Bugfix][Spec Decode][Structured Output] DSpark: fix the grammar bitmask mapping when the draft budget is zero

[oops-oom]

Purpose

Fixes a crash when DSpark adaptive verification (enable_adaptive_verification, added in #47808) is combined with structured outputs and the chosen draft budget is zero.

assert num_masks == len(mapping)  // fails when draft_budget == 0

Scenario: serving a structured-outputs request (JSON schema / grammar) with DSpark adaptive verification enabled crashes on an assert as soon as the drafter's confidence drops far enough that adaptive verification decides to verify zero drafts this step.

When is draft_budget == 0? (not an edge case)

That is not an exotic corner. Zero is the routine choice whenever drafts stop paying for their verification cost. Per step, the controller picks draft_budget = argmax_b (est_accepted_tokens(b)/cost(b)) over b ∈ [0, max_draft_budget] (adaptive_verification.py:329). b = 0 ("verify no drafts, sample only the bonus token") is a normal output of that argmax, selected whenever verifying even the most-confident draft is net-negative for throughput — i.e. whenever the drafter's confidence is low. That happens routinely (high-entropy tokens, warm-up, grammar-constrained steps), so a sustained DSpark + structured-outputs workload hits it repeatedly and the crash reproduces under ordinary load. This is exactly the behavior #47808 introduced.

Root cause

The scheduler and the worker disagree on how many grammar bitmask rows exist:

• The scheduler sizes the bitmask from the scheduled drafts — len(drafts) + 1 rows per request (vllm/v1/structured_output/__init__.py), and it has no knowledge of the device-side budget trimming.
• compact_batch()'s zero-budget branch rewrites cu_num_logits_np to the exact bonus-only layout (adaptive_verification.py:361). That rewrite exists to serve _iter_request_chunks, which slices the compacted logits with these offsets.

apply_grammar_bitmask derived its bitmask → logits mapping from those rewritten offsets, so it produced one row per request instead of one per scheduled draft and tripped assert num_masks == len(mapping) (structured_outputs.py:102).

For a batch of 2 verification requests with 2 drafts each plus 1 bonus token, the scheduler emits 6 rows while the old expression yields 2:

bitmask rows = 6,  mapping = [0, 3]  (len 2)  ->  assert fails

Why the other two budget regimes were unaffected

Only the zero-budget branch rewrites cu_num_logits_np; the other two return it untouched, still holding the scheduled layout that the scheduler used. Exercising all three branches through the real compact_batch:

regime	cu_num_logits_np	rewritten	old mapping len	bitmask rows
budget == num_drafts	[0, 3, 6]	no	6	6	OK
0 < budget < num_drafts	[0, 3, 6]	no	6	6	OK
budget == 0	[0, 1, 2]	yes	2	6	assert fails

So the old code was correct in two regimes only because that field happened to carry scheduled-layout semantics there.

Fix

Size the mapping from num_draft_tokens_per_req + num_bonus_tokens instead. That array comes from the same scheduled_spec_decode_tokens the scheduler used (model_runner.py:1132) and is never touched by compaction, so all three regimes agree and the row count no longer depends on a field whose meaning shifts.

Rows that the compacted device layout has no room for are already masked by position_is_active in the kernel (structured_outputs.py:144), which resolves true per-request offsets from the GPU-side cu_num_logits. This is the mechanism #47808 introduced; the fix just stops the host side from second-guessing it.

Test Plan

Added test_zero_budget_keeps_one_grammar_row_per_scheduled_draft to tests/v1/spec_decode/test_adaptive_verification.py.

.venv/bin/python -m pytest tests/v1/spec_decode/test_adaptive_verification.py -v
# related suites touching the same cu_num_logits_np contract:
.venv/bin/python -m pytest tests/v1/worker/test_gpu_rejection_sampler_chunking.py \
                           tests/v1/worker/test_gpu_batch_ordering.py -v
pre-commit run --all-files

Test Result

tests/v1/spec_decode/test_adaptive_verification.py  6 passed
+ test_gpu_rejection_sampler_chunking.py, test_gpu_batch_ordering.py
                                                  17 passed, 4 skipped
pre-commit                                        all hooks passed (incl. mypy, ruff format)

Verified the new test actually fails without the fix by reverting only the helper body while keeping its signature:

without fix:  E   assert 3 == 7
                  where 3 = len([0, 3, 6])
with fix:     6 passed

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[LucasWilkinson]

Thanks for the PR; great catch! 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)

[github-actions]

✅ @oops-oom, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[LucasWilkinson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84063 for commit b58c7640567e.

[mergify]

Hi @oops-oom, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[oops-oom]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #84063.

[oops-oom]

@LucasWilkinson Thanks for the review, and for the ready label. Let me read through #52477 before answering the num_draft_tokens_per_req question properly.

[oops-oom]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84068 for commit cb8b32647aca.

[oops-oom]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #84068.

[oops-oom]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #84068.

[oops-oom]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #84068.

[oops-oom]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84079 for commit af118ba37f12.

[mergify]

Hi @oops-oom, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[oops-oom]

/ci run

[github-actions]

✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/84079

[oops-oom]

/ci cancel

[github-actions]

✅ Requested cancellation of 1 CI build for bugfix/adaptive-zero-budget-structured-output: #84079.

[oops-oom]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84080 for commit af118ba37f12.

[oops-oom]

/ci cancel

[github-actions]

✅ Requested cancellation of 1 CI build for bugfix/adaptive-zero-budget-structured-output: #84080.

[oops-oom]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84082 for commit af118ba37f12.

[oops-oom]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84084 for commit 625673438f84.