| 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] |