[Bench] Add fixed-prompt mode and per-request spec accept length metrics (#30615)
This commit is contained in:
@@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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 ""
|
||||
|
||||
Reference in New Issue
Block a user