[Bugfix][Spec Decode][Structured Output] DSpark: fix the grammar bitmask mapping when the draft budget is zero (#52436)
Signed-off-by: oops-oom <73481342@qq.com> Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
oops-oom committed
Aug 16, 2026 at 19:09 UTC
4d2a68d64d9e05921ed5c4099146e768a92d71d5
3 files changed
+89
-15
tests/v1/spec_decode/test_adaptive_verification.py
+49
@@ -9,6 +9,7 @@ from vllm.v1.worker.gpu.async_utils import StepTimingSample
9
from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
10
AdaptiveVerificationManager,
11
)
12
+from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping
13
14
15
def make_manager(
@@ -168,3 +169,51 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits():
169
assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype
170
# The prefill keeps its scheduled tokens; only drafts are dropped.
171
assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32))
172
+
173
+
174
+def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft():
175
+ # The scheduler sizes the grammar bitmask from the *scheduled* drafts
176
+ # (len(drafts) + 1 rows per request), but a zero budget rewrites
177
+ # cu_num_logits_np to bonus-only. Deriving the bitmask -> logits mapping
178
+ # from those rewritten offsets drops rows and trips the
179
+ # `num_masks == len(mapping)` assert in apply_grammar_bitmask.
180
+ manager = make_manager(
181
+ np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32),
182
+ np.ones(64),
183
+ )
184
+ manager.req_states.req_id_to_index["prefill"] = 2
185
+ manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32)
186
+ manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32)
187
+ manager._max_total_logits = 2 # < 3 requests * 1 bonus token
188
+
189
+ scheduled_spec_decode_tokens = {"low": [1, 2], "high": [3, 4]}
190
+ manager.get_num_tokens(
191
+ {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens
192
+ )
193
+ assert manager._batch_budget[2] == 0
194
+
195
+ req_ids = ["low", "high", "prefill"]
196
+ num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32)
197
+ _, cu_num_logits_np = manager.compact_batch(
198
+ num_draft_tokens_per_req,
199
+ np.array([3, 3, 40], dtype=np.int32),
200
+ np.array([0, 3, 6, 7], dtype=np.int32),
201
+ )
202
+
203
+ mask_stride = manager.num_speculative_steps + manager.num_bonus_tokens
204
+ mapping = _build_grammar_mapping(
205
+ req_ids,
206
+ req_ids,
207
+ cu_num_logits_np,
208
+ num_draft_tokens_per_req,
209
+ manager.num_bonus_tokens,
210
+ mask_stride,
211
+ )
212
+
213
+ num_bitmask_rows = sum(
214
+ len(scheduled_spec_decode_tokens.get(req_id, ())) + 1 for req_id in req_ids
215
+ )
216
+ assert len(mapping) == num_bitmask_rows
217
+ # (request, position) keys, so the kernel can mask rows the compacted
218
+ # device layout no longer has room for.
219
+ assert mapping == [0, 1, 2, 3, 4, 5, 6]
vllm/v1/worker/gpu/model_runner.py
+1
@@ -425,6 +425,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
425
vocab_size=self.vocab_size,
426
device=self.device,
427
mask_stride=self.decode_query_len,
428
+ num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
429
)
430
431
if self.is_pooling_model and self.is_last_pp_rank:
vllm/v1/worker/gpu/structured_outputs.py
+39
-15
@@ -10,6 +10,32 @@ from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
10
from vllm.v1.worker.gpu.input_batch import InputBatch
11
12
13
+def _build_grammar_mapping(
14
+ req_ids: list[str],
15
+ grammar_req_ids: list[str],
16
+ cu_num_logits_np: np.ndarray,
17
+ num_draft_tokens_per_req: np.ndarray | None,
18
+ num_bonus_tokens: int,
19
+ mask_stride: int,
20
+) -> list[int]:
21
+ mapping: list[int] = []
22
+ req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)}
23
+ for grammar_req_id in grammar_req_ids:
24
+ req_idx = req_id_to_idx[grammar_req_id]
25
+ if num_draft_tokens_per_req is None:
26
+ num_positions = int(
27
+ cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx]
28
+ )
29
+ else:
30
+ # Grammar masks follow the scheduled layout even when adaptive
31
+ # verification compacts the actual CPU logit offsets to bonus-only.
32
+ num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens
33
+ mapping.extend(
34
+ req_idx * mask_stride + position for position in range(num_positions)
35
+ )
36
+ return mapping
37
+
38
+
39
class StructuredOutputsWorker:
40
def __init__(
41
self,
@@ -17,6 +43,7 @@ class StructuredOutputsWorker:
43
vocab_size: int,
44
device: torch.device,
45
mask_stride: int,
46
+ num_bonus_tokens: int,
47
):
48
self.logits_indices = torch.zeros(
49
max_num_logits, dtype=torch.int32, device=device
@@ -27,6 +54,7 @@ class StructuredOutputsWorker:
54
self.device = device
55
self.copy_stream = torch.cuda.Stream()
56
self.mask_stride = mask_stride
57
+ self.num_bonus_tokens = num_bonus_tokens
58
59
def apply_grammar_bitmask(
60
self,
@@ -45,21 +73,17 @@ class StructuredOutputsWorker:
73
)
74
75
# Construct bitmask -> logits mapping
48
- mapping: list[int] = []
49
- req_ids = input_batch.req_ids
50
- cu_num_logits = input_batch.cu_num_logits_np.tolist()
51
- req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)}
52
- for grammar_req_id in grammar_req_ids:
53
- req_idx = req_id_to_idx[grammar_req_id]
54
- logits_start_idx = cu_num_logits[req_idx]
55
- logits_end_idx = cu_num_logits[req_idx + 1]
56
- # Key by (request, position) rather than absolute logit index:
57
- # adaptive verification finalizes per-request logit offsets on
58
- # device, so the kernel resolves them from the GPU cu_num_logits.
59
- mapping.extend(
60
- req_idx * self.mask_stride + position
61
- for position in range(logits_end_idx - logits_start_idx)
62
- )
76
+ # Key by (request, position) rather than absolute logit index:
77
+ # adaptive verification finalizes per-request logit offsets on
78
+ # device, so the kernel resolves them from the GPU cu_num_logits.
79
+ mapping = _build_grammar_mapping(
80
+ input_batch.req_ids,
81
+ grammar_req_ids,
82
+ input_batch.cu_num_logits_np,
83
+ input_batch.num_draft_tokens_per_req,
84
+ self.num_bonus_tokens,
85
+ self.mask_stride,
86
+ )
87
88
# Asynchronously copy the mapping to GPU.
89
with torch.cuda.stream(self.copy_stream):