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

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

docs/serving/online_serving/derenderer.md

@@ -98,7 +98,7 @@ The server keeps no state between calls. Everything the next call needs is in `s
 - **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`.
 - **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk.
 
-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.
+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.
 
 ## Streaming cost
 
@@ -115,6 +115,15 @@ When a tool or reasoning parser is configured, parser internal state (buffered m
 
 `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.
 
+## Streaming state and logprobs
+
+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:
+
+- `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
+- `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0
+
+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`.
+
 ## Example
 
 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

@@ -169,6 +169,7 @@ def _make_stream_chunk(
     request_id: str = "test-req",
     usage: dict | None = None,
     metrics: PerRequestMetrics | None = None,
+    logprobs: dict | None = None,
 ) -> GenerateStreamResponse:
     """Build a GenerateStreamResponse SSE chunk."""
     return GenerateStreamResponse(
@@ -178,6 +179,7 @@ def _make_stream_chunk(
                 index=index,
                 token_ids=token_ids,
                 finish_reason=finish_reason,
+                logprobs=logprobs,
             )
         ],
         usage=UsageInfo(**usage) if usage else None,
@@ -185,6 +187,22 @@ def _make_stream_chunk(
     )
 
 
+def _placeholder_logprobs(token_ids: list[int]) -> dict:
+    """Per-token logprob entries using token_id:N placeholders, as sent by
+    the generate worker."""
+    return {
+        "content": [
+            {
+                "token": f"token_id:{tid}",
+                "logprob": -0.5,
+                "bytes": None,
+                "top_logprobs": [],
+            }
+            for tid in token_ids
+        ]
+    }
+
+
 def _make_usage_chunk(
     completion_tokens: int,
     prompt_tokens: int = 0,
@@ -536,6 +554,115 @@ class TestDerenderCompletionStream:
         assert chunk.choices[0].finish_reason == "length"
 
 
+class TestStreamLogprobs:
+    """Streaming derender must carry per-chunk logprobs with placeholders
+    resolved, matching what the generate streaming path emits."""
+
+    @pytest.mark.asyncio
+    async def test_chat_stream_logprobs_resolved_per_chunk(self, derenderer, tokenizer):
+        """Each streamed chunk carries logprobs with token_id:N resolved."""
+        token_ids = tokenizer.encode("hello world")[:6]
+        mid = len(token_ids) // 2
+
+        state = None
+        for part in (token_ids[:mid], token_ids[mid:]):
+            chunk, state = await derenderer.derender_chat_stream(
+                model=MODEL_NAME,
+                generate_chunk=_make_stream_chunk(
+                    part, logprobs=_placeholder_logprobs(part)
+                ),
+                state=state,
+            )
+            logprobs = chunk.choices[0].logprobs
+            assert logprobs is not None and logprobs.content is not None
+            assert len(logprobs.content) == len(part)
+            for entry in logprobs.content:
+                assert not entry.token.startswith("token_id:"), (
+                    f"placeholder not resolved: {entry.token!r}"
+                )
+
+    @pytest.mark.asyncio
+    async def test_chat_stream_logprobs_multibyte_across_chunks(
+        self, derenderer, tokenizer
+    ):
+        """Byte-fallback correction works when a multi-byte character's
+        tokens are split across chunks (context carried in stream_state)."""
+        token_ids = tokenizer.encode("👍", add_special_tokens=False)
+        if len(token_ids) < 2 or "�" not in tokenizer.decode([token_ids[-1]]):
+            pytest.skip("Tokenizer does not byte-split this character")
+
+        state = None
+        last_chunk = None
+        # One token per chunk: every entry after the first needs cross-chunk
+        # context to resolve without U+FFFD.
+        for i, tid in enumerate(token_ids):
+            last_chunk, state = await derenderer.derender_chat_stream(
+                model=MODEL_NAME,
+                generate_chunk=_make_stream_chunk(
+                    [tid], logprobs=_placeholder_logprobs([tid])
+                ),
+                state=state,
+            )
+
+        assert last_chunk is not None
+        final_entry = last_chunk.choices[0].logprobs.content[0]
+        assert not final_entry.token.endswith("�"), (
+            "byte-fallback correction failed across chunk boundary: "
+            f"{final_entry.token!r}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_completion_stream_logprobs_text_offset_absolute(
+        self, derenderer, tokenizer
+    ):
+        """text_offset continues across chunks instead of restarting at 0."""
+        token_ids = tokenizer.encode("hello world")[:6]
+        mid = len(token_ids) // 2
+
+        chunk1, state = await derenderer.derender_completion_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                token_ids[:mid], logprobs=_placeholder_logprobs(token_ids[:mid])
+            ),
+        )
+        chunk2, _ = await derenderer.derender_completion_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(
+                token_ids[mid:], logprobs=_placeholder_logprobs(token_ids[mid:])
+            ),
+            state=state,
+        )
+
+        assert chunk1.choices[0].logprobs.text_offset[0] == 0
+        assert chunk2.choices[0].logprobs.text_offset[0] == len(chunk1.choices[0].text)
+
+    @pytest.mark.asyncio
+    async def test_stream_logprob_state_stays_bounded(self, derenderer, tokenizer):
+        """Carried logprob context never exceeds the 4-token window."""
+        token_ids = tokenizer.encode("the quick brown fox jumps over the lazy dog")
+        state = None
+        for tid in token_ids:
+            _, state = await derenderer.derender_chat_stream(
+                model=MODEL_NAME,
+                generate_chunk=_make_stream_chunk(
+                    [tid], logprobs=_placeholder_logprobs([tid])
+                ),
+                state=state,
+            )
+        assert state is not None
+        assert len(state.logprob_context_token_ids) <= 4
+
+    @pytest.mark.asyncio
+    async def test_stream_without_logprobs_unchanged(self, derenderer, tokenizer):
+        """Chunks without logprobs keep logprobs=None on the output choice."""
+        token_ids = tokenizer.encode("hello")[:3]
+        chunk, _ = await derenderer.derender_chat_stream(
+            model=MODEL_NAME,
+            generate_chunk=_make_stream_chunk(token_ids),
+        )
+        assert chunk.choices[0].logprobs is None
+
+
 class TestDerenderChatStream:
     """derender_chat_stream: plain detok branch (no parser)."""
 

vllm/entrypoints/scale_out/token_in_token_out/protocol.py

@@ -546,6 +546,17 @@ class DerenderStreamState(BaseModel):
     window is transiently empty (e.g. usage only final chunk).
     """
 
+    logprob_context_token_ids: list[int] = Field(default_factory=list, max_length=4)
+    """Trailing sampled token IDs carried across chunks so byte-fallback
+    (U+FFFD) correction during logprob placeholder resolution has context
+    at chunk boundaries. Bounded to the 4-token window that
+    ``_correct_decoded_token`` reads."""
+
+    logprob_text_offset: int = Field(default=0, ge=0)
+    """Cumulative emitted text length. Seeds ``text_offset`` for completion
+    streaming logprobs so offsets stay absolute across chunks, mirroring
+    ``initial_text_offset`` in the generate streaming path."""
+
     output_token_ids: list[int] = Field(default_factory=list)
     """All output tokens seen so far. Parser path only.
 

vllm/renderers/online_derenderer.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Sequence
 from typing import Any
 
 from vllm.config import ModelConfig
@@ -368,9 +369,30 @@ class OnlineDerenderer:
                 tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special
             )
 
+            # NOTE: parser-configured servers dispatch to
+            # _derender_chat_stream_parsed above and never reach this plain
+            # path. That parsed path does not resolve logprobs yet; when it
+            # does, mirror the generate chat streaming path, which suppresses
+            # logprobs entirely when a parser is configured and reasoning is
+            # hidden, because decoded logprob token text would leak hidden
+            # reasoning.
+            resolved_logprobs = None
+            if choice.logprobs is not None:
+                resolved_logprobs = _resolve_logprobs(
+                    choice.logprobs,
+                    tokenizer,
+                    initial_context_token_ids=state.logprob_context_token_ids,
+                )
+
             include_role = not updated_state.role_sent
-            if include_role:
-                updated_state = updated_state.model_copy(update={"role_sent": True})
+            updated_state = updated_state.model_copy(
+                update={
+                    "role_sent": True,
+                    "logprob_context_token_ids": _logprob_context_tail(
+                        state.logprob_context_token_ids, delta_tids
+                    ),
+                }
+            )
 
             delta = DeltaMessage(
                 role="assistant" if include_role else None,
@@ -380,6 +402,7 @@ class OnlineDerenderer:
                 ChatCompletionResponseStreamChoice(
                     index=choice.index,
                     delta=delta,
+                    logprobs=resolved_logprobs,
                     finish_reason=choice.finish_reason,
                 )
             )
@@ -721,10 +744,32 @@ class OnlineDerenderer:
             new_text, updated_state = await self._detokenize_delta_async(
                 tokenizer, delta_tids, updated_state, skip_special_tokens=skip_special
             )
+
+            completion_logprobs = None
+            if choice.logprobs is not None:
+                resolved = _resolve_logprobs(
+                    choice.logprobs,
+                    tokenizer,
+                    initial_context_token_ids=state.logprob_context_token_ids,
+                )
+                completion_logprobs = _convert_chat_logprobs_to_completion_logprobs(
+                    resolved, initial_text_offset=state.logprob_text_offset
+                )
+
+            updated_state = updated_state.model_copy(
+                update={
+                    "logprob_context_token_ids": _logprob_context_tail(
+                        state.logprob_context_token_ids, delta_tids
+                    ),
+                    "logprob_text_offset": state.logprob_text_offset + len(new_text),
+                }
+            )
+
             stream_choices.append(
                 CompletionResponseStreamChoice(
                     index=choice.index,
                     text=new_text,
+                    logprobs=completion_logprobs,
                     finish_reason=choice.finish_reason,
                 )
             )
@@ -750,6 +795,18 @@ class OnlineDerenderer:
         return chunk, updated_state
 
 
+# Number of preceding sampled token IDs `_correct_decoded_token` reads to
+# repair U+FFFD from byte-fallback tokenization.
+_LOGPROB_CONTEXT_WINDOW = 4
+
+
+def _logprob_context_tail(
+    context_token_ids: list[int], delta_token_ids: list[int]
+) -> list[int]:
+    """Advance the carried logprob context by this chunk's sampled tokens."""
+    return (list(context_token_ids) + list(delta_token_ids))[-_LOGPROB_CONTEXT_WINDOW:]
+
+
 def _parse_token_id_placeholder(token: str) -> int | None:
     """Extract token ID from a 'token_id:N' placeholder string."""
     if not token.startswith("token_id:"):
@@ -799,13 +856,20 @@ def _correct_decoded_token(
 
 
 def _resolve_logprobs(
-    logprobs: ChatCompletionLogProbs, tokenizer: TokenizerLike
+    logprobs: ChatCompletionLogProbs,
+    tokenizer: TokenizerLike,
+    initial_context_token_ids: Sequence[int] = (),
 ) -> ChatCompletionLogProbs:
-    """Resolve token_id:N placeholders in a ChatCompletionLogProbs object."""
+    """Resolve token_id:N placeholders in a ChatCompletionLogProbs object.
+
+    ``initial_context_token_ids`` seeds the byte-fallback correction context
+    with sampled IDs from preceding chunks (streaming), so multi-byte
+    characters split across chunk boundaries still resolve.
+    """
     if logprobs.content is None:
         return logprobs
 
-    context_token_ids: list[int] = []
+    context_token_ids: list[int] = list(initial_context_token_ids)
     resolved_content = []
 
     for entry in logprobs.content:
@@ -845,9 +909,14 @@ def _resolve_logprobs(
 
 def _convert_chat_logprobs_to_completion_logprobs(
     logprobs: ChatCompletionLogProbs,
+    initial_text_offset: int = 0,
 ) -> CompletionLogProbs:
     """Convert ChatCompletionLogProbs (per-token objects) to CompletionLogProbs
-    (parallel flat lists) as required by the /v1/completions response schema."""
+    (parallel flat lists) as required by the /v1/completions response schema.
+
+    ``initial_text_offset`` keeps ``text_offset`` absolute across streaming
+    chunks, mirroring the generate streaming path.
+    """
     if logprobs.content is None:
         return CompletionLogProbs()
 
@@ -856,7 +925,7 @@ def _convert_chat_logprobs_to_completion_logprobs(
     top_logprobs_list: list[dict[str, float] | None] = []
     text_offset: list[int] = []
 
-    offset = 0
+    offset = initial_text_offset
     for entry in logprobs.content:
         text_offset.append(offset)
         tokens.append(entry.token)