chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
193
third_party/sglang/benchmark/benchmark_batch/benchmark_batch.py
vendored
Normal file
193
third_party/sglang/benchmark/benchmark_batch/benchmark_batch.py
vendored
Normal file
@@ -0,0 +1,193 @@
|
||||
import concurrent.futures
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
from statistics import mean
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
|
||||
|
||||
###############################################################################
|
||||
# CONFIG
|
||||
###############################################################################
|
||||
ENDPOINT_URL = "http://127.0.0.1:30000"
|
||||
TOKENIZER_DIR = "/models/meta-llama/Llama-3.2-3B"
|
||||
|
||||
# Benchmark configurations
|
||||
NUM_REQUESTS = 10 # Total number of requests (each with BATCH_SIZE prompts)
|
||||
NUM_TOKENS = 32000 # Tokens per prompt
|
||||
BATCH_SIZE = 8 # Number of prompts per request
|
||||
GEN_TOKENS = 0 # Tokens to generate per prompt
|
||||
|
||||
|
||||
###############################################################################
|
||||
# REQUEST GENERATION (in parallel)
|
||||
###############################################################################
|
||||
def generate_random_prompt(index, tokenizer_dir, num_tokens):
|
||||
"""Generate a single random prompt with specified token count."""
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
|
||||
vocab_size = tokenizer.vocab_size
|
||||
|
||||
def generate_random_text(num_toks):
|
||||
random_token_ids = [random.randint(0, vocab_size - 1) for _ in range(num_toks)]
|
||||
return tokenizer.decode(random_token_ids, clean_up_tokenization_spaces=True)
|
||||
|
||||
random_text = generate_random_text(num_tokens)
|
||||
return f"Prompt {index}: {random_text}"
|
||||
|
||||
|
||||
def prepare_all_prompts(num_requests, batch_size, num_tokens, tokenizer_dir):
|
||||
"""Generate prompts for all requests in parallel."""
|
||||
total_prompts = num_requests * batch_size
|
||||
all_prompts = [None] * total_prompts
|
||||
max_workers = min(os.cpu_count() or 1, total_prompts)
|
||||
|
||||
with ProcessPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(generate_random_prompt, i, tokenizer_dir, num_tokens)
|
||||
for i in range(total_prompts)
|
||||
]
|
||||
for future in tqdm(
|
||||
concurrent.futures.as_completed(futures),
|
||||
total=total_prompts,
|
||||
desc="Generating prompts",
|
||||
):
|
||||
index = futures.index(future)
|
||||
all_prompts[index] = future.result()
|
||||
|
||||
batched_prompts = [
|
||||
all_prompts[i * batch_size : (i + 1) * batch_size] for i in range(num_requests)
|
||||
]
|
||||
|
||||
print(
|
||||
f"Generated {total_prompts} prompts with {num_tokens} tokens each, grouped into {num_requests} requests of {batch_size} prompts.\n"
|
||||
)
|
||||
return batched_prompts
|
||||
|
||||
|
||||
###############################################################################
|
||||
# HTTP CALLS
|
||||
###############################################################################
|
||||
def send_batch_request(endpoint, prompts, gen_tokens, request_id):
|
||||
"""Send a batch of prompts to the /generate endpoint synchronously."""
|
||||
sampling_params = {
|
||||
"max_new_tokens": gen_tokens,
|
||||
"temperature": 0.7,
|
||||
"stop": "\n",
|
||||
}
|
||||
data = {"text": prompts, "sampling_params": sampling_params}
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
response = requests.post(
|
||||
endpoint.base_url + "/generate", json=data, timeout=3600
|
||||
)
|
||||
if response.status_code != 200:
|
||||
error = response.json()
|
||||
raise RuntimeError(f"Request {request_id} failed: {error}")
|
||||
result = response.json()
|
||||
elapsed_time = (time.perf_counter() - start_time) * 1000 # Convert to ms
|
||||
avg_per_prompt = elapsed_time / len(prompts) if prompts else 0
|
||||
return request_id, elapsed_time, avg_per_prompt, True, len(prompts)
|
||||
except Exception as e:
|
||||
print(f"[Request] Error for request {request_id}: {e}")
|
||||
return request_id, 0, 0, False, len(prompts)
|
||||
|
||||
|
||||
def run_benchmark(endpoint, batched_prompts, batch_size, gen_tokens):
|
||||
"""Run the benchmark sequentially."""
|
||||
results = []
|
||||
num_requests = len(batched_prompts)
|
||||
|
||||
# Record start time for total latency
|
||||
benchmark_start_time = time.perf_counter()
|
||||
|
||||
for i, batch_prompts in enumerate(batched_prompts):
|
||||
request_id = i + 1
|
||||
assert (
|
||||
len(batch_prompts) == batch_size
|
||||
), f"Request {request_id} should have {batch_size} prompts, got {len(batch_prompts)}"
|
||||
|
||||
print(
|
||||
f"[Request] Sending request {request_id}/{num_requests} with {len(batch_prompts)} prompts at {int(time.time()*1000)}"
|
||||
)
|
||||
result = send_batch_request(endpoint, batch_prompts, gen_tokens, request_id)
|
||||
results.append(result)
|
||||
|
||||
# Calculate total latency
|
||||
total_latency = (time.perf_counter() - benchmark_start_time) * 1000 # Convert to ms
|
||||
|
||||
return results, total_latency
|
||||
|
||||
|
||||
###############################################################################
|
||||
# RESULTS
|
||||
###############################################################################
|
||||
def process_results(results, total_latency, num_requests):
|
||||
"""Process and display benchmark results."""
|
||||
total_time = 0
|
||||
successful_requests = 0
|
||||
failed_requests = 0
|
||||
request_latencies = []
|
||||
per_prompt_latencies = []
|
||||
total_prompts = 0
|
||||
|
||||
for request_id, elapsed_time, avg_per_prompt, success, batch_size in results:
|
||||
if success:
|
||||
successful_requests += 1
|
||||
total_prompts += batch_size
|
||||
request_latencies.append(elapsed_time)
|
||||
per_prompt_latencies.append(avg_per_prompt)
|
||||
total_time += elapsed_time / 1000 # Convert to seconds
|
||||
else:
|
||||
failed_requests += 1
|
||||
|
||||
avg_request_latency = mean(request_latencies) if request_latencies else 0
|
||||
avg_per_prompt_latency = mean(per_prompt_latencies) if per_prompt_latencies else 0
|
||||
throughput = total_prompts / total_time if total_time > 0 else 0
|
||||
|
||||
print("\nBenchmark Summary:")
|
||||
print(f" Total requests sent: {len(results)}")
|
||||
print(f" Total prompts sent: {total_prompts}")
|
||||
print(f" Successful requests: {successful_requests}")
|
||||
print(f" Failed requests: {failed_requests}")
|
||||
print(f" Total latency (all requests): {total_latency:.2f} ms")
|
||||
print(f" Avg per request latency: {avg_request_latency:.2f} ms")
|
||||
print(f" Avg per prompt latency: {avg_per_prompt_latency:.2f} ms")
|
||||
print(f" Throughput: {throughput:.2f} prompts/second\n")
|
||||
|
||||
|
||||
###############################################################################
|
||||
# MAIN
|
||||
###############################################################################
|
||||
def main():
|
||||
# Initialize endpoint
|
||||
endpoint = RuntimeEndpoint(ENDPOINT_URL)
|
||||
|
||||
# Generate prompts
|
||||
batched_prompts = prepare_all_prompts(
|
||||
NUM_REQUESTS, BATCH_SIZE, NUM_TOKENS, TOKENIZER_DIR
|
||||
)
|
||||
|
||||
# Flush cache before benchmark
|
||||
# endpoint.flush_cache()
|
||||
|
||||
# Run benchmark
|
||||
print(
|
||||
f"Starting benchmark: NUM_TOKENS={NUM_TOKENS}, BATCH_SIZE={BATCH_SIZE}, NUM_REQUESTS={NUM_REQUESTS}\n"
|
||||
)
|
||||
results, total_latency = run_benchmark(
|
||||
endpoint, batched_prompts, BATCH_SIZE, GEN_TOKENS
|
||||
)
|
||||
|
||||
# Process and display results
|
||||
process_results(results, total_latency, NUM_REQUESTS)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
random.seed(0)
|
||||
main()
|
||||
237
third_party/sglang/benchmark/benchmark_batch/benchmark_tokenizer.py
vendored
Normal file
237
third_party/sglang/benchmark/benchmark_batch/benchmark_tokenizer.py
vendored
Normal file
@@ -0,0 +1,237 @@
|
||||
import argparse
|
||||
import random
|
||||
import time
|
||||
from statistics import mean
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from sglang.srt.utils.patch_tokenizer import patch_tokenizer
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
print("Tokenizer Benchmark: Sequential vs Batch Processing")
|
||||
print("-" * 60)
|
||||
print(f"Tokenizer: {args.tokenizer}")
|
||||
print(f"Functions: {', '.join(args.function)}")
|
||||
print(f"Tokens per prompt: {args.num_tokens}")
|
||||
print(f"Number of runs per batch size: {args.num_runs}")
|
||||
print(f"Batch mode: {', '.join(args.batch_mode)}")
|
||||
print("-" * 60)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
|
||||
tokenizer = patch_tokenizer(tokenizer)
|
||||
max_batch_size = max(args.batch_sizes)
|
||||
|
||||
token_ids = generate_random_token_ids(
|
||||
num_prompts=max_batch_size, num_tokens=args.num_tokens, tokenizer=tokenizer
|
||||
)
|
||||
|
||||
if "encode" in args.function:
|
||||
prompts = [
|
||||
tokenizer.decode(ids, clean_up_tokenization_spaces=True)
|
||||
for ids in token_ids
|
||||
]
|
||||
run_benchmark(
|
||||
name="encode",
|
||||
data=prompts,
|
||||
sequential_fn=lambda batch: [tokenizer.encode(p) for p in batch],
|
||||
batch_fn=lambda batch: tokenizer(batch),
|
||||
batch_sizes=args.batch_sizes,
|
||||
num_runs=args.num_runs,
|
||||
batch_mode=args.batch_mode,
|
||||
)
|
||||
|
||||
if "decode" in args.function:
|
||||
# mimic DetokenizerManager's usual case
|
||||
decode_kwargs = dict(
|
||||
skip_special_tokens=True,
|
||||
spaces_between_special_tokens=True,
|
||||
)
|
||||
run_benchmark(
|
||||
name="decode",
|
||||
data=token_ids,
|
||||
sequential_fn=lambda batch: [
|
||||
tokenizer.decode(ids, **decode_kwargs) for ids in batch
|
||||
],
|
||||
batch_fn=lambda batch: tokenizer.batch_decode(batch, **decode_kwargs),
|
||||
batch_sizes=args.batch_sizes,
|
||||
num_runs=args.num_runs,
|
||||
batch_mode=args.batch_mode,
|
||||
)
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
*, name, data, sequential_fn, batch_fn, batch_sizes, num_runs, batch_mode
|
||||
):
|
||||
print("\n" + "=" * 60)
|
||||
print(f"{name.upper()} BENCHMARK")
|
||||
print("=" * 60)
|
||||
|
||||
results = [
|
||||
benchmark(
|
||||
data=data,
|
||||
batch_size=bs,
|
||||
sequential_fn=sequential_fn,
|
||||
batch_fn=batch_fn,
|
||||
num_runs=num_runs,
|
||||
batch_mode=batch_mode,
|
||||
)
|
||||
for bs in batch_sizes
|
||||
]
|
||||
print_results(results=results, func_name=name, batch_mode=batch_mode)
|
||||
|
||||
|
||||
def benchmark(*, data, batch_size, sequential_fn, batch_fn, num_runs, batch_mode):
|
||||
batch_data = data[:batch_size]
|
||||
run_single = "single" in batch_mode
|
||||
run_batch = "batch" in batch_mode
|
||||
|
||||
out = {"batch_size": batch_size}
|
||||
|
||||
if run_single:
|
||||
sequential_times = measure_times(
|
||||
fn=lambda: sequential_fn(batch_data), num_runs=num_runs
|
||||
)
|
||||
out |= {
|
||||
"avg_sequential_ms": mean(sequential_times),
|
||||
"sequential_runs": sequential_times,
|
||||
}
|
||||
|
||||
if run_batch:
|
||||
batch_times = measure_times(fn=lambda: batch_fn(batch_data), num_runs=num_runs)
|
||||
out |= {
|
||||
"avg_batch_ms": mean(batch_times),
|
||||
"batch_runs": batch_times,
|
||||
}
|
||||
|
||||
if run_single and run_batch:
|
||||
out["speedup_factor"] = (
|
||||
out["avg_sequential_ms"] / out["avg_batch_ms"]
|
||||
if out["avg_batch_ms"] > 0
|
||||
else 0
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def print_results(*, results, func_name, batch_mode):
|
||||
run_single = "single" in batch_mode
|
||||
run_batch = "batch" in batch_mode
|
||||
|
||||
for r in results:
|
||||
print(f"\nBatch size: {r['batch_size']}")
|
||||
if run_single:
|
||||
print_runs(
|
||||
label=f"Sequential {func_name}",
|
||||
runs=r["sequential_runs"],
|
||||
avg=r["avg_sequential_ms"],
|
||||
)
|
||||
if run_batch:
|
||||
print_runs(
|
||||
label=f"Batch {func_name}", runs=r["batch_runs"], avg=r["avg_batch_ms"]
|
||||
)
|
||||
if run_single and run_batch:
|
||||
print(f" Speedup factor: {r['speedup_factor']:.2f}x")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"SUMMARY: {func_name.upper()}")
|
||||
print("=" * 60)
|
||||
|
||||
headers = ["Batch Size"]
|
||||
if run_single:
|
||||
headers.append("Sequential (ms)")
|
||||
if run_batch:
|
||||
headers.append("Batch (ms)")
|
||||
if run_single and run_batch:
|
||||
headers.append("Speedup")
|
||||
print("".join(f"{h:<18}" for h in headers))
|
||||
print("-" * (18 * len(headers)))
|
||||
|
||||
for r in results:
|
||||
row = [f"{r['batch_size']}"]
|
||||
if run_single:
|
||||
row.append(f"{r['avg_sequential_ms']:.2f} ms")
|
||||
if run_batch:
|
||||
row.append(f"{r['avg_batch_ms']:.2f} ms")
|
||||
if run_single and run_batch:
|
||||
row.append(f"{r['speedup_factor']:.2f}x")
|
||||
print("".join(f"{v:<18}" for v in row))
|
||||
|
||||
|
||||
def print_runs(*, label, runs, avg):
|
||||
print(f" {label}:")
|
||||
for i, t in enumerate(runs):
|
||||
print(f" Run {i+1}: {t:.2f} ms")
|
||||
print(f" Average: {avg:.2f} ms")
|
||||
|
||||
|
||||
def measure_times(*, fn, num_runs):
|
||||
times = []
|
||||
for _ in range(num_runs):
|
||||
start = time.perf_counter()
|
||||
fn()
|
||||
times.append((time.perf_counter() - start) * 1000)
|
||||
return times
|
||||
|
||||
|
||||
def generate_random_token_ids(*, num_prompts, num_tokens, tokenizer):
|
||||
vocab_size = tokenizer.vocab_size
|
||||
print(f"Generating {num_prompts} random sequences with {num_tokens} tokens each...")
|
||||
return [
|
||||
[random.randint(0, vocab_size - 1) for _ in range(num_tokens)]
|
||||
for _ in range(num_prompts)
|
||||
]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Tokenizer Benchmark: Sequential vs Batch Processing"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tokenizer",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Tokenizer name or path (e.g. nvidia/Kimi-K2-Thinking-NVFP4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--function",
|
||||
type=str,
|
||||
nargs="+",
|
||||
choices=["encode", "decode"],
|
||||
default=["encode", "decode"],
|
||||
help="Functions to benchmark (default: encode decode)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-tokens",
|
||||
type=int,
|
||||
default=20000,
|
||||
help="Number of tokens per prompt (default: 20000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 2, 4, 8],
|
||||
help="Batch sizes to test (default: 1 2 4 8)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-mode",
|
||||
nargs="+",
|
||||
choices=["single", "batch"],
|
||||
default=["single", "batch"],
|
||||
help="Benchmark modes to run (default: single batch)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-runs",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Number of runs per batch size (default: 5)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
random.seed(0)
|
||||
main()
|
||||
Reference in New Issue
Block a user