sgl-project--sglang/benchmark/hicache/bench_warm_cache.py

# 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())