Update benchmarks

Author: jmanning-stackav

benchmarks/paged_attention_benchmark.py

@@ -26,7 +26,7 @@
     "--head-dim",
     required=True,
     type=int,
-    default=256,
+    default=128,
     help="Head dimension",
 )
 @click.option(
@@ -47,14 +47,14 @@
     "--batch-size",
     required=False,
     type=int,
-    default=4,
+    default=128,
     help="Batch size",
 )
 @click.option(
     "--num-query-heads",
     required=False,
     type=int,
-    default=8,
+    default=32,
     help="Number of query heads",
 )
 @click.option(

benchmarks/paged_attention_vs_flash_benchmark.py

@@ -26,7 +26,7 @@
     "--head-dim",
     required=True,
     type=int,
-    default=256,
+    default=128,
     help="Head dimension",
 )
 @click.option(
@@ -47,14 +47,14 @@
     "--batch-size",
     required=False,
     type=int,
-    default=4,
+    default=128,
     help="Batch size",
 )
 @click.option(
     "--num-query-heads",
     required=False,
     type=int,
-    default=8,
+    default=32,
     help="Number of query heads",
 )
 @click.option(

benchmarks/varlen_attention_benchmark.py

@@ -27,14 +27,14 @@
     "--head-dim",
     required=True,
     type=int,
-    default=256,
+    default=128,
     help="Head dimension",
 )
 @click.option(
     "--seq-len",
     required=True,
     type=int,
-    default=1024,
+    default=512,
     help="Sequence length (for k/v)",
 )
 @click.option(
@@ -48,21 +48,21 @@
     "--batch-size",
     required=False,
     type=int,
-    default=10,
+    default=64,
     help="Batch size",
 )
 @click.option(
     "--num-query-heads",
     required=False,
     type=int,
-    default=8,
+    default=32,
     help="Number of query heads",
 )
 @click.option(
     "--num-kv-heads",
     required=False,
     type=int,
-    default=4,
+    default=8,
     help="Number of kv heads",
 )
 @click.option(

tools/create_benchmark_results_table.py

@@ -24,6 +24,7 @@
     "GeLU, Tanh, and Mul": "gelu_tanh_and_mul_benchmark",
     "SiLU and Mul": "silu_and_mul_benchmark",
     "Paged Attention": "paged_attention_vs_flash_benchmark",
+    "Varlen Attention": "varlen_attention_benchmark",
     "Rotary Embedding": "rotary_embedding_benchmark",
     "RMS Norm (Gemma-style)": "gemma_rms_norm_benchmark",
     "RMS Norm (Llama-style)": "rms_norm_benchmark",
@@ -45,6 +46,11 @@
     "unknown": [],
 }
 
+# Add any extra flags for each benchmark here
+_EXTRA_BENCHMARK_FLAGS: Final = {
+    "varlen_attention_benchmark": ["--causal"],
+}
+
 
 @click.command()
 @click.option(
@@ -90,9 +96,12 @@ def main(results_directory: Path, use_cached_results: bool) -> None:
             # Run benchmark and redirect output
             print(f"Running benchmark for {op_name}...")
 
+            # Some benchmark args are flags to enable things that default false, so we add any per-benchmark here
+            extra_flags = _EXTRA_BENCHMARK_FLAGS[benchmark_name] if benchmark_name in _EXTRA_BENCHMARK_FLAGS else []
+
             with results_csv.open("w") as results_file:
                 run(
-                    ["python", f"benchmarks/{benchmark_name}.py", "--csv"],
+                    ["python", f"benchmarks/{benchmark_name}.py", "--csv"] + extra_flags,
                     check=True,
                     stdout=results_file,
                     env=os.environ,