Support multi-round conversations in bench_serving (#6135)

python/sglang/bench_serving.py

@@ -12,6 +12,7 @@ python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-pro
 
 import argparse
 import asyncio
+import copy
 import importlib.util
 import io
 import json
@@ -27,12 +28,12 @@ import uuid
 import warnings
 from argparse import ArgumentParser
 from copy import deepcopy
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime
 from functools import lru_cache
 from json import JSONDecodeError
 from pathlib import Path
-from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, Union
 
 import aiohttp
 import numpy as np
@@ -80,7 +81,7 @@ def _create_bench_client_session():
 
 @dataclass
 class RequestFuncInput:
-    prompt: str
+    prompt: Union[str, List[str], List[Dict[str, str]]]
     api_url: str
     prompt_len: int
     output_len: int
@@ -339,7 +340,9 @@ async def async_request_openai_chat_completions(
             f'rid={rid} time={request_start_time} message="request start" request_func_input="{str(input_partial)}"'
         )
 
-    if request_func_input.image_data:
+    if isinstance(request_func_input.prompt, list):
+        messages = request_func_input.prompt
+    elif request_func_input.image_data:
         # Build multi-image content: a list of image_url entries followed by the text
         content_items = [
             {
@@ -1768,9 +1771,10 @@ def sample_generated_shared_prefix_requests(
 ) -> List[DatasetRow]:
     """Generate benchmark requests with shared system prompts using random tokens and caching."""
     send_routing_key = getattr(args, "gsp_send_routing_key", False)
+    num_turns = getattr(args, "gsp_num_turns", 1)
 
     cache_path = get_gen_prefix_cache_path(args, tokenizer)
-    should_cache = (range_ratio == 1) and not send_routing_key
+    should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1
 
     # Try to load from cache first
     if cache_path.exists() and should_cache:
@@ -1780,7 +1784,7 @@ def sample_generated_shared_prefix_requests(
 
     print(
         f"\nGenerating new input data... "
-        f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=})"
+        f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})"
     )
 
     run_random_str = uuid.uuid4().hex[:8]
@@ -1794,26 +1798,31 @@ def sample_generated_shared_prefix_requests(
     question_lens = compute_random_lens(
         full_len=question_len,
         range_ratio=range_ratio,
-        num=num_groups * prompts_per_group,
+        num=num_groups * prompts_per_group * num_turns,
-    )
+    ).reshape(num_groups, prompts_per_group, num_turns)
     output_lens = compute_random_lens(
         full_len=output_len,
         range_ratio=range_ratio,
         num=num_groups * prompts_per_group,
-    )
+    ).reshape(num_groups, prompts_per_group)
     del system_prompt_len, question_len, output_len
 
     # Generate system prompts for each group
-    system_prompts = []
+    system_prompts = [
-    for i in range(num_groups):
+        gen_prompt(tokenizer, system_prompt_lens[i].item()) for i in range(num_groups)
-        system_prompt = gen_prompt(tokenizer, system_prompt_lens[i].item())
+    ]
-        system_prompts.append(system_prompt)
 
-    # Generate questions
+    # Generate questions: shape (num_groups, prompts_per_group, num_turns)
-    questions = []
+    questions = [
-    for i in range(num_groups * prompts_per_group):
+        [
-        question = gen_prompt(tokenizer, question_lens[i].item())
+            [
-        questions.append(question)
+                gen_prompt(tokenizer, question_lens[g, p, t].item())
+                for t in range(num_turns)
+            ]
+            for p in range(prompts_per_group)
+        ]
+        for g in range(num_groups)
+    ]
 
     # Combine system prompts with questions
     input_requests = []
@@ -1830,25 +1839,28 @@ def sample_generated_shared_prefix_requests(
         for prompt_idx in tqdm(
             range(prompts_per_group), desc="Generating questions", leave=False
         ):
-            flat_index = group_idx * prompts_per_group + prompt_idx
+            turn_questions = questions[group_idx][prompt_idx]
-            question = questions[flat_index]
+            turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[
-            full_prompt = f"{system_prompt}\n\n{question}"
+                1:
+            ]
+            full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts
             prompt_len = (
                 1
                 if getattr(args, "gsp_fast_prepare", False)
-                else len(tokenizer.encode(full_prompt))
+                else len(tokenizer.encode(turn_prompts[0]))
             )
+            output_len_val = output_lens[group_idx, prompt_idx].item()
 
             input_requests.append(
                 DatasetRow(
                     prompt=full_prompt,
                     prompt_len=prompt_len,
-                    output_len=output_lens[flat_index].item(),
+                    output_len=output_len_val,
                     routing_key=routing_key,
                 )
             )
             total_input_tokens += prompt_len
-            total_output_tokens += output_lens[flat_index].item()
+            total_output_tokens += output_len_val
 
     # Shuffle questions
     random.shuffle(input_requests)
@@ -1857,6 +1869,7 @@ def sample_generated_shared_prefix_requests(
     print(f"\nGenerated shared prefix dataset statistics:")
     print(f"Number of groups: {num_groups}")
     print(f"Prompts per group: {prompts_per_group}")
+    print(f"Number of turns: {num_turns}")
     print(f"Total prompts: {len(input_requests)}")
     if not getattr(args, "gsp_fast_prepare", False):
         print(f"Total input tokens: {total_input_tokens}")
@@ -1864,8 +1877,9 @@ def sample_generated_shared_prefix_requests(
         print(
             f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens"
         )
+        all_questions = [q for group in questions for conv in group for q in conv]
         print(
-            f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n"
+            f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n"
         )
 
     # Save to cache
@@ -1919,7 +1933,7 @@ async def get_request(
 
 
 def calculate_metrics(
-    input_requests: List[DatasetRow],
+    input_requests: Optional[List[DatasetRow]],
     outputs: List[RequestFuncOutput],
     dur_s: float,
     tokenizer: PreTrainedTokenizerBase,
@@ -1953,9 +1967,10 @@ def calculate_metrics(
                 tokenizer.encode(outputs[i].generated_text, add_special_tokens=False)
             )
             retokenized_output_lens.append(retokenized_output_len)
-            total_input += input_requests[i].prompt_len
+            if input_requests is not None:
-            total_input_text += input_requests[i].text_prompt_len
+                total_input += input_requests[i].prompt_len
-            total_input_vision += input_requests[i].vision_prompt_len
+                total_input_text += input_requests[i].text_prompt_len
+                total_input_vision += input_requests[i].vision_prompt_len
             if output_len > 1:
                 tpots.append((outputs[i].latency - outputs[i].ttft) / (output_len - 1))
             if use_retokenized_itl:
@@ -2092,6 +2107,42 @@ def calculate_metrics(
     return metrics, output_lens
 
 
+MULTI_TURN_BACKENDS = {"sglang-oai-chat", "vllm-chat", "lmdeploy-chat"}
+
+
+def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable:
+    assert (
+        backend in MULTI_TURN_BACKENDS
+    ), f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}"
+
+    async def f(
+        request_func_input: RequestFuncInput,
+        pbar: Optional[tqdm] = None,
+    ) -> List[RequestFuncOutput]:
+        prompts: List[str] = request_func_input.prompt
+        prev_messages: List[Dict[str, str]] = []
+        outputs = []
+
+        for round_index in range(len(prompts)):
+            prev_messages.append({"role": "user", "content": prompts[round_index]})
+
+            inner_input = replace(
+                copy.deepcopy(request_func_input), prompt=copy.deepcopy(prev_messages)
+            )
+            output = await request_func(
+                inner_input, pbar=pbar if round_index == len(prompts) - 1 else None
+            )
+            outputs.append(output)
+
+            prev_messages.append(
+                {"role": "assistant", "content": output.generated_text}
+            )
+
+        return outputs
+
+    return f
+
+
 async def benchmark(
     backend: str,
     api_url: str,
@@ -2121,6 +2172,10 @@ async def benchmark(
     else:
         raise ValueError(f"Unknown backend: {backend}")
 
+    is_multi_turn = isinstance(input_requests[0].prompt, list)
+    if is_multi_turn:
+        request_func = wrap_multi_turn_request_func(request_func, backend=backend)
+
     # Limit concurrency
     # From https://github.com/vllm-project/vllm/pull/9390
     semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None
@@ -2186,6 +2241,8 @@ async def benchmark(
         )
 
     warmup_outputs = await asyncio.gather(*warmup_tasks)
+    if is_multi_turn:
+        warmup_outputs = [x for output in warmup_outputs for x in output]
 
     # Check if at least one warmup request succeeded
     if warmup_requests > 0 and not any(output.success for output in warmup_outputs):
@@ -2291,6 +2348,8 @@ async def benchmark(
             )
         )
     outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks)
+    if is_multi_turn:
+        outputs = [x for output in outputs for x in output]
 
     # Stop profiler
     if profile:
@@ -2334,7 +2393,7 @@ async def benchmark(
     # Compute metrics and print results
     benchmark_duration = time.perf_counter() - benchmark_start_time
     metrics, output_lens = calculate_metrics(
-        input_requests=input_requests,
+        input_requests=None if is_multi_turn else input_requests,
         outputs=outputs,
         dur_s=benchmark_duration,
         tokenizer=tokenizer,
@@ -3124,6 +3183,12 @@ if __name__ == "__main__":
         action="store_true",
         help="Send routing key in requests via X-SMG-Routing-Key header. Requests with the same prefix share the same routing key.",
     )
+    group.add_argument(
+        "--gsp-num-turns",
+        type=int,
+        default=1,
+        help="Number of turns for multi-turn conversations. If > 1, each prompt becomes a list of questions sharing the same system prefix.",
+    )
     mooncake_group = parser.add_argument_group("mooncake dataset arguments")
     mooncake_group.add_argument(
         "--mooncake-slowdown-factor",

python/sglang/test/test_utils.py

@@ -752,6 +752,7 @@ def get_similarities(vec1, vec2):
 
 def get_benchmark_args(
     base_url="",
+    backend="sglang",
     dataset_name="",
     dataset_path="",
     tokenizer="",
@@ -769,9 +770,15 @@ def get_benchmark_args(
     lora_name=None,
     lora_request_distribution="uniform",
     lora_zipf_alpha=1.5,
+    gsp_num_groups=4,
+    gsp_prompts_per_group=4,
+    gsp_system_prompt_len=128,
+    gsp_question_len=32,
+    gsp_output_len=32,
+    gsp_num_turns=1,
 ):
     return SimpleNamespace(
-        backend="sglang",
+        backend=backend,
         base_url=base_url,
         host=None,
         port=None,
@@ -803,6 +810,12 @@ def get_benchmark_args(
         prompt_suffix="",
         device=device,
         pd_separated=pd_separated,
+        gsp_num_groups=gsp_num_groups,
+        gsp_prompts_per_group=gsp_prompts_per_group,
+        gsp_system_prompt_len=gsp_system_prompt_len,
+        gsp_question_len=gsp_question_len,
+        gsp_output_len=gsp_output_len,
+        gsp_num_turns=gsp_num_turns,
     )
 
 

test/registered/bench_fn/test_bench_serving_functionality.py

@@ -0,0 +1,100 @@
+import json
+import tempfile
+import time
+import unittest
+from pathlib import Path
+
+from sglang.bench_serving import run_benchmark
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    get_benchmark_args,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=300, suite="nightly-1-gpu", nightly=True)
+
+MODEL = "Qwen/Qwen3-0.6B"
+NUM_CONVERSATIONS, NUM_TURNS = 4, 3
+
+
+class TestBenchServingFunctionality(CustomTestCase):
+    def test_gsp_multi_turn(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            process = popen_launch_server(
+                MODEL,
+                DEFAULT_URL_FOR_TEST,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--mem-fraction-static",
+                    "0.7",
+                    "--log-requests",
+                    "--log-requests-level",
+                    "3",
+                    "--log-requests-format",
+                    "json",
+                    "--log-requests-target",
+                    "stdout",
+                    temp_dir,
+                ],
+            )
+            try:
+                args = get_benchmark_args(
+                    base_url=DEFAULT_URL_FOR_TEST,
+                    backend="sglang-oai-chat",
+                    tokenizer=MODEL,
+                    dataset_name="generated-shared-prefix",
+                    num_prompts=NUM_CONVERSATIONS,
+                    request_rate=float("inf"),
+                    gsp_num_groups=2,
+                    gsp_prompts_per_group=2,
+                    gsp_system_prompt_len=64,
+                    gsp_question_len=16,
+                    gsp_output_len=16,
+                    gsp_num_turns=NUM_TURNS,
+                )
+                args.warmup_requests = 0
+                res = run_benchmark(args)
+                self.assertEqual(res["completed"], NUM_CONVERSATIONS * NUM_TURNS)
+
+                time.sleep(1)
+                logs = "".join(f.read_text() for f in Path(temp_dir).glob("*.log"))
+                self._verify_multi_turn_logs(logs)
+            finally:
+                kill_process_tree(process.pid)
+
+    def _verify_multi_turn_logs(self, content: str):
+        reqs = []
+        for line in content.splitlines():
+            if not line.startswith("{"):
+                continue
+            obj = json.loads(line)
+            if obj.get("event") != "request.finished":
+                continue
+            text = obj.get("obj", {}).get("text")
+            rid = obj.get("rid", "")
+            if text and not rid.startswith("HEALTH_CHECK"):
+                reqs.append(text)
+
+        self.assertGreaterEqual(len(reqs), NUM_CONVERSATIONS * NUM_TURNS)
+
+        # Verify prefix relationships
+        reqs_sorted = sorted(reqs, key=len)
+        prefix_count = 0
+        for i, text in enumerate(reqs_sorted):
+            for j in range(i + 1, len(reqs_sorted)):
+                if reqs_sorted[j].startswith(text):
+                    prefix_count += 1
+                    break
+
+        expected = NUM_CONVERSATIONS * (NUM_TURNS - 1)
+        self.assertGreaterEqual(
+            prefix_count, expected, f"Expected at least {expected} prefix pairs"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()