[Bench] Add fixed-prompt mode and per-request spec accept length metrics (#30615)

This commit is contained in:
Liangsheng Yin
2026-07-09 02:06:04 -07:00
committed by GitHub
parent bd7e54d737
commit bc5d376c2c
4 changed files with 180 additions and 45 deletions
+126 -43
View File
@@ -19,6 +19,7 @@ import json
import random
import re
import time
from functools import lru_cache
from types import SimpleNamespace
from typing import Callable, List, Optional, Tuple
@@ -121,6 +122,8 @@ class BenchArgs:
profile_output_dir: Optional[str] = None
dataset_path: str = ""
dataset_name: str = "random"
fixed_prompt_file: str = ""
apply_chat_template: bool = False
gsp_num_groups: int = 1
gsp_system_prompt_len: int = 2048
gsp_question_len: int = 128
@@ -218,6 +221,19 @@ class BenchArgs:
choices=["mmmu", "random", "random-ids", "generated-shared-prefix"],
help="Name of the dataset to benchmark on.",
)
parser.add_argument(
"--fixed-prompt-file",
type=str,
default=BenchArgs.fixed_prompt_file,
help="Use this file's prompt for every request in the batch, "
"bypassing --dataset-name.",
)
parser.add_argument(
"--apply-chat-template",
action="store_true",
help="Encode the prompt as a single user message through the "
"model's chat template. Requires --fixed-prompt-file.",
)
parser.add_argument(
"--gsp-num-groups",
type=int,
@@ -467,6 +483,51 @@ def _flush_cache_with_retry(url: str, endpoint: str, max_retries: int = 3):
time.sleep(2)
# FIXME: this mirrors the chat-encoding dispatch in
# serving_chat._resolve_chat_encoding_spec (DeepSeek-V4 custom encoding vs HF
# chat template) so the benchmark reproduces the serving token stream. Unify
# the dispatch into a shared resolver instead of duplicating it client-side.
@lru_cache(maxsize=None)
def _is_deepseek_v4_model(name_or_path: str) -> bool:
from transformers import AutoConfig
from sglang.srt.configs.model_config import is_deepseek_v4
try:
hf_config = AutoConfig.from_pretrained(name_or_path, trust_remote_code=True)
except Exception as e:
print(
f"Warning: could not load config for {name_or_path!r} ({e}); "
"assuming a non-DeepSeek-V4 model for --apply-chat-template."
)
return False
return is_deepseek_v4(hf_config)
def _encode_fixed_prompt(
tok_inner, prompt_text: str, apply_chat_template: bool
) -> List[int]:
if not apply_chat_template:
return tok_inner.encode(prompt_text)
messages = [{"role": "user", "content": prompt_text}]
# DeepSeek-V4 chat encoding does not go through the HF chat template; use
# its own encoder so the token stream matches /v1/chat/completions.
if _is_deepseek_v4_model(getattr(tok_inner, "name_or_path", "") or ""):
from sglang.srt.entrypoints.openai import encoding_dsv4
real_input = encoding_dsv4.encode_messages(messages, thinking_mode="chat")
return tok_inner.encode(real_input)
if getattr(tok_inner, "chat_template", None) is None:
raise ValueError(
"--apply-chat-template requires a tokenizer with a chat template, "
f"but {getattr(tok_inner, 'name_or_path', tok_inner)!r} has none."
)
return tok_inner.apply_chat_template(
messages, add_generation_prompt=True, tokenize=True
)
def run_one_case(
url: str,
batch_size: int,
@@ -500,6 +561,8 @@ def run_one_case(
lora_name: Optional[List[str]] = None,
lora_request_distribution: str = BenchArgs.lora_request_distribution,
lora_zipf_alpha: float = BenchArgs.lora_zipf_alpha,
fixed_prompt_file: str = "",
apply_chat_template: bool = False,
):
if backend == "vllm":
# You need to have export VLLM_SERVER_DEV_MODE=1 in your environment to use this endpoint.
@@ -507,51 +570,60 @@ def run_one_case(
else:
_flush_cache_with_retry(url, "/flush_cache")
# Load input token ids via benchmark.datasets.get_dataset
supported_datasets = ("random", "random-ids", "mmmu", "generated-shared-prefix")
if dataset_name not in supported_datasets:
raise ValueError(
f"Unsupported dataset for batch benchmark: {dataset_name}. "
f"Supported: {supported_datasets}"
)
actual_gsp_groups = min(gsp_num_groups, batch_size)
dataset_args = SimpleNamespace(
dataset_name=dataset_name,
num_prompts=batch_size,
random_input_len=input_len,
random_output_len=output_len,
random_range_ratio=1.0,
dataset_path=dataset_path,
tokenize_prompt=dataset_name not in ("mmmu", "generated-shared-prefix"),
backend=backend,
seed=BenchArgs.seed,
gsp_num_groups=actual_gsp_groups,
gsp_prompts_per_group=(batch_size + actual_gsp_groups - 1) // actual_gsp_groups,
gsp_system_prompt_len=gsp_system_prompt_len,
gsp_question_len=gsp_question_len,
gsp_output_len=gsp_output_len,
# The generated-shared-prefix dataset's from_args requires these; the
# batch-bench path only ever uses the uniform group distribution.
gsp_group_distribution="uniform",
gsp_zipf_alpha=None,
)
tok_inner = getattr(tokenizer, "tokenizer", tokenizer)
dataset_model_id = model_name or getattr(tok_inner, "name_or_path", None)
input_requests = get_dataset(dataset_args, tokenizer, model_id=dataset_model_id)
if dataset_name == "generated-shared-prefix":
input_requests = input_requests[:batch_size]
input_ids = [tokenizer.encode(req.prompt) for req in input_requests]
input_len = sum(len(ids) for ids in input_ids) // len(input_ids)
output_len = gsp_output_len
if fixed_prompt_file:
tok_inner = getattr(tokenizer, "tokenizer", tokenizer)
with open(fixed_prompt_file) as f:
prompt_ids = _encode_fixed_prompt(tok_inner, f.read(), apply_chat_template)
input_ids = [list(prompt_ids) for _ in range(batch_size)]
input_len = len(prompt_ids)
image_data = None
elif dataset_name == "mmmu":
input_ids = [tok_inner.encode(req.prompt) for req in input_requests]
image_data = [req.image_data for req in input_requests]
else:
input_ids = [req.prompt for req in input_requests]
image_data = None
# Load input token ids via benchmark.datasets.get_dataset
supported_datasets = ("random", "random-ids", "mmmu", "generated-shared-prefix")
if dataset_name not in supported_datasets:
raise ValueError(
f"Unsupported dataset for batch benchmark: {dataset_name}. "
f"Supported: {supported_datasets}"
)
actual_gsp_groups = min(gsp_num_groups, batch_size)
dataset_args = SimpleNamespace(
dataset_name=dataset_name,
num_prompts=batch_size,
random_input_len=input_len,
random_output_len=output_len,
random_range_ratio=1.0,
dataset_path=dataset_path,
tokenize_prompt=dataset_name not in ("mmmu", "generated-shared-prefix"),
backend=backend,
seed=BenchArgs.seed,
gsp_num_groups=actual_gsp_groups,
gsp_prompts_per_group=(batch_size + actual_gsp_groups - 1)
// actual_gsp_groups,
gsp_system_prompt_len=gsp_system_prompt_len,
gsp_question_len=gsp_question_len,
gsp_output_len=gsp_output_len,
# The generated-shared-prefix dataset's from_args requires these; the
# batch-bench path only ever uses the uniform group distribution.
gsp_group_distribution="uniform",
gsp_zipf_alpha=None,
)
tok_inner = getattr(tokenizer, "tokenizer", tokenizer)
dataset_model_id = model_name or getattr(tok_inner, "name_or_path", None)
input_requests = get_dataset(dataset_args, tokenizer, model_id=dataset_model_id)
if dataset_name == "generated-shared-prefix":
input_requests = input_requests[:batch_size]
input_ids = [tokenizer.encode(req.prompt) for req in input_requests]
input_len = sum(len(ids) for ids in input_ids) // len(input_ids)
output_len = gsp_output_len
image_data = None
elif dataset_name == "mmmu":
input_ids = [tok_inner.encode(req.prompt) for req in input_requests]
image_data = [req.image_data for req in input_requests]
else:
input_ids = [req.prompt for req in input_requests]
image_data = None
# Build payload based on backend
if backend == "vllm":
@@ -979,6 +1051,13 @@ def run_benchmark_internal(
bench_args.lora_zipf_alpha > 1
), f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}"
if bench_args.apply_chat_template and not bench_args.fixed_prompt_file:
raise ValueError(
"--apply-chat-template requires --fixed-prompt-file: the other "
"datasets generate token ids directly, so there is no prompt text "
"to run through a chat template."
)
gsp_kwargs = dict(
gsp_num_groups=bench_args.gsp_num_groups,
gsp_system_prompt_len=bench_args.gsp_system_prompt_len,
@@ -1013,6 +1092,8 @@ def run_benchmark_internal(
lora_name=bench_args.lora_name,
lora_request_distribution=bench_args.lora_request_distribution,
lora_zipf_alpha=bench_args.lora_zipf_alpha,
fixed_prompt_file=bench_args.fixed_prompt_file,
apply_chat_template=bench_args.apply_chat_template,
**gsp_kwargs,
)
print("=" * 8 + " Warmup End " + "=" * 8 + "\n")
@@ -1056,6 +1137,8 @@ def run_benchmark_internal(
lora_name=bench_args.lora_name,
lora_request_distribution=bench_args.lora_request_distribution,
lora_zipf_alpha=bench_args.lora_zipf_alpha,
fixed_prompt_file=bench_args.fixed_prompt_file,
apply_chat_template=bench_args.apply_chat_template,
**gsp_kwargs,
)
)
+11
View File
@@ -108,6 +108,7 @@ class RequestFuncOutput:
start_time: float = 0.0
cached_tokens: int = 0
cached_tokens_details: Optional[Dict[str, Any]] = None
spec_accept_length: float = 0.0
@staticmethod
def init_new(request_func_input: RequestFuncInput):
@@ -473,6 +474,10 @@ async def async_request_openai_chat_completions(
output.output_len = response_json.get("usage", {}).get(
"completion_tokens", output_len
)
_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
)
if getattr(args, "cache_report", False):
_extract_cache_from_sglext(response_json, output)
else:
@@ -695,6 +700,12 @@ async def async_request_sglang_generate(
else:
data = json.loads(chunk)
_meta_info = data.get("meta_info") or {}
if _meta_info.get("spec_accept_length") is not None:
output.spec_accept_length = _meta_info[
"spec_accept_length"
]
# NOTE: Some completion API might have a last
# usage summary response without a token so we
# want to check a token was generated
+31
View File
@@ -6,6 +6,7 @@ python3 -m sglang.test.run_eval --port 30000 --eval-name mmlu --num-examples 10
import argparse
import json
import os
import statistics
import time
from sglang.test.simple_eval_common import (
@@ -82,6 +83,7 @@ def run_eval_once(args, base_url: str, eval_obj: Eval) -> dict:
**common_kwargs,
reasoning_effort=getattr(args, "reasoning_effort", None),
extra_body=extra_body if extra_body else None,
record_meta_info=True,
)
# Run eval
@@ -92,6 +94,30 @@ def run_eval_once(args, base_url: str, eval_obj: Eval) -> dict:
return result, latency, sampler
def print_accept_length_summary(samplers: list) -> None:
accept_lengths = [
m["spec_accept_length"]
for sampler in samplers
for m in getattr(sampler, "_meta_infos", [])
if m.get("spec_accept_length") is not None
]
print("=" * 20)
if not accept_lengths:
print(
"Speculative decoding: no per-request spec_accept_length in responses "
"(non-speculative server, or --api completion which lacks return_meta_info)."
)
else:
print(
f"Speculative accept length (per-request, from meta_info): "
f"n={len(accept_lengths)} "
f"mean={statistics.fmean(accept_lengths):.4f} "
f"min={min(accept_lengths):.4f} "
f"max={max(accept_lengths):.4f}"
)
print("=" * 20)
def run_eval(args):
# Lazy import to avoid circular dependency with test_utils
from sglang.test.test_utils import dump_metric
@@ -194,6 +220,7 @@ def run_eval(args):
if getattr(args, "repeat", 1) == 1:
result, latency, sampler = run_eval_once(args, base_url, eval_obj)
samplers = [sampler]
metrics = result.metrics | {"score": result.score}
metrics["latency"] = latency
print(f"Total latency: {latency:.3f} s")
@@ -229,9 +256,11 @@ def run_eval(args):
scores_repeat = []
latencies = []
total_completion_tokens = 0
samplers = []
for f in futures:
result, latency, sampler = f.result()
samplers.append(sampler)
scores_repeat.append(result.score)
latencies.append(latency)
total_completion_tokens += sum(sampler._completion_tokens)
@@ -266,6 +295,8 @@ def run_eval(args):
executor.shutdown()
print_accept_length_summary(samplers)
# Dump reports
file_stem = f"{args.eval_name}_{sampler.model.replace('/', '_')}"
report_filename = f"/tmp/{file_stem}.html"
+12 -2
View File
@@ -95,6 +95,7 @@ class ChatCompletionSampler(SamplerBase):
reasoning_effort: Optional[str] = None,
max_tokens: int = 2048,
extra_body: Optional[Dict[str, Any]] = None,
record_meta_info: bool = False,
):
self.client = OpenAI(base_url=base_url, http_client=LargerHttpxClient())
@@ -110,8 +111,10 @@ class ChatCompletionSampler(SamplerBase):
self.extra_body = extra_body
self.image_format = "url"
self._completion_tokens: list[int] = []
self.record_meta_info = record_meta_info
self._meta_infos: List[Dict[str, Any]] = []
print(
f"ChatCompletionSampler initialized with {self.system_message=} {self.temperature=} {self.max_tokens=} {self.reasoning_effort=} {self.extra_body=}"
f"ChatCompletionSampler initialized with {self.system_message=} {self.temperature=} {self.max_tokens=} {self.reasoning_effort=} {self.extra_body=} {self.record_meta_info=}"
)
def _handle_image(
@@ -140,6 +143,9 @@ class ChatCompletionSampler(SamplerBase):
message_list = [
self._pack_message("system", self.system_message)
] + message_list
extra_body = self.extra_body
if self.record_meta_info:
extra_body = {**(self.extra_body or {}), "return_meta_info": True}
trial = 0
while trial < 6: # 126 seconds in total
try:
@@ -150,8 +156,12 @@ class ChatCompletionSampler(SamplerBase):
top_p=self.top_p,
max_tokens=self.max_tokens,
reasoning_effort=self.reasoning_effort,
extra_body=self.extra_body,
extra_body=extra_body,
)
if self.record_meta_info:
meta_info = getattr(response.choices[0], "meta_info", None)
if meta_info:
self._meta_infos.append(meta_info)
if response.usage and response.usage.completion_tokens is not None:
self._completion_tokens.append(response.usage.completion_tokens)
return response.choices[0].message.content or ""