@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4dc57d44c4

[Frontend] Attach resolved logprobs to streaming derender chunks (#55029)

Signed-off-by: Shimi Bandiel <shimib@google.com> Co-authored-by: Claude <noreply@anthropic.com>

Shimi Bandiel committed Sep 29, 2026 at 11:43 UTC 4dc57d44c4a790a210ad17cde2e1af98146f8bca

4 files changed +224 -8

docs/serving/online_serving/derenderer.md

+10 -1

@@ -98,7 +98,7 @@ The server keeps no state between calls. Everything the next call needs is in `s

98 - **Send the same context on every call.** `chat_request` and `prompt_token_ids` aren't kept between calls, so they go with every chunk including the usage chunk. Both are required when a tool or reasoning parser is configured. `prompt_token_ids` is the `token_ids` of the `GenerateRequest` returned by `/render`.

99 - **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk.

100

101 -Streaming chunks don't carry logprobs yet, so `logprobs` on a generate chunk are dropped. The non-streaming endpoints do resolve them, including `token_id:N` placeholders.

101 +Streaming derender resolves `logprobs` on the plain detokenization path, including `token_id:N` placeholders, same as the non-streaming endpoints (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them.

102

103 ## Streaming cost

104

@@ -115,6 +115,15 @@ When a tool or reasoning parser is configured, parser internal state (buffered m

115

116 `output_token_ids` and `prompt_token_ids` are both bounded by `max_model_len` but callers streaming long reasoning traces through a parser configured model should expect materially more state transport and CPU cost than the plain detokenization path.

117

118 +## Streaming state and logprobs

119 +

120 +Streaming derender is stateless on the server side: all mutable state lives in the client-carried `stream_state` (`DerenderStreamState`), passed back on every per-chunk call. Besides the bounded incremental detokenization window (`prev_tokens`, `prefix_offset`, `read_offset`), `role_sent`, and the parser path's replay fields (see [Streaming cost](#streaming-cost)), the state carries two fields for logprob handling:

121 +

122 +- `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings

123 +- `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0

124 +

125 +When a streamed `GenerateResponseStreamChoice` carries `logprobs`, the `token_id:N` placeholders are resolved per chunk and the resolved logprobs are attached to the corresponding streamed choice — for chat as `ChatCompletionLogProbs`, for completions converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`.

126 +

127 ## Example

128

129 The example below drives the full `render → generate → derender` round trip for a chat request against a GPU less render server (`/render`, `/derender`) and a token-in / token-out engine (`/inference/v1/generate`).

tests/entrypoints/scale_out/derender/test_derender_stream.py

+127

@@ -169,6 +169,7 @@ def _make_stream_chunk(

169 request_id: str = "test-req",

170 usage: dict | None = None,

171 metrics: PerRequestMetrics | None = None,

172 + logprobs: dict | None = None,

173 ) -> GenerateStreamResponse:

174 """Build a GenerateStreamResponse SSE chunk."""

175 return GenerateStreamResponse(

@@ -178,6 +179,7 @@ def _make_stream_chunk(

179 index=index,

180 token_ids=token_ids,

181 finish_reason=finish_reason,

182 + logprobs=logprobs,

183 )

184 ],

185 usage=UsageInfo(**usage) if usage else None,

@@ -185,6 +187,22 @@ def _make_stream_chunk(

187 )

188

189

190 +def _placeholder_logprobs(token_ids: list[int]) -> dict:

191 + """Per-token logprob entries using token_id:N placeholders, as sent by

192 + the generate worker."""

193 + return {

194 + "content": [

195 + {

196 + "token": f"token_id:{tid}",

197 + "logprob": -0.5,

198 + "bytes": None,

199 + "top_logprobs": [],

200 + }

201 + for tid in token_ids

202 + ]

203 + }

204 +

205 +

206 def _make_usage_chunk(

207 completion_tokens: int,

208 prompt_tokens: int = 0,

@@ -536,6 +554,115 @@ class TestDerenderCompletionStream:

554 assert chunk.choices[0].finish_reason == "length"

555

556

557 +class TestStreamLogprobs:

558 + """Streaming derender must carry per-chunk logprobs with placeholders

559 + resolved, matching what the generate streaming path emits."""

560 +

561 + @pytest.mark.asyncio

562 + async def test_chat_stream_logprobs_resolved_per_chunk(self, derenderer, tokenizer):

563 + """Each streamed chunk carries logprobs with token_id:N resolved."""

564 + token_ids = tokenizer.encode("hello world")[:6]

565 + mid = len(token_ids) // 2

566 +

567 + state = None

568 + for part in (token_ids[:mid], token_ids[mid:]):

569 + chunk, state = await derenderer.derender_chat_stream(

570 + model=MODEL_NAME,

571 + generate_chunk=_make_stream_chunk(

572 + part, logprobs=_placeholder_logprobs(part)

573 + ),

574 + state=state,

575 + )

576 + logprobs = chunk.choices[0].logprobs

577 + assert logprobs is not None and logprobs.content is not None

578 + assert len(logprobs.content) == len(part)

579 + for entry in logprobs.content:

580 + assert not entry.token.startswith("token_id:"), (

581 + f"placeholder not resolved: {entry.token!r}"

582 + )

583 +

584 + @pytest.mark.asyncio

585 + async def test_chat_stream_logprobs_multibyte_across_chunks(

586 + self, derenderer, tokenizer

587 + ):

588 + """Byte-fallback correction works when a multi-byte character's

589 + tokens are split across chunks (context carried in stream_state)."""

590 + token_ids = tokenizer.encode("👍", add_special_tokens=False)

591 + if len(token_ids) < 2 or "�" not in tokenizer.decode([token_ids[-1]]):

592 + pytest.skip("Tokenizer does not byte-split this character")

593 +

594 + state = None

595 + last_chunk = None

596 + # One token per chunk: every entry after the first needs cross-chunk

597 + # context to resolve without U+FFFD.

598 + for i, tid in enumerate(token_ids):

599 + last_chunk, state = await derenderer.derender_chat_stream(

600 + model=MODEL_NAME,

601 + generate_chunk=_make_stream_chunk(

602 + [tid], logprobs=_placeholder_logprobs([tid])

603 + ),

604 + state=state,

605 + )

606 +

607 + assert last_chunk is not None

608 + final_entry = last_chunk.choices[0].logprobs.content[0]

609 + assert not final_entry.token.endswith("�"), (

610 + "byte-fallback correction failed across chunk boundary: "

611 + f"{final_entry.token!r}"

612 + )

613 +

614 + @pytest.mark.asyncio

615 + async def test_completion_stream_logprobs_text_offset_absolute(

616 + self, derenderer, tokenizer

617 + ):

618 + """text_offset continues across chunks instead of restarting at 0."""

619 + token_ids = tokenizer.encode("hello world")[:6]

620 + mid = len(token_ids) // 2

621 +

622 + chunk1, state = await derenderer.derender_completion_stream(

623 + model=MODEL_NAME,

624 + generate_chunk=_make_stream_chunk(

625 + token_ids[:mid], logprobs=_placeholder_logprobs(token_ids[:mid])

626 + ),

627 + )

628 + chunk2, _ = await derenderer.derender_completion_stream(

629 + model=MODEL_NAME,

630 + generate_chunk=_make_stream_chunk(

631 + token_ids[mid:], logprobs=_placeholder_logprobs(token_ids[mid:])

632 + ),

633 + state=state,

634 + )

635 +

636 + assert chunk1.choices[0].logprobs.text_offset[0] == 0

637 + assert chunk2.choices[0].logprobs.text_offset[0] == len(chunk1.choices[0].text)

638 +

639 + @pytest.mark.asyncio

640 + async def test_stream_logprob_state_stays_bounded(self, derenderer, tokenizer):

641 + """Carried logprob context never exceeds the 4-token window."""

642 + token_ids = tokenizer.encode("the quick brown fox jumps over the lazy dog")

643 + state = None

644 + for tid in token_ids:

645 + _, state = await derenderer.derender_chat_stream(

646 + model=MODEL_NAME,

647 + generate_chunk=_make_stream_chunk(

648 + [tid], logprobs=_placeholder_logprobs([tid])

649 + ),

650 + state=state,

651 + )

652 + assert state is not None

653 + assert len(state.logprob_context_token_ids) <= 4

654 +

655 + @pytest.mark.asyncio

656 + async def test_stream_without_logprobs_unchanged(self, derenderer, tokenizer):

657 + """Chunks without logprobs keep logprobs=None on the output choice."""

658 + token_ids = tokenizer.encode("hello")[:3]

659 + chunk, _ = await derenderer.derender_chat_stream(

660 + model=MODEL_NAME,

661 + generate_chunk=_make_stream_chunk(token_ids),

662 + )

663 + assert chunk.choices[0].logprobs is None

664 +

665 +

666 class TestDerenderChatStream:

667 """derender_chat_stream: plain detok branch (no parser)."""

668

vllm/entrypoints/scale_out/token_in_token_out/protocol.py

+11

@@ -546,6 +546,17 @@ class DerenderStreamState(BaseModel):

546 window is transiently empty (e.g. usage only final chunk).

547 """

548

549 + logprob_context_token_ids: list[int] = Field(default_factory=list, max_length=4)

550 + """Trailing sampled token IDs carried across chunks so byte-fallback

551 + (U+FFFD) correction during logprob placeholder resolution has context

552 + at chunk boundaries. Bounded to the 4-token window that

553 + ``_correct_decoded_token`` reads."""

554 +

555 + logprob_text_offset: int = Field(default=0, ge=0)

556 + """Cumulative emitted text length. Seeds ``text_offset`` for completion

557 + streaming logprobs so offsets stay absolute across chunks, mirroring

558 + ``initial_text_offset`` in the generate streaming path."""

559 +

560 output_token_ids: list[int] = Field(default_factory=list)

561 """All output tokens seen so far. Parser path only.

562

vllm/renderers/online_derenderer.py

+76 -7

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Sequence

4 from typing import Any

6 from vllm.config import ModelConfig

@@ -368,9 +369,30 @@ class OnlineDerenderer:

369 tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special

370 )

371

372 + # NOTE: parser-configured servers dispatch to

373 + # _derender_chat_stream_parsed above and never reach this plain

374 + # path. That parsed path does not resolve logprobs yet; when it

375 + # does, mirror the generate chat streaming path, which suppresses

376 + # logprobs entirely when a parser is configured and reasoning is

377 + # hidden, because decoded logprob token text would leak hidden

378 + # reasoning.

379 + resolved_logprobs = None

380 + if choice.logprobs is not None:

381 + resolved_logprobs = _resolve_logprobs(

382 + choice.logprobs,

383 + tokenizer,

384 + initial_context_token_ids=state.logprob_context_token_ids,

385 + )

386 +

387 include_role = not updated_state.role_sent

372 - if include_role:

373 - updated_state = updated_state.model_copy(update={"role_sent": True})

388 + updated_state = updated_state.model_copy(

389 + update={

390 + "role_sent": True,

391 + "logprob_context_token_ids": _logprob_context_tail(

392 + state.logprob_context_token_ids, delta_tids

393 + ),

394 + }

395 + )

396

397 delta = DeltaMessage(

398 role="assistant" if include_role else None,

@@ -380,6 +402,7 @@ class OnlineDerenderer:

402 ChatCompletionResponseStreamChoice(

403 index=choice.index,

404 delta=delta,

405 + logprobs=resolved_logprobs,

406 finish_reason=choice.finish_reason,

407 )

408 )

@@ -721,10 +744,32 @@ class OnlineDerenderer:

744 new_text, updated_state = await self._detokenize_delta_async(

745 tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special

746 )

747 +

748 + completion_logprobs = None

749 + if choice.logprobs is not None:

750 + resolved = _resolve_logprobs(

751 + choice.logprobs,

752 + tokenizer,

753 + initial_context_token_ids=state.logprob_context_token_ids,

754 + )

755 + completion_logprobs = _convert_chat_logprobs_to_completion_logprobs(

756 + resolved, initial_text_offset=state.logprob_text_offset

757 + )

758 +

759 + updated_state = updated_state.model_copy(

760 + update={

761 + "logprob_context_token_ids": _logprob_context_tail(

762 + state.logprob_context_token_ids, delta_tids

763 + ),

764 + "logprob_text_offset": state.logprob_text_offset + len(new_text),

765 + }

766 + )

767 +

768 stream_choices.append(

769 CompletionResponseStreamChoice(

770 index=choice.index,

771 text=new_text,

772 + logprobs=completion_logprobs,

773 finish_reason=choice.finish_reason,

774 )

775 )

@@ -750,6 +795,18 @@ class OnlineDerenderer:

795 return chunk, updated_state

796

797

798 +# Number of preceding sampled token IDs `_correct_decoded_token` reads to

799 +# repair U+FFFD from byte-fallback tokenization.

800 +_LOGPROB_CONTEXT_WINDOW = 4

801 +

802 +

803 +def _logprob_context_tail(

804 + context_token_ids: list[int], delta_token_ids: list[int]

805 +) -> list[int]:

806 + """Advance the carried logprob context by this chunk's sampled tokens."""

807 + return (list(context_token_ids) + list(delta_token_ids))[-_LOGPROB_CONTEXT_WINDOW:]

808 +

809 +

810 def _parse_token_id_placeholder(token: str) -> int | None:

811 """Extract token ID from a 'token_id:N' placeholder string."""

812 if not token.startswith("token_id:"):

@@ -799,13 +856,20 @@ def _correct_decoded_token(

856

857

858 def _resolve_logprobs(

802 - logprobs: ChatCompletionLogProbs, tokenizer: TokenizerLike

859 + logprobs: ChatCompletionLogProbs,

860 + tokenizer: TokenizerLike,

861 + initial_context_token_ids: Sequence[int] = (),

862 ) -> ChatCompletionLogProbs:

804 - """Resolve token_id:N placeholders in a ChatCompletionLogProbs object."""

863 + """Resolve token_id:N placeholders in a ChatCompletionLogProbs object.

864 +

865 + ``initial_context_token_ids`` seeds the byte-fallback correction context

866 + with sampled IDs from preceding chunks (streaming), so multi-byte

867 + characters split across chunk boundaries still resolve.

868 + """

869 if logprobs.content is None:

870 return logprobs

871

808 - context_token_ids: list[int] = []

872 + context_token_ids: list[int] = list(initial_context_token_ids)

873 resolved_content = []

874

875 for entry in logprobs.content:

@@ -845,9 +909,14 @@ def _resolve_logprobs(

909

910 def _convert_chat_logprobs_to_completion_logprobs(

911 logprobs: ChatCompletionLogProbs,

912 + initial_text_offset: int = 0,

913 ) -> CompletionLogProbs:

914 """Convert ChatCompletionLogProbs (per-token objects) to CompletionLogProbs

850 - (parallel flat lists) as required by the /v1/completions response schema."""

915 + (parallel flat lists) as required by the /v1/completions response schema.

916 +

917 + ``initial_text_offset`` keeps ``text_offset`` absolute across streaming

918 + chunks, mirroring the generate streaming path.

919 + """

920 if logprobs.content is None:

921 return CompletionLogProbs()

922

@@ -856,7 +925,7 @@ def _convert_chat_logprobs_to_completion_logprobs(

925 top_logprobs_list: list[dict[str, float] | None] = []

926 text_offset: list[int] = []

927

859 - offset = 0

928 + offset = initial_text_offset

929 for entry in logprobs.content:

930 text_offset.append(offset)

931 tokens.append(entry.token)