fix(sglang): preserve unary output-flat logprobs requests

components/src/dynamo/sglang/engine_generate.py

@@ -43,6 +43,7 @@ def build_native_generate_request(
     request_id: str,
     priority: int | None,
     sampling_overrides: Mapping[str, Any] | None = None,
+    is_prefill: bool = False,
     bootstrap_host: str | None = None,
     bootstrap_port: int | None = None,
     bootstrap_room: int | None = None,
@@ -79,7 +80,15 @@ def build_native_generate_request(
         payload["cache_salt"] = salt
     payload["input_ids"] = input_ids
     payload["rid"] = request_id
-    payload["stream"] = True
+    if is_prefill and payload.get("return_flat_raw_output_top_logprobs"):
+        # Prefill must remain incremental so decode can rendezvous for KV transfer.
+        payload.pop("return_flat_raw_output_top_logprobs")
+        if not payload.get("return_flat_raw_top_logprobs"):
+            payload.pop("return_flat_raw_top_logprobs_b64", None)
+    # Output-flat logprobs are only emitted on SGLang's unary response.
+    # Keep ordinary requests incremental, even for unary frontend clients.
+    if not payload.get("return_flat_raw_output_top_logprobs"):
+        payload["stream"] = True
     if priority is None:
         payload.pop("priority", None)
     else:

components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py

@@ -212,6 +212,7 @@ async def generate(
                 request_id=sglang_request_id,
                 priority=priority_kwargs.get("priority"),
                 sampling_overrides={"n": 1, "max_new_tokens": 1},
+                is_prefill=True,
                 bootstrap_host=bootstrap_host,
                 bootstrap_port=bootstrap_port,
                 bootstrap_room=bootstrap_room,

components/src/dynamo/sglang/tests/test_sglang_decode_handler.py

@@ -746,6 +746,61 @@ def test_engine_generate_preserves_native_fields_and_overrides_worker_state(
     }
 
 
+@pytest.mark.parametrize(
+    "output_flat,streaming", [(False, False), (True, False), (True, True)]
+)
+def test_engine_generate_output_flat_logprobs_preserves_unary_mode(
+    monkeypatch, output_flat, streaming
+):
+    if not hasattr(GenerateReqInput, "return_flat_raw_output_top_logprobs"):
+        pytest.skip("installed SGLang does not support output-flat logprobs")
+    monkeypatch.setenv("DYN_SGL_ALLOW_TOP_LOGPROBS", "1")
+    native = build_native_generate_request(
+        {
+            "stream": streaming,
+            "return_logprob": True,
+            "top_logprobs_num": 2,
+            "return_flat_raw_output_top_logprobs": output_flat,
+        },
+        input_ids=[1, 2],
+        request_id="request",
+        priority=None,
+    )
+    if output_flat and streaming:
+        with pytest.raises(ValueError, match="does not support stream=True"):
+            native.normalize_batch_and_arguments()
+    else:
+        native.normalize_batch_and_arguments()
+        assert native.stream is (not output_flat)
+        assert native.return_flat_raw_output_top_logprobs is output_flat
+
+
+@pytest.mark.parametrize("prompt_flat", [False, True])
+def test_engine_generate_prefill_output_flat_keeps_streaming(monkeypatch, prompt_flat):
+    if not hasattr(GenerateReqInput, "return_flat_raw_output_top_logprobs"):
+        pytest.skip("installed SGLang does not support output-flat logprobs")
+    monkeypatch.setenv("DYN_SGL_ALLOW_TOP_LOGPROBS", "1")
+    native = build_native_generate_request(
+        {
+            "stream": False,
+            "return_logprob": True,
+            "top_logprobs_num": 2,
+            "return_flat_raw_output_top_logprobs": True,
+            "return_flat_raw_top_logprobs": prompt_flat,
+            "return_flat_raw_top_logprobs_b64": True,
+        },
+        input_ids=[1, 2],
+        request_id="request",
+        priority=None,
+        is_prefill=True,
+        sampling_overrides={"n": 1, "max_new_tokens": 1},
+    )
+    native.normalize_batch_and_arguments()
+    assert native.stream is True
+    assert native.return_flat_raw_output_top_logprobs is False
+    assert native.return_flat_raw_top_logprobs_b64 is prompt_flat
+
+
 def test_native_generate_rejects_salt_without_engine_support(monkeypatch):
     monkeypatch.delitem(GenerateReqInput.__dataclass_fields__, "cache_salt")
     kwargs = {"input_ids": [1], "request_id": "request", "priority": None}

lib/llm/src/protocols/sglang/generate.rs

@@ -43,6 +43,14 @@ pub struct SglangGenerateRequest {
 
 impl SglangGenerateRequest {
     pub fn validate(&self) -> Result<(), String> {
+        if self.stream
+            && self.passthrough.get("return_flat_raw_output_top_logprobs")
+                == Some(&Value::Bool(true))
+        {
+            return Err(
+                "return_flat_raw_output_top_logprobs does not support stream=true.".to_string(),
+            );
+        }
         if self.input_ids.is_empty() {
             return Err("input_ids cannot be empty.".to_string());
         }
@@ -103,7 +111,7 @@ impl SglangGenerateRequest {
                 .unwrap_or(Value::Null),
         );
         envelope.insert("rid".to_string(), Value::String(request_id.to_string()));
-        envelope.insert("stream".to_string(), Value::Bool(true));
+        envelope.insert("stream".to_string(), Value::Bool(self.stream));
         (self.input_ids, Value::Object(envelope))
     }
 }
@@ -146,7 +154,7 @@ mod tests {
         let (input_ids, envelope) = request.into_worker_envelope("resolved-request");
         assert_eq!(input_ids, [1, 2, 3]);
         assert_eq!(envelope["rid"], "resolved-request");
-        assert_eq!(envelope["stream"], true);
+        assert_eq!(envelope["stream"], false);
         assert_eq!(envelope["session_id"], "session-1");
         assert_eq!(envelope["return_text_in_logprobs"], true);
         assert_eq!(envelope["token_ids_logprob"], serde_json::json!([17]));
@@ -159,6 +167,24 @@ mod tests {
         assert!(envelope.get("priority").is_none());
     }
 
+    #[test]
+    fn output_flat_logprobs_require_and_preserve_unary_mode() {
+        for streaming in [false, true] {
+            let request: SglangGenerateRequest = serde_json::from_value(serde_json::json!({
+                "input_ids": [1],
+                "stream": streaming,
+                "return_flat_raw_output_top_logprobs": true
+            }))
+            .unwrap();
+            assert_eq!(request.validate().is_ok(), !streaming);
+            if !streaming {
+                let (_, envelope) = request.into_worker_envelope("request");
+                assert_eq!(envelope["stream"], false);
+                assert_eq!(envelope["return_flat_raw_output_top_logprobs"], true);
+            }
+        }
+    }
+
     #[test]
     fn omitted_sampling_defaults_stay_omitted() {
         let request: SglangGenerateRequest = serde_json::from_value(serde_json::json!({

lib/sidecar/sglang/src/native_http.rs

@@ -1,7 +1,7 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-//! Opaque transport for SGLang's native streaming `/generate` API.
+//! Opaque transport for SGLang's native `/generate` API.
 
 use std::{collections::HashMap, io, time::Duration};
 
@@ -67,7 +67,19 @@ pub(crate) fn request(
 
     body.insert("input_ids".into(), serde_json::json!(request.token_ids));
     body.insert("rid".into(), Value::String(request_id.to_string()));
-    body.insert("stream".into(), Value::Bool(true));
+    if mode.is_prefill()
+        && body.get("return_flat_raw_output_top_logprobs") == Some(&Value::Bool(true))
+    {
+        // The HTTP headers must arrive before publishing the KV handoff.
+        body.remove("return_flat_raw_output_top_logprobs");
+        if body.get("return_flat_raw_top_logprobs") != Some(&Value::Bool(true)) {
+            body.remove("return_flat_raw_top_logprobs_b64");
+        }
+    }
+    // SGLang only supports output-flat logprobs on its unary response.
+    if body.get("return_flat_raw_output_top_logprobs") != Some(&Value::Bool(true)) {
+        body.insert("stream".into(), Value::Bool(true));
+    }
 
     let routing = request.routing.as_ref();
     if let Some(priority) = routing.and_then(|routing| routing.priority) {
@@ -331,10 +343,15 @@ impl NativeHttp {
     }
 
     async fn open(self, body: Value) -> Result<Response, DynamoError> {
+        let accept = if body.get("stream") == Some(&Value::Bool(true)) {
+            "text/event-stream"
+        } else {
+            "application/json"
+        };
         let response = self
             .client
             .post(self.endpoint.with_path("/generate"))
-            .header(header::ACCEPT, "text/event-stream")
+            .header(header::ACCEPT, accept)
             .json(&body)
             .send()
             .await
@@ -500,6 +517,36 @@ fn response_stream(
     response: Response,
     max_event_bytes: usize,
 ) -> BoxStream<'static, Result<String, LinesCodecError>> {
+    if response
+        .headers()
+        .get(header::CONTENT_TYPE)
+        .and_then(|value| value.to_str().ok())
+        .is_some_and(|value| value.split(';').next() == Some("application/json"))
+    {
+        // Adapt the bounded unary body to the same response decoder as SSE.
+        return Box::pin(async_stream::stream! {
+            let chunks = response.bytes_stream();
+            futures::pin_mut!(chunks);
+            let mut body = Vec::new();
+            while let Some(chunk) = chunks.next().await {
+                let chunk = match chunk {
+                    Ok(chunk) => chunk,
+                    Err(error) => {
+                        yield Err(LinesCodecError::Io(io::Error::other(error)));
+                        return;
+                    }
+                };
+                if chunk.len() > max_event_bytes.saturating_sub(body.len()) {
+                    yield Err(LinesCodecError::MaxLineLengthExceeded);
+                    return;
+                }
+                body.extend_from_slice(&chunk);
+            }
+            yield String::from_utf8(body)
+                .map(|body| format!("data: {body}"))
+                .map_err(|error| LinesCodecError::Io(io::Error::other(error)));
+        });
+    }
     let bytes = response.bytes_stream().map_err(io::Error::other);
     let reader = StreamReader::new(bytes);
     Box::pin(FramedRead::new(
@@ -723,6 +770,15 @@ mod tests {
     }
 
     async fn serve_once(body: String, status: &str) -> (u16, tokio::task::JoinHandle<()>) {
+        serve_once_with_content_type(body, status, "text/event-stream").await
+    }
+
+    async fn serve_once_with_content_type(
+        body: String,
+        status: &str,
+        content_type: &str,
+    ) -> (u16, tokio::task::JoinHandle<()>) {
+        let content_type = content_type.to_string();
         let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
         let port = listener.local_addr().unwrap().port();
         let status = status.to_string();
@@ -731,7 +787,7 @@ mod tests {
             let mut request = [0_u8; 4096];
             let _ = socket.read(&mut request).await.unwrap();
             let response = format!(
-                "HTTP/1.1 {status}\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
+                "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
                 body.len()
             );
             socket.write_all(response.as_bytes()).await.unwrap();
@@ -924,6 +980,42 @@ mod tests {
         );
     }
 
+    #[test]
+    fn prefill_output_flat_keeps_streaming_and_prompt_encoding() {
+        for prompt_flat in [false, true] {
+            let mut canonical = canonical_request();
+            canonical.extra_args = Some(json!({"sglang_tito": {
+                "stream": false,
+                "return_flat_raw_output_top_logprobs": true,
+                "return_flat_raw_top_logprobs": prompt_flat,
+                "return_flat_raw_top_logprobs_b64": true
+            }}));
+            let native = request(
+                &canonical,
+                "request",
+                DisaggregationMode::Prefill,
+                Some("prefill"),
+                Some(5000),
+            )
+            .unwrap()
+            .unwrap();
+            assert_eq!(native.body["stream"], true);
+            assert!(
+                native
+                    .body
+                    .get("return_flat_raw_output_top_logprobs")
+                    .is_none()
+            );
+            assert_eq!(
+                native
+                    .body
+                    .get("return_flat_raw_top_logprobs_b64")
+                    .and_then(serde_json::Value::as_bool),
+                prompt_flat.then_some(true),
+            );
+        }
+    }
+
     #[test]
     fn both_native_request_paths_forward_the_current_trace_context() {
         let _subscriber = tracing::subscriber::set_default(
@@ -1199,6 +1291,67 @@ mod tests {
         server.await.unwrap();
     }
 
+    #[tokio::test]
+    async fn native_unary_output_flat_logprobs_preserve_response_and_size_limit() {
+        let response = json!({
+            "output_ids": [101, 102],
+            "meta_info": {
+                "finish_reason": {"type": "length", "length": 2},
+                "output_top_logprobs_val_flat": [-0.1, -0.2, -0.3, -0.4],
+                "output_top_logprobs_idx_flat": [101, 201, 102, 202],
+                "output_top_logprobs_shape": [2, 2]
+            }
+        });
+        for limit in [MAX_EVENT_BYTES, 16] {
+            let mut canonical = canonical_request();
+            canonical.extra_args = Some(json!({"sglang_tito": {
+                "stream": false,
+                "return_logprob": true,
+                "top_logprobs_num": 2,
+                "return_flat_raw_output_top_logprobs": true
+            }}));
+            let native = request(
+                &canonical,
+                "request",
+                DisaggregationMode::Aggregated,
+                None,
+                None,
+            )
+            .unwrap()
+            .unwrap();
+            assert_eq!(native.body["stream"], false);
+            assert_eq!(native.body["return_flat_raw_output_top_logprobs"], true);
+            let (port, server) = serve_once_with_content_type(
+                serde_json::to_string_pretty(&response).unwrap(),
+                "200 OK",
+                "application/json",
+            )
+            .await;
+            let mut http = native_http(port);
+            http.max_event_bytes = limit;
+            let mut stream = http.generate(
+                native,
+                GenerateContext::new(dynamo_backend_common::testing::mock_context(), None),
+                CancellationToken::new(),
+            );
+            let output = stream.next().await.unwrap();
+            if limit == MAX_EVENT_BYTES {
+                let output = output.unwrap();
+                assert_eq!(output.token_ids, [101, 102]);
+                assert_eq!(output.engine_data.unwrap()["sglang_response"], response);
+                assert!(stream.next().await.is_none());
+            } else {
+                assert!(
+                    output
+                        .unwrap_err()
+                        .to_string()
+                        .contains("exceeded 16 bytes")
+                );
+            }
+            server.await.unwrap();
+        }
+    }
+
     #[tokio::test]
     async fn native_stream_rejects_invalid_output_ids() {
         let body = "data: {\"output_ids\":[-1],\"meta_info\":{\"finish_reason\":null}}\n\n";