feat: logprob / top_logprobs=1 handling

messiaen · messiaen · commit 406ef5e5445f · 2025-08-13T11:27:19.000-04:00
Signed-off-by: Greg Clark &lt;grclark@nvidia.com&gt;
diff --git a/components/backends/sglang/src/dynamo/sglang/worker/main.py b/components/backends/sglang/src/dynamo/sglang/worker/main.py
@@ -7,7 +7,8 @@
 import signal
 import socket
 import sys
-from typing import Any, Dict, Optional, Union
+from operator import itemgetter
+from typing import Any, Optional
 
 import sglang as sgl
 import uvloop
@@ -210,6 +211,7 @@ async def generate(self, request: dict):
                 else request["batch_token_ids"],
                 sampling_params=sampling_params,
                 stream=True,
+                return_logprob=True,
                 bootstrap_host=bootstrap_host,
                 bootstrap_port=bootstrap_port,
                 bootstrap_room=bootstrap_room,
@@ -231,54 +233,49 @@ async def generate(self, request: dict):
                 else request["batch_token_ids"],
                 sampling_params=sampling_params,
                 stream=True,
+                return_logprob=True,
             )
 
             async for out in self._process_stream(g, unpack=False, is_batch=is_batch):
                 yield out
 
     async def _process_stream(self, stream_source, unpack: bool, is_batch: bool):
-        # Initialize based on batch mode
-        num_output_tokens_so_far: Union[Dict[int, int], int]
-        if is_batch:
-            num_output_tokens_so_far = {}
-        else:
-            num_output_tokens_so_far = 0
+        assert not is_batch, "Batch processing is not supported."
+        num_output_tokens_so_far = 0
 
         async for res in stream_source:
             data = res.data() if unpack else res
             finish_reason = data["meta_info"]["finish_reason"]
 
-            if is_batch:
-                # Handle batch response
-                assert isinstance(num_output_tokens_so_far, dict)
-                index = data.get("index", 0)
-                if index not in num_output_tokens_so_far:
-                    num_output_tokens_so_far[index] = 0
-
-                if finish_reason:
-                    out = {
-                        "token_ids": [],
-                        "finish_reason": finish_reason["type"],
-                        "index": index,
-                    }
-                else:
-                    next_total_toks = len(data["output_ids"])
-                    new_tokens = data["output_ids"][num_output_tokens_so_far[index] :]
-                    out = {
-                        "token_ids": new_tokens,
-                        "index": index,
-                    }
-                    num_output_tokens_so_far[index] = next_total_toks
+            # Handle single response
+            assert isinstance(num_output_tokens_so_far, int)
+            if finish_reason:
+                out = {"token_ids": [], "finish_reason": finish_reason["type"]}
             else:
-                # Handle single response
-                assert isinstance(num_output_tokens_so_far, int)
-                if finish_reason:
-                    out = {"token_ids": [], "finish_reason": finish_reason["type"]}
-                else:
-                    next_total_toks = len(data["output_ids"])
-                    out = {"token_ids": data["output_ids"][num_output_tokens_so_far:]}
-                    num_output_tokens_so_far = next_total_toks
+                next_total_toks = len(res["meta_info"]["output_token_logprobs"])
+                new_tokens = list(
+                    map(
+                        itemgetter(1),
+                        res["meta_info"]["output_token_logprobs"][
+                            num_output_tokens_so_far:
+                        ],
+                    )
+                )
+                new_logprobs = list(
+                    map(
+                        itemgetter(0),
+                        res["meta_info"]["output_token_logprobs"][
+                            num_output_tokens_so_far:
+                        ],
+                    )
+                )
+                out = {
+                    "token_ids": new_tokens,
+                    "log_probs": new_logprobs,
+                }
+                num_output_tokens_so_far = next_total_toks
 
+            logging.debug(f"Generated output: {out}")
             yield out
 
     async def _prefill_generator(self, prefill):
diff --git a/lib/engines/llamacpp/src/lib.rs b/lib/engines/llamacpp/src/lib.rs
@@ -268,6 +268,7 @@ fn run_request(
             //text: if output.text.is_empty() { None } else { Some(output.text) },
             cum_log_probs: None, // TODO output.cumulative_logprob.map(|v| v as f64),
             log_probs: None,     // TODO  output.logprobs
+            top_logprobs: None,
             finish_reason: None,
             index: None,
         };
diff --git a/lib/engines/mistralrs/src/lib.rs b/lib/engines/mistralrs/src/lib.rs
@@ -590,7 +590,7 @@ impl
                             None => None,
                         };
                         #[allow(deprecated)]
-                        let inner = response_generator.create_choice(0, Some(from_assistant), None);
+                        let inner = response_generator.create_choice(0, Some(from_assistant), None, None);
                         let ann = Annotated{
                             id: None,
                             data: Some(inner),
diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs
@@ -218,12 +218,14 @@ impl
         //let mdcsum = self.mdcsum.clone();
         let stream = processed_stream.map(move |output| {
             output.map_data(|data| {
+                log::info!("data: {:?}", data);
                 Ok(BackendOutput {
                     token_ids: data.token_ids,
                     tokens: data.tokens.unwrap_or_default(),
                     text: data.text,
                     cum_log_probs: data.cum_log_probs,
                     log_probs: data.log_probs,
+                    top_logprobs: data.top_logprobs,
                     finish_reason: data.finish_reason,
                     //mdcsum: mdcsum.clone(),
                     index: data.index,
diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs
@@ -102,6 +102,7 @@ fn delta_core(tok: u32) -> Annotated<LLMEngineOutput> {
         text: None,
         cum_log_probs: None,
         log_probs: None,
+        top_logprobs: None,
         finish_reason: None,
         index: None,
     };
@@ -242,11 +243,11 @@ impl
             let mut id = 1;
             for c in chars_string.chars() {
                 tokio::time::sleep(*TOKEN_ECHO_DELAY).await;
-                let response = deltas.create_choice(0, Some(c.to_string()), None);
+                let response = deltas.create_choice(0, Some(c.to_string()), None, None);
                 yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None };
                 id += 1;
             }
-            let response = deltas.create_choice(0, None, Some(async_openai::types::CompletionFinishReason::Stop));
+            let response = deltas.create_choice(0, None, Some(async_openai::types::CompletionFinishReason::Stop), None);
             yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None };
 
         };
diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs
@@ -166,7 +166,7 @@ impl RetryManager {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::protocols::common::{SamplingOptions, StopConditions};
+    use crate::protocols::common::{OutputOptions, SamplingOptions, StopConditions};
     use dynamo_runtime::pipeline::context::Controller;
     use dynamo_runtime::pipeline::AsyncEngine;
     use std::sync::atomic::{AtomicU32, Ordering};
@@ -183,6 +183,7 @@ mod tests {
                 ..Default::default()
             },
             sampling_options: SamplingOptions::default(),
+            output_options: OutputOptions::default(),
             eos_token_ids: vec![],
             mdc_sum: None,
             annotations: vec![],
@@ -198,6 +199,7 @@ mod tests {
             text: Some(format!("token_{}", token_id)),
             cum_log_probs: None,
             log_probs: None,
+            top_logprobs: None,
             finish_reason: None,
             index: None,
         })
diff --git a/lib/llm/src/mocker/engine.rs b/lib/llm/src/mocker/engine.rs
@@ -405,6 +405,7 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
                             text: None,
                             cum_log_probs: None,
                             log_probs: None,
+                            top_logprobs: None,
                             finish_reason: None,
                             index: None,
                         };
@@ -525,7 +526,7 @@ mod integration_tests {
     use super::*;
     use crate::kv_router::indexer::RouterEvent;
     use crate::kv_router::KV_EVENT_SUBJECT;
-    use crate::protocols::common::{SamplingOptions, StopConditions};
+    use crate::protocols::common::{OutputOptions, SamplingOptions, StopConditions};
     use dynamo_runtime::{
         pipeline::Context,
         pipeline::{network::Ingress, PushRouter},
@@ -641,6 +642,7 @@ mod integration_tests {
                 ..Default::default()
             },
             sampling_options: SamplingOptions::default(),
+            output_options: OutputOptions::default(),
             eos_token_ids: vec![],
             mdc_sum: None,
             annotations: vec![format!("dp_rank:{dp_rank}")],
diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs
@@ -33,7 +33,7 @@ use dynamo_runtime::pipeline::{
 use dynamo_runtime::protocols::annotated::{Annotated, AnnotationsProvider};
 
 use crate::protocols::{
-    common::{SamplingOptionsProvider, StopConditionsProvider},
+    common::{OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider},
     openai::{
         chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse},
         completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
@@ -146,6 +146,7 @@ impl OpenAIPreprocessor {
             + AnnotationsProvider
             + SamplingOptionsProvider
             + StopConditionsProvider
+            + OutputOptionsProvider
             + NvExtProvider,
     >(
         &self,
@@ -249,6 +250,7 @@ impl OpenAIPreprocessor {
 
         builder.stop_conditions(stop_conditions);
         builder.sampling_options(request.extract_sampling_options()?);
+        builder.output_options(request.extract_output_options()?);
         builder.annotations(request.annotations().unwrap_or_default());
         builder.mdc_sum(Some(self.mdcsum.clone()));
         builder.estimated_prefix_hit_num_blocks(None);
diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs
@@ -45,6 +45,10 @@ pub trait StopConditionsProvider {
     fn extract_stop_conditions(&self) -> Result<StopConditions>;
 }
 
+pub trait OutputOptionsProvider {
+    fn extract_output_options(&self) -> Result<OutputOptions>;
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
 pub enum FinishReason {
     #[serde(rename = "eos")]
@@ -179,6 +183,9 @@ pub struct CompletionRequest {
     /// are needed.
     pub sampling_options: SamplingOptions,
 
+    #[builder(default)]
+    pub output_options: OutputOptions,
+
     /// The computed checksum of the Model Deployment Card (MDC).
     #[builder(default)]
     pub mdc_sum: Option<String>,
diff --git a/lib/llm/src/protocols/common/llm_backend.rs b/lib/llm/src/protocols/common/llm_backend.rs
@@ -23,6 +23,15 @@ use dynamo_runtime::protocols::maybe_error::MaybeError;
 pub type TokenType = Option<String>;
 pub type LogProbs = Vec<f64>;
 
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+pub struct TopLogprob {
+    pub rank: u32,
+    pub token_id: TokenIdType,
+    pub token: TokenType,
+    pub logprob: f64,
+}
+pub type TopLogprobs = Vec<Vec<TopLogprob>>; // num_tokens x top_logprobs
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct BackendOutput {
     /// New token_ids generated from the LLM Engine
@@ -41,6 +50,8 @@ pub struct BackendOutput {
     /// Optional log probabilities
     pub log_probs: Option<LogProbs>,
 
+    pub top_logprobs: Option<TopLogprobs>,
+
     // TODO: Enrich this with more information as can apply our first-level postprocessing
     // logic and return more detailed information
     pub finish_reason: Option<FinishReason>,
@@ -77,6 +88,8 @@ pub struct LLMEngineOutput {
     /// Optional log probabilities
     pub log_probs: Option<LogProbs>,
 
+    pub top_logprobs: Option<TopLogprobs>,
+
     // TODO: Enrich this with more information as can apply our first-level postprocessing
     // logic and return more detailed information
     pub finish_reason: Option<FinishReason>,
@@ -93,6 +106,7 @@ impl LLMEngineOutput {
             text: None,
             cum_log_probs: None,
             log_probs: None,
+            top_logprobs: None,
             finish_reason: Some(FinishReason::Cancelled),
             index: None,
         }
@@ -106,6 +120,7 @@ impl LLMEngineOutput {
             cum_log_probs: None,
             log_probs: None,
             finish_reason: Some(FinishReason::Stop),
+            top_logprobs: None,
             index: None,
         }
     }
@@ -117,6 +132,7 @@ impl LLMEngineOutput {
             text: None,
             cum_log_probs: None,
             log_probs: None,
+            top_logprobs: None,
             finish_reason: Some(FinishReason::Length),
             index: None,
         }
@@ -129,6 +145,7 @@ impl LLMEngineOutput {
             text: None,
             cum_log_probs: None,
             log_probs: None,
+            top_logprobs: None,
             finish_reason: Some(FinishReason::Error(err_msg)),
             index: None,
         }
diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs
@@ -4,7 +4,7 @@
 use derive_builder::Builder;
 use serde::{Deserialize, Serialize};
 
-use super::{SamplingOptions, StopConditions};
+use super::{OutputOptions, SamplingOptions, StopConditions};
 use crate::protocols::TokenIdType;
 
 /// [`PreprocessedRequest`] is the internal representation of an LLM request. The [`dynamo.llm-preprocessor`]
@@ -29,6 +29,10 @@ pub struct PreprocessedRequest {
     /// are needed.
     pub sampling_options: SamplingOptions,
 
+    /// OutputOptions are options that control the output of the inference engine such as whether
+    /// to return log probabilities, or whether to skip special tokens in output.
+    pub output_options: OutputOptions,
+
     /// The EOS token ID(s) for the Model
     /// Not every backend needs this, but those that do can find it here.
     /// TODO - refactor this to a better location
diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs
@@ -17,7 +17,7 @@ use anyhow::Result;
 use serde::{Deserialize, Serialize};
 
 use super::{
-    common::{self, SamplingOptionsProvider, StopConditionsProvider},
+    common::{self, OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider},
     ContentProvider,
 };
 use crate::protocols::openai::common_ext::CommonExtProvider;
@@ -79,6 +79,16 @@ trait OpenAIStopConditionsProvider {
     }
 }
 
+trait OpenAIOutputOptionsProvider {
+    fn get_logprobs(&self) -> Option<u32>;
+
+    fn get_prompt_logprobs(&self) -> Option<u32>;
+
+    fn get_skip_special_tokens(&self) -> Option<bool>;
+
+    fn get_formatted_prompt(&self) -> Option<bool>;
+}
+
 impl<T: OpenAISamplingOptionsProvider + CommonExtProvider> SamplingOptionsProvider for T {
     fn extract_sampling_options(&self) -> Result<common::SamplingOptions> {
         // let result = self.validate();
@@ -168,6 +178,22 @@ impl<T: OpenAIStopConditionsProvider> StopConditionsProvider for T {
     }
 }
 
+impl<T: OpenAIOutputOptionsProvider> OutputOptionsProvider for T {
+    fn extract_output_options(&self) -> Result<common::OutputOptions> {
+        let logprobs = self.get_logprobs();
+        let prompt_logprobs = self.get_prompt_logprobs();
+        let skip_special_tokens = self.get_skip_special_tokens();
+        let formatted_prompt = self.get_formatted_prompt();
+
+        Ok(common::OutputOptions {
+            logprobs,
+            prompt_logprobs,
+            skip_special_tokens,
+            formatted_prompt,
+        })
+    }
+}
+
 pub trait DeltaGeneratorExt<ResponseType: Send + Sync + 'static + std::fmt::Debug>:
     Send + Sync + 'static
 {
diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs
diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs
diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs
diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs
