[Bugfix] Don't let a structured-output request sample from an unmasked row (#54442)
Signed-off-by: ArcheyChen <45146080+ArcheyChen@users.noreply.github.com> Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: GPT-5.6 <noreply@openai.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Ausar committed
Oct 1, 2026 at 10:11 UTC
34c458d27f512bb888d92acced70163b013ed1a5
8 files changed
+142
-3
tests/v1/spec_decode/test_rejection_sampler_utils.py
+1
@@ -909,6 +909,7 @@ def test_verify_rejects_unproposed_drafts():
909
inputs["expanded_idx_mapping"],
910
inputs["expanded_local_pos"],
911
None,
912
+ draft_sampled,
913
)
914
return sampled, num_sampled
915
tests/v1/worker/test_grammar_invalid_drafts.py
new
+56
@@ -0,0 +1,56 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Which speculative drafts a structured-output request may accept.
4
+
5
+`StructuredOutputManager.grammar_bitmask` fills row i before it inspects
6
+`req_tokens[i]`, so with the first -1 placeholder at index j rows 0..j carry a
7
+real mask and rows j+1.. carry the all-permissive `_full_mask`. Drafts j..K-1
8
+must therefore be rejected, or the request samples with no grammar constraint
9
+at a position the model really does sample.
10
+"""
11
+
12
+from types import SimpleNamespace
13
+
14
+import numpy as np
15
+import torch
16
+
17
+from vllm.v1.worker.gpu.structured_outputs import grammar_invalid_drafts
18
+
19
+
20
+def _input_batch(req_ids, num_drafts, num_admitted=None):
21
+ """`num_admitted` is what adaptive verification kept on device."""
22
+ num_admitted = num_drafts if num_admitted is None else num_admitted
23
+ rows = [n + 1 for n in num_admitted]
24
+ return SimpleNamespace(
25
+ req_ids=list(req_ids),
26
+ num_reqs=len(req_ids),
27
+ num_draft_tokens=sum(num_drafts),
28
+ num_draft_tokens_per_req=np.array(num_drafts, dtype=np.int32),
29
+ cu_num_logits=torch.tensor([0, *np.cumsum(rows)], dtype=torch.int32),
30
+ expanded_local_pos=torch.cat(
31
+ [torch.arange(n, dtype=torch.int32) for n in rows]
32
+ ),
33
+ )
34
+
35
+
36
+def _invalid_rows(batch, grammar_req_ids, num_acceptable):
37
+ mask = grammar_invalid_drafts(batch, grammar_req_ids, num_acceptable)
38
+ return None if mask is None else mask.nonzero().flatten().tolist()
39
+
40
+
41
+def test_reject_whole_window_when_nothing_was_backfilled():
42
+ # Request "g" is at rows 0..3 (1 + 3 drafts), so its drafts are rows 1..3.
43
+ batch = _input_batch(["g", "p"], [3, 3])
44
+ assert _invalid_rows(batch, ["g"], [0]) == [1, 2, 3]
45
+
46
+
47
+def test_keep_the_drafts_the_bitmask_could_see():
48
+ batch = _input_batch(["g", "p"], [3, 3])
49
+ assert _invalid_rows(batch, ["g"], [2]) == [3]
50
+
51
+
52
+def test_follow_the_device_layout_under_adaptive_verification():
53
+ # Scheduled 3 drafts each, but adaptive verification admitted 2 for "a" and
54
+ # 3 for "b": the real rows are a=0..2, b=3..6, not the scheduled 0..3, 4..7.
55
+ batch = _input_batch(["a", "b"], [3, 3], num_admitted=[2, 3])
56
+ assert _invalid_rows(batch, ["a", "b"], [1, 0]) == [2, 4, 5, 6]
vllm/v1/core/sched/output.py
+5
@@ -335,3 +335,8 @@ class GrammarOutput:
335
structured_output_request_ids: list[str]
336
# Bitmask ordered as structured_output_request_ids.
337
grammar_bitmask: "npt.NDArray[np.int32]"
338
+ # Per request, ordered as structured_output_request_ids: how many leading
339
+ # drafts the bitmask constrained. `grammar_bitmask` fills every row after
340
+ # the first -1 placeholder with the all-permissive `_full_mask`, so drafts
341
+ # from this index on must not be accepted. None invalidates every draft.
342
+ num_acceptable_drafts: list[int] | None = None
vllm/v1/core/sched/scheduler.py
+9
-1
@@ -71,6 +71,7 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
71
from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
72
from vllm.v1.spec_decode.metrics import SpecDecodingStats
73
from vllm.v1.structured_output import StructuredOutputManager
74
+from vllm.v1.structured_output.utils import strip_speculative_padding
75
from vllm.v1.utils import record_function_or_nullcontext
76
77
logger = init_logger(__name__)
@@ -1963,7 +1964,14 @@ class Scheduler(SchedulerInterface):
1964
structured_output_request_ids,
1965
scheduler_output.scheduled_spec_decode_tokens,
1966
)
1966
- return GrammarOutput(structured_output_request_ids, bitmask)
1967
+ spec_tokens = scheduler_output.scheduled_spec_decode_tokens
1968
+ num_acceptable_drafts = [
1969
+ len(strip_speculative_padding(spec_tokens.get(req_id, [])))
1970
+ for req_id in structured_output_request_ids
1971
+ ]
1972
+ return GrammarOutput(
1973
+ structured_output_request_ids, bitmask, num_acceptable_drafts
1974
+ )
1975
1976
def update_from_output(
1977
self,
vllm/v1/worker/gpu/model_runner.py
+11
-1
@@ -164,7 +164,10 @@ from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
164
from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
165
from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
166
from vllm.v1.worker.gpu.states import RequestState
167
-from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
167
+from vllm.v1.worker.gpu.structured_outputs import (
168
+ StructuredOutputsWorker,
169
+ grammar_invalid_drafts,
170
+)
171
from vllm.v1.worker.gpu.ubatch_utils import (
172
UBatchRunner,
173
UBatchState,
@@ -1576,6 +1579,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1579
sample_hidden_states = hidden_states[input_batch.logits_indices]
1580
logits = self.model.compute_logits(sample_hidden_states)
1581
1582
+ invalid_drafts = None
1583
# A diffusion prefill has no logit rows even when a bitmask row
1584
# arrived for it.
1585
if grammar_output is not None and logits.shape[0] > 0:
@@ -1587,6 +1591,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1591
grammar_output.structured_output_request_ids,
1592
grammar_output.grammar_bitmask,
1593
)
1594
+ invalid_drafts = grammar_invalid_drafts(
1595
+ input_batch,
1596
+ grammar_output.structured_output_request_ids,
1597
+ grammar_output.num_acceptable_drafts,
1598
+ )
1599
1600
sampler_output: SamplerOutput | None
1601
if input_batch.num_reqs == 0:
@@ -1605,6 +1614,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1614
input_batch,
1615
# Draft logits are needed for probabilistic rejection sampling.
1616
self.speculator.draft_logits,
1617
+ invalid_drafts,
1618
)
1619
1620
if shard_metadata is not None:
vllm/v1/worker/gpu/sample/batch_shard.py
+7
@@ -339,9 +339,16 @@ def _shard_grammar_output(
339
cursor += num_req_logits
340
if not local_ids:
341
return None
342
+ # num_acceptable_drafts is ordered with structured_output_request_ids.
343
+ num_acceptable = grammar_output.num_acceptable_drafts
344
+ if num_acceptable is not None:
345
+ owned = set(local_ids)
346
+ ids = grammar_output.structured_output_request_ids
347
+ num_acceptable = [n for i, n in zip(ids, num_acceptable) if i in owned]
348
return GrammarOutput(
349
structured_output_request_ids=local_ids,
350
grammar_bitmask=grammar_output.grammar_bitmask[keep_indices],
351
+ num_acceptable_drafts=num_acceptable,
352
)
353
354
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+16
-1
@@ -253,6 +253,7 @@ class RejectionSampler:
253
expanded_idx_mapping: torch.Tensor,
254
expanded_local_pos: torch.Tensor,
255
seq_lens_upper_bound_np: np.ndarray,
256
+ verify_draft_sampled: torch.Tensor,
257
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
258
processed_logits = self.sampler.apply_sampling_params(
259
logits,
@@ -267,7 +268,7 @@ class RejectionSampler:
268
sampled, num_sampled = rejection_sample(
269
processed_logits,
270
draft_logits,
270
- draft_sampled,
271
+ verify_draft_sampled,
272
cu_num_logits,
273
pos,
274
idx_mapping,
@@ -294,7 +295,10 @@ class RejectionSampler:
295
pos: torch.Tensor,
296
max_chunk_logits: int,
297
max_num_logprobs: int,
298
+ verify_draft_sampled: torch.Tensor | None = None,
299
) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]:
300
+ if verify_draft_sampled is None:
301
+ verify_draft_sampled = draft_sampled
302
cu_num_logits_np = input_batch.cu_num_logits_np
303
use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES
304
num_reqs = input_batch.num_reqs
@@ -330,6 +334,7 @@ class RejectionSampler:
334
input_batch.expanded_idx_mapping[lo:hi],
335
input_batch.expanded_local_pos[lo:hi],
336
input_batch.seq_lens_cpu_upper_bound.numpy()[start:end],
337
+ verify_draft_sampled[lo:hi],
338
)
339
chunk_logprobs = self._get_logprobs_tensors(
340
sampled,
@@ -368,6 +373,7 @@ class RejectionSampler:
373
logits: torch.Tensor,
374
input_batch: InputBatch,
375
draft_logits: torch.Tensor | None = None,
376
+ invalid_drafts: torch.Tensor | None = None,
377
) -> SamplerOutput:
378
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
379
# that num_nans is computed before applying penalties and temperature.
@@ -382,6 +388,14 @@ class RejectionSampler:
388
self.sampler.req_states.prefill_len.gpu,
389
)
390
391
+ # Marking a draft invalid (`is_valid_draft = draft_sampled >= 0`) pins the
392
+ # accepted length before it reaches a permissive bitmask row. Only
393
+ # verification sees the copy: `apply_sampling_params` and watermarking
394
+ # keep the real draft ids.
395
+ verify_draft_sampled = draft_sampled
396
+ if invalid_drafts is not None:
397
+ verify_draft_sampled = draft_sampled.masked_fill(invalid_drafts, -1)
398
+
399
max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
400
input_batch.idx_mapping_np
401
)
@@ -394,6 +408,7 @@ class RejectionSampler:
408
pos,
409
chunk_logit_limit,
410
max_num_logprobs,
411
+ verify_draft_sampled,
412
)
413
414
num_sampled, num_rejected = get_num_sampled_and_rejected(
vllm/v1/worker/gpu/structured_outputs.py
+37
@@ -35,6 +35,43 @@ def _build_grammar_mapping(
35
return mapping
36
37
38
+def grammar_invalid_drafts(
39
+ input_batch: InputBatch,
40
+ grammar_req_ids: list[str],
41
+ num_acceptable_drafts: list[int] | None,
42
+) -> torch.Tensor | None:
43
+ """Mask over logit rows whose draft was verified against a permissive row.
44
+
45
+ Draft i sits at local position i + 1, and drafts from `num_acceptable_drafts`
46
+ on met `_full_mask`. Local positions come from the device `cu_num_logits`, so
47
+ this stays right when adaptive verification trims drafts on device.
48
+ Returns None when there is nothing to invalidate.
49
+ """
50
+ if not grammar_req_ids or input_batch.num_draft_tokens == 0:
51
+ return None
52
+ req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
53
+ limit = np.full(input_batch.num_reqs, np.iinfo(np.int32).max, dtype=np.int32)
54
+ for i, req_id in enumerate(grammar_req_ids):
55
+ req_idx = req_id_to_idx.get(req_id)
56
+ if req_idx is not None:
57
+ # None (an older scheduler, or warmup): invalidate the whole window.
58
+ limit[req_idx] = (
59
+ num_acceptable_drafts[i] if num_acceptable_drafts is not None else 0
60
+ )
61
+ # Adaptive verification only trims drafts, so this bound holds either way.
62
+ num_drafts = input_batch.num_draft_tokens_per_req
63
+ if num_drafts is not None and (limit >= num_drafts).all():
64
+ return None
65
+ local_pos = input_batch.expanded_local_pos
66
+ cu_num_logits = input_batch.cu_num_logits
67
+ row_limit = torch.repeat_interleave(
68
+ async_tensor_h2d(limit, device=local_pos.device),
69
+ cu_num_logits[1:] - cu_num_logits[:-1],
70
+ output_size=local_pos.shape[0],
71
+ )
72
+ return local_pos > row_limit
73
+
74
+
75
class StructuredOutputsWorker:
76
def __init__(
77
self,