diff --git a/benchmark/benchmark_batch/benchmark_tokenizer.py b/benchmark/benchmark_batch/benchmark_tokenizer.py index 4439f3a05..8296b7e64 100644 --- a/benchmark/benchmark_batch/benchmark_tokenizer.py +++ b/benchmark/benchmark_batch/benchmark_tokenizer.py @@ -5,6 +5,8 @@ from statistics import mean from transformers import AutoTokenizer +from sglang.srt.utils.patch_tokenizer import patch_tokenizer + def main(): args = parse_args() @@ -15,13 +17,16 @@ def main(): 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"Skip batch: {args.no_batch}") + 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(max_batch_size, args.num_tokens, tokenizer) + token_ids = generate_random_token_ids( + num_prompts=max_batch_size, num_tokens=args.num_tokens, tokenizer=tokenizer + ) if "encode" in args.function: prompts = [ @@ -35,7 +40,7 @@ def main(): batch_fn=lambda batch: tokenizer(batch), batch_sizes=args.batch_sizes, num_runs=args.num_runs, - skip_batch=args.no_batch, + batch_mode=args.batch_mode, ) if "decode" in args.function: @@ -53,82 +58,116 @@ def main(): batch_fn=lambda batch: tokenizer.batch_decode(batch, **decode_kwargs), batch_sizes=args.batch_sizes, num_runs=args.num_runs, - skip_batch=args.no_batch, + batch_mode=args.batch_mode, ) def run_benchmark( - name, data, sequential_fn, batch_fn, batch_sizes, num_runs, skip_batch + *, 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, bs, sequential_fn, batch_fn, num_runs, skip_batch) + 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, name, skip_batch) + print_results(results=results, func_name=name, batch_mode=batch_mode) -def benchmark(data, batch_size, sequential_fn, batch_fn, num_runs, skip_batch): +def benchmark(*, data, batch_size, sequential_fn, batch_fn, num_runs, batch_mode): batch_data = data[:batch_size] - sequential_times = measure_times(lambda: sequential_fn(batch_data), num_runs) - avg_seq = mean(sequential_times) + run_single = "single" in batch_mode + run_batch = "batch" in batch_mode - out = { - "batch_size": batch_size, - "avg_sequential_ms": avg_seq, - "sequential_runs": sequential_times, - } + out = {"batch_size": batch_size} - if not skip_batch: - batch_times = measure_times(lambda: batch_fn(batch_data), num_runs) - avg_batch = mean(batch_times) + if run_single: + sequential_times = measure_times( + fn=lambda: sequential_fn(batch_data), num_runs=num_runs + ) out |= { - "avg_batch_ms": avg_batch, - "speedup_factor": avg_seq / avg_batch if avg_batch > 0 else 0, + "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, skip_batch): +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']}") - print_runs( - f"Sequential {func_name}", r["sequential_runs"], r["avg_sequential_ms"] - ) - if not skip_batch: - print_runs(f"Batch {func_name}", r["batch_runs"], r["avg_batch_ms"]) + 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", "Sequential (ms)"] - if not skip_batch: - headers += ["Batch (ms)", "Speedup"] + 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']}", f"{r['avg_sequential_ms']:.2f} ms"] - if not skip_batch: - row += [f"{r['avg_batch_ms']:.2f} ms", f"{r['speedup_factor']:.2f}x"] + 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): +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): +def measure_times(*, fn, num_runs): times = [] for _ in range(num_runs): start = time.perf_counter() @@ -137,7 +176,7 @@ def measure_times(fn, num_runs): return times -def generate_random_token_ids(num_prompts, num_tokens, tokenizer): +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 [ @@ -178,9 +217,11 @@ def parse_args(): help="Batch sizes to test (default: 1 2 4 8)", ) parser.add_argument( - "--no-batch", - action="store_true", - help="Skip batch benchmark, only run sequential", + "--batch-mode", + nargs="+", + choices=["single", "batch"], + default=["single", "batch"], + help="Benchmark modes to run (default: single batch)", ) parser.add_argument( "--num-runs",