bench: Add comprehensive VLM benchmark + fix CUDA graph timing methodology

- Add bench_kbit_vlm.py: sweeps all kernel variants (scalar GEMV, MMA,
  dequant+cuBLAS) with and without Hadamard rotation across VLM-relevant
  M values (1-1024) on Qwen3-Coder-Next 70B shapes, k=2..5.

- Rewrite bench_hadamard.py to use batched graph replay: replay the
  captured graph N times within one event pair, then divide. This
  amortizes the ~14 us per-replay timing floor to <0.03 us, revealing
  true kernel execution times that were previously masked.

- Update CLAUDE.md with benchmarking section: where to find scripts,
  how to run them, and the batched replay methodology.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

CLAUDE.md

@@ -23,3 +23,32 @@ pytest tests/ -v --tb=short -n 4
 ```
 
 Best practices, benchmark data, and known architecture-specific issues: `agents/testing_guide.md`
+
+# Benchmarking
+
+Benchmark scripts live in `benchmarks/`. The two kbit-specific ones:
+
+- `bench_hadamard.py` — Hadamard rotation kernel + M=1 pipeline (rotation + scalar GEMV) vs cuBLAS FP16. Quick focused benchmark for the decode path.
+- `bench_kbit_vlm.py` — Comprehensive sweep across all VLM-relevant M values (1 to 1024), all kernel variants (scalar GEMV, MMA, dequant+cuBLAS), all k values (2-5), with and without Hadamard rotation. Qwen3-Coder-Next 70B shapes.
+
+```bash
+# Quick M=1 decode benchmark
+python benchmarks/bench_hadamard.py
+
+# Full VLM sweep (all M, all k)
+python benchmarks/bench_kbit_vlm.py
+
+# Single k value, subset of M
+python benchmarks/bench_kbit_vlm.py --k 4 --m 1,4,16,256,1024
+
+# Higher accuracy (more iterations)
+python benchmarks/bench_kbit_vlm.py --inner 1000 --outer 30
+```
+
+## CUDA graph benchmarking methodology
+
+Single graph replay has a ~14 us timing floor (on RTX 4090) that masks sub-14 us kernel differences. The benchmarks use **batched graph replay**: replay the graph N times within one event-timed region, then divide. This amortizes the per-replay overhead to ~14/N us per iteration.
+
+The `--inner` flag controls N (replays per measurement). Default 500 gives ~0.03 us amortized overhead. Use `--inner 1000` for the highest accuracy when comparing kernels that differ by < 1 us.
+
+`--outer` controls the number of measurements (default 15). The median is reported to reject outliers.

benchmarks/bench_hadamard.py

@@ -1,14 +1,23 @@
-"""Benchmark for Hadamard rotation kernel and full kbit pipeline.
+"""Benchmark for Hadamard rotation kernel and kbit M=1 pipeline.
 
 Measures:
-1. Rotation standalone: all block sizes × Qwen3 K values × M=1,4
-2. Full pipeline (rotate + kbit_scalar_gemv_tiled): Qwen3 dense shapes at M=1, k=2,3,4
+1. Rotation standalone: all block sizes x Qwen3 K values x M=1,4
+2. Full pipeline (rotate + kbit_scalar_gemv_tiled): Qwen3 dense shapes at M=1, k=2..5
 3. cuBLAS FP16 baseline: same shapes
 4. Speedup table: pipeline vs cuBLAS
 
-All timing via CUDA graph capture + replay for clean kernel-only measurements.
+Timing methodology:
+  CUDA graph capture + batched replay. Each measurement replays the graph
+  INNER times within a single event-timed region, then divides. This
+  amortizes the ~14 us per-replay overhead down to negligible levels,
+  revealing true kernel execution times. Median of OUTER measurements.
+
+Usage:
+  python benchmarks/bench_hadamard.py
+  python benchmarks/bench_hadamard.py --inner 1000 --outer 30   # higher accuracy
 """
 
+import argparse
 import sys
 
 import torch
@@ -22,9 +31,7 @@
     quantize_kbit,
 )
 
-BLOCKSIZE = 32
-WARMUP = 50
-ITERS = 200
+ROTATION_BLOCK_SIZE = 64
 
 
 def create_normal_float_codebook(k: int) -> torch.Tensor:
@@ -35,14 +42,18 @@ def create_normal_float_codebook(k: int) -> torch.Tensor:
     return values.cuda()
 
 
-def bench_graph(fn, warmup=WARMUP, iters=ITERS):
-    """Time a function using CUDA graph capture + replay. Returns median time in us."""
-    # Warm up on default stream
-    for _ in range(warmup):
+def bench(fn, inner: int, outer: int) -> float:
+    """Batched CUDA graph replay timing. Returns median us per iteration.
+
+    Captures fn into a CUDA graph, then replays it `inner` times within a
+    single CUDA event pair. The per-replay overhead (~14 us on RTX 4090)
+    is amortized to ~14/inner us per iteration. Takes the median of `outer`
+    such measurements.
+    """
+    for _ in range(30):
         fn()
     torch.cuda.synchronize()
 
-    # Capture graph
     s = torch.cuda.Stream()
     s.wait_stream(torch.cuda.current_stream())
     with torch.cuda.stream(s):
@@ -55,27 +66,34 @@ def bench_graph(fn, warmup=WARMUP, iters=ITERS):
         fn()
     torch.cuda.synchronize()
 
-    # Warm up replay
-    for _ in range(10):
+    for _ in range(50):
         g.replay()
     torch.cuda.synchronize()
 
-    # Time replay
     times = []
-    for _ in range(iters):
+    for _ in range(outer):
         start = torch.cuda.Event(enable_timing=True)
         end = torch.cuda.Event(enable_timing=True)
         start.record()
-        g.replay()
+        for _ in range(inner):
+            g.replay()
         end.record()
         torch.cuda.synchronize()
-        times.append(start.elapsed_time(end) * 1000)  # ms -> us
-
+        times.append(start.elapsed_time(end) * 1000 / inner)  # ms -> us/iter
     times.sort()
-    return times[len(times) // 2]  # median
+    return times[len(times) // 2]
+
+
+def prepare_kbit_weights(K_dim, N, k):
+    """Quantize random weights and repack for tiled access."""
+    W = torch.randn(N, K_dim, dtype=torch.float16, device="cuda")
+    codebook = create_normal_float_codebook(k)
+    packed, absmax, _ = quantize_kbit(W, k=k, codebook=codebook)
+    packed_tiled, absmax_tiled = torch.ops.bitsandbytes.repack_kbit(packed, absmax, K_dim, N, k)
+    return packed_tiled, absmax_tiled, codebook
 
 
-def bench_rotation_standalone():
+def bench_rotation_standalone(inner, outer):
     """Benchmark rotation kernel standalone across block sizes and shapes."""
     print("=" * 70)
     print("1. ROTATION STANDALONE")
@@ -91,31 +109,20 @@ def bench_rotation_standalone():
         for K in k_values:
             for bs in block_sizes:
                 A = torch.randn(M, K, dtype=torch.float16, device="cuda")
-                t = bench_graph(lambda: hadamard_rotate(A, block_size=bs))
-                # BW: read + write = 2 * numel * 2 bytes (fp16)
+                t = bench(lambda: hadamard_rotate(A, block_size=bs), inner, outer)
                 bw = 2 * A.numel() * 2 / (t / 1e6) / 1e9
-                print(f"{M:>4} {K:>6} {bs:>4} {t:>10.2f} {bw:>10.1f}")
+                print(f"{M:>4} {K:>6} {bs:>4} {t:>10.3f} {bw:>10.1f}")
         print()
 
 
-def prepare_kbit_weights(K_dim, N, k):
-    """Quantize random weights and repack for tiled access."""
-    W = torch.randn(N, K_dim, dtype=torch.float16, device="cuda")
-    codebook = create_normal_float_codebook(k)
-    packed, absmax, _ = quantize_kbit(W, k=k, codebook=codebook)
-    packed_tiled, absmax_tiled = torch.ops.bitsandbytes.repack_kbit(packed, absmax, K_dim, N, k)
-    return packed_tiled, absmax_tiled, codebook
-
-
-def bench_pipeline():
+def bench_pipeline(inner, outer):
     """Benchmark full pipeline: rotate(A) + kbit_scalar_gemv."""
     print("=" * 70)
     print("2. FULL PIPELINE: rotate + kbit_scalar_gemv_tiled")
     print("=" * 70)
     print(f"{'M':>4} {'K':>6} {'N':>6} {'k':>2} {'Rotate(us)':>11} {'GEMV(us)':>9} {'Total(us)':>10} {'TFLOPS':>7}")
     print("-" * 65)
 
-    # Qwen3-Coder-Next 70B dense shapes
     shapes = [
         (1, 2048, 5120, "gate/up"),
         (1, 5120, 2048, "down"),
@@ -126,42 +133,41 @@ def bench_pipeline():
         (4, 5120, 2048, "down M=4"),
     ]
 
-    for k in [2, 3, 4]:
+    for k in [2, 3, 4, 5]:
         print(f"\n--- k={k} ---")
         for M, K_dim, N, label in shapes:
             packed_tiled, absmax_tiled, codebook = prepare_kbit_weights(K_dim, N, k)
             A = torch.randn(M, K_dim, dtype=torch.float16, device="cuda")
 
-            # Benchmark rotation alone
             A_copy = A.clone()
-            t_rot = bench_graph(lambda: hadamard_rotate(A_copy, block_size=64))
+            t_rot = bench(lambda: hadamard_rotate(A_copy, block_size=ROTATION_BLOCK_SIZE), inner, outer)
 
-            # Benchmark GEMV alone (tiled layout, pre-allocated output)
             out = torch.zeros(M, N, dtype=torch.float16, device="cuda")
-            t_gemv = bench_graph(
+            t_gemv = bench(
                 lambda: torch.ops.bitsandbytes.kbit_scalar_gemv_tiled_(
                     A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out
-                )
+                ),
+                inner,
+                outer,
             )
 
-            # Benchmark combined
             def pipeline():
-                hadamard_rotate(A_copy, block_size=64)
+                hadamard_rotate(A_copy, block_size=ROTATION_BLOCK_SIZE)
                 torch.ops.bitsandbytes.kbit_scalar_gemv_tiled_(
                     A_copy, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out
                 )
 
-            t_total = bench_graph(pipeline)
+            t_total = bench(pipeline, inner, outer)
 
             flops = 2 * M * K_dim * N
             tflops = flops / (t_total / 1e6) / 1e12
             print(
-                f"{M:>4} {K_dim:>6} {N:>6} {k:>2} {t_rot:>11.2f} {t_gemv:>9.2f} "
-                f"{t_total:>10.2f} {tflops:>7.3f}  {label}"
+                f"{M:>4} {K_dim:>6} {N:>6} {k:>2} {t_rot:>11.3f} {t_gemv:>9.3f} "
+                f"{t_total:>10.3f} {tflops:>7.3f}  {label}"
             )
 
 
-def bench_cublas_baseline():
+def bench_cublas_baseline(inner, outer):
     """Benchmark cuBLAS FP16 GEMM for the same shapes."""
     print("\n" + "=" * 70)
     print("3. cuBLAS FP16 BASELINE")
@@ -184,16 +190,16 @@ def bench_cublas_baseline():
         W = torch.randn(N, K_dim, dtype=torch.float16, device="cuda")
         out = torch.empty(M, N, dtype=torch.float16, device="cuda")
 
-        t = bench_graph(lambda: torch.mm(A, W.t(), out=out))
+        t = bench(lambda: torch.mm(A, W.t(), out=out), inner, outer)
         flops = 2 * M * K_dim * N
         tflops = flops / (t / 1e6) / 1e12
-        print(f"{M:>4} {K_dim:>6} {N:>6} {t:>9.2f} {tflops:>7.3f}")
+        print(f"{M:>4} {K_dim:>6} {N:>6} {t:>9.3f} {tflops:>7.3f}")
 
 
-def bench_speedup_table():
+def bench_speedup_table(inner, outer):
     """Print a speedup comparison table: pipeline vs cuBLAS."""
     print("\n" + "=" * 70)
-    print("4. SPEEDUP TABLE: kbit pipeline vs cuBLAS FP16")
+    print("4. SPEEDUP TABLE: Rot + kbit GEMV vs cuBLAS FP16")
     print("=" * 70)
 
     shapes = [
@@ -208,7 +214,7 @@ def bench_speedup_table():
     print(f"{'Shape':>20} {'k':>2} {'Pipeline(us)':>13} {'cuBLAS(us)':>11} {'Speedup':>8}")
     print("-" * 65)
 
-    for k in [2, 3, 4]:
+    for k in [2, 3, 4, 5]:
         print(f"\n--- k={k} ---")
         for M, K_dim, N, label in shapes:
             packed_tiled, absmax_tiled, codebook = prepare_kbit_weights(K_dim, N, k)
@@ -217,40 +223,36 @@ def bench_speedup_table():
             out = torch.zeros(M, N, dtype=torch.float16, device="cuda")
             A_copy = A.clone()
 
-            # Pipeline: rotate + GEMV
             def pipeline():
-                hadamard_rotate(A_copy, block_size=64)
+                hadamard_rotate(A_copy, block_size=ROTATION_BLOCK_SIZE)
                 torch.ops.bitsandbytes.kbit_scalar_gemv_tiled_(
                     A_copy, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out
                 )
 
-            t_pipe = bench_graph(pipeline)
-
-            # cuBLAS baseline
-            t_cublas = bench_graph(lambda: torch.mm(A, W.t(), out=out))
+            t_pipe = bench(pipeline, inner, outer)
+            t_cublas = bench(lambda: torch.mm(A, W.t(), out=out), inner, outer)
 
             speedup = t_cublas / t_pipe
             shape_str = f"{M}x{K_dim}x{N}"
-            print(f"{shape_str:>20} {k:>2} {t_pipe:>13.2f} {t_cublas:>11.2f} {speedup:>7.2f}x  {label}")
+            print(f"{shape_str:>20} {k:>2} {t_pipe:>13.3f} {t_cublas:>11.3f} {speedup:>7.2f}x  {label}")
 
 
-def bench_cuda_graph_capture():
-    """Verify that all benchmarks above were graph-captured (implicit from bench_graph).
-    This just confirms the pipeline captures as a single graph explicitly."""
-    print("\n" + "=" * 70)
-    print("5. CUDA GRAPH CAPTURE VERIFICATION")
-    print("=" * 70)
-    print("All benchmarks above used CUDA graph capture + replay for timing.")
-    print("If they produced numbers, graph capture succeeded for all operations.")
+def main():
+    parser = argparse.ArgumentParser(description="Hadamard rotation + kbit M=1 pipeline benchmark")
+    parser.add_argument("--inner", type=int, default=500, help="Graph replays per measurement (default: 500)")
+    parser.add_argument("--outer", type=int, default=15, help="Measurements per benchmark (default: 15)")
+    args = parser.parse_args()
 
-
-if __name__ == "__main__":
     print(f"GPU: {torch.cuda.get_device_name(0)}")
     print(f"CUDA: {torch.version.cuda}")
+    print(f"Timing: batched graph replay ({args.inner} replays/measurement, median of {args.outer})")
     print()
 
-    bench_rotation_standalone()
-    bench_pipeline()
-    bench_cublas_baseline()
-    bench_speedup_table()
-    bench_cuda_graph_capture()
+    bench_rotation_standalone(args.inner, args.outer)
+    bench_pipeline(args.inner, args.outer)
+    bench_cublas_baseline(args.inner, args.outer)
+    bench_speedup_table(args.inner, args.outer)
+
+
+if __name__ == "__main__":
+    main()