| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Round trip parity CI: render -> generate -> derender -> render == render. |
| 5 | |
| 6 | Pins the coupled parsing path (`/v1/chat/completions` on a normal GPU |
| 7 | server which parses generated tokens incrementally as they stream out) |
| 8 | against the disaggregated path (`OnlineDerenderer.derender_chat` via |
| 9 | `/v1/chat/completions/derender` which parses the same tokens all at once, |
| 10 | out of process). A standard `vllm serve` GPU server mounts both, so one |
| 11 | real generation lets both parsers run on identical input. |
| 12 | |
| 13 | Both paths consume the same generated token IDs (extracted from the coupled |
| 14 | response's `token_ids` via `return_token_ids=True`), so generation |
| 15 | nondeterminism is irrelevant. The only variable under test is whether the |
| 16 | two parsing code paths agree. Parity is asserted unconditionally as only the |
| 17 | stronger per case assertions (e.g. "a tool call was produced") are gated |
| 18 | behind the marker actually having been emitted since a 1.5B model is not |
| 19 | guaranteed to emit `<think>` / `<tool_call>`. |
| 20 | |
| 21 | The streaming cases do the same for `stream=true`. Each coupled chunk's |
| 22 | `token_ids` becomes one `/inference/v1/generate` stream chunk, fed |
| 23 | through the chunked derender path with `stream_state` threaded across calls |
| 24 | the way a client would, so both parsers see the same chunk boundaries. |
| 25 | """ |
| 26 | |
| 27 | import json |
| 28 | |
| 29 | import httpx |
| 30 | import pytest |
| 31 | import pytest_asyncio |
| 32 | import regex as re |
| 33 | |
| 34 | from tests.entrypoints.scale_out.derender.utils import ( |
| 35 | assemble_stream, |
| 36 | stream_chat_derender, |
| 37 | ) |
| 38 | from tests.utils import RemoteOpenAIServer |
| 39 | |
| 40 | MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" |
| 41 | ARGS = [ |
| 42 | "--enable-auto-tool-choice", |
| 43 | "--tool-call-parser", |
| 44 | "hermes", |
| 45 | "--reasoning-parser", |
| 46 | "deepseek_r1", |
| 47 | "--enable-scale-out", |
| 48 | ] |
| 49 | |
| 50 | TOOLS = [ |
| 51 | { |
| 52 | "type": "function", |
| 53 | "function": { |
| 54 | "name": "get_weather", |
| 55 | "description": "Get weather for a city", |
| 56 | "parameters": { |
| 57 | "type": "object", |
| 58 | "properties": {"city": {"type": "string"}}, |
| 59 | }, |
| 60 | }, |
| 61 | } |
| 62 | ] |
| 63 | FORCE_WEATHER_TOOL = {"type": "function", "function": {"name": "get_weather"}} |
| 64 | |
| 65 | |
| 66 | @pytest.fixture(scope="module") |
| 67 | def server(): |
| 68 | with RemoteOpenAIServer(MODEL, ARGS) as remote_server: |
| 69 | yield remote_server |
| 70 | |
| 71 | |
| 72 | @pytest_asyncio.fixture |
| 73 | async def client(server): |
| 74 | async with httpx.AsyncClient( |
| 75 | base_url=server.url_for(""), timeout=60.0 |
| 76 | ) as http_client: |
| 77 | yield http_client |
| 78 | |
| 79 | |
| 80 | # --------------------------------------------------------------------------- |
| 81 | # Helpers |
| 82 | # --------------------------------------------------------------------------- |
| 83 | |
| 84 | |
| 85 | async def _coupled(client: httpx.AsyncClient, messages: list[dict], **extra) -> dict: |
| 86 | resp = await client.post( |
| 87 | "/v1/chat/completions", |
| 88 | json={ |
| 89 | "model": MODEL, |
| 90 | "messages": messages, |
| 91 | "temperature": 0, |
| 92 | "max_tokens": 128, |
| 93 | "return_token_ids": True, |
| 94 | **extra, |
| 95 | }, |
| 96 | ) |
| 97 | assert resp.status_code == 200, resp.text |
| 98 | return resp.json() |
| 99 | |
| 100 | |
| 101 | async def _disagg( |
| 102 | client: httpx.AsyncClient, |
| 103 | output_ids: list[int], |
| 104 | prompt_tokens: int, |
| 105 | finish_reason: str, |
| 106 | chat_request: dict, |
| 107 | logprobs: dict | None = None, |
| 108 | ) -> dict: |
| 109 | resp = await client.post( |
| 110 | "/v1/chat/completions/derender", |
| 111 | json={ |
| 112 | "model": MODEL, |
| 113 | "generate_response": { |
| 114 | "request_id": "parity", |
| 115 | "choices": [ |
| 116 | { |
| 117 | "index": 0, |
| 118 | "token_ids": output_ids, |
| 119 | "finish_reason": finish_reason, |
| 120 | "logprobs": logprobs, |
| 121 | } |
| 122 | ], |
| 123 | }, |
| 124 | "prompt_tokens": prompt_tokens, |
| 125 | "chat_request": chat_request, |
| 126 | }, |
| 127 | ) |
| 128 | assert resp.status_code == 200, resp.text |
| 129 | return resp.json() |
| 130 | |
| 131 | |
| 132 | def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]: |
| 133 | """[(name, json normalized args)] so key ordering / whitespace don't |
| 134 | cause false negatives.""" |
| 135 | return [ |
| 136 | (tc["function"]["name"], json.loads(tc["function"]["arguments"])) |
| 137 | for tc in (response_choice["message"].get("tool_calls") or []) |
| 138 | ] |
| 139 | |
| 140 | |
| 141 | def _assert_parity(coupled: dict, disagg: dict) -> None: |
| 142 | """Both paths saw the same tokens, so they must agree unconditionally.""" |
| 143 | c, d = coupled["choices"][0], disagg["choices"][0] |
| 144 | assert d["message"]["content"] == c["message"]["content"] |
| 145 | assert d["message"].get("reasoning") == c["message"].get("reasoning") |
| 146 | assert _tool_sig(d) == _tool_sig(c) |
| 147 | assert d["finish_reason"] == c["finish_reason"] |
| 148 | assert disagg["usage"]["prompt_tokens"] == coupled["usage"]["prompt_tokens"] |
| 149 | assert disagg["usage"]["completion_tokens"] == len(c["token_ids"]) |
| 150 | |
| 151 | |
| 152 | async def _run_parity_case( |
| 153 | client: httpx.AsyncClient, messages: list[dict], **extra |
| 154 | ) -> tuple[dict, dict]: |
| 155 | """Run the coupled request then feed its generated tokens into the |
| 156 | disaggregated derender endpoint. Returns (coupled, disagg).""" |
| 157 | coupled = await _coupled(client, messages, **extra) |
| 158 | ch = coupled["choices"][0] |
| 159 | chat_request = {"model": MODEL, "messages": messages, **extra} |
| 160 | disagg = await _disagg( |
| 161 | client, |
| 162 | ch["token_ids"], |
| 163 | coupled["usage"]["prompt_tokens"], |
| 164 | ch["finish_reason"], |
| 165 | chat_request, |
| 166 | ) |
| 167 | return coupled, disagg |
| 168 | |
| 169 | |
| 170 | # --------------------------------------------------------------------------- |
| 171 | # Parity cases |
| 172 | # --------------------------------------------------------------------------- |
| 173 | |
| 174 | |
| 175 | @pytest.mark.asyncio |
| 176 | async def test_parity_plain(client): |
| 177 | """Plain detokenization parity. No reasoning/tool markers involved.""" |
| 178 | messages = [ |
| 179 | {"role": "user", "content": "What is 2+2? Answer in one short sentence."} |
| 180 | ] |
| 181 | coupled, disagg = await _run_parity_case(client, messages) |
| 182 | _assert_parity(coupled, disagg) |
| 183 | |
| 184 | |
| 185 | @pytest.mark.asyncio |
| 186 | async def test_parity_reasoning(client): |
| 187 | """Reasoning/content split parity for <think>...</think> outputs.""" |
| 188 | messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}] |
| 189 | coupled, disagg = await _run_parity_case( |
| 190 | client, messages, include_reasoning=True, max_tokens=256 |
| 191 | ) |
| 192 | _assert_parity(coupled, disagg) |
| 193 | |
| 194 | if not coupled["choices"][0]["message"].get("reasoning"): |
| 195 | pytest.skip("Model did not emit a <think> block") |
| 196 | assert disagg["choices"][0]["message"]["reasoning"] |
| 197 | |
| 198 | |
| 199 | @pytest.mark.asyncio |
| 200 | async def test_parity_tool_call(client): |
| 201 | """Tool call name+args parity.""" |
| 202 | messages = [{"role": "user", "content": "What's the weather in Paris?"}] |
| 203 | coupled, disagg = await _run_parity_case( |
| 204 | client, messages, tools=TOOLS, tool_choice=FORCE_WEATHER_TOOL, max_tokens=1024 |
| 205 | ) |
| 206 | _assert_parity(coupled, disagg) |
| 207 | |
| 208 | if not _tool_sig(coupled["choices"][0]): |
| 209 | pytest.skip("Model did not emit a tool call") |
| 210 | assert _tool_sig(disagg["choices"][0]) |
| 211 | |
| 212 | |
| 213 | @pytest.mark.asyncio |
| 214 | async def test_parity_reasoning_and_tool_call(client): |
| 215 | """Combined reasoning + tool call parity means the highest drift risk |
| 216 | since it exercises both parser branches on the same output.""" |
| 217 | messages = [{"role": "user", "content": "What's the weather in Paris?"}] |
| 218 | coupled, disagg = await _run_parity_case( |
| 219 | client, |
| 220 | messages, |
| 221 | tools=TOOLS, |
| 222 | tool_choice=FORCE_WEATHER_TOOL, |
| 223 | include_reasoning=True, |
| 224 | max_tokens=1024, |
| 225 | ) |
| 226 | _assert_parity(coupled, disagg) |
| 227 | |
| 228 | c_msg = coupled["choices"][0]["message"] |
| 229 | if not (c_msg.get("reasoning") and _tool_sig(coupled["choices"][0])): |
| 230 | pytest.skip("Model did not emit both a <think> block and a tool call") |
| 231 | d_msg = disagg["choices"][0]["message"] |
| 232 | assert d_msg["reasoning"] |
| 233 | assert _tool_sig(disagg["choices"][0]) |
| 234 | |
| 235 | |
| 236 | @pytest.mark.asyncio |
| 237 | async def test_parity_logprobs(client): |
| 238 | """token_id:N resolution parity vs. the coupled server's real strings. |
| 239 | |
| 240 | A real disaggregated worker only has token IDs so it emits logprobs |
| 241 | with `token_id:N` placeholders (`return_tokens_as_token_ids=True` |
| 242 | reproduces that shape here). `/derender` must resolve those |
| 243 | placeholders to the same token strings/bytes the coupled server |
| 244 | resolves them to directly. |
| 245 | """ |
| 246 | messages = [{"role": "user", "content": "What is 2+2?"}] |
| 247 | extra = {"logprobs": True, "top_logprobs": 3} |
| 248 | |
| 249 | # What a real GPU less worker would hand to /derender is token IDs plus |
| 250 | # logprobs still in token_id:N placeholder form |
| 251 | placeholder = await _coupled( |
| 252 | client, messages, return_tokens_as_token_ids=True, **extra |
| 253 | ) |
| 254 | ch = placeholder["choices"][0] |
| 255 | chat_request = {"model": MODEL, "messages": messages, **extra} |
| 256 | disagg = await _disagg( |
| 257 | client, |
| 258 | ch["token_ids"], |
| 259 | placeholder["usage"]["prompt_tokens"], |
| 260 | ch["finish_reason"], |
| 261 | chat_request, |
| 262 | logprobs=ch["logprobs"], |
| 263 | ) |
| 264 | |
| 265 | # The coupled server resolving the same greedy generation |
| 266 | # to real token strings itself. |
| 267 | resolved = await _coupled(client, messages, **extra) |
| 268 | assert resolved["choices"][0]["token_ids"] == ch["token_ids"], ( |
| 269 | "greedy (temperature=0) generation was expected to be deterministic " |
| 270 | "across the two coupled calls used to build this test's fixtures" |
| 271 | ) |
| 272 | _assert_parity(resolved, disagg) |
| 273 | |
| 274 | r_content = resolved["choices"][0]["logprobs"]["content"] |
| 275 | d_content = disagg["choices"][0]["logprobs"]["content"] |
| 276 | assert len(d_content) == len(r_content) |
| 277 | for d_entry, r_entry in zip(d_content, r_content): |
| 278 | assert d_entry["token"] == r_entry["token"] |
| 279 | assert d_entry["bytes"] == r_entry["bytes"] |
| 280 | |
| 281 | |
| 282 | # --------------------------------------------------------------------------- |
| 283 | # Streaming helpers |
| 284 | # --------------------------------------------------------------------------- |
| 285 | |
| 286 | |
| 287 | async def _coupled_stream( |
| 288 | client: httpx.AsyncClient, messages: list[dict], **extra |
| 289 | ) -> tuple[list[dict], list[int]]: |
| 290 | """Stream `/v1/chat/completions` and return the choices that carry |
| 291 | `token_ids` along with the prompt token IDs. |
| 292 | |
| 293 | The leading role only chunk has no `token_ids` and no counterpart on |
| 294 | `/inference/v1/generate`, so it is dropped. |
| 295 | """ |
| 296 | choices: list[dict] = [] |
| 297 | prompt_token_ids: list[int] | None = None |
| 298 | async with client.stream( |
| 299 | "POST", |
| 300 | "/v1/chat/completions", |
| 301 | json={ |
| 302 | "model": MODEL, |
| 303 | "messages": messages, |
| 304 | "temperature": 0, |
| 305 | "max_tokens": 128, |
| 306 | "return_token_ids": True, |
| 307 | "stream": True, |
| 308 | **extra, |
| 309 | }, |
| 310 | ) as resp: |
| 311 | assert resp.status_code == 200, await resp.aread() |
| 312 | async for line in resp.aiter_lines(): |
| 313 | if not line.startswith("data: "): |
| 314 | continue |
| 315 | data = line[len("data: ") :] |
| 316 | if data == "[DONE]": |
| 317 | break |
| 318 | chunk = json.loads(data) |
| 319 | assert "error" not in chunk, chunk |
| 320 | if chunk.get("prompt_token_ids") is not None: |
| 321 | prompt_token_ids = chunk["prompt_token_ids"] |
| 322 | choices.extend( |
| 323 | ch for ch in chunk["choices"] if ch.get("token_ids") is not None |
| 324 | ) |
| 325 | assert prompt_token_ids is not None |
| 326 | return choices, prompt_token_ids |
| 327 | |
| 328 | |
| 329 | def _tool_call_id_shape(tool_call_id: str | None) -> str | None: |
| 330 | """Drop the random suffix so IDs compare by format, e.g. `chatcmpl-tool-`. |
| 331 | Deterministic IDs such as kimi_k2's `functions.<name>:<idx>` are kept |
| 332 | whole.""" |
| 333 | if tool_call_id is None: |
| 334 | return None |
| 335 | return re.sub(r"[0-9a-f]{16}$", "", tool_call_id) |
| 336 | |
| 337 | |
| 338 | async def _run_stream_parity_case( |
| 339 | client: httpx.AsyncClient, |
| 340 | messages: list[dict], |
| 341 | finish_only_tail: bool = False, |
| 342 | **extra, |
| 343 | ) -> tuple[dict, list[dict]]: |
| 344 | """Stream the coupled request, replay its chunks through the streaming |
| 345 | derender endpoint and assert the assembled messages agree. Returns the |
| 346 | assembled coupled message and the coupled choices. |
| 347 | |
| 348 | Derender gets the engine's finish reason, since `"tool_calls"` is a chat |
| 349 | level rewrite of `"stop"` that derender has to reapply itself. With |
| 350 | `finish_only_tail` it goes on a trailing chunk with no tokens. |
| 351 | """ |
| 352 | coupled_choices, prompt_token_ids = await _coupled_stream(client, messages, **extra) |
| 353 | finish_reason = coupled_choices[-1]["finish_reason"] |
| 354 | disagg = await stream_chat_derender( |
| 355 | client, |
| 356 | [tid for ch in coupled_choices for tid in ch["token_ids"]], |
| 357 | [len(ch["token_ids"]) for ch in coupled_choices] |
| 358 | + ([0] if finish_only_tail else []), |
| 359 | {"model": MODEL, "messages": messages, **extra}, |
| 360 | len(prompt_token_ids), |
| 361 | prompt_token_ids, |
| 362 | finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, |
| 363 | ) |
| 364 | coupled = assemble_stream(coupled_choices) |
| 365 | for message in (coupled, disagg): |
| 366 | for tc in message["tool_calls"]: |
| 367 | tc["id"] = _tool_call_id_shape(tc["id"]) |
| 368 | assert disagg == coupled |
| 369 | return coupled, coupled_choices |
| 370 | |
| 371 | |
| 372 | # --------------------------------------------------------------------------- |
| 373 | # Streaming parity cases |
| 374 | # --------------------------------------------------------------------------- |
| 375 | |
| 376 | |
| 377 | @pytest.mark.asyncio |
| 378 | async def test_stream_parity_plain(client): |
| 379 | """Plain streamed detokenization parity.""" |
| 380 | messages = [ |
| 381 | {"role": "user", "content": "What is 2+2? Answer in one short sentence."} |
| 382 | ] |
| 383 | await _run_stream_parity_case(client, messages) |
| 384 | |
| 385 | |
| 386 | @pytest.mark.asyncio |
| 387 | async def test_stream_parity_reasoning(client): |
| 388 | """Streamed reasoning/content split parity.""" |
| 389 | messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}] |
| 390 | coupled, _ = await _run_stream_parity_case( |
| 391 | client, messages, include_reasoning=True, max_tokens=256 |
| 392 | ) |
| 393 | |
| 394 | if not coupled["reasoning"]: |
| 395 | pytest.skip("Model did not emit a <think> block") |
| 396 | |
| 397 | |
| 398 | @pytest.mark.asyncio |
| 399 | async def test_stream_parity_tool_call(client): |
| 400 | """Streamed tool call parity under `tool_choice="auto"`, where a tool |
| 401 | call also rewrites `finish_reason` to `"tool_calls"`.""" |
| 402 | messages = [{"role": "user", "content": "What's the weather in Paris?"}] |
| 403 | coupled, _ = await _run_stream_parity_case( |
| 404 | client, messages, tools=TOOLS, tool_choice="auto", max_tokens=512 |
| 405 | ) |
| 406 | |
| 407 | if not coupled["tool_calls"]: |
| 408 | pytest.skip("Model did not emit a tool call") |
| 409 | |
| 410 | |
| 411 | @pytest.mark.asyncio |
| 412 | @pytest.mark.parametrize( |
| 413 | "chunking", ["engine_steps", "multi_token", "finish_only_tail"] |
| 414 | ) |
| 415 | async def test_stream_parity_reasoning_and_tool_call(client, chunking): |
| 416 | """Streamed reasoning + forced tool call parity across chunk shapes. |
| 417 | |
| 418 | `multi_token` batches tokens per chunk with `stream_interval` the way |
| 419 | speculative decoding does, so replay has to reuse the original chunk |
| 420 | boundaries. `finish_only_tail` sends the finish reason on a chunk with no |
| 421 | tokens, which must still flush buffered tool call arguments. |
| 422 | """ |
| 423 | messages = [{"role": "user", "content": "What's the weather in Paris?"}] |
| 424 | extra: dict = { |
| 425 | "tools": TOOLS, |
| 426 | "tool_choice": FORCE_WEATHER_TOOL, |
| 427 | "include_reasoning": True, |
| 428 | "max_tokens": 512, |
| 429 | } |
| 430 | if chunking == "multi_token": |
| 431 | extra["stream_interval"] = 4 |
| 432 | coupled, coupled_choices = await _run_stream_parity_case( |
| 433 | client, |
| 434 | messages, |
| 435 | finish_only_tail=chunking == "finish_only_tail", |
| 436 | **extra, |
| 437 | ) |
| 438 | |
| 439 | if chunking == "multi_token": |
| 440 | assert any(len(ch["token_ids"]) > 1 for ch in coupled_choices) |
| 441 | if not (coupled["reasoning"] and coupled["tool_calls"]): |
| 442 | pytest.skip("Model did not emit both a <think> block and a tool call") |