From c105a3124b9d840742c6101fd8216a3c119018d4 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 6 Jan 2026 11:59:39 +0800 Subject: [PATCH] Support multi-round conversations in bench_serving (#6135) --- python/sglang/bench_serving.py | 125 +++++++++++++----- python/sglang/test/test_utils.py | 15 ++- .../test_bench_serving_functionality.py | 100 ++++++++++++++ 3 files changed, 209 insertions(+), 31 deletions(-) create mode 100644 test/registered/bench_fn/test_bench_serving_functionality.py diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index 1b9fa5ae7..906b45a62 100644 --- a/python/sglang/bench_serving.py +++ b/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 = [] - for i in range(num_groups): - system_prompt = gen_prompt(tokenizer, system_prompt_lens[i].item()) - system_prompts.append(system_prompt) + system_prompts = [ + gen_prompt(tokenizer, system_prompt_lens[i].item()) for i in range(num_groups) + ] - # Generate questions - questions = [] - for i in range(num_groups * prompts_per_group): - question = gen_prompt(tokenizer, question_lens[i].item()) - questions.append(question) + # Generate questions: shape (num_groups, prompts_per_group, num_turns) + questions = [ + [ + [ + 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 - question = questions[flat_index] - full_prompt = f"{system_prompt}\n\n{question}" + turn_questions = questions[group_idx][prompt_idx] + turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[ + 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 - total_input_text += input_requests[i].text_prompt_len - total_input_vision += input_requests[i].vision_prompt_len + if input_requests is not None: + total_input += input_requests[i].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", diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 387edfa7c..73f04b0e8 100644 --- a/python/sglang/test/test_utils.py +++ b/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, ) diff --git a/test/registered/bench_fn/test_bench_serving_functionality.py b/test/registered/bench_fn/test_bench_serving_functionality.py new file mode 100644 index 000000000..c2ca47634 --- /dev/null +++ b/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()