From 97aa930ea143f36896db9c7c3599128682948459 Mon Sep 17 00:00:00 2001 From: papertager <2567587994@qq.com> Date: Sat, 6 Jun 2026 11:03:44 +0800 Subject: [PATCH] feat: add controllable FlashInfer benchmark runs --- .../flashinfer_baseline/bench_batch_decode.py | 86 ++++++++++++----- .../flashinfer_baseline/bench_batch_mla.py | 83 ++++++++++++----- .../bench_batch_prefill_paged.py | 83 ++++++++++------- .../bench_batch_prefill_ragged.py | 92 ++++++++++++------- .../flashinfer_baseline/bench_common.py | 84 ++++++++++++++++- 5 files changed, 318 insertions(+), 110 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py index 32861ca..49982ff 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py @@ -4,11 +4,22 @@ seq_len_q=1 (decode mode), seq_len_kv from 1K to 16K """ import itertools -import pandas as pd import torch import flashinfer -from bench_common import dtype, page_block_size, setup_workspace, setup_paged_kv_indptr, run_with_profiler, get_csv_path, compute_reps +from bench_common import ( + compute_reps, + dtype, + get_device_info, + limited_cases, + page_block_size, + parse_benchmark_args, + record_failed_case, + run_with_profiler, + setup_paged_kv_indptr, + setup_workspace, + write_csv, +) target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] @@ -20,6 +31,8 @@ def bench_batch_decode( num_kv_heads, head_dim, page_block_size, + warmup, + base_reps, ): """Benchmark BatchDecodeWithPagedKVCacheWrapper""" seq_lens = [seq_len_kv] * batch_size @@ -45,15 +58,20 @@ def bench_batch_decode( q_data_type=dtype, ) - reps = compute_reps(batch_size, seq_len_kv, head_dim, base_reps=100) - ms = run_with_profiler(lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps) + reps = compute_reps(batch_size, seq_len_kv, head_dim, base_reps=base_reps) + ms = run_with_profiler( + lambda: wrapper.run(q, kv_data), + target_kernels=target_kernels, + warmup=warmup, + reps=reps, + ) io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() flops = 2 * batch_size * seq_len_kv * num_qo_heads * num_kv_heads * head_dim - return ms, io, flops + return ms, io, flops, reps -def run_benchmark(): +def run_benchmark(args): records = [] batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] @@ -61,7 +79,9 @@ def run_benchmark(): seq_lens_kv = [512, 1024, 2048, 4096, 8192, 16384] api_name = "BatchDecodeWithPagedKVCacheWrapper" - test_cases = list(itertools.product(batch_sizes, seq_lens_kv, head_dims)) + test_cases = limited_cases( + list(itertools.product(batch_sizes, seq_lens_kv, head_dims)), args.max_cases + ) total_cases = len(test_cases) print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") @@ -69,22 +89,45 @@ def run_benchmark(): for idx, (bs, sl_kv, hd) in enumerate(test_cases, 1): num_qo_heads = 32 num_kv_heads = 8 if hd == 64 else 4 - ms, io, flops = bench_batch_decode(bs, sl_kv, num_qo_heads, num_kv_heads, hd, page_block_size) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append({ - "api": api_name, + case = { "batch_size": bs, "seq_len_q": 1, "seq_len_kv": sl_kv, "num_qo_heads": num_qo_heads, "num_kv_heads": num_kv_heads, "head_dim": hd, - "time_ms": ms, - "bandwidth_GB_s": bw, - "tflops": tflops, - }) - print(f" [{idx}/{total_cases}] bs={bs}, kv_len={sl_kv}, hd={hd}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + } + try: + ms, io, flops, reps = bench_batch_decode( + bs, + sl_kv, + num_qo_heads, + num_kv_heads, + hd, + page_block_size, + args.warmup, + args.base_reps, + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + record = { + "api": api_name, + "status": "ok", + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + "warmup": args.warmup, + "reps": reps, + } + record.update(case) + record.update(get_device_info()) + records.append(record) + print(f" [{idx}/{total_cases}] bs={bs}, kv_len={sl_kv}, hd={hd}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + except Exception as exc: + if not args.continue_on_error: + raise + records.append(record_failed_case(api_name, exc, **case)) + print(f" [{idx}/{total_cases}] bs={bs}, kv_len={sl_kv}, hd={hd}: failed: {exc}") return records @@ -94,8 +137,7 @@ if __name__ == "__main__": np.random.seed(42) torch.random.manual_seed(42) - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchDecodeWithPagedKVCacheWrapper") - df.to_csv(csv_path, index=False) - print(f"\nResults saved to {csv_path}") \ No newline at end of file + args = parse_benchmark_args(__doc__) + records = run_benchmark(args) + csv_path = write_csv("BatchDecodeWithPagedKVCacheWrapper", records, args.output_dir) + print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py index 5c48bc9..e2424e2 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py @@ -4,11 +4,20 @@ headdim: ckv=512, kpe=64 (DeepSeek MLA configuration) """ import itertools -import pandas as pd import torch import flashinfer -from bench_common import dtype, page_block_size, setup_workspace, run_with_profiler, get_csv_path, compute_reps +from bench_common import ( + compute_reps, + dtype, + get_device_info, + limited_cases, + parse_benchmark_args, + record_failed_case, + run_with_profiler, + setup_workspace, + write_csv, +) target_kernels = ["BatchMLAPagedAttentionKernel"] @@ -19,6 +28,8 @@ def bench_batch_mla_paged_attention( num_heads, head_dim_ckv, head_dim_kpe, + warmup, + base_reps, ): """Benchmark BatchMLAPagedAttentionWrapper for DeepSeek MLA""" # MLA decode mode: q has length 1, not seq_len @@ -54,16 +65,21 @@ def bench_batch_mla_paged_attention( ckv.dtype, ) - reps = compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe, base_reps=100) - ms = run_with_profiler(lambda: wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False), target_kernels=target_kernels, reps=reps) + reps = compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe, base_reps=base_reps) + ms = run_with_profiler( + lambda: wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False), + target_kernels=target_kernels, + warmup=warmup, + reps=reps, + ) io = sum([t.numel() * t.element_size() for t in [q_nope, q_pe, ckv, kpe]]) # MLA FLOPs: 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len flops = 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len - return ms, io, flops + return ms, io, flops, reps -def run_benchmark(): +def run_benchmark(args): records = [] # MLA configuration - same as DeepSeek @@ -74,26 +90,50 @@ def run_benchmark(): num_heads_list = [64, 128] api_name = "BatchMLAPagedAttentionWrapper" - test_cases = list(itertools.product(num_heads_list, batch_sizes, seq_lens)) + test_cases = limited_cases( + list(itertools.product(num_heads_list, batch_sizes, seq_lens)), args.max_cases + ) total_cases = len(test_cases) print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") for idx, (num_heads, bs, sl) in enumerate(test_cases, 1): - ms, io, flops = bench_batch_mla_paged_attention(bs, sl, num_heads, head_dim_ckv, head_dim_kpe) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append({ - "api": api_name, + case = { "batch_size": bs, "seq_len": sl, "num_heads": num_heads, "head_dim_ckv": head_dim_ckv, "head_dim_kpe": head_dim_kpe, - "time_ms": ms, - "bandwidth_GB_s": bw, - "tflops": tflops, - }) - print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, num_heads={num_heads}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + } + try: + ms, io, flops, reps = bench_batch_mla_paged_attention( + bs, + sl, + num_heads, + head_dim_ckv, + head_dim_kpe, + args.warmup, + args.base_reps, + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + record = { + "api": api_name, + "status": "ok", + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + "warmup": args.warmup, + "reps": reps, + } + record.update(case) + record.update(get_device_info()) + records.append(record) + print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, num_heads={num_heads}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + except Exception as exc: + if not args.continue_on_error: + raise + records.append(record_failed_case(api_name, exc, **case)) + print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, num_heads={num_heads}: failed: {exc}") return records @@ -103,8 +143,7 @@ if __name__ == "__main__": np.random.seed(42) torch.random.manual_seed(42) - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchMLAPagedAttentionWrapper") - df.to_csv(csv_path, index=False) - print(f"\nResults saved to {csv_path}") \ No newline at end of file + args = parse_benchmark_args(__doc__) + records = run_benchmark(args) + csv_path = write_csv("BatchMLAPagedAttentionWrapper", records, args.output_dir) + print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py index 1dee3d0..2e04769 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py @@ -4,17 +4,20 @@ headdim: 64/128/256 """ import itertools -import pandas as pd import torch import flashinfer from bench_common import ( - dtype, - setup_workspace, - setup_paged_kv_indptr, - run_with_profiler, - get_csv_path, compute_reps, + dtype, + get_device_info, + limited_cases, + parse_benchmark_args, + record_failed_case, + run_with_profiler, + setup_paged_kv_indptr, + setup_workspace, + write_csv, ) target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] @@ -26,6 +29,8 @@ def bench_batch_prefill_with_paged_kv_cache( num_qo_heads, num_kv_heads, head_dim, + warmup, + base_reps, causal=True, ): """Benchmark BatchPrefillWithPagedKVCacheWrapper""" @@ -59,9 +64,12 @@ def bench_batch_prefill_with_paged_kv_cache( kv_data_type=dtype, ) - reps = compute_reps(batch_size, seq_len, head_dim, base_reps=100) + reps = compute_reps(batch_size, seq_len, head_dim, base_reps=base_reps) ms = run_with_profiler( - lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps + lambda: wrapper.run(q, kv_data), + target_kernels=target_kernels, + warmup=warmup, + reps=reps, ) io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() @@ -77,10 +85,10 @@ def bench_batch_prefill_with_paged_kv_cache( * head_dim * (1 if causal else 2) ) - return ms, io, flops + return ms, io, flops, reps -def run_benchmark(): +def run_benchmark(args): records = [] batch_sizes = [1, 4, 16, 64] @@ -88,34 +96,48 @@ def run_benchmark(): head_dims = [128, 256] api_name = "BatchPrefillWithPagedKVCacheWrapper" - test_cases = list(itertools.product(head_dims, batch_sizes, seq_lens)) + test_cases = limited_cases( + list(itertools.product(head_dims, batch_sizes, seq_lens)), args.max_cases + ) total_cases = len(test_cases) print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") for idx, (head_dim, bs, sl) in enumerate(test_cases, 1): num_qo_heads = 32 num_kv_heads = 8 if head_dim == 64 else 4 - ms, io, flops = bench_batch_prefill_with_paged_kv_cache( - bs, sl, num_qo_heads, num_kv_heads, head_dim - ) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append( - { + case = { + "batch_size": bs, + "seq_len": sl, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + } + try: + ms, io, flops, reps = bench_batch_prefill_with_paged_kv_cache( + bs, sl, num_qo_heads, num_kv_heads, head_dim, args.warmup, args.base_reps + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + record = { "api": api_name, - "batch_size": bs, - "seq_len": sl, - "num_qo_heads": num_qo_heads, - "num_kv_heads": num_kv_heads, - "head_dim": head_dim, + "status": "ok", "time_ms": ms, "bandwidth_GB_s": bw, "tflops": tflops, + "warmup": args.warmup, + "reps": reps, } - ) - print( - f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd={head_dim}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" - ) + record.update(case) + record.update(get_device_info()) + records.append(record) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd={head_dim}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" + ) + except Exception as exc: + if not args.continue_on_error: + raise + records.append(record_failed_case(api_name, exc, **case)) + print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd={head_dim}: failed: {exc}") return records @@ -126,8 +148,7 @@ if __name__ == "__main__": np.random.seed(42) torch.random.manual_seed(42) - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchPrefillWithPagedKVCacheWrapper") - df.to_csv(csv_path, index=False) + args = parse_benchmark_args(__doc__) + records = run_benchmark(args) + csv_path = write_csv("BatchPrefillWithPagedKVCacheWrapper", records, args.output_dir) print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py index 43aa4e0..868aef7 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py @@ -4,16 +4,19 @@ headdim configurations: [64,64], [128,128], [192,128], [256,256] """ import itertools -import pandas as pd import torch import flashinfer from bench_common import ( - dtype, - setup_workspace, - run_with_profiler, - get_csv_path, compute_reps, + dtype, + get_device_info, + limited_cases, + parse_benchmark_args, + record_failed_case, + run_with_profiler, + setup_workspace, + write_csv, ) target_kernels = [ @@ -29,6 +32,8 @@ def bench_batch_prefill_with_ragged_kv_cache( num_kv_heads, head_dim_qk, head_dim_vo, + warmup, + base_reps, causal=True, ): """Benchmark BatchPrefillWithRaggedKVCacheWrapper for MLA""" @@ -58,9 +63,12 @@ def bench_batch_prefill_with_ragged_kv_cache( kv_data_type=dtype, ) - reps = compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo, base_reps=100) + reps = compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo, base_reps=base_reps) ms = run_with_profiler( - lambda: wrapper.run(q, k, v), target_kernels=target_kernels, reps=reps + lambda: wrapper.run(q, k, v), + target_kernels=target_kernels, + warmup=warmup, + reps=reps, ) io = ( @@ -78,10 +86,10 @@ def bench_batch_prefill_with_ragged_kv_cache( * (1 if causal else 2) ) - return ms, io, flops + return ms, io, flops, reps -def run_benchmark(): +def run_benchmark(args): records = [] # headdim combinations: [qk, vo] @@ -90,35 +98,58 @@ def run_benchmark(): seq_lens = [1024, 4096, 8192, 16384] api_name = "BatchPrefillWithRaggedKVCacheWrapper" - test_cases = list(itertools.product(head_dim_configs, batch_sizes, seq_lens)) + test_cases = limited_cases( + list(itertools.product(head_dim_configs, batch_sizes, seq_lens)), args.max_cases + ) total_cases = len(test_cases) print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") for idx, ((head_dim_qk, head_dim_vo), bs, sl) in enumerate(test_cases, 1): num_qo_heads = 32 num_kv_heads = 4 - ms, io, flops = bench_batch_prefill_with_ragged_kv_cache( - bs, sl, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo - ) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append( - { + case = { + "batch_size": bs, + "seq_len": sl, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim_qk": head_dim_qk, + "head_dim_vo": head_dim_vo, + } + try: + ms, io, flops, reps = bench_batch_prefill_with_ragged_kv_cache( + bs, + sl, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + args.warmup, + args.base_reps, + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + record = { "api": api_name, - "batch_size": bs, - "seq_len": sl, - "num_qo_heads": num_qo_heads, - "num_kv_heads": num_kv_heads, - "head_dim_qk": head_dim_qk, - "head_dim_vo": head_dim_vo, + "status": "ok", "time_ms": ms, "bandwidth_GB_s": bw, "tflops": tflops, + "warmup": args.warmup, + "reps": reps, } - ) - print( - f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd=[{head_dim_qk},{head_dim_vo}]: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" - ) + record.update(case) + record.update(get_device_info()) + records.append(record) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd=[{head_dim_qk},{head_dim_vo}]: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" + ) + except Exception as exc: + if not args.continue_on_error: + raise + records.append(record_failed_case(api_name, exc, **case)) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd=[{head_dim_qk},{head_dim_vo}]: failed: {exc}" + ) return records @@ -129,8 +160,7 @@ if __name__ == "__main__": np.random.seed(42) torch.random.manual_seed(42) - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchPrefillWithRaggedKVCacheWrapper") - df.to_csv(csv_path, index=False) + args = parse_benchmark_args(__doc__) + records = run_benchmark(args) + csv_path = write_csv("BatchPrefillWithRaggedKVCacheWrapper", records, args.output_dir) print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py index 3a0563a..8a534f4 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py @@ -1,10 +1,11 @@ -""" -Common utilities for FlashInfer benchmarks -""" +"""Common utilities for FlashInfer benchmarks.""" +import argparse +import csv import os import random from datetime import datetime +from pathlib import Path import numpy as np import torch @@ -21,6 +22,81 @@ def get_csv_path(prefix): return f"{prefix}_{get_timestamp()}.csv" +def parse_benchmark_args(description): + parser = argparse.ArgumentParser(description=description) + parser.add_argument( + "--output-dir", + default=".", + help="Directory used for CSV benchmark output.", + ) + parser.add_argument( + "--max-cases", + type=int, + default=0, + help="Run at most this many cases. 0 means run all cases.", + ) + parser.add_argument( + "--warmup", + type=int, + default=10, + help="Profiler warmup iterations for each case.", + ) + parser.add_argument( + "--base-reps", + type=int, + default=100, + help="Base profiler repetitions before workload-based scaling.", + ) + parser.add_argument( + "--continue-on-error", + action="store_true", + help="Record failed cases and continue instead of aborting.", + ) + return parser.parse_args() + + +def limited_cases(test_cases, max_cases): + if max_cases and max_cases > 0: + return test_cases[:max_cases] + return test_cases + + +def write_csv(prefix, records, output_dir): + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + csv_path = output_path / get_csv_path(prefix) + fieldnames = sorted({key for record in records for key in record.keys()}) + with csv_path.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(records) + return csv_path + + +def get_device_info(): + if not torch.cuda.is_available(): + return {"device_name": "unavailable", "device_memory_gib": 0.0} + + device = torch.cuda.current_device() + props = torch.cuda.get_device_properties(device) + return { + "device_name": props.name, + "device_memory_gib": round(props.total_memory / 1024**3, 2), + } + + +def record_failed_case(api_name, error, **case): + row = { + "api": api_name, + "status": "failed", + "error_type": type(error).__name__, + "error": str(error), + } + row.update(case) + row.update(get_device_info()) + return row + + def generate_random_seqlens(batch_size, min_len=1024, max_len=16384): """Generate random sequence lengths simulating real LLM workloads""" return [random.randint(min_len, max_len) for _ in range(batch_size)] @@ -88,4 +164,4 @@ def compute_reps(batch_size, seq_len, head_dim, base_reps=100): elif workload < 1e9: # very large workload return base_reps // 16 else: # huge workload - return base_reps // 32 \ No newline at end of file + return max(1, base_reps // 32) -- 2.34.1