From bc5d376c2cf1a18c172dd752da5fbf51ad6a814f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 9 Jul 2026 02:06:04 -0700 Subject: [PATCH] [Bench] Add fixed-prompt mode and per-request spec accept length metrics (#30615) --- python/sglang/benchmark/one_batch_server.py | 169 +++++++++++++++----- python/sglang/benchmark/serving.py | 11 ++ python/sglang/test/run_eval.py | 31 ++++ python/sglang/test/simple_eval_common.py | 14 +- 4 files changed, 180 insertions(+), 45 deletions(-) diff --git a/python/sglang/benchmark/one_batch_server.py b/python/sglang/benchmark/one_batch_server.py index df896106d..4af30b6e8 100644 --- a/python/sglang/benchmark/one_batch_server.py +++ b/python/sglang/benchmark/one_batch_server.py @@ -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, ) ) diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 869c190ef..2a26b6809 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -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 diff --git a/python/sglang/test/run_eval.py b/python/sglang/test/run_eval.py index c80b09540..e5dcffeb8 100644 --- a/python/sglang/test/run_eval.py +++ b/python/sglang/test/run_eval.py @@ -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" diff --git a/python/sglang/test/simple_eval_common.py b/python/sglang/test/simple_eval_common.py index 56f8b5cd3..e010e1194 100644 --- a/python/sglang/test/simple_eval_common.py +++ b/python/sglang/test/simple_eval_common.py @@ -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 ""