@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 160 lines 5.81 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Backend and engine-reasoner checks that are not manager-flow tests."""
4
5 from typing import Literal
6
7 import pytest
8 from transformers import AutoTokenizer
9
10 from vllm.config import StructuredOutputsConfig, VllmConfig
11 from vllm.config.model import ModelConfig
12 from vllm.config.speculative import SpeculativeConfig
13 from vllm.parser.engine.adapters import ParserEngineReasoningAdapter
14 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
15 from vllm.v1.request import Request
16 from vllm.v1.structured_output import StructuredOutputManager
17
18 TOKENIZER = "gpt2"
19 NUM_SPEC_TOKENS = 4
20
21
22 def _make_manager_and_request(
23 backend: Literal["xgrammar", "guidance"], prompt_str: str = '{"a": "b"}'
24 ):
25 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
26 prompt = tokenizer.encode(prompt_str)
27
28 vllm_config = VllmConfig(
29 model_config=ModelConfig(tokenizer=TOKENIZER),
30 structured_outputs_config=StructuredOutputsConfig(backend=backend),
31 speculative_config=SpeculativeConfig(
32 model="[ngram]", num_speculative_tokens=NUM_SPEC_TOKENS
33 ),
34 )
35 manager = StructuredOutputManager(vllm_config)
36
37 sampling_params = SamplingParams(
38 structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
39 )
40 assert sampling_params.structured_outputs is not None
41 sampling_params.structured_outputs._backend = backend
42 sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)
43
44 request = Request(
45 "mtp_req",
46 prompt_token_ids=prompt,
47 sampling_params=sampling_params,
48 pooling_params=None,
49 )
50 manager.grammar_init(request)
51 assert request.structured_output_request is not None
52 while not request.structured_output_request._check_grammar_completion():
53 continue
54
55 return tokenizer, manager, request, prompt
56
57
58 def test_xgrammar_accept_tokens_stops_at_termination(capfd):
59 """Tokens after a terminating EOS do not reach the matcher."""
60 tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
61 grammar = request.structured_output_request.grammar
62
63 assert grammar.accept_tokens(request.request_id, prompt)
64
65 eos = tokenizer.eos_token_id
66 trailing = tokenizer.encode("\n")[0]
67 processed_before = grammar.num_processed_tokens
68
69 assert grammar.accept_tokens(request.request_id, [eos, trailing])
70 assert grammar.is_terminated()
71 assert grammar.num_processed_tokens == processed_before + 1
72 assert "trying to accept new token" not in capfd.readouterr().err
73
74 processed_after_eos = grammar.num_processed_tokens
75 assert grammar.accept_tokens(request.request_id, [trailing])
76 assert grammar.num_processed_tokens == processed_after_eos
77 assert "trying to accept new token" not in capfd.readouterr().err
78
79 grammar.reset()
80 assert not grammar.is_terminated()
81 assert grammar.num_processed_tokens == 0
82
83
84 def test_xgrammar_validate_tokens_stops_at_termination(capfd):
85 """Validation rolls back after reaching a terminating EOS."""
86 tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
87 grammar = request.structured_output_request.grammar
88
89 assert grammar.accept_tokens(request.request_id, prompt)
90
91 eos = tokenizer.eos_token_id
92 trailing = tokenizer.encode("\n")[0]
93
94 assert grammar.validate_tokens([eos, trailing]) == [eos]
95 assert "trying to accept new token" not in capfd.readouterr().err
96 assert not grammar.matcher.is_terminated()
97
98 assert grammar.accept_tokens(request.request_id, [eos])
99 assert grammar.is_terminated()
100
101 assert grammar.validate_tokens([trailing]) == []
102 assert "trying to accept new token" not in capfd.readouterr().err
103
104
105 class _EngineReasonerStub(ParserEngineReasoningAdapter):
106 """Adapter-typed reasoner with a fixed end-token set and no real engine."""
107
108 def __init__(self, end_token_ids):
109 self._end_token_ids = frozenset(end_token_ids)
110 self.windows: list[list[int]] = []
111
112 @property
113 def reasoning_end_token_ids(self):
114 return self._end_token_ids
115
116 def find_reasoning_end_offset(self, token_ids):
117 self.windows.append(list(token_ids))
118 for offset, token in enumerate(token_ids):
119 if token in self._end_token_ids:
120 return offset
121 return len(token_ids)
122
123 def is_reasoning_end(self, input_ids):
124 return any(token in self._end_token_ids for token in input_ids)
125
126 def is_reasoning_end_streaming(self, input_ids, delta_ids):
127 raise AssertionError("engine path must not rescan draft prefixes")
128
129
130 @pytest.mark.parametrize("backend", ["xgrammar", "guidance"])
131 def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend):
132 """Engine reasoners see the draft window once, without -1 padding."""
133 tokenizer, manager, request, prompt = _make_manager_and_request(backend)
134 grammar = request.structured_output_request.grammar
135
136 assert grammar.accept_tokens(request.request_id, prompt)
137
138 marker = tokenizer.encode("\n")[0]
139 reasoner = _EngineReasonerStub({marker})
140 manager.reasoner_cls = _EngineReasonerStub
141 request.structured_output_request.reasoner = reasoner
142 request.structured_output_request.reasoning_ended = False
143
144 pre = tokenizer.encode(" ")[0]
145 post = tokenizer.encode(",")[0]
146 drafts = [pre, marker, post, -1]
147
148 bitmask = manager.grammar_bitmask(
149 requests={request.request_id: request},
150 structured_output_request_ids=[request.request_id],
151 scheduled_spec_decode_tokens={request.request_id: drafts},
152 )
153
154 assert bitmask is not None
155 assert bitmask.shape[0] == len(drafts) + 1
156 assert (bitmask[0] == -1).all()
157 assert (bitmask[1] == -1).all()
158 assert not (bitmask[2] == -1).all()
159 assert reasoner.windows == [[pre, marker, post]]
160 assert not grammar.is_terminated()