# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py # Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py """ Benchmark online serving with dynamic requests. Usage: python3 -m sglang.benchmark.serving --backend sglang --num-prompt 10 python3 -m sglang.benchmark.serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 """ import argparse import asyncio import copy import importlib.util import json import math import os import random import shutil import sys import time import traceback import uuid import warnings from argparse import ArgumentParser from copy import deepcopy from dataclasses import dataclass, field, replace from datetime import datetime from pathlib import Path from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, Union import aiohttp import numpy as np import requests from tqdm.asyncio import tqdm from transformers import AutoTokenizer, PreTrainedTokenizerBase from sglang.benchmark.datasets import DatasetRow, get_dataset from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time from sglang.benchmark.utils import ( get_tokenizer, parse_custom_headers, remove_prefix, set_ulimit, ) from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST from sglang.srt.utils.network import resolve_base_url, resolve_host_port _ROUTING_KEY_HEADER = "X-SMG-Routing-Key" _EMBEDDING_UNSUPPORTED_DATASETS = {"image", "mmmu", "mooncake"} TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and ( shutil.which("gnuplot") is not None ) global args # don't want to import sglang package here def _get_bool_env_var(name: str, default: str = "false") -> bool: value = os.getenv(name, default) return value.lower() in ("true", "1") def _create_bench_client_session(): # When the pressure is big, the read buffer could be full before aio thread read # the content. We increase the read_bufsize from 64K to 10M. # Define constants for timeout and buffer size for clarity and maintainability BENCH_AIOHTTP_TIMEOUT_SECONDS = 6 * 60 * 60 # 6 hours BENCH_AIOHTTP_READ_BUFSIZE_BYTES = 10 * 1024**2 # 10 MB aiohttp_timeout = aiohttp.ClientTimeout(total=BENCH_AIOHTTP_TIMEOUT_SECONDS) return aiohttp.ClientSession( timeout=aiohttp_timeout, read_bufsize=BENCH_AIOHTTP_READ_BUFSIZE_BYTES ) @dataclass class RequestFuncInput: prompt: Union[str, List[str], List[Dict[str, str]]] api_url: str prompt_len: int output_len: int model: str lora_name: str image_data: Optional[List[str]] extra_request_body: Dict[str, Any] timestamp: Optional[float] = None routing_key: Optional[str] = None @dataclass class RequestFuncOutput: generated_text: str = "" success: bool = False latency: float = 0.0 ttft: float = 0.0 # Time to first token itl: List[float] = field(default_factory=list) # List of inter-token latencies text_chunks: List[str] = field(default_factory=list) prompt_len: int = 0 error: str = "" output_len: int = 0 start_time: float = 0.0 cached_tokens: int = 0 cached_tokens_details: Optional[Dict[str, Any]] = None spec_accept_length: float = 0.0 spec_cap_length: float = 0.0 spec_block_accept_length: float = 0.0 spec_cap_lens_histogram: List[int] = field(default_factory=list) @staticmethod def init_new(request_func_input: RequestFuncInput): output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len return output def get_auth_headers() -> Dict[str, str]: openai_api_key = os.environ.get("OPENAI_API_KEY") if openai_api_key: return {"Authorization": f"Bearer {openai_api_key}"} else: api_key = os.environ.get("API_KEY") if api_key: return {"Authorization": f"{api_key}"} return {} def get_request_headers() -> Dict[str, str]: headers = get_auth_headers() if h := getattr(args, "header", None): headers.update(parse_custom_headers(h)) return headers def _combine_openai_chat_content(message: Dict[str, Any]) -> str: # Most OpenAI-compatible servers use ``reasoning_content``. vLLM's Kimi # parser instead streams its reasoning in ``reasoning``. Prefer the # standard field when both are present to avoid counting the same tokens # twice on servers that expose aliases. reasoning = message.get("reasoning_content") or message.get("reasoning") or "" return reasoning + (message.get("content") or "") def wait_for_endpoint(url: str, timeout_sec: int = 60) -> bool: """Wait for the server to become ready by polling the given URL.""" print(f"Waiting up to {timeout_sec}s for {url} to become ready...") start_time = time.perf_counter() headers = get_auth_headers() while True: try: response = requests.get(url, headers=headers, timeout=5) if response.status_code == 200: elapsed = time.perf_counter() - start_time print(f"Server ready in {elapsed:.1f}s.") return True except requests.exceptions.RequestException: pass elapsed = time.perf_counter() - start_time if elapsed >= timeout_sec: print(f"Server did not become ready within {timeout_sec}s timeout.") return False time.sleep(1) # trt llm does not support ignore_eos # https://github.com/triton-inference-server/tensorrtllm_backend/issues/505 async def async_request_trt_llm( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url assert api_url.endswith("generate_stream") async with _create_bench_client_session() as session: payload = { "accumulate_tokens": True, "text_input": request_func_input.prompt, "temperature": 0.000001, "top_p": 1.0, "max_tokens": request_func_input.output_len, "stream": True, "min_length": request_func_input.output_len, "end_id": 1048576, **request_func_input.extra_request_body, } if args.disable_ignore_eos: del payload["min_length"] del payload["end_id"] output = RequestFuncOutput.init_new(request_func_input) ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st try: async with session.post(url=api_url, json=payload) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:") data = json.loads(chunk) output.generated_text += data["text_output"] timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = timestamp - st output.ttft = ttft # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp output.latency = most_recent_timestamp - st output.success = True output.output_len = request_func_input.output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output def _extract_cache_from_sglext(data, output): """Extract cache hit details from sglext in OAI-compatible responses.""" sglext = data.get("sglext") or {} details = sglext.get("cached_tokens_details") if details: output.cached_tokens = ( (details.get("device") or 0) + (details.get("host") or 0) + (details.get("storage") or 0) ) output.cached_tokens_details = details # set ignore_eos True by default async def async_request_openai_completions( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url assert api_url.endswith( "completions" ), "OpenAI Completions API URL must end with 'completions'." prompt = request_func_input.prompt async with _create_bench_client_session() as session: # Build payload with defaults that can be overridden by extra_request_body payload = { "model": request_func_input.model, "prompt": prompt, "best_of": 1, "max_tokens": request_func_input.output_len, "stream": not args.disable_stream, } # Add temperature default only if not specified in extra_request_body if "temperature" not in request_func_input.extra_request_body: payload["temperature"] = 0.0 # Add ignore_eos default only if not specified in extra_request_body if "ignore_eos" not in request_func_input.extra_request_body: payload["ignore_eos"] = not args.disable_ignore_eos if args.return_logprob and args.top_logprobs_num > 0: payload["logprobs"] = args.top_logprobs_num # Merge in extra parameters - these will override defaults if present payload.update(request_func_input.extra_request_body) # hack to accommodate different LoRA conventions between SGLang and vLLM. if request_func_input.lora_name: payload["model"] = request_func_input.lora_name payload["lora_path"] = request_func_input.lora_name if request_func_input.image_data: payload.update({"image_data": request_func_input.image_data}) headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) generated_text = "" output_len = request_func_input.output_len ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) if getattr(args, "cache_report", False): _extract_cache_from_sglext(data, output) # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if data["choices"][0]["text"]: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: output.text_chunks.append( data["choices"][0]["text"] ) output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp generated_text += data["choices"][0]["text"] output_len = (data.get("usage") or {}).get( "completion_tokens", output_len ) output.generated_text = generated_text output.success = True output.latency = latency output.output_len = output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_openai_chat_completions( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: """Makes a request to the OpenAI Chat Completions API. Handles both streaming and non-streaming responses, including support for image data in messages. Calculates and returns various performance metrics. Args: request_func_input: Input parameters for the request. pbar: Optional tqdm progress bar to update. Returns: RequestFuncOutput: Output of the request, including generated text, latency, TTFT, ITL, and success status. """ api_url = request_func_input.api_url assert api_url.endswith( "chat/completions" ), "OpenAI Chat Completions API URL must end with 'chat/completions'." # TODO put it to other functions when `pbar` logic is refactored if getattr(args, "print_requests", False): rid = str(uuid.uuid4()) input_partial = deepcopy(request_func_input) input_partial.prompt = "..." request_start_time = time.time() print( f'rid={rid} time={request_start_time} message="request start" request_func_input="{str(input_partial)}"' ) 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 = [ { "type": "image_url", "image_url": {"url": img_url}, } for img_url in request_func_input.image_data ] content_items.append({"type": "text", "text": request_func_input.prompt}) messages = [ { "role": "user", "content": content_items, }, ] else: messages = [{"role": "user", "content": request_func_input.prompt}] async with _create_bench_client_session() as session: # Build payload with defaults that can be overridden by extra_request_body payload = { "model": request_func_input.model, "messages": messages, "max_completion_tokens": request_func_input.output_len, "stream": not args.disable_stream, } # Add temperature default only if not specified in extra_request_body if "temperature" not in request_func_input.extra_request_body: payload["temperature"] = 0.0 # Add ignore_eos default only if not specified in extra_request_body # Default to False for more realistic behavior (respect EOS tokens) if "ignore_eos" not in request_func_input.extra_request_body: payload["ignore_eos"] = not args.disable_ignore_eos # Merge in extra parameters (tools, temperature, top_p, etc.) # These will override defaults if present payload.update(request_func_input.extra_request_body) # hack to accommodate different LoRA conventions between SGLang and vLLM. if request_func_input.lora_name: payload["model"] = request_func_input.lora_name payload["lora_path"] = request_func_input.lora_name headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) generated_text = "" output_len = request_func_input.output_len ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: if args.disable_stream: # Non-streaming response response_json = await response.json() message = response_json["choices"][0]["message"] output.generated_text = _combine_openai_chat_content(message) output.success = True output.latency = time.perf_counter() - st output.ttft = ( output.latency ) # For non-streaming, TTFT = total latency 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 ) output.spec_cap_length = ( _meta_info.get("spec_cap_length", 0.0) or 0.0 ) output.spec_block_accept_length = ( _meta_info.get("spec_block_accept_length", 0.0) or 0.0 ) output.spec_cap_lens_histogram = ( _meta_info.get("spec_cap_lens_histogram", []) or [] ) if getattr(args, "cache_report", False): _extract_cache_from_sglext(response_json, output) else: # Streaming response async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # Check for usage info in final chunks. OpenAI-compatible # servers may emit usage-only chunks with choices=[]. output_len = (data.get("usage") or {}).get( "completion_tokens", output_len ) if getattr(args, "cache_report", False): _extract_cache_from_sglext(data, output) choices = data.get("choices") or [] if not choices: continue # Reasoning models stream thoughts via # `reasoning_content`; count them like content. delta = choices[0].get("delta") or {} content = _combine_openai_chat_content(delta) if content: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = timestamp - st output.ttft = ttft # Decoding phase else: output.text_chunks.append(content) output.itl.append( timestamp - most_recent_timestamp ) most_recent_timestamp = timestamp generated_text += content output.generated_text = generated_text output.success = True output.latency = latency output.output_len = output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) # TODO put it to other functions when `pbar` logic is refactored if getattr(args, "print_requests", False): curr_t = time.time() output_partial = deepcopy(output) output_partial.generated_text = "..." print( f'rid={rid} time={curr_t} time_delta={curr_t - request_start_time} message="request end" output="{str(output_partial)}"' ) if pbar: pbar.update(1) return output async def async_request_truss( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url prompt = request_func_input.prompt async with _create_bench_client_session() as session: payload = { "model": request_func_input.model, "prompt": prompt, "temperature": 0.0, "best_of": 1, "max_tokens": request_func_input.output_len, "stream": not args.disable_stream, "ignore_eos": not args.disable_ignore_eos, **request_func_input.extra_request_body, } headers = get_request_headers() output = RequestFuncOutput.init_new(request_func_input) generated_text = "" ttft = 0.0 st = time.perf_counter() most_recent_timestamp = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass else: data = json.loads(chunk) # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated if data["choices"][0]["text"]: timestamp = time.perf_counter() # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: output.itl.append(timestamp - most_recent_timestamp) most_recent_timestamp = timestamp generated_text += data["choices"][0]["text"] output.generated_text = generated_text output.success = True output.latency = latency output.output_len = request_func_input.output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_sglang_generate( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url prompt = request_func_input.prompt async with _create_bench_client_session() as session: sampling_params = { "temperature": args.temperature, "max_new_tokens": request_func_input.output_len, "ignore_eos": not args.disable_ignore_eos, } if args.top_p < 1.0: sampling_params["top_p"] = args.top_p payload = { ("text" if isinstance(prompt, str) else "input_ids"): prompt, "sampling_params": sampling_params, "stream": not args.disable_stream, "lora_path": request_func_input.lora_name, "return_logprob": args.return_logprob, "return_routed_experts": args.return_routed_experts, "logprob_start_len": args.logprob_start_len, **request_func_input.extra_request_body, } if args.top_logprobs_num > 0: payload["top_logprobs_num"] = args.top_logprobs_num if args.token_ids_logprob is not None: payload["token_ids_logprob"] = args.token_ids_logprob # Add image data if available (list of image urls/base64) if request_func_input.image_data: payload["image_data"] = request_func_input.image_data headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) generated_text = "" output_len = request_func_input.output_len ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st last_output_len = 0 try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: async for chunk_bytes in response.content: chunk_bytes = chunk_bytes.strip() if not chunk_bytes: continue chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") latency = time.perf_counter() - st if chunk == "[DONE]": pass 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 if getattr(args, "cache_report", False): _meta = data.get("meta_info") or {} output.cached_tokens = _meta.get("cached_tokens", 0) output.cached_tokens_details = _meta.get( "cached_tokens_details" ) if "text" in data and data["text"]: timestamp = time.perf_counter() generated_text = data["text"] output_len = data["meta_info"]["completion_tokens"] # First token if ttft == 0.0: ttft = time.perf_counter() - st output.ttft = ttft # Decoding phase else: num_new_tokens = output_len - last_output_len if num_new_tokens == 0: continue chunk_gap = timestamp - most_recent_timestamp adjust_itl = chunk_gap / num_new_tokens output.itl.extend([adjust_itl] * num_new_tokens) most_recent_timestamp = timestamp last_output_len = output_len output.generated_text = generated_text output.success = True output.latency = latency output.output_len = output_len else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) print(f"{output.error=}") if pbar: pbar.update(1) return output async def async_request_openai_embeddings( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url async with _create_bench_client_session() as session: payload = { "input": request_func_input.prompt, "model": request_func_input.model, } if request_func_input.lora_name: payload["model"] = request_func_input.lora_name payload["lora_path"] = request_func_input.lora_name payload.update(request_func_input.extra_request_body) headers = get_request_headers() if request_func_input.routing_key: headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key output = RequestFuncOutput.init_new(request_func_input) st = time.perf_counter() output.start_time = st try: async with session.post( url=api_url, json=payload, headers=headers ) as response: if response.status == 200: await response.json() output.latency = time.perf_counter() - st output.success = True output.output_len = 0 else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) if pbar: pbar.update(1) return output async def async_request_gserver( request_func_input: RequestFuncInput, pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: raise NotImplementedError() async def async_request_profile(api_url: str) -> RequestFuncOutput: async with _create_bench_client_session() as session: output = RequestFuncOutput() try: if api_url.endswith("/start_profile"): num_steps = getattr(args, "profile_num_steps", None) profile_by_stage = getattr(args, "profile_by_stage", None) if profile_by_stage and num_steps is None: num_steps = 5 output_dir = getattr(args, "profile_output_dir", None) if output_dir is None: output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") output_dir = Path(os.path.abspath(os.path.normpath(output_dir))) / str( time.time() ) output_dir.mkdir(exist_ok=True, parents=True) output_dir = str(output_dir) body = { "activities": getattr(args, "profile_activities", []), "num_steps": num_steps, "profile_by_stage": profile_by_stage, "profile_stages": getattr(args, "profile_stages", None), "output_dir": output_dir, "profile_prefix": getattr(args, "profile_prefix", None), } else: # stop_profile doesn't need any parameters body = {} print(f"async_request_profile {api_url=} {body=}") # Add optional profiling parameters if provided if ( hasattr(args, "profile_start_step") and args.profile_start_step is not None ): body["start_step"] = str(args.profile_start_step) if hasattr(args, "profile_steps") and args.profile_steps is not None: body["num_steps"] = str(args.profile_steps) async with session.post(url=api_url, json=body) as response: if response.status == 200: output.success = True else: output.error = ( (response.reason or "") + ": " + (await response.text()) ) output.success = False except Exception: output.success = False exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) return output def _build_profile_urls( profile_prefill_url: Optional[List[str]], profile_decode_url: Optional[List[str]], ) -> List[Tuple[str, str]]: """Build profile URLs list from prefill/decode URL arguments. Returns: List of (worker_type, url) tuples. e.g., [("Prefill-0", "http://..."), ("Decode-0", "http://...")] """ profile_urls = [] if profile_prefill_url: for idx, url in enumerate(profile_prefill_url): profile_urls.append((f"Prefill-{idx}", url)) if profile_decode_url: for idx, url in enumerate(profile_decode_url): profile_urls.append((f"Decode-{idx}", url)) return profile_urls async def _call_profile_pd(profile_urls: List[Tuple[str, str]], mode: str) -> None: """Call profile endpoint (start/stop) on PD separated workers. Args: profile_urls: List of (worker_type, url) tuples mode: "start" or "stop" """ endpoint = "/start_profile" if mode == "start" else "/stop_profile" action = "Starting" if mode == "start" else "Stopping" action_past = "started" if mode == "start" else "stopped" print(f"{action} profiler...") for worker_type, url in profile_urls: profile_output = await async_request_profile(api_url=url + endpoint) if profile_output.success: print(f"Profiler {action_past} for {worker_type} worker at {url}") else: print( f"Failed to {mode} profiler for {worker_type} worker at {url}: {profile_output.error}" ) ASYNC_REQUEST_FUNCS = { "sglang": async_request_sglang_generate, "sglang-native": async_request_sglang_generate, "sglang-oai": async_request_openai_completions, "sglang-oai-chat": async_request_openai_chat_completions, "sglang-embedding": async_request_openai_embeddings, "vllm": async_request_openai_completions, "vllm-chat": async_request_openai_chat_completions, "lmdeploy": async_request_openai_completions, "lmdeploy-chat": async_request_openai_chat_completions, "trt": async_request_trt_llm, "gserver": async_request_gserver, "truss": async_request_truss, } # API path appended to the base URL per backend. gserver is special (bare # host:port, no path) and is handled separately, so it is not listed here. _BACKEND_API_PATHS = { "sglang": "/generate", "sglang-native": "/generate", "sglang-oai": "/v1/completions", "sglang-oai-chat": "/v1/chat/completions", "sglang-embedding": "/v1/embeddings", "vllm": "/v1/completions", "vllm-chat": "/v1/chat/completions", "lmdeploy": "/v1/completions", "lmdeploy-chat": "/v1/chat/completions", "trt": "/v2/models/ensemble/generate_stream", "truss": "/v1/models/model:predict", } @dataclass class BenchmarkMetrics: # Request counts and token totals completed: int total_input: int total_input_text: int total_input_vision: int total_output: int total_output_retokenized: int # Throughput (req/s and tok/s) request_throughput: float input_throughput: float output_throughput: float output_throughput_retokenized: float total_throughput: float total_throughput_retokenized: float # TTFT - Time to First Token (ms) mean_ttft_ms: float median_ttft_ms: float std_ttft_ms: float p90_ttft_ms: float p95_ttft_ms: float p99_ttft_ms: float # TPOT - Time per Output Token, excluding the first token (ms) mean_tpot_ms: float median_tpot_ms: float std_tpot_ms: float p90_tpot_ms: float p95_tpot_ms: float p99_tpot_ms: float # ITL - Inter-Token Latency (ms) mean_itl_ms: float median_itl_ms: float std_itl_ms: float p90_itl_ms: float p95_itl_ms: float p99_itl_ms: float max_itl_ms: float # E2E - End-to-End request latency (ms) mean_e2e_latency_ms: float median_e2e_latency_ms: float std_e2e_latency_ms: float p90_e2e_latency_ms: float p95_e2e_latency_ms: float p99_e2e_latency_ms: float # Concurrency and peak metrics concurrency: float max_output_tokens_per_s: float = 0.0 max_concurrent_requests: int = 0 async def get_request( input_requests: List[DatasetRow], request_rate: float, use_trace_timestamps: bool = False, slowdown_factor: float = 1.0, ) -> AsyncGenerator[DatasetRow, None]: if use_trace_timestamps: print( f"Using trace timestamps for request generation with slowdown factor {slowdown_factor}." ) # Sort requests by timestamp for correct replay input_requests.sort(key=lambda r: r.timestamp) start_time = time.perf_counter() trace_start_time_ms = input_requests[0].timestamp if input_requests else 0 for request in input_requests: trace_time_s = (request.timestamp - trace_start_time_ms) / 1000.0 target_arrival_time = start_time + (trace_time_s * slowdown_factor) sleep_duration = target_arrival_time - time.perf_counter() if sleep_duration > 0: await asyncio.sleep(sleep_duration) yield request else: input_requests_iter = iter(input_requests) for request in input_requests_iter: yield request if request_rate == float("inf"): # If the request rate is infinity, then we don't need to wait. continue # Sample the request interval from the exponential distribution. interval = np.random.exponential(1.0 / request_rate) # The next request will be sent after the interval. await asyncio.sleep(interval) def calculate_metrics( input_requests: Optional[List[DatasetRow]], outputs: List[RequestFuncOutput], dur_s: float, tokenizer: PreTrainedTokenizerBase, backend: str, accept_length: Optional[float] = None, plot_throughput: bool = False, ) -> Tuple[BenchmarkMetrics, List[int]]: output_lens: List[int] = [] retokenized_output_lens: List[int] = [] total_input = 0 total_input_text = 0 total_input_vision = 0 completed = 0 itls: List[float] = [] tpots: List[float] = [] ttfts: List[float] = [] e2e_latencies: List[float] = [] retokenized_itls: List[float] = [] use_retokenized_itl = ( accept_length is not None and accept_length > 0 and backend in ("sglang-oai", "sglang-oai-chat") ) for i in range(len(outputs)): if outputs[i].success: output_len = outputs[i].output_len output_lens.append(output_len) retokenized_output_len = len( tokenizer.encode(outputs[i].generated_text, add_special_tokens=False) ) retokenized_output_lens.append(retokenized_output_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: for k, itl in enumerate(outputs[i].itl): num_tokens = len( tokenizer.encode( outputs[i].text_chunks[k], add_special_tokens=False ) ) adjusted_itl = itl / num_tokens retokenized_itls.extend([adjusted_itl] * num_tokens) else: itls += outputs[i].itl ttfts.append(outputs[i].ttft) e2e_latencies.append(outputs[i].latency) completed += 1 else: output_lens.append(0) retokenized_output_lens.append(0) if completed == 0: warnings.warn( "All requests failed. This is likely due to a misconfiguration " "on the benchmark arguments.", stacklevel=2, ) max_output_tokens_per_s = 0.0 max_concurrent_requests = 0 successful_outputs = [output for output in outputs if output.success] if successful_outputs: min_start_time = min(output.start_time for output in successful_outputs) max_end_time = max( output.start_time + output.latency for output in successful_outputs ) duration_seconds = int(np.ceil(max_end_time - min_start_time)) + 1 tokens_per_second = np.zeros(duration_seconds) concurrent_requests_per_second = np.zeros(duration_seconds) for output in outputs: if not output.success: continue token_times = [output.start_time + output.ttft] current_time = token_times[0] for itl_value in output.itl: current_time += itl_value token_times.append(current_time) for token_time in token_times: second_bucket = int(token_time - min_start_time) if 0 <= second_bucket < duration_seconds: tokens_per_second[second_bucket] += 1 request_start_second = int(output.start_time - min_start_time) request_end_second = int( (output.start_time + output.latency) - min_start_time ) for second in range( request_start_second, min(request_end_second + 1, duration_seconds) ): concurrent_requests_per_second[second] += 1 if len(tokens_per_second) > 0: max_output_tokens_per_s = float(np.max(tokens_per_second)) max_concurrent_requests = int(np.max(concurrent_requests_per_second)) if plot_throughput: if TERM_PLOTLIB_AVAILABLE: import termplotlib as tpl fig = tpl.figure() fig.plot( np.arange(len(tokens_per_second)), tokens_per_second, title="Output tokens per second", xlabel="Time (s)", ) fig.plot( np.arange(len(concurrent_requests_per_second)), concurrent_requests_per_second, title="Concurrent requests per second", xlabel="Time (s)", ) fig.show() else: print("tip: install termplotlib and gnuplot to plot the metrics") itls = retokenized_itls if use_retokenized_itl else itls metrics = BenchmarkMetrics( completed=completed, total_input=total_input, total_input_text=total_input_text, total_input_vision=total_input_vision, total_output=sum(output_lens), total_output_retokenized=sum(retokenized_output_lens), request_throughput=completed / dur_s, input_throughput=total_input / dur_s, output_throughput=sum(output_lens) / dur_s, output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, total_throughput=(total_input + sum(output_lens)) / dur_s, total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) / dur_s, mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by backend median_ttft_ms=np.median(ttfts or 0) * 1000, std_ttft_ms=np.std(ttfts or 0) * 1000, p90_ttft_ms=np.percentile(ttfts or 0, 90) * 1000, p95_ttft_ms=np.percentile(ttfts or 0, 95) * 1000, p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, mean_tpot_ms=np.mean(tpots or 0) * 1000, median_tpot_ms=np.median(tpots or 0) * 1000, std_tpot_ms=np.std(tpots or 0) * 1000, p90_tpot_ms=np.percentile(tpots or 0, 90) * 1000, p95_tpot_ms=np.percentile(tpots or 0, 95) * 1000, p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, mean_itl_ms=np.mean(itls or 0) * 1000, median_itl_ms=np.median(itls or 0) * 1000, std_itl_ms=np.std(itls or 0) * 1000, p90_itl_ms=np.percentile(itls or 0, 90) * 1000, p95_itl_ms=np.percentile(itls or 0, 95) * 1000, p99_itl_ms=np.percentile(itls or 0, 99) * 1000, max_itl_ms=np.max(itls or 0) * 1000, mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, median_e2e_latency_ms=np.median(e2e_latencies) * 1000, std_e2e_latency_ms=np.std(e2e_latencies) * 1000, p90_e2e_latency_ms=np.percentile(e2e_latencies, 90) * 1000, p95_e2e_latency_ms=np.percentile(e2e_latencies, 95) * 1000, p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, concurrency=np.sum(e2e_latencies) / dur_s, max_output_tokens_per_s=max_output_tokens_per_s, max_concurrent_requests=max_concurrent_requests, ) return metrics, output_lens MULTI_TURN_BACKENDS = {"sglang-oai-chat", "vllm-chat", "lmdeploy-chat"} def _normalize_round_messages(turn: Any) -> Optional[List[Dict[str, str]]]: """Normalize a multi-turn round to a list of message dicts. Accepts ``str`` (single user message) or ``List[Dict]`` with role/content (e.g. multiple tool observations bundled into one round). Returns ``None`` on any other shape so callers can also use it as a predicate. """ if isinstance(turn, str): return [{"role": "user", "content": turn}] if ( isinstance(turn, list) and turn and all(isinstance(m, dict) and "role" in m and "content" in m for m in turn) ): return [{"role": m["role"], "content": m["content"]} for m in turn] return None 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 = request_func_input.prompt prev_messages: List[Dict[str, str]] = [] outputs = [] for round_index in range(len(prompts)): normalized = _normalize_round_messages(prompts[round_index]) if normalized is None: raise ValueError( f"Multi-turn round {round_index} must be a str or a " "non-empty List[Dict] of role/content messages, got: " f"{type(prompts[round_index]).__name__}" ) prev_messages.extend(normalized) 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, base_url: str, model_id: str, tokenizer: PreTrainedTokenizerBase, input_requests: List[DatasetRow], request_rate: float, max_concurrency: Optional[int], disable_tqdm: bool, lora_names: List[str], lora_request_distribution: Optional[str], lora_zipf_alpha: Optional[float], extra_request_body: Dict[str, Any], profile: bool, pd_separated: bool = False, flush_cache: bool = False, warmup_requests: int = 1, use_trace_timestamps: bool = False, mooncake_slowdown_factor=1.0, mooncake_num_rounds=1, profile_prefill_url: Optional[List[str]] = None, profile_decode_url: Optional[List[str]] = None, ): if backend in ASYNC_REQUEST_FUNCS: request_func = ASYNC_REQUEST_FUNCS[backend] else: raise ValueError(f"Unknown backend: {backend}") # Multi-turn iff prompt[0] is a valid per-round payload. Single-shot # OpenAI messages (List[Dict]) is excluded since its first element is a dict. first_prompt = input_requests[0].prompt is_multi_turn = ( isinstance(first_prompt, list) and bool(first_prompt) and _normalize_round_messages(first_prompt[0]) is not None ) 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 async def limited_request_func(request_func_input, pbar): if semaphore is None: return await request_func(request_func_input=request_func_input, pbar=pbar) async with semaphore: return await request_func(request_func_input=request_func_input, pbar=pbar) # Warmup print(f"Starting warmup with {warmup_requests} sequences...") # Handle the data structure difference for the warmup request if args.dataset_name == "mooncake": # For mooncake, input_requests is a list of dicts. # We need to build a temporary DatasetRow for the warmup phase. warmup_record = input_requests[0] # Build prompt from hash_ids, just like in the async generator hash_ids = warmup_record.get("hash_ids", []) prompt_text = "" for hash_id in hash_ids: prompt_text += f"{hash_id}" + " ".join(["hi"] * 512) prompt_text += "Can you tell me a detailed story in 1000 words?" output_len = warmup_record.get("output_length", 32) prompt_len = len(tokenizer.encode(prompt_text)) # Create a temporary DatasetRow object for warmup test_request = DatasetRow( prompt=prompt_text, prompt_len=prompt_len, output_len=output_len, image_data=None, # Mooncake doesn't have image data ) else: # For all other datasets, input_requests is a list of DatasetRow objects test_request = input_requests[0] if lora_names is not None and len(lora_names) != 0: lora_name = lora_names[0] else: lora_name = None # Create the test input once test_input = RequestFuncInput( model=model_id, prompt=test_request.prompt, api_url=api_url, prompt_len=test_request.prompt_len, output_len=min(test_request.output_len, 32), lora_name=lora_name, image_data=test_request.image_data, extra_request_body=extra_request_body, ) # Run warmup requests warmup_tasks = [] for _ in range(warmup_requests): warmup_tasks.append( asyncio.create_task(request_func(request_func_input=test_input)) ) 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): raise ValueError( "Warmup failed - Please make sure benchmark arguments " f"are correctly specified. Error: {warmup_outputs[0].error}" ) else: print( f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..." ) # Flush cache if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache: requests.post(base_url + "/flush_cache", headers=get_auth_headers()) time.sleep(1.0) # Build profile URLs for PD separated mode (do this once at the beginning) pd_profile_urls = [] if profile and pd_separated: pd_profile_urls = _build_profile_urls(profile_prefill_url, profile_decode_url) if not pd_profile_urls: print( "Warning: PD separated mode requires --profile-prefill-url or --profile-decode-url" ) print("Skipping profiler start. Please specify worker URLs for profiling.") # Start profiler if profile: if pd_separated: if pd_profile_urls: await _call_profile_pd(pd_profile_urls, "start") else: print("Starting profiler...") profile_output = await async_request_profile( api_url=base_url + "/start_profile" ) if profile_output.success: print("Profiler started") # Run all requests benchmark_start_time = time.perf_counter() tasks: List[asyncio.Task] = [] pbar_total = len(input_requests) if ( backend == "sglang" and args.dataset_name == "mooncake" ): # Assuming mooncake is mainly for sglang or similar backends print("Using time-based Mooncake request scheduler, ignoring --request-rate.") request_generator = get_mooncake_request_over_time( input_requests, tokenizer, mooncake_slowdown_factor, mooncake_num_rounds ) print( f"Starting Mooncake trace replay. Sessions: {len(input_requests)}, Rounds per session: {mooncake_num_rounds}. Slowdown factor: {mooncake_slowdown_factor}" ) pbar_total *= args.mooncake_num_rounds else: request_generator = get_request(input_requests, request_rate) # Prepare LoRA request distribution parameters if lora_request_distribution == "distinct": lora_idx = 0 elif lora_request_distribution == "skewed": weights = np.array([lora_zipf_alpha**-i for i in range(len(lora_names))]) lora_probs = weights / np.sum(weights) else: lora_idx = None lora_probs = None pbar = None if disable_tqdm else tqdm(total=pbar_total) async for request in request_generator: if lora_names is not None and len(lora_names) != 0: if lora_request_distribution == "uniform": lora_name = random.choice(lora_names) elif lora_request_distribution == "distinct": lora_name = lora_names[lora_idx] lora_idx = (lora_idx + 1) % len(lora_names) else: assert ( lora_request_distribution == "skewed" ), f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." lora_name = np.random.choice(lora_names, p=lora_probs) else: lora_name = None # Merge global extra_request_body with per-request extras # Per-request parameters take precedence over global ones merged_extra_body = {**extra_request_body, **request.extra_request_body} request_func_input = RequestFuncInput( model=model_id, prompt=request.prompt, api_url=api_url, prompt_len=request.prompt_len, output_len=request.output_len, lora_name=lora_name, image_data=request.image_data, extra_request_body=merged_extra_body, timestamp=request.timestamp, routing_key=request.routing_key, ) tasks.append( asyncio.create_task( limited_request_func(request_func_input=request_func_input, pbar=pbar) ) ) outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) if is_multi_turn: outputs = [x for output in outputs for x in output] # Stop profiler (only if profile_steps was not provided, as it auto-stops) if profile and not ( hasattr(args, "profile_steps") and args.profile_steps is not None ): if pd_separated: if pd_profile_urls: await _call_profile_pd(pd_profile_urls, "stop") else: if getattr(args, "profile_num_steps", None) is None: print("Stopping profiler...") profile_output = await async_request_profile( api_url=base_url + "/stop_profile" ) if profile_output.success: print("Profiler stopped") if pbar is not None: pbar.close() if "sglang" in backend: server_info = requests.get( base_url + "/server_info", headers=get_auth_headers() ) if server_info.status_code == 200: server_info_json = server_info.json() if "decode" in server_info_json: server_info_json = server_info_json["decode"][0] if ( "internal_states" in server_info_json and server_info_json["internal_states"] ): accept_length = server_info_json["internal_states"][0].get( "avg_spec_accept_length", None ) else: accept_length = None else: accept_length = None else: accept_length = None # Compute metrics and print results benchmark_duration = time.perf_counter() - benchmark_start_time metrics, output_lens = calculate_metrics( input_requests=None if is_multi_turn else input_requests, outputs=outputs, dur_s=benchmark_duration, tokenizer=tokenizer, backend=backend, accept_length=accept_length, plot_throughput=args.plot_throughput, ) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) print("{:<40} {:<10}".format("Backend:", backend)) print( "{:<40} {:<10}".format( "Traffic request rate:", "trace" if use_trace_timestamps else request_rate ) ) print( "{:<40} {:<10}".format( "Max request concurrency:", max_concurrency if max_concurrency else "not set", ) ) print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) print("{:<40} {:<10}".format("Total input text tokens:", metrics.total_input_text)) if args.dataset_name in ["image", "mmmu"]: print( "{:<40} {:<10}".format( "Total input vision tokens:", metrics.total_input_vision ) ) is_embedding = backend == "sglang-embedding" if not is_embedding: print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) print( "{:<40} {:<10}".format( "Total generated tokens (retokenized):", metrics.total_output_retokenized, ) ) print( "{:<40} {:<10.2f}".format( "Request throughput (req/s):", metrics.request_throughput ) ) print( "{:<40} {:<10.2f}".format( "Input token throughput (tok/s):", metrics.input_throughput ) ) if not is_embedding: print( "{:<40} {:<10.2f}".format( "Output token throughput (tok/s):", metrics.output_throughput ) ) print( "{:<40} {:<10.2f}".format( "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s, ) ) print( "{:<40} {:<10}".format( "Peak concurrent requests:", metrics.max_concurrent_requests ) ) if not is_embedding: print( "{:<40} {:<10.2f}".format( "Total token throughput (tok/s):", metrics.total_throughput ) ) print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) if accept_length: print("{:<40} {:<10.2f}".format("Accept length:", accept_length)) print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) print( "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format( "Median E2E Latency (ms):", metrics.median_e2e_latency_ms ) ) print( "{:<40} {:<10.2f}".format("P90 E2E Latency (ms):", metrics.p90_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format("P95 E2E Latency (ms):", metrics.p95_e2e_latency_ms) ) print( "{:<40} {:<10.2f}".format("P99 E2E Latency (ms):", metrics.p99_e2e_latency_ms) ) if not is_embedding: print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) print("{:<40} {:<10.2f}".format("P90 TTFT (ms):", metrics.p90_ttft_ms)) print("{:<40} {:<10.2f}".format("P95 TTFT (ms):", metrics.p95_ttft_ms)) print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) print( "{s:{c}^{n}}".format( s="Time per Output Token (excl. 1st token)", n=50, c="-" ) ) print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) print("{:<40} {:<10.2f}".format("P90 TPOT (ms):", metrics.p90_tpot_ms)) print("{:<40} {:<10.2f}".format("P95 TPOT (ms):", metrics.p95_tpot_ms)) print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) print("{s:{c}^{n}}".format(s="Inter-Token Latency", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) print("{:<40} {:<10.2f}".format("P90 ITL (ms):", metrics.p90_itl_ms)) print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms)) print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms)) if args.cache_report: total_prompt_tokens = 0 total_cached = 0 total_device = total_host = total_storage = 0 storage_backend_name = None has_details = False for o in outputs: if not o.success: continue total_prompt_tokens += o.prompt_len total_cached += o.cached_tokens if o.cached_tokens_details: has_details = True total_device += o.cached_tokens_details.get("device") or 0 total_host += o.cached_tokens_details.get("host") or 0 s = o.cached_tokens_details.get("storage") or 0 if s: total_storage += s storage_backend_name = o.cached_tokens_details.get( "storage_backend" ) hit_rate = ( total_cached / total_prompt_tokens * 100 if total_prompt_tokens > 0 else 0.0 ) print("{s:{c}^{n}}".format(s="Cache Hit Details", n=50, c="-")) print("{:<40} {:<10}".format("Total prompt tokens:", total_prompt_tokens)) print("{:<40} {:<10}".format("Total cached tokens:", total_cached)) if has_details and total_cached > 0: print("{:<40} {:<10}".format(" Device:", total_device)) print("{:<40} {:<10}".format(" Host:", total_host)) if total_storage > 0: label = ( f" Storage ({storage_backend_name}):" if storage_backend_name else " Storage:" ) print("{:<40} {:<10}".format(label, total_storage)) print("{:<40} {:.1f}%".format("Cache hit rate:", hit_rate)) if has_details and total_cached > 0: device_pct = total_device / total_cached * 100 host_pct = total_host / total_cached * 100 print("{:<40} {:.1f}%".format(" Device:", device_pct)) print("{:<40} {:.1f}%".format(" Host:", host_pct)) if total_storage > 0: storage_pct = total_storage / total_cached * 100 label = ( f" Storage ({storage_backend_name}):" if storage_backend_name else " Storage:" ) print("{:<40} {:.1f}%".format(label, storage_pct)) print("=" * 50) resp = requests.get(base_url + "/server_info", headers=get_auth_headers()) server_info = resp.json() if resp.status_code == 200 else None if ( metrics.median_ttft_ms is not None and metrics.mean_itl_ms is not None and metrics.output_throughput is not None ): result = { # Arguments "tag": getattr(args, "tag", None), "backend": args.backend, "dataset_name": args.dataset_name, "request_rate": "trace" if use_trace_timestamps else request_rate, "max_concurrency": max_concurrency, "sharegpt_output_len": args.sharegpt_output_len, "random_input_len": args.random_input_len, "random_output_len": args.random_output_len, "random_range_ratio": args.random_range_ratio, # Information "server_info": server_info, # Results "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, "total_input_text_tokens": metrics.total_input_text, "total_input_vision_tokens": metrics.total_input_vision, "total_output_tokens": metrics.total_output, "total_output_tokens_retokenized": metrics.total_output_retokenized, "request_throughput": metrics.request_throughput, "input_throughput": metrics.input_throughput, "output_throughput": metrics.output_throughput, "total_throughput": metrics.total_throughput, "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, "median_e2e_latency_ms": metrics.median_e2e_latency_ms, "std_e2e_latency_ms": metrics.std_e2e_latency_ms, "p90_e2e_latency_ms": metrics.p90_e2e_latency_ms, "p95_e2e_latency_ms": metrics.p95_e2e_latency_ms, "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, "mean_ttft_ms": metrics.mean_ttft_ms, "median_ttft_ms": metrics.median_ttft_ms, "std_ttft_ms": metrics.std_ttft_ms, "p90_ttft_ms": metrics.p90_ttft_ms, "p95_ttft_ms": metrics.p95_ttft_ms, "p99_ttft_ms": metrics.p99_ttft_ms, "mean_tpot_ms": metrics.mean_tpot_ms, "median_tpot_ms": metrics.median_tpot_ms, "std_tpot_ms": metrics.std_tpot_ms, "p90_tpot_ms": metrics.p90_tpot_ms, "p95_tpot_ms": metrics.p95_tpot_ms, "p99_tpot_ms": metrics.p99_tpot_ms, "mean_itl_ms": metrics.mean_itl_ms, "median_itl_ms": metrics.median_itl_ms, "std_itl_ms": metrics.std_itl_ms, "p90_itl_ms": metrics.p90_itl_ms, "p95_itl_ms": metrics.p95_itl_ms, "p99_itl_ms": metrics.p99_itl_ms, "concurrency": metrics.concurrency, "accept_length": accept_length, "max_output_tokens_per_s": metrics.max_output_tokens_per_s, "max_concurrent_requests": metrics.max_concurrent_requests, } if args.cache_report: result["cache_report"] = { "total_prompt_tokens": total_prompt_tokens, "total_cached_tokens": total_cached, "cache_hit_rate_pct": round(hit_rate, 2), "device_cached_tokens": total_device if has_details else None, "host_cached_tokens": total_host if has_details else None, "storage_cached_tokens": (total_storage if total_storage > 0 else None), "storage_backend": storage_backend_name, } else: print(f"Error running benchmark for request rate: {request_rate}") print("-" * 30) # Determine output file name if args.output_file: output_file_name = args.output_file else: now = datetime.now().strftime("%m%d") if args.dataset_name == "image": output_file_name = ( f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_" f"{args.random_output_len}_{args.image_count}imgs_" f"{args.image_resolution}.jsonl" ) elif args.dataset_name.startswith("random"): output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" else: output_file_name = ( f"{args.backend}_{now}_{args.num_prompts}_{args.dataset_name}.jsonl" ) result_details = { "input_lens": [output.prompt_len for output in outputs], "output_lens": output_lens, "ttfts": [output.ttft for output in outputs], "itls": [output.itl for output in outputs], "generated_texts": [output.generated_text for output in outputs], "errors": [output.error for output in outputs], } if args.cache_report: result_details["cached_tokens"] = [o.cached_tokens for o in outputs] result_details["cached_tokens_details"] = [ o.cached_tokens_details for o in outputs ] # Append results to a JSONL file with open(output_file_name, "a") as file: if args.output_details: result_for_dump = result | result_details else: result_for_dump = result file.write(json.dumps(result_for_dump) + "\n") return result | result_details def check_chat_template(model_path): try: tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) return "chat_template" in tokenizer.init_kwargs except Exception as e: print(f"Fail to load tokenizer config with error={e}") return False def set_global_args(args_: argparse.Namespace): """Set the global args.""" global args args = args_ def run_benchmark(args_: argparse.Namespace): global args args = args_ # Set default value for max_concurrency if not present if not hasattr(args, "max_concurrency"): args.max_concurrency = None # Set default value for warmup_requests if not present if not hasattr(args, "warmup_requests"): args.warmup_requests = 1 if not hasattr(args, "output_details"): args.output_details = False if not hasattr(args, "tokenize_prompt"): args.tokenize_prompt = False if not hasattr(args, "plot_throughput"): args.plot_throughput = False if not hasattr(args, "top_logprobs_num"): args.top_logprobs_num = 0 if not hasattr(args, "token_ids_logprob"): args.token_ids_logprob = None if not hasattr(args, "logprob_start_len"): args.logprob_start_len = -1 if not hasattr(args, "return_logprob"): args.return_logprob = False if not hasattr(args, "temperature"): args.temperature = 0.0 if not hasattr(args, "top_p"): args.top_p = 1.0 if not hasattr(args, "use_trace_timestamps"): args.use_trace_timestamps = False if not hasattr(args, "mooncake_slowdown_factor"): args.mooncake_slowdown_factor = 1.0 if not hasattr(args, "mooncake_slowdown_factor"): args.mooncake_slowdown_factor = 1.0 if not hasattr(args, "mooncake_num_rounds"): args.mooncake_num_rounds = 1 if not hasattr(args, "served_model_name"): args.served_model_name = None if not hasattr(args, "cache_report"): args.cache_report = False if getattr(args, "print_requests", False): assert args.backend == "sglang-oai-chat" # only support this now print(f"benchmark_args={args}") # Set global environments set_ulimit() random.seed(args.seed) np.random.seed(args.seed) extra_request_body = {} if args.extra_request_body: extra_request_body = json.loads(args.extra_request_body) if args.cache_report: sglang_backends = ("sglang", "sglang-native", "sglang-oai", "sglang-oai-chat") if args.backend not in sglang_backends: print("WARNING: --cache-report is only supported with sglang backends.") elif args.backend in ("sglang-oai", "sglang-oai-chat"): extra_request_body["return_cached_tokens_details"] = True # Inject bootstrap fields for fake decode benchmarking if getattr(args, "fake_prefill", False): extra_request_body["bootstrap_host"] = FAKE_BOOTSTRAP_HOST extra_request_body["bootstrap_room"] = 0 if args.tokenize_prompt: assert ( args.backend == "sglang" ), "`--tokenize-prompt` only compatible with `--backend sglang` currently" # Set url if args.port is None: args.port = { "sglang": 30000, "sglang-native": 30000, "sglang-oai": 30000, "lmdeploy": 23333, "vllm": 8000, "trt": 8000, "gserver": 9988, "truss": 8080, }.get(args.backend, 30000) # Base URL the client sends to: --base-url if given, else http://host:port # (IPv6-correct). gserver uses the scheme-less host:port form instead. base_url = resolve_base_url(args.base_url, args.host, args.port) model_url = f"{base_url}/v1/models" if args.backend == "gserver": # gRPC server takes a bare host:port, not an http URL. api_url = resolve_host_port(args.base_url, args.host, args.port) args.model = args.model or "default" else: api_url = f"{base_url}{_BACKEND_API_PATHS[args.backend]}" if args.backend == "trt" and args.model is None: print("Please provide a model using `--model` when using `trt` backend.") sys.exit(1) # Wait for server to be ready if args.ready_check_timeout_sec > 0: health_url = model_url if args.backend not in ("trt", "gserver") else base_url if not wait_for_endpoint(health_url, args.ready_check_timeout_sec): print(f"Server at {health_url} is not ready. Exiting.") sys.exit(1) # Get model name if args.model is None: if args.backend == "truss": print( "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" ) sys.exit(1) try: response = requests.get(model_url, headers=get_auth_headers()) model_list = response.json().get("data", []) args.model = model_list[0]["id"] if model_list else None except Exception as e: print(f"Failed to fetch model from {model_url}. Error: {e}") print( "Please specify the correct host and port using `--host` and `--port`." ) sys.exit(1) if args.model is None: print("No model specified or found. Please provide a model using `--model`.") sys.exit(1) if args.backend != "sglang-embedding" and not check_chat_template(args.model): print( "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n" "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n" ) if ( args.backend == "sglang-embedding" and args.dataset_name in _EMBEDDING_UNSUPPORTED_DATASETS ): print(f"{args.dataset_name} dataset is unsupported for embeddings benchmark") sys.exit(1) if args.dataset_name in ["image", "mmmu"]: args.apply_chat_template = True assert ( not args.tokenize_prompt ), "`--tokenize-prompt` not compatible with image dataset" if args.lora_request_distribution in ["distinct", "skewed"]: assert ( args.lora_name is not None and len(args.lora_name) > 1 ), "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." assert ( args.lora_zipf_alpha > 1 ), f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." print(f"{args}\n") # Read dataset backend = args.backend model_id = args.served_model_name or args.model tokenizer_id = args.tokenizer if tokenizer_id is None: try: resp = requests.get( base_url + "/model_info", headers=get_auth_headers(), timeout=5 ) if resp.status_code == 200: info = resp.json() tokenizer_id = info.get("tokenizer_path") or info.get("model_path") except Exception: pass if tokenizer_id is None: tokenizer_id = args.model tokenizer = get_tokenizer(tokenizer_id) input_requests = get_dataset(args, tokenizer, model_id) # compatible with SimpleNamespace if not hasattr(args, "flush_cache"): args.flush_cache = False # Prepare LoRA arguments lora_request_distribution = ( args.lora_request_distribution if args.lora_name is not None else None ) lora_zipf_alpha = ( args.lora_zipf_alpha if args.lora_name is not None and args.lora_request_distribution == "skewed" else None ) return asyncio.run( benchmark( backend=backend, api_url=api_url, base_url=base_url, model_id=model_id, tokenizer=tokenizer, input_requests=input_requests, request_rate=args.request_rate, max_concurrency=args.max_concurrency, disable_tqdm=args.disable_tqdm, lora_names=args.lora_name, lora_request_distribution=lora_request_distribution, lora_zipf_alpha=lora_zipf_alpha, extra_request_body=extra_request_body, profile=args.profile, pd_separated=args.pd_separated, flush_cache=args.flush_cache, warmup_requests=args.warmup_requests, use_trace_timestamps=args.use_trace_timestamps, mooncake_slowdown_factor=args.mooncake_slowdown_factor, mooncake_num_rounds=args.mooncake_num_rounds, profile_prefill_url=getattr(args, "profile_prefill_url", None), profile_decode_url=getattr(args, "profile_decode_url", None), ) ) def _finite_positive_float(value) -> float: """argparse type for a finite, strictly positive float.""" try: parsed = float(value) except (TypeError, ValueError) as exc: raise argparse.ArgumentTypeError( f"expected a finite float > 0, got {value!r}" ) from exc if not math.isfinite(parsed) or parsed <= 0: raise argparse.ArgumentTypeError(f"expected a finite float > 0, got {value!r}") return parsed def _validate_parsed_gsp_args( parser: argparse.ArgumentParser, args: argparse.Namespace ) -> None: """Reject malformed GSP distribution/alpha combinations at parse time. Invoked from the CLI entry point right after ``parser.parse_args()`` so users see a clear argparse-style error before any server, model, or tokenizer setup runs and masks the real cause with an unrelated network failure. """ distribution = getattr(args, "gsp_group_distribution", None) alpha = getattr(args, "gsp_zipf_alpha", None) if distribution == "zipf" and alpha is None: parser.error( "--gsp-group-distribution=zipf requires --gsp-zipf-alpha " "(a finite float > 0)" ) if distribution == "uniform" and alpha is not None: parser.error( "--gsp-zipf-alpha is only meaningful with " "--gsp-group-distribution=zipf; remove --gsp-zipf-alpha " "or set --gsp-group-distribution=zipf" ) class LoRAPathAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, []) for lora_name in values: getattr(namespace, self.dest).append(lora_name) def cli_main(): parser = ArgumentParser(description="Benchmark the online serving throughput.") parser.add_argument( "--backend", type=str, choices=list(ASYNC_REQUEST_FUNCS.keys()), default="sglang", help="Must specify a backend, depending on the LLM Inference Engine.", ) parser.add_argument( "--base-url", type=str, default=None, help="Server or API base url if not using http host and port.", ) parser.add_argument( "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." ) parser.add_argument( "--port", type=int, help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", ) parser.add_argument( "--ready-check-timeout-sec", type=int, default=60, help="Maximum time in seconds to wait for the server to be ready before benchmarking. Set to 0 to skip. Default: 60.", ) parser.add_argument( "--dataset-name", type=str, default="sharegpt", choices=[ "agentic-trace", "autobench", "sharegpt", "custom", "openai", "random", "random-ids", "generated-shared-prefix", "mmmu", "image", "mooncake", "longbench_v2", "speed-bench", ], help="Name of the dataset to benchmark on.", ) parser.add_argument( "--dataset-path", type=str, default="", help="Path to the dataset." ) parser.add_argument( "--dataset-offset", type=int, default=0, help="Rotate the conversation list by this many entries before sampling " "(agentic-trace dataset), so successive sweep steps start on fresh " "conversations.", ) parser.add_argument( "--agentic-max-turns", type=int, default=None, help="Cap each conversation to at most this many turns (agentic-trace " "dataset). Default: use all turns in the trace.", ) parser.add_argument( "--speed-bench-category", type=str, default=None, choices=["low_entropy", "mixed", "high_entropy"], help="Category filter for the speed-bench dataset.", ) parser.add_argument( "--speed-bench-output-len", type=int, default=512, help="Fixed output length for speed-bench requests (default: 512).", ) parser.add_argument( "--model", type=str, help="Name or path of the model. If not set, the default model will request /v1/models for conf.", ) parser.add_argument( "--served-model-name", type=str, help="The name of the model as served by the serving service. If not set, this defaults to the value of --model.", ) parser.add_argument( "--tokenizer", type=str, help="Name or path of the tokenizer. If not set, using the model conf.", ) parser.add_argument( "--num-prompts", type=int, default=1000, help="Number of prompts to process. Default is 1000.", ) parser.add_argument( "--sharegpt-output-len", type=int, default=None, help="Output length for each request. Overrides the output length from the ShareGPT dataset.", ) parser.add_argument( "--sharegpt-context-len", type=int, default=None, help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", ) parser.add_argument( "--random-input-len", type=int, default=1024, help="Number of input tokens per request, used only for random and image dataset.", ) parser.add_argument( "--random-output-len", default=1024, type=int, help="Number of output tokens per request, used only for random and image dataset.", ) parser.add_argument( "--random-range-ratio", type=float, default=0.0, help="Range of sampled ratio of input/output length, " "used only for random and image dataset.", ) # image dataset args parser.add_argument( "--image-count", type=int, default=1, help="Number of images per request (only available with the image dataset)", ) parser.add_argument( "--image-resolution", type=str, default="1080p", help=( "Resolution of images for image dataset. " "Supports presets 4k/1080p/720p/360p, custom 'heightxwidth' " "(e.g., 1080x1920), or random 'random:x-x' " "bounds (e.g., random:256x256-1024x1024)." ), ) parser.add_argument( "--random-image-count", action="store_true", help="Enable Random Image Count", ) parser.add_argument( "--image-format", type=str, default="jpeg", help=("Format of images for image dataset. " "Supports jpeg and png."), ) parser.add_argument( "--image-content", type=str, default="random", help=("Content for images for image dataset. " "Supports random and blank."), ) parser.add_argument( "--request-rate", type=float, default=float("inf"), help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", ) parser.add_argument( "--use-trace-timestamps", action="store_true", help="Use timestamps from the trace file for request scheduling. Only valid for 'mooncake' dataset.", ) parser.add_argument( "--max-concurrency", type=int, default=None, help="Maximum number of concurrent requests. This can be used " "to help simulate an environment where a higher level component " "is enforcing a maximum number of concurrent requests. While the " "--request-rate argument controls the rate at which requests are " "initiated, this argument will control how many are actually allowed " "to execute at a time. This means that when used in combination, the " "actual request rate may be lower than specified with --request-rate, " "if the server is not processing requests fast enough to keep up.", ) parser.add_argument("--output-file", type=str, help="Output JSONL file name.") parser.add_argument( "--output-details", action="store_true", help="Output details of benchmarking." ) parser.add_argument( "--print-requests", action="store_true", help="Print requests immediately during benchmarking. Useful to quickly realize issues.", ) parser.add_argument( "--disable-tqdm", action="store_true", help="Specify to disable tqdm progress bar.", ) parser.add_argument( "--disable-stream", action="store_true", help="Disable streaming mode.", ) parser.add_argument( "--return-logprob", action="store_true", help="Return logprob.", ) parser.add_argument( "--top-logprobs-num", type=int, default=0, help="Number of top logprobs to return per token. Only used with --return-logprob.", ) parser.add_argument( "--token-ids-logprob", type=int, nargs="+", default=None, help="Token IDs to probe logprobs for. E.g. --token-ids-logprob 1 2 10 100 1000. Only used with --return-logprob.", ) parser.add_argument( "--logprob-start-len", type=int, default=-1, help="Start position for returning input logprobs. -1 means no input logprobs, 0 means all. Only used with --return-logprob.", ) parser.add_argument( "--return-routed-experts", action="store_true", help="Return routed experts.", ) parser.add_argument( "--cache-report", action="store_true", help="Collect and display cache hit statistics after the benchmark. " "Supported with sglang backends (native, oai, oai-chat).", ) parser.add_argument("--seed", type=int, default=42, help="The random seed.") parser.add_argument( "--disable-ignore-eos", action="store_true", help="Disable ignoring EOS.", ) parser.add_argument( "--temperature", type=float, default=0.0, help="Sampling temperature.", ) parser.add_argument( "--top-p", type=float, default=1.0, help="Nucleus sampling parameter.", ) parser.add_argument( "--extra-request-body", metavar='{"key1": "value1", "key2": "value2"}', type=str, help="Append given JSON object to the request payload. You can use this to specify" "additional generate params like sampling params.", ) parser.add_argument( "--apply-chat-template", action="store_true", help="Apply chat template", ) parser.add_argument( "--profile", action="store_true", help="Use Torch Profiler. The endpoint must be launched with " "SGLANG_TORCH_PROFILER_DIR to enable profiler.", ) parser.add_argument( "--plot-throughput", action="store_true", help="Plot throughput and concurrent requests over time. Requires termplotlib and gnuplot.", ) # TODO unify all these parser.add_argument( "--profile-activities", type=str, nargs="+", default=["CPU", "GPU"], choices=["CPU", "GPU", "CUDA_PROFILER", "XPU", "MEM"], help="Profiler activities to capture: CPU, GPU, XPU, CUDA_PROFILER, MEM " "(MEM dumps a torch.cuda.memory snapshot, viewable at https://pytorch.org/memory_viz).", ) parser.add_argument( "--profile-start-step", type=int, default=None, help="Start profiling after this many forward steps. Useful for warmup.", ) parser.add_argument( "--profile-steps", type=int, default=None, help="Number of steps to profile. If specified, profiling stops automatically after this many steps.", ) parser.add_argument("--profile-num-steps", type=int, default=None) parser.add_argument("--profile-by-stage", action="store_true", default=False) parser.add_argument("--profile-stages", nargs="+", default=None) parser.add_argument( "--profile-output-dir", type=str, default=None, help="Output directory for profile traces.", ) parser.add_argument( "--profile-prefix", type=str, default=None, help="Prefix for profile trace filenames.", ) parser.add_argument( "--lora-name", type=str, nargs="*", default=None, action=LoRAPathAction, help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...", ) parser.add_argument( "--lora-request-distribution", type=str, default="uniform", choices=[ "uniform", "distinct", "skewed", ], help="What distribution to sample the LoRA adapters specified in --lora-name. Borrowed from the Punica paper. " "'distinct' distribution means selecting a new LoRA adapter for every request. " "'skewed' distribution follows the Zipf distribution, where the number of requests " "to model i specified in --lora-name is α times the number of requests for model i+1, " "where α > 1.", ) parser.add_argument( "--lora-zipf-alpha", type=float, default=1.5, help="The parameter to use for the Zipf distribution when --lora-request-distribution='skewed'.", ) parser.add_argument( "--prompt-suffix", type=str, default="", help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", ) parser.add_argument( "--pd-separated", action="store_true", help="Benchmark PD disaggregation server", ) # Create a mutually exclusive group for profiling URLs # In PD separated mode, prefill and decode workers must be profiled separately profile_url_group = parser.add_mutually_exclusive_group() profile_url_group.add_argument( "--profile-prefill-url", type=str, nargs="*", default=None, help="URL(s) of the prefill worker(s) for profiling in PD separated mode. " "Can specify multiple URLs: --profile-prefill-url http://localhost:30000 http://localhost:30001. " "NOTE: Cannot be used together with --profile-decode-url. " "In PD separated mode, prefill and decode workers must be profiled separately.", ) profile_url_group.add_argument( "--profile-decode-url", type=str, nargs="*", default=None, help="URL(s) of the decode worker(s) for profiling in PD separated mode. " "Can specify multiple URLs: --profile-decode-url http://localhost:30010 http://localhost:30011. " "NOTE: Cannot be used together with --profile-prefill-url. " "In PD separated mode, prefill and decode workers must be profiled separately.", ) parser.add_argument( "--flush-cache", action="store_true", help="Flush the cache before running the benchmark", ) parser.add_argument( "--warmup-requests", type=int, default=1, help="Number of warmup requests to run before the benchmark", ) parser.add_argument( "--tokenize-prompt", action="store_true", help="Use integer ids instead of string for inputs. Useful to control prompt lengths accurately", ) group = parser.add_argument_group("generated-shared-prefix dataset arguments") group.add_argument( "--gsp-num-groups", type=int, default=64, help="Number of system prompt groups for generated-shared-prefix dataset", ) group.add_argument( "--gsp-prompts-per-group", type=int, default=16, help="Number of prompts per system prompt group for generated-shared-prefix dataset", ) group.add_argument( "--gsp-system-prompt-len", type=int, default=2048, help="Target length in tokens for system prompts in generated-shared-prefix dataset", ) group.add_argument( "--gsp-question-len", type=int, default=128, help="Target length in tokens for questions in generated-shared-prefix dataset", ) group.add_argument( "--gsp-output-len", type=int, default=256, help="Target length in tokens for outputs in generated-shared-prefix dataset", ) parser.add_argument( "--gsp-range-ratio", type=float, # WARN: The default 1.0 is for backward compatibility, and is different from the default 0.0 for random dataset default=1.0, help="Range of sampled ratio of input/output length, used only for gsp dataset.", ) group.add_argument( "--gsp-fast-prepare", action="store_true", help="Speedup preparing by removing statistics computation, which will make some output statistics inaccurate but suitable for pressure tests.", ) group.add_argument( "--gsp-send-routing-key", 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.", ) group.add_argument( "--gsp-ordered", action="store_true", help="Keep requests in order without shuffling. By default, requests are shuffled randomly.", ) group.add_argument( "--gsp-group-distribution", type=str, choices=["uniform", "zipf"], default="uniform", help=( "Prefix-group sampling distribution for generated-shared-prefix. " "'uniform' (default) assigns each group an equal number of requests. " "'zipf' samples each request's group by rank with " "p(rank) = (1/rank**alpha) / sum_k(1/k**alpha); rank starts at 1 " "and group index 0 is the hottest. Requires --gsp-zipf-alpha " "(a finite float > 0) when set to 'zipf'. Total request count is " "still num_groups * prompts_per_group, identical to uniform mode; " "only the per-request group assignment changes. The on-disk " "dataset cache uses a distinct key per (group_distribution, " "zipf_alpha), so uniform-mode caches are never mixed with " "zipf-mode caches and zipf runs with different alpha use " "separate files." ), ) group.add_argument( "--gsp-zipf-alpha", type=_finite_positive_float, default=None, help=( "Zipf exponent alpha for --gsp-group-distribution=zipf, with " "p(rank) = (1/rank**alpha) / sum_k(1/k**alpha) and rank starting " "at 1. Must be a finite float strictly greater than 0; larger " "values concentrate requests on lower-ranked (hotter) groups." ), ) mooncake_group = parser.add_argument_group("mooncake dataset arguments") mooncake_group.add_argument( "--mooncake-slowdown-factor", type=float, default=1.0, help="Slowdown factor for replaying the mooncake trace. " "A value of 2.0 means the replay is twice as slow. " "NOTE: --request-rate is IGNORED in mooncake mode.", ) mooncake_group.add_argument( "--mooncake-num-rounds", type=int, default=1, help="Number of conversation rounds for each session in the mooncake dataset. " "A value > 1 will enable true multi-turn session benchmarking.", ) mooncake_group.add_argument( "--mooncake-workload", type=str, default="conversation", choices=[ "mooncake", "conversation", "synthetic", "toolagent", ], help="Underlying workload for the mooncake dataset.", ) parser.add_argument( "--fake-prefill", action="store_true", default=False, help="Enable fake prefill mode for decode-only benchmarking. " "Use with a decode server running --disaggregation-transfer-backend fake " "to benchmark pure decode performance without a real prefill node.", ) parser.add_argument( "--tag", type=str, default=None, help="The tag to be dumped to output." ) parser.add_argument( "--header", type=str, nargs="+", default=None, help="Custom HTTP headers in Key=Value format. Example: --header MyHeader=MY_VALUE MyAnotherHeader=myanothervalue", ) args = parser.parse_args() _validate_parsed_gsp_args(parser, args) run_benchmark(args) if __name__ == "__main__": cli_main()