[bench] Take each request's prompt length from the server

python/sglang/benchmark/serving.py

@@ -255,6 +255,31 @@ def _extract_cache_from_sglext(data, output):
         output.cached_tokens_details = details
 
 
+def _record_server_prompt_len(data, output):
+    """Take the prompt length from the server, the only side that knows it.
+
+    ``RequestFuncOutput.prompt_len`` is seeded from the dataset row by
+    ``RequestFuncOutput.init_new``. That is correct for a single-turn row, where
+    one row is one request. A multi-turn row is instead replayed as one request
+    per round, and every round's output inherits the row's single value -- so
+    summing ``prompt_len`` across outputs counts one number once per round
+    rather than adding up each request's own prompt.
+
+    Two consumers divide by that sum: the ``--cache-report`` hit rate, and
+    ``input_lens`` in the JSON output.
+
+    The server reports the real figure as ``usage.prompt_tokens`` on the
+    OpenAI-compatible routes and ``meta_info.prompt_tokens`` on the native one.
+    It is the same quantity for a single-turn row, so preferring it needs no
+    per-dataset branch, and leaving the seeded value in place when the server
+    reports nothing keeps behaviour unchanged for any backend that does not.
+    """
+    reported = data.get("usage") or data.get("meta_info") or {}
+    prompt_tokens = reported.get("prompt_tokens")
+    if prompt_tokens:
+        output.prompt_len = prompt_tokens
+
+
 # set ignore_eos True by default
 async def async_request_openai_completions(
     request_func_input: RequestFuncInput,
@@ -327,6 +352,7 @@ async def async_request_openai_completions(
                             pass
                         else:
                             data = json.loads(chunk)
+                            _record_server_prompt_len(data, output)
 
                             if getattr(args, "cache_report", False):
                                 _extract_cache_from_sglext(data, output)
@@ -484,6 +510,7 @@ async def async_request_openai_chat_completions(
                         output.output_len = response_json.get("usage", {}).get(
                             "completion_tokens", output_len
                         )
+                        _record_server_prompt_len(response_json, output)
                         _meta_info = response_json["choices"][0].get("meta_info") or {}
                         output.spec_accept_length = (
                             _meta_info.get("spec_accept_length", 0.0) or 0.0
@@ -517,6 +544,7 @@ async def async_request_openai_chat_completions(
                                 output_len = (data.get("usage") or {}).get(
                                     "completion_tokens", output_len
                                 )
+                                _record_server_prompt_len(data, output)
 
                                 if getattr(args, "cache_report", False):
                                     _extract_cache_from_sglext(data, output)
@@ -734,6 +762,8 @@ async def async_request_sglang_generate(
                             # NOTE: Some completion API might have a last
                             # usage summary response without a token so we
                             # want to check a token was generated
+                            _record_server_prompt_len(data, output)
+
                             if getattr(args, "cache_report", False):
                                 _meta = data.get("meta_info") or {}
                                 output.cached_tokens = _meta.get("cached_tokens", 0)
@@ -1321,6 +1351,15 @@ async def f(
             inner_input = replace(
                 copy.deepcopy(request_func_input), prompt=copy.deepcopy(prev_messages)
             )
+            # Each round's prompt length comes from the server's usage block,
+            # which a streamed response only carries when asked for. Default the
+            # key rather than the object, so a user's stream_options still gets
+            # it, while an explicit include_usage=false is kept.
+            if not args.disable_stream:
+                body = inner_input.extra_request_body
+                options = body.get("stream_options") or {}
+                options.setdefault("include_usage", True)
+                body["stream_options"] = options
             output = await request_func(
                 inner_input, pbar=pbar if round_index == len(prompts) - 1 else None
             )

test/registered/unit/bench/test_bench_serving_prompt_len.py

@@ -0,0 +1,139 @@
+"""Unit tests for the per-request prompt length bench_serving records."""
+
+import asyncio
+import json
+import threading
+import unittest
+from argparse import Namespace
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+from sglang.benchmark.serving import (
+    RequestFuncInput,
+    async_request_openai_chat_completions,
+    set_global_args,
+    wrap_multi_turn_request_func,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+# The dataset row's value; never what the stub reports.
+ROW_PROMPT_LEN = 9999
+
+
+class _ChatHandler(BaseHTTPRequestHandler):
+    """Reports server.prompt_tokens[i] for the i-th request. Like a default SGLang
+    server, a streamed response carries usage only when include_usage is set."""
+
+    def do_POST(self):  # noqa: N802
+        body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+        self.server.bodies.append(body)
+        usage = {
+            "completion_tokens": 1,
+            "prompt_tokens": self.server.prompt_tokens[len(self.server.bodies) - 1],
+        }
+        if body.get("stream"):
+            chunks = [{"choices": [{"index": 0, "delta": {"content": "hi"}}]}]
+            if (body.get("stream_options") or {}).get("include_usage"):
+                chunks.append({"choices": [], "usage": usage})
+            content_type = "text/event-stream"
+            payload = b"".join(f"data: {json.dumps(c)}\n\n".encode() for c in chunks)
+            payload += b"data: [DONE]\n\n"
+        else:
+            message = {"role": "assistant", "content": "hi"}
+            choice = {"index": 0, "message": message, "finish_reason": "length"}
+            content_type = "application/json"
+            payload = json.dumps({"choices": [choice], "usage": usage}).encode()
+        self.send_response(200)
+        self.send_header("Content-Type", content_type)
+        self.end_headers()
+        self.wfile.write(payload)
+
+    def log_message(self, fmt, *args):
+        pass
+
+
+class TestBenchServingPromptLen(CustomTestCase):
+    def setUp(self):
+        self.server = HTTPServer(("127.0.0.1", 0), _ChatHandler)
+        threading.Thread(target=self.server.serve_forever, daemon=True).start()
+        self.addCleanup(self.server.server_close)
+        self.addCleanup(self.server.shutdown)
+
+    def _run(self, request_func, prompt, prompt_tokens, stream=True, extra=None):
+        self.server.bodies = []
+        self.server.prompt_tokens = prompt_tokens
+        set_global_args(
+            Namespace(
+                disable_stream=not stream,
+                disable_ignore_eos=True,
+                print_requests=False,
+                tokenizer="",
+                header=None,
+                cache_report=False,
+            )
+        )
+        request = RequestFuncInput(
+            prompt=prompt,
+            api_url=f"http://127.0.0.1:{self.server.server_port}/v1/chat/completions",
+            prompt_len=ROW_PROMPT_LEN,
+            output_len=8,
+            model="dummy-model",
+            lora_name="",
+            image_data=None,
+            extra_request_body=extra or {},
+        )
+        out = asyncio.run(request_func(request))
+        outputs = out if isinstance(out, list) else [out]
+        for o in outputs:
+            self.assertTrue(o.success, o.error)
+        return [o.prompt_len for o in outputs]
+
+    def test_single_request(self):
+        usage_on = {"stream_options": {"include_usage": True}}
+        cases = [
+            ("streaming with usage", True, usage_on, 123),
+            ("non-streaming", False, None, 123),
+            ("no usage keeps row value", True, None, ROW_PROMPT_LEN),
+        ]
+        for name, stream, extra, expected in cases:
+            with self.subTest(name):
+                lens = self._run(
+                    async_request_openai_chat_completions,
+                    "hello",
+                    [123],
+                    stream=stream,
+                    extra=extra,
+                )
+                self.assertEqual(lens, [expected])
+
+    def test_multi_turn_requests_usage_per_round(self):
+        per_round = [100, 250, 400]
+        multi_turn = wrap_multi_turn_request_func(
+            async_request_openai_chat_completions, backend="sglang-oai-chat"
+        )
+        on = {"include_usage": True}
+        other = {"continuous_usage_stats": False}
+        off = {"include_usage": False}
+        # (user's extra_request_body, stream_options sent, prompt_len per round)
+        cases = [
+            (None, on, per_round),
+            ({"stream_options": {}}, on, per_round),
+            ({"stream_options": None}, on, per_round),
+            ({"stream_options": other}, {**other, **on}, per_round),
+            ({"stream_options": off}, off, [ROW_PROMPT_LEN] * 3),
+        ]
+        for extra, sent, expected in cases:
+            with self.subTest(extra=extra):
+                lens = self._run(
+                    multi_turn, ["one", "two", "three"], per_round, extra=extra
+                )
+                self.assertEqual(lens, expected)
+                self.assertEqual(
+                    [b["stream_options"] for b in self.server.bodies], [sent] * 3
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()