Add benchmark/hicache/bench_warm_cache.py for exact warm-cache shared-prefix benchmarking (#24083)
This commit is contained in:
Clint
@@ -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())
|
||||
Reference in New Issue
Block a user