[Fix] Guard conditional top-logprob keys in the completions echo path (#34776)

Co-authored-by: James Liu <jamesl@modal.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

python/sglang/srt/entrypoints/openai/serving_completions.py

@@ -74,7 +74,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
     ) -> tuple[GenerateReqInput, CompletionRequest]:
         """Convert OpenAI completion request to internal format"""
         # NOTE: with openai API, the prompt's logprobs are always not computed
-        if request.echo and request.logprobs:
+        if request.echo and request.logprobs is not None:
             logger.warning(
                 "Echo is not compatible with logprobs. "
                 "To compute logprobs of input prompt, please use the native /generate API."
@@ -274,6 +274,9 @@ class OpenAIServingCompletion(OpenAIServingBase):
                         content["meta_info"]
                     )
 
+                finish_reason = content["meta_info"].get("finish_reason", None)
+                finish_reason_type = finish_reason["type"] if finish_reason else None
+
                 is_first_chunk = index not in stream_offsets
                 offset = stream_offsets.get(index, 0)
                 # Handle echo for first chunk
@@ -286,11 +289,13 @@ class OpenAIServingCompletion(OpenAIServingBase):
                 logprobs = None
                 if request.logprobs is not None:
                     # The first chunk and echo is enabled.
-                    if is_first_chunk and request.echo:
+                    if is_first_chunk and request.echo and request.logprobs:
                         input_token_logprobs = content["meta_info"][
                             "input_token_logprobs"
                         ]
-                        input_top_logprobs = content["meta_info"]["input_top_logprobs"]
+                        input_top_logprobs = content["meta_info"].get(
+                            "input_top_logprobs", None
+                        )
                     else:
                         input_token_logprobs = None
                         input_top_logprobs = None
@@ -343,8 +348,6 @@ class OpenAIServingCompletion(OpenAIServingBase):
                 else:
                     delta = text[offset:]
                 stream_offsets[index] = len(content["text"])
-                finish_reason = content["meta_info"].get("finish_reason", None)
-                finish_reason_type = finish_reason["type"] if finish_reason else None
 
                 # Abort with an explicit error status_code is a system error
                 # (timeout, OOM, validation): emit a streaming error chunk.
@@ -575,9 +578,11 @@ class OpenAIServingCompletion(OpenAIServingBase):
             # Handle logprobs
             logprobs = None
             if request.logprobs is not None:
-                if echo:
+                if echo and request.logprobs:
                     input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"]
-                    input_top_logprobs = ret_item["meta_info"]["input_top_logprobs"]
+                    input_top_logprobs = ret_item["meta_info"].get(
+                        "input_top_logprobs", None
+                    )
                 else:
                     input_token_logprobs = None
                     input_top_logprobs = None

test/registered/unit/entrypoints/openai/test_serving_chat.py

@@ -2677,8 +2677,9 @@ class ServingChatTestCase(unittest.TestCase):
                         "status_code": err_code,
                         "message": err_msg,
                     },
-                    "output_token_logprobs": None,
-                    "output_top_logprobs": None,
+                    "output_token_logprobs": [],
+                    "output_token_logprobs_length": 0,
+                    "output_top_logprobs": [],
                 },
                 "index": 0,
             }
@@ -2691,6 +2692,8 @@ class ServingChatTestCase(unittest.TestCase):
             temperature=0.7,
             max_tokens=100,
             stream=True,
+            logprobs=True,
+            top_logprobs=5,
         )
 
         with patch(

test/registered/unit/entrypoints/openai/test_serving_completions.py

@@ -270,8 +270,9 @@ class ServingCompletionTestCase(unittest.TestCase):
                         "status_code": err_code,
                         "message": err_msg,
                     },
-                    "output_token_logprobs": None,
-                    "output_top_logprobs": None,
+                    "output_token_logprobs": [],
+                    "output_token_logprobs_length": 0,
+                    "output_top_logprobs": [],
                 },
                 "index": 0,
             }
@@ -283,6 +284,7 @@ class ServingCompletionTestCase(unittest.TestCase):
             prompt="Hello world",
             max_tokens=100,
             stream=True,
+            logprobs=5,
         )
 
         adapted_request, _ = self.sc._convert_to_internal_request(req)
@@ -318,6 +320,100 @@ class ServingCompletionTestCase(unittest.TestCase):
         self.assertGreaterEqual(len(chunks), 2)
         self.assertIn("error", chunks[0])
 
+    def test_echo_with_zero_logprobs_streaming(self):
+        """logprobs=0 requests token logprobs without top-logprobs, so the
+        scheduler never fills the top-logprob keys. The echo branch must not
+        assume they are present."""
+
+        async def _mock_generate(*args, **kwargs):
+            yield {
+                "text": "Hello world",
+                "meta_info": {
+                    "id": "cmpl-test",
+                    "prompt_tokens": 2,
+                    "completion_tokens": 2,
+                    "cached_tokens": 0,
+                    "finish_reason": {"type": "stop"},
+                    # top_logprobs_num == 0, so no input/output top-logprob keys.
+                    "input_token_logprobs": [],
+                    "output_token_logprobs": [
+                        (-0.1, 3, "Hello"),
+                        (-0.2, 4, " world"),
+                    ],
+                    "output_token_logprobs_length": 2,
+                },
+                "index": 0,
+            }
+
+        self.sc.tokenizer_manager.generate_request = _mock_generate
+
+        req = CompletionRequest(
+            model="x",
+            prompt="Hi",
+            max_tokens=100,
+            stream=True,
+            echo=True,
+            logprobs=0,
+        )
+        adapted_request, _ = self.sc._convert_to_internal_request(req)
+
+        async def run_stream():
+            return [
+                chunk
+                async for chunk in self.sc._generate_completion_stream(
+                    adapted_request, req, self.fastapi_request
+                )
+            ]
+
+        loop = get_or_create_event_loop()
+        chunks = loop.run_until_complete(run_stream())
+
+        # Assert on the payload, not just termination: a regression that
+        # silently drops logprobs still produces a well-formed stream.
+        self.assertNotIn("error", chunks[0])
+        self.assertEqual(chunks[-1], "data: [DONE]\n\n")
+
+        choice = json.loads(chunks[0][len("data: ") :])["choices"][0]
+        self.assertTrue(choice["text"].startswith("Hi"))
+        logprobs = choice["logprobs"]
+        # logprobs=0 asks for token logprobs but no top-logprobs, and the echoed
+        # prompt contributes none because input logprobs were never requested.
+        self.assertEqual(logprobs["tokens"], ["Hello", " world"])
+        self.assertEqual(logprobs["token_logprobs"], [-0.1, -0.2])
+        self.assertEqual(logprobs["top_logprobs"], [])
+
+    def test_echo_with_zero_logprobs_non_streaming(self):
+        """Same contract on the non-streaming path."""
+        req = CompletionRequest(
+            model="x",
+            prompt="Hi",
+            max_tokens=100,
+            echo=True,
+            logprobs=0,
+        )
+        ret = [
+            {
+                "text": " world",
+                "meta_info": {
+                    "id": "cmpl-test",
+                    "prompt_tokens": 2,
+                    "completion_tokens": 1,
+                    "cached_tokens": 0,
+                    "finish_reason": {"type": "stop"},
+                    "weight_version": "v1",
+                    "input_token_logprobs": [],
+                    "output_token_logprobs": [(-0.1, 3, " world")],
+                    "output_token_logprobs_length": 1,
+                },
+            }
+        ]
+
+        response = self.sc._build_completion_response(req, ret, 1234567890)
+
+        self.assertEqual(len(response.choices), 1)
+        self.assertEqual(response.choices[0].logprobs.token_logprobs, [-0.1])
+        self.assertEqual(response.choices[0].logprobs.top_logprobs, [])
+
     def test_streaming_token_ids_deltas_cover_output_exactly(self):
         req = CompletionRequest(
             model="x",