diff --git a/benchmark/hicache/bench_warm_cache.py b/benchmark/hicache/bench_warm_cache.py new file mode 100644 index 000000000..f0f266633 --- /dev/null +++ b/benchmark/hicache/bench_warm_cache.py @@ -0,0 +1,673 @@ +# Adapted from benchmark/hicache/bench_serving.py and python/sglang/bench_serving.py + +""" +Benchmark warm-cache serving with exact shared-prefix control. + +This benchmark is designed for cache-focused studies where each request has a +fixed total input length and an exactly controlled shared-prefix ratio. For each +shared-prefix percentage, the benchmark: + +1. Flushes the server KV cache. +2. Builds prompts with an identical shared prefix and random unique suffixes. +3. Warms only the shared prefix once. +4. Benchmarks the full prompts through SGLang's native /generate endpoint. + +Compared with the existing hicache shared-prefix benchmarks, this benchmark +provides direct control over total length, shared-prefix length, and suffix +length at the token-id level. +""" + +import argparse +import asyncio +import json +import random +import time +import warnings +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import aiohttp +import numpy as np +import requests +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.utils import get_tokenizer, remove_prefix, set_ulimit + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) +AIOHTTP_READ_BUFSIZE = 10 * 1024**2 + + +global args + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + ttft: float = 0.0 + itl: List[float] = field(default_factory=list) + prompt_len: int = 0 + error: str = "" + output_len: int = 0 + start_time: float = 0.0 + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + total_output_retokenized: int + request_throughput: float + input_throughput: float + output_throughput: float + output_throughput_retokenized: float + total_throughput: float + total_throughput_retokenized: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + p90_ttft_ms: float + p99_ttft_ms: float + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + p90_tpot_ms: float + p99_tpot_ms: float + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + p90_itl_ms: float + p99_itl_ms: float + mean_e2e_latency_ms: float + median_e2e_latency_ms: float + std_e2e_latency_ms: float + p99_e2e_latency_ms: float + concurrency: float + + +def _create_bench_client_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession( + timeout=AIOHTTP_TIMEOUT, + read_bufsize=AIOHTTP_READ_BUFSIZE, + ) + + +async def async_request_sglang_generate( + api_url: str, + input_ids: List[int], + prompt_len: int, + output_len: int, + pbar: Optional[Any] = None, +) -> RequestFuncOutput: + async with _create_bench_client_session() as session: + payload = { + "input_ids": input_ids, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": output_len, + "ignore_eos": not args.disable_ignore_eos, + }, + "stream": True, + **args.extra_request_body, + } + + output = RequestFuncOutput(prompt_len=prompt_len) + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + output.start_time = st + most_recent_timestamp = st + last_output_len = 0 + latency = 0.0 + + 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: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + continue + + data = json.loads(chunk) + + if "text" in data and data["text"]: + timestamp = time.perf_counter() + generated_text = data["text"] + current_output_len = data["meta_info"]["completion_tokens"] + + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + else: + num_new_tokens = current_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 = current_output_len + output.output_len = current_output_len + + output.generated_text = generated_text + output.success = True + output.latency = latency + else: + output.error = ( + (response.reason or "") + ": " + (await response.text()) + ) + output.success = False + except Exception as exc: + output.success = False + output.error = str(exc) + + if pbar: + pbar.update(1) + return output + + +async def run_batch( + api_url: str, + prompts: List[Dict[str, Any]], + output_len: int, + max_concurrency: Optional[int], + pbar: Optional[Any] = None, +) -> List[RequestFuncOutput]: + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def limited_request(prompt: Dict[str, Any]) -> RequestFuncOutput: + if semaphore is None: + return await async_request_sglang_generate( + api_url=api_url, + input_ids=prompt["input_ids"], + prompt_len=prompt["prompt_len"], + output_len=output_len, + pbar=pbar, + ) + async with semaphore: + return await async_request_sglang_generate( + api_url=api_url, + input_ids=prompt["input_ids"], + prompt_len=prompt["prompt_len"], + output_len=output_len, + pbar=pbar, + ) + + tasks = [asyncio.create_task(limited_request(prompt)) for prompt in prompts] + return await asyncio.gather(*tasks) + + +def flush_cache(base_url: str) -> None: + response = requests.post(f"{base_url}/flush_cache", timeout=30) + response.raise_for_status() + + +def gen_token_ids( + vocab_ids: List[int], + token_num: int, + rng: random.Random, +) -> List[int]: + if token_num <= 0: + return [] + return rng.choices(vocab_ids, k=token_num) + + +def build_prompts( + vocab_ids: List[int], + total_tokens: int, + shared_pct: int, + num_prompts: int, + rng: random.Random, +) -> List[Dict[str, Any]]: + prefix_len = total_tokens * shared_pct // 100 + suffix_len = total_tokens - prefix_len + + shared_prefix = gen_token_ids(vocab_ids, prefix_len, rng) + + prompts: List[Dict[str, Any]] = [] + for _ in range(num_prompts): + suffix = gen_token_ids(vocab_ids, suffix_len, rng) + input_ids = shared_prefix + suffix + prompts.append({"input_ids": input_ids, "prompt_len": len(input_ids)}) + + return prompts + + +async def warm_shared_prefix(api_url: str, shared_prefix_ids: List[int]) -> None: + if not shared_prefix_ids: + return + + warmup = await async_request_sglang_generate( + api_url=api_url, + input_ids=shared_prefix_ids, + prompt_len=len(shared_prefix_ids), + output_len=1, + ) + if not warmup.success: + raise RuntimeError( + "Warmup failed - Please make sure benchmark arguments are correctly " + f"specified. Error: {warmup.error}" + ) + + +def calculate_metrics( + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, +) -> Tuple[BenchmarkMetrics, List[int]]: + output_lens: List[int] = [] + retokenized_output_lens: List[int] = [] + total_input = 0 + completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + ttfts: List[float] = [] + e2e_latencies: List[float] = [] + + for output in outputs: + if output.success: + output_len = output.output_len + output_lens.append(output_len) + retokenized_output_len = len( + tokenizer.encode(output.generated_text, add_special_tokens=False) + ) + retokenized_output_lens.append(retokenized_output_len) + total_input += output.prompt_len + if output_len > 1: + tpots.append((output.latency - output.ttft) / (output_len - 1)) + itls += output.itl + ttfts.append(output.ttft) + e2e_latencies.append(output.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, + ) + + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + 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, + 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, + 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, + 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, + p99_itl_ms=np.percentile(itls or 0, 99) * 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, + p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, + concurrency=np.sum(e2e_latencies) / dur_s, + ) + return metrics, output_lens + + +def print_benchmark_result( + metrics: BenchmarkMetrics, + benchmark_duration: float, + backend: str, + request_rate: float, + max_concurrency: Optional[int], +) -> None: + 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:", 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 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 + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", metrics.total_throughput + ) + ) + print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) + 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("{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("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("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("P99 ITL (ms):", metrics.p99_itl_ms)) + print("=" * 50) + + +def maybe_write_summary_jsonl( + pct: int, + prefix_len: int, + suffix_len: int, + metrics: BenchmarkMetrics, + output_file: Optional[str], + benchmark_duration: float, +) -> None: + if not output_file: + return + + result = { + "backend": args.backend, + "dataset_name": "warm-cache", + "request_rate": float("inf"), + "max_concurrency": args.max_concurrency, + "shared_prefix_pct": pct, + "prefix_len": prefix_len, + "suffix_len": suffix_len, + "total_tokens": args.total_tokens, + "num_prompts": args.num_prompts, + "output_len": args.output_len, + "completed": metrics.completed, + "benchmark_duration": benchmark_duration, + "total_input": metrics.total_input, + "total_output": metrics.total_output, + "total_output_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "output_throughput_retokenized": metrics.output_throughput_retokenized, + "total_throughput": metrics.total_throughput, + "total_throughput_retokenized": metrics.total_throughput_retokenized, + "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, + "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, + "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, + "p99_itl_ms": metrics.p99_itl_ms, + "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, + "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, + "concurrency": metrics.concurrency, + } + + with open(output_file, "a", encoding="utf-8") as fout: + fout.write(json.dumps(result) + "\n") + + +async def benchmark_shared_prefix_pct( + api_url: str, + base_url: str, + tokenizer: PreTrainedTokenizerBase, + vocab_ids: List[int], + rng: random.Random, + pct: int, +) -> Tuple[BenchmarkMetrics, float, int, int, int]: + prefix_len = args.total_tokens * pct // 100 + suffix_len = args.total_tokens - prefix_len + + print(f"\n{'=' * 70}") + print( + f"shared_prefix={pct}% prefix_len={prefix_len} " + f"suffix_len={suffix_len} total={prefix_len + suffix_len}" + ) + print(f"{'=' * 70}") + + print("Flushing KV cache ...") + flush_cache(base_url) + time.sleep(1) + + print(f"Building {args.num_prompts} prompts ...") + prompts = build_prompts( + vocab_ids=vocab_ids, + total_tokens=args.total_tokens, + shared_pct=pct, + num_prompts=args.num_prompts, + rng=rng, + ) + + if prefix_len > 0: + print(f"Warming shared prefix only ({prefix_len} tokens) ...") + await warm_shared_prefix( + api_url=api_url, shared_prefix_ids=prompts[0]["input_ids"][:prefix_len] + ) + + print(f"Sending requests (max_concurrency={args.max_concurrency}) ...") + benchmark_start_time = time.perf_counter() + outputs = await run_batch( + api_url=api_url, + prompts=prompts, + output_len=args.output_len, + max_concurrency=args.max_concurrency, + pbar=None, + ) + benchmark_duration = time.perf_counter() - benchmark_start_time + + failed_outputs = [output for output in outputs if not output.success] + if failed_outputs: + print(f"WARNING: {len(failed_outputs)}/{len(outputs)} requests failed") + for output in failed_outputs[:5]: + print(f" {output.error[:160]}") + + metrics, _ = calculate_metrics( + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + ) + + if metrics.completed == 0: + raise RuntimeError("All requests failed for this shared-prefix percentage.") + + print_benchmark_result( + metrics=metrics, + benchmark_duration=benchmark_duration, + backend=args.backend, + request_rate=float("inf"), + max_concurrency=args.max_concurrency, + ) + + maybe_write_summary_jsonl( + pct=pct, + prefix_len=prefix_len, + suffix_len=suffix_len, + metrics=metrics, + output_file=args.output_file, + benchmark_duration=benchmark_duration, + ) + + return metrics, benchmark_duration, prefix_len, suffix_len, len(outputs) + + +async def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--backend", + type=str, + default="sglang", + choices=["sglang"], + help="Warm-cache benchmark currently supports the native SGLang /generate endpoint.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server base url if not using host and port.", + ) + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + parser.add_argument( + "--model", + type=str, + required=True, + help="Name or path of the model. Used to load the tokenizer and vocab ids.", + ) + parser.add_argument( + "--tokenizer", + type=str, + default=None, + help="Name or path of the tokenizer. Defaults to --model.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=64, + help="Number of prompts to process per shared-prefix percentage.", + ) + parser.add_argument( + "--total-tokens", + type=int, + default=70000, + help="Total input tokens per request (shared prefix + unique suffix).", + ) + parser.add_argument( + "--output-len", + type=int, + default=200, + help="Output length for each request.", + ) + parser.add_argument( + "--max-concurrency", + type=int, + default=4, + help="Maximum number of concurrent requests.", + ) + parser.add_argument( + "--pcts", + type=str, + default="0,10,20,30,40,50,60,70,80,90,92,95,97,99", + help="Comma-separated shared-prefix percentages to sweep.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for synthetic prompt generation.", + ) + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignoring EOS.", + ) + parser.add_argument( + "--output-file", + type=str, + default=None, + help="Optional JSONL file to append one result object per shared-prefix percentage.", + ) + 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.", + ) + global args + args = parser.parse_args() + + args.extra_request_body = ( + json.loads(args.extra_request_body) if args.extra_request_body else {} + ) + + base_url = args.base_url or f"http://{args.host}:{args.port}" + api_url = f"{base_url}/generate" + pcts = [int(p.strip()) for p in args.pcts.split(",") if p.strip()] + rng = random.Random(args.seed) + + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer = get_tokenizer(tokenizer_id) + vocab_ids = list(tokenizer.get_vocab().values()) + + print(f"{args}\n") + print(f"Loading tokenizer from {tokenizer_id} ...") + print(f"Tokenizer loaded (vocab_size={len(vocab_ids)})") + + for pct in pcts: + await benchmark_shared_prefix_pct( + api_url=api_url, + base_url=base_url, + tokenizer=tokenizer, + vocab_ids=vocab_ids, + rng=rng, + pct=pct, + ) + + if args.output_file: + print(f"JSONL results saved to {args.output_file}") + + +if __name__ == "__main__": + set_ulimit() + asyncio.run(main())