chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support multi-round conversations in bench_serving (#6135)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-06 11:59:39 +08:00
committed by GitHub
parent 4cf2bbd084
commit c105a3124b
3 changed files with 209 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
python/sglang/bench_serving.py
View File

@@ -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",

15
python/sglang/test/test_utils.py
View File

@@ -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,
)