| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import openai # use the official client for correctness check |
| 5 | import pytest |
| 6 | import pytest_asyncio |
| 7 | |
| 8 | from tests.utils import RemoteOpenAIServer |
| 9 | |
| 10 | # any model with a chat template defined in tokenizer_config should work here |
| 11 | MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" |
| 12 | |
| 13 | |
| 14 | @pytest.fixture(scope="module") |
| 15 | def default_server_args(): |
| 16 | return [ |
| 17 | # use half precision for speed and memory savings in CI environment |
| 18 | "--max-model-len", |
| 19 | "2048", |
| 20 | "--max-num-seqs", |
| 21 | "128", |
| 22 | "--enforce-eager", |
| 23 | ] |
| 24 | |
| 25 | |
| 26 | @pytest.fixture(scope="module") |
| 27 | def server(default_server_args): |
| 28 | with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server: |
| 29 | yield remote_server |
| 30 | |
| 31 | |
| 32 | @pytest_asyncio.fixture |
| 33 | async def client(server): |
| 34 | async with server.get_async_client() as async_client: |
| 35 | yield async_client |
| 36 | |
| 37 | |
| 38 | @pytest.mark.asyncio |
| 39 | @pytest.mark.parametrize( |
| 40 | "model_name", |
| 41 | [MODEL_NAME], |
| 42 | ) |
| 43 | async def test_invalid_json_schema(client: openai.AsyncOpenAI, model_name: str) -> None: |
| 44 | invalid_json_schema = { |
| 45 | "$defs": { |
| 46 | "CarType": { |
| 47 | "enum": ["sedan", "SUV", "Truck", "Coupe"], |
| 48 | "title": "CarType", |
| 49 | "type": "string", |
| 50 | } |
| 51 | }, |
| 52 | "properties": { |
| 53 | "brand": {"title": "Brand", "type": "string"}, |
| 54 | "model": {"title": "Model", "type": "string"}, |
| 55 | "car_type": {"$ref": "#/$defs/CarType"}, |
| 56 | "foo": "bar", |
| 57 | }, |
| 58 | "required": ["brand", "model", "car_type"], |
| 59 | "title": "CarDescription", |
| 60 | "type": "object", |
| 61 | } |
| 62 | prompt = ( |
| 63 | "Generate a JSON with the brand, model and car_type of" |
| 64 | "the most iconic car from the 90's" |
| 65 | ) |
| 66 | with pytest.raises((openai.BadRequestError, openai.APIError)): |
| 67 | await client.chat.completions.create( |
| 68 | model=model_name, |
| 69 | messages=[ |
| 70 | { |
| 71 | "role": "user", |
| 72 | "content": prompt, |
| 73 | } |
| 74 | ], |
| 75 | extra_body={"structured_outputs": {"json": invalid_json_schema}}, |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | @pytest.mark.asyncio |
| 80 | @pytest.mark.parametrize( |
| 81 | "model_name", |
| 82 | [MODEL_NAME], |
| 83 | ) |
| 84 | async def test_invalid_regex(client: openai.AsyncOpenAI, model_name: str): |
| 85 | prompt = ( |
| 86 | "Generate an email address for Alan Turing, who works in Enigma." |
| 87 | "End in .com and new line. Example result:" |
| 88 | "alan.turing@enigma.com\n" |
| 89 | ) |
| 90 | |
| 91 | with pytest.raises((openai.BadRequestError, openai.APIError)): |
| 92 | await client.chat.completions.create( |
| 93 | model=model_name, |
| 94 | messages=[ |
| 95 | { |
| 96 | "role": "user", |
| 97 | "content": prompt, |
| 98 | } |
| 99 | ], |
| 100 | extra_body={"structured_outputs": {"regex": r"[.*"}, "stop": ["\n"]}, |
| 101 | ) |
| 102 | |
| 103 | |
| 104 | @pytest.mark.asyncio |
| 105 | @pytest.mark.parametrize( |
| 106 | "model_name", |
| 107 | [MODEL_NAME], |
| 108 | ) |
| 109 | async def test_invalid_grammar(client: openai.AsyncOpenAI, model_name: str): |
| 110 | invalid_simplified_sql_grammar = """ |
| 111 | root ::= select_statementinvalidsyntax |
| 112 | |
| 113 | select_statement ::= "SELECT " column " from " table " where " condition |
| 114 | |
| 115 | column ::= "col_1 " | "col_2 " |
| 116 | |
| 117 | table ::= "table_1 " | "table_2 " |
| 118 | |
| 119 | condition ::= column "= " number |
| 120 | |
| 121 | number ::= "1 " | "2 " |
| 122 | """ |
| 123 | |
| 124 | prompt = ( |
| 125 | "Generate an SQL query to show the 'username' and 'email'" |
| 126 | "from the 'users' table." |
| 127 | ) |
| 128 | with pytest.raises((openai.BadRequestError, openai.APIError)): |
| 129 | await client.chat.completions.create( |
| 130 | model=model_name, |
| 131 | messages=[ |
| 132 | { |
| 133 | "role": "user", |
| 134 | "content": prompt, |
| 135 | } |
| 136 | ], |
| 137 | extra_body={ |
| 138 | "structured_outputs": {"grammar": invalid_simplified_sql_grammar} |
| 139 | }, |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | @pytest.mark.asyncio |
| 144 | @pytest.mark.parametrize( |
| 145 | "model_name", |
| 146 | [MODEL_NAME], |
| 147 | ) |
| 148 | async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> None: |
| 149 | prompt = "Say hello" |
| 150 | with pytest.raises((openai.BadRequestError, openai.APIError)): |
| 151 | await client.chat.completions.create( |
| 152 | model=model_name, |
| 153 | messages=[ |
| 154 | { |
| 155 | "role": "user", |
| 156 | "content": prompt, |
| 157 | } |
| 158 | ], |
| 159 | extra_body={"structured_outputs": {"grammar": ""}}, |
| 160 | ) |
| 161 | |
| 162 | |
| 163 | # Decode-side token reuse for disaggregated serving. The router forwards the |
| 164 | # prefill stage's prompt token ids in kv_transfer_params so the decode stage |
| 165 | # skips re-tokenizing. |
| 166 | |
| 167 | TOKEN_IN_MESSAGES = [{"role": "user", "content": "Hello, how are you today?"}] |
| 168 | DECODE_MESSAGES = [{"role": "user", "content": "unrelated decode-side text"}] |
| 169 | |
| 170 | |
| 171 | @pytest.mark.asyncio |
| 172 | async def test_kv_transfer_prompt_token_ids_round_trip(client: openai.AsyncOpenAI): |
| 173 | """Ids forwarded in kv_transfer_params are used verbatim, skipping tokenize. |
| 174 | |
| 175 | The decode request carries different messages, so a response whose |
| 176 | prompt_token_ids match the forwarded ids proves the ids were used rather |
| 177 | than the request's own messages. Generated text is not compared across |
| 178 | requests because vLLM greedy decoding is not bitwise-reproducible. |
| 179 | """ |
| 180 | baseline = await client.chat.completions.create( |
| 181 | model=MODEL_NAME, |
| 182 | messages=TOKEN_IN_MESSAGES, |
| 183 | max_completion_tokens=16, |
| 184 | temperature=0, |
| 185 | extra_body={"return_token_ids": True}, |
| 186 | ) |
| 187 | reused_ids = baseline.prompt_token_ids |
| 188 | assert reused_ids |
| 189 | |
| 190 | decode = await client.chat.completions.create( |
| 191 | model=MODEL_NAME, |
| 192 | messages=DECODE_MESSAGES, |
| 193 | max_completion_tokens=16, |
| 194 | temperature=0, |
| 195 | extra_body={ |
| 196 | "kv_transfer_params": {"prompt_token_ids": reused_ids}, |
| 197 | "return_token_ids": True, |
| 198 | }, |
| 199 | ) |
| 200 | |
| 201 | # The engine saw the forwarded ids, not the decode request's own messages. |
| 202 | assert decode.prompt_token_ids == reused_ids |
| 203 | # text-out: reuse still yields a detokenized message. |
| 204 | assert decode.choices[0].message.content |
| 205 | |
| 206 | |
| 207 | @pytest.mark.asyncio |
| 208 | async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI): |
| 209 | """Decode-side token reuse streams chat-formatted text-out.""" |
| 210 | baseline = await client.chat.completions.create( |
| 211 | model=MODEL_NAME, |
| 212 | messages=TOKEN_IN_MESSAGES, |
| 213 | max_completion_tokens=16, |
| 214 | temperature=0, |
| 215 | extra_body={"return_token_ids": True}, |
| 216 | ) |
| 217 | reused_ids = baseline.prompt_token_ids |
| 218 | assert reused_ids |
| 219 | |
| 220 | stream = await client.chat.completions.create( |
| 221 | model=MODEL_NAME, |
| 222 | messages=DECODE_MESSAGES, |
| 223 | max_completion_tokens=16, |
| 224 | temperature=0, |
| 225 | stream=True, |
| 226 | extra_body={ |
| 227 | "kv_transfer_params": {"prompt_token_ids": reused_ids}, |
| 228 | "return_token_ids": True, |
| 229 | }, |
| 230 | ) |
| 231 | |
| 232 | content = "" |
| 233 | delta_token_ids: list[int] = [] |
| 234 | first_chunk = True |
| 235 | async for chunk in stream: |
| 236 | if first_chunk: |
| 237 | # prompt_token_ids arrives once, on the first chunk. |
| 238 | assert chunk.prompt_token_ids == reused_ids |
| 239 | first_chunk = False |
| 240 | if not chunk.choices: |
| 241 | continue |
| 242 | if chunk.choices[0].delta.content: |
| 243 | content += chunk.choices[0].delta.content |
| 244 | if tids := getattr(chunk.choices[0], "token_ids", None): |
| 245 | delta_token_ids.extend(tids) |
| 246 | |
| 247 | # streamed text-out, reconstructed from deltas, with generated token ids. |
| 248 | assert content |
| 249 | assert delta_token_ids |