fix(frontend): return SGLang chat logprobs

components/src/dynamo/frontend/sglang_prepost.py

@@ -979,6 +979,8 @@ def __init__(
         # incomplete byte-fallback sequence.
         self._decode_context_ids = list((prompt_token_ids or [])[-5:])
         self._pending_decode_ids: list[int] = []
+        self._logprob_context_ids: list[int] = []
+        self._pending_logprobs_content: list[dict[str, Any]] = []
         self._has_emitted_role: bool = False
         # Tool call accumulation.  SGLang's streaming parser returns
         # deltas (name in one chunk, argument fragments across subsequent
@@ -1044,6 +1046,108 @@ def _incremental_decode(
         self._pending_decode_ids = []
         return delta_text
 
+    def _build_openai_logprobs(
+        self,
+        log_probs: list[float],
+        top_logprobs: list[list[dict[str, Any]]] | None,
+        token_ids: list[int],
+    ) -> dict[str, Any] | None:
+        if len(log_probs) != len(token_ids):
+            return None
+
+        content: list[dict[str, Any]] = []
+        for index, (token_id, logprob) in enumerate(zip(token_ids, log_probs)):
+            context_token_ids = (self._logprob_context_ids + token_ids[:index])[-4:]
+            token = self._decode_logprob_token(token_id, None, context_token_ids)
+            candidates = top_logprobs[index] if top_logprobs else []
+            openai_top_logprobs = []
+            for candidate in candidates:
+                candidate_token = self._decode_logprob_token(
+                    candidate.get("token_id"),
+                    candidate.get("token"),
+                    context_token_ids,
+                )
+                candidate_bytes = candidate.get("bytes")
+                if candidate_bytes is None:
+                    candidate_bytes = (
+                        list(candidate_token.encode("utf-8"))
+                        if candidate_token
+                        else None
+                    )
+                openai_top_logprobs.append(
+                    {
+                        "token": candidate_token,
+                        "logprob": float(candidate["logprob"]),
+                        "bytes": candidate_bytes,
+                    }
+                )
+            content.append(
+                {
+                    "token": token,
+                    "logprob": float(logprob),
+                    "bytes": list(token.encode("utf-8")) if token else None,
+                    "top_logprobs": openai_top_logprobs,
+                }
+            )
+
+        return {"content": content, "refusal": None} if content else None
+
+    def _decode_logprob_token(
+        self,
+        token_id: int | None,
+        token: str | None,
+        context_token_ids: list[int],
+    ) -> str:
+        if token is None:
+            if token_id is None:
+                return ""
+            token = self.tokenizer.decode([token_id], skip_special_tokens=False)
+
+        if not token.endswith("\ufffd") or token_id is None:
+            return token
+
+        for context_size in range(1, min(len(context_token_ids), 4) + 1):
+            context = context_token_ids[-context_size:]
+            decoded = self.tokenizer.decode(
+                context + [token_id], skip_special_tokens=False
+            )
+            if decoded.endswith("\ufffd"):
+                continue
+
+            clean_end = len(context)
+            for context_index in range(len(context) - 1, -1, -1):
+                context_token = self.tokenizer.decode(
+                    [context[context_index]], skip_special_tokens=False
+                )
+                if context_token.endswith("\ufffd"):
+                    clean_end = context_index
+                else:
+                    break
+
+            clean_prefix = (
+                self.tokenizer.decode(context[:clean_end], skip_special_tokens=False)
+                if clean_end
+                else ""
+            )
+            if decoded.startswith(clean_prefix):
+                return decoded[len(clean_prefix) :]
+
+            common_prefix_length = 0
+            for prefix_char, decoded_char in zip(clean_prefix, decoded):
+                if prefix_char != decoded_char:
+                    break
+                common_prefix_length += 1
+            return decoded[common_prefix_length:]
+
+        return ""
+
+    def _take_pending_logprobs(self) -> dict[str, Any] | None:
+        if not self._pending_logprobs_content:
+            return None
+        content = self._pending_logprobs_content
+        self._pending_logprobs_content = []
+        return {"content": content, "refusal": None}
+
     def _parse_reasoning_delta(
         self, delta_text: str, finish_reason: str | None
     ) -> tuple[str | None, str]:
@@ -1106,29 +1210,45 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No
         raw_ids = engine_response.get("token_ids")
         token_ids = raw_ids if isinstance(raw_ids, list) else list(raw_ids or [])
         finish_reason = engine_response.get("finish_reason")
+        log_probs = engine_response.get("log_probs")
+        top_logprobs = engine_response.get("top_logprobs")
         if finish_reason is not None:
+            raw_token_count = len(token_ids)
             token_ids = self._strip_trailing_eos_token_ids(list(token_ids))
+            retained_token_count = len(token_ids)
+            if log_probs is not None and len(log_probs) == raw_token_count:
+                log_probs = log_probs[:retained_token_count]
+            if top_logprobs is not None and len(top_logprobs) == raw_token_count:
+                top_logprobs = top_logprobs[:retained_token_count]
 
         delta_text = (
             self._incremental_decode(token_ids, flush=finish_reason is not None)
             if token_ids or finish_reason is not None
             else ""
         )
+        openai_logprobs = None
+        if log_probs is not None:
+            openai_logprobs = self._build_openai_logprobs(
+                log_probs, top_logprobs, token_ids
+            )
+            if openai_logprobs is not None:
+                self._pending_logprobs_content.extend(openai_logprobs["content"])
+        self._logprob_context_ids = (self._logprob_context_ids + token_ids)[-4:]
 
         if self._fast_plain_text:
             if delta_text:
                 return {
                     "index": 0,
                     "delta": self._with_initial_role({"content": delta_text}),
                     "finish_reason": finish_reason,
-                    "logprobs": None,
+                    "logprobs": self._take_pending_logprobs(),
                 }
             elif finish_reason:
                 return {
                     "index": 0,
                     "delta": self._with_initial_role({}),
                     "finish_reason": finish_reason,
-                    "logprobs": None,
+                    "logprobs": self._take_pending_logprobs(),
                 }
             return None
 
@@ -1345,7 +1465,7 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No
                 "index": 0,
                 "delta": self._with_initial_role(delta),
                 "finish_reason": effective_finish,
-                "logprobs": None,
+                "logprobs": self._take_pending_logprobs(),
             }
 
         return None

components/src/dynamo/frontend/sglang_processor.py

@@ -625,6 +625,8 @@ async def _generate_and_stream(
             # finish_reason.  Use si=1 for the first chunk to minimize
             # TTFT, then switch to the configured interval.
             pending_token_ids: list[int] = []
+            pending_log_probs: list[float] | None = None
+            pending_top_logprobs: list[list[dict[str, Any]]] | None = None
             pending_usage: dict[str, Any] | None = None
             first_chunk = True
             input_tokens = len(tokens)
@@ -637,6 +639,91 @@ async def _generate_and_stream(
             video_count = len(_mm_counts.get("video_url", []))
             audio_count = len(_mm_counts.get("audio_url", []))
 
+            def flush_pending(
+                *,
+                finish_reason: str | None,
+                stop_reason: Any | None,
+                engine_data: Any | None,
+            ) -> dict[str, Any]:
+                nonlocal pending_token_ids
+                nonlocal pending_log_probs
+                nonlocal pending_top_logprobs
+                nonlocal pending_usage
+                nonlocal first_chunk
+                nonlocal post_proc_total_ms
+                nonlocal token_count
+
+                chunk_token_count = len(pending_token_ids)
+                usage_for_metrics = pending_usage
+                mapped_response: dict[str, Any] = {
+                    "token_ids": pending_token_ids,
+                    "finish_reason": finish_reason,
+                }
+                if pending_log_probs is not None:
+                    mapped_response["log_probs"] = pending_log_probs
+                if pending_top_logprobs is not None:
+                    mapped_response["top_logprobs"] = pending_top_logprobs
+
+                if self.debug_perf:
+                    t_pp0 = time.monotonic()
+
+                choice = post.process_output(mapped_response)
+
+                if self.debug_perf:
+                    t_pp1 = time.monotonic()
+                    post_proc_total_ms += (t_pp1 - t_pp0) * 1000.0
+                    token_count += chunk_token_count
+
+                envelope: dict[str, Any] = {"_dynamo_annotated": True}
+                if choice:
+                    dynamo_out: dict[str, Any] = {
+                        "id": request_id,
+                        "choices": [choice],
+                        "created": created_ts,
+                        "model": request["model"],
+                        "object": "chat.completion.chunk",
+                    }
+                    if pending_usage:
+                        dynamo_out["usage"] = pending_usage
+                    response_nvext: dict[str, Any] = {}
+                    if stop_reason is not None and nvext_extra_field_requested(
+                        request, "stop_reason"
+                    ):
+                        response_nvext["stop_reason"] = stop_reason
+                    if engine_data is not None and nvext_extra_field_requested(
+                        request, "engine_data"
+                    ):
+                        response_nvext["engine_data"] = engine_data
+                    if response_nvext:
+                        dynamo_out["nvext"] = response_nvext
+
+                    envelope["data"] = dynamo_out
+
+                metrics: dict[str, Any] = {
+                    "input_tokens": input_tokens,
+                    "output_tokens": cumulative_output_tokens,
+                    "chunk_tokens": chunk_token_count,
+                }
+                # Include nonzero counts on every frame (text-only carries nothing).
+                if image_count:
+                    metrics["image_count"] = image_count
+                if video_count:
+                    metrics["video_count"] = video_count
+                if audio_count:
+                    metrics["audio_count"] = audio_count
+                cached_tokens = _cached_tokens_from_usage(usage_for_metrics)
+                if cached_tokens is not None:
+                    metrics["cached_tokens"] = cached_tokens
+                envelope["event"] = "llm_metrics"
+                envelope["comment"] = [json.dumps(metrics)]
+
+                pending_token_ids = []
+                pending_log_probs = None
+                pending_top_logprobs = None
+                pending_usage = None
+                first_chunk = False
+                return envelope
+
             async for dynamo_response in dynamo_stream:
                 if dynamo_response.is_error():
                     comments = dynamo_response.comments() or []
@@ -663,6 +750,25 @@ async def _generate_and_stream(
                     break
 
                 new_ids = engine_response["token_ids"]
+                log_probs = engine_response.get("log_probs")
+                top_logprobs = engine_response.get("top_logprobs")
+
+                if new_ids and pending_token_ids:
+                    pending_logprob_shape = (
+                        pending_log_probs is not None,
+                        pending_top_logprobs is not None,
+                    )
+                    chunk_logprob_shape = (
+                        log_probs is not None,
+                        top_logprobs is not None,
+                    )
+                    if pending_logprob_shape != chunk_logprob_shape:
+                        yield flush_pending(
+                            finish_reason=None,
+                            stop_reason=None,
+                            engine_data=None,
+                        )
+
                 chunk_tokens = len(new_ids)
                 cumulative_output_tokens += chunk_tokens
                 raw_finish = engine_response.get("finish_reason")
@@ -674,76 +780,24 @@ async def _generate_and_stream(
                 engine_data = engine_response.get("engine_data")
 
                 pending_token_ids.extend(new_ids)
+                if log_probs is not None:
+                    if pending_log_probs is None:
+                        pending_log_probs = []
+                    pending_log_probs.extend(log_probs)
+                if top_logprobs is not None:
+                    if pending_top_logprobs is None:
+                        pending_top_logprobs = []
+                    pending_top_logprobs.extend(top_logprobs)
 
                 # Flush on finish or when we've accumulated enough tokens.
                 # First chunk flushes immediately (si=1) to minimize TTFT.
                 flush_threshold = 1 if first_chunk else stream_interval
                 if finish_reason or len(pending_token_ids) >= flush_threshold:
-                    usage_for_metrics = pending_usage
-                    mapped_response = {
-                        "token_ids": pending_token_ids,
-                        "finish_reason": finish_reason,
-                    }
-
-                    if self.debug_perf:
-                        t_pp0 = time.monotonic()
-
-                    choice = post.process_output(mapped_response)
-
-                    if self.debug_perf:
-                        t_pp1 = time.monotonic()
-                        post_proc_total_ms += (t_pp1 - t_pp0) * 1000.0
-                        token_count += len(pending_token_ids)
-
-                    envelope: dict[str, Any] = {"_dynamo_annotated": True}
-                    if choice:
-                        dynamo_out: dict[str, Any] = {
-                            "id": request_id,
-                            "choices": [choice],
-                            "created": created_ts,
-                            "model": request["model"],
-                            "object": "chat.completion.chunk",
-                        }
-                        if pending_usage:
-                            dynamo_out["usage"] = pending_usage
-                            pending_usage = None
-                        response_nvext: dict[str, Any] = {}
-                        if stop_reason is not None and nvext_extra_field_requested(
-                            request, "stop_reason"
-                        ):
-                            response_nvext["stop_reason"] = stop_reason
-                        if engine_data is not None and (
-                            nvext_extra_field_requested(request, "engine_data")
-                        ):
-                            response_nvext["engine_data"] = engine_data
-                        if response_nvext:
-                            dynamo_out["nvext"] = response_nvext
-
-                        envelope["data"] = dynamo_out
-
-                    metrics: dict[str, Any] = {
-                        "input_tokens": input_tokens,
-                        "output_tokens": cumulative_output_tokens,
-                        "chunk_tokens": len(pending_token_ids),
-                    }
-                    # Include nonzero counts on every frame (text-only carries nothing).
-                    if image_count:
-                        metrics["image_count"] = image_count
-                    if video_count:
-                        metrics["video_count"] = video_count
-                    if audio_count:
-                        metrics["audio_count"] = audio_count
-                    cached_tokens = _cached_tokens_from_usage(usage_for_metrics)
-                    if cached_tokens is not None:
-                        metrics["cached_tokens"] = cached_tokens
-                    envelope["event"] = "llm_metrics"
-                    envelope["comment"] = [json.dumps(metrics)]
-
-                    yield envelope
-
-                    pending_token_ids = []
-                    pending_usage = None
-                    first_chunk = False
+                    yield flush_pending(
+                        finish_reason=finish_reason,
+                        stop_reason=stop_reason,
+                        engine_data=engine_data,
+                    )
         except Unknown:
             raise
         except Exception as e: