Add CUDA graph mode, stddev, and O shape to tiled vs flat benchmark

- Add --graph flag for CUDA graph replay timing (kernel-only, no dispatch overhead)
- Add --trials flag with stddev reporting across multiple trials
- Add missing O (4096x2048) shape to match full Qwen 72B shape set
- Results: tiled layout 5-30% slower than flat on large shapes, neutral on KV

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

Author: TimDettmers

benchmarks/bench_tiled_vs_flat.py

@@ -23,21 +23,28 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--ncu", action="store_true", help="NCU mode: single iteration, no timing")
+parser.add_argument("--graph", action="store_true", help="Use CUDA graph replay for accurate kernel timing")
 parser.add_argument("--warmup", type=int, default=20, help="Warmup iterations")
-parser.add_argument("--iters", type=int, default=100, help="Timed iterations")
+parser.add_argument("--iters", type=int, default=100, help="Timed iterations per trial")
+parser.add_argument("--trials", type=int, default=5, help="Number of trials for stddev (graph mode)")
 args = parser.parse_args()
 
 SHAPES = [
     ("gateup", 2048, 5120),
     ("down", 5120, 2048),
     ("Q", 2048, 4096),
+    ("O", 4096, 2048),
     ("KV", 2048, 512),
 ]
 K_VALUES = [2, 3, 4, 5]
 M_VALUES = [1, 2, 4]
 
-print(f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}  {'flat_us':>8} {'tiled_us':>8} {'diff%':>7}")
-print("-" * 60)
+if args.graph:
+    print(f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}  {'flat_us':>8} {'±flat':>6} {'tiled_us':>8} {'±tiled':>6} {'diff%':>7}")
+    print("-" * 76)
+else:
+    print(f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}  {'flat_us':>8} {'tiled_us':>8} {'diff%':>7}")
+    print("-" * 60)
 
 for name, K_dim, N in SHAPES:
     for k in K_VALUES:
@@ -68,44 +75,78 @@
                 print(f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}  {'ncu':>8} {'ncu':>8} {'ncu':>7}")
                 continue
 
-            # CUDA events timing
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
 
-            # --- Flat ---
-            for _ in range(args.warmup):
+            def call_flat():
                 torch.ops.bitsandbytes.kbit_scalar_gemv.out(
                     A, packed_flat, absmax_flat, codebook, K_dim, N, k, out_flat
                 )
-            torch.cuda.synchronize()
 
-            start.record()
-            for _ in range(args.iters):
-                torch.ops.bitsandbytes.kbit_scalar_gemv.out(
-                    A, packed_flat, absmax_flat, codebook, K_dim, N, k, out_flat
-                )
-            end.record()
-            torch.cuda.synchronize()
-            flat_us = start.elapsed_time(end) * 1000 / args.iters  # ms -> us
-
-            # --- Tiled ---
-            for _ in range(args.warmup):
+            def call_tiled():
                 torch.ops.bitsandbytes.kbit_scalar_gemv_tiled_(
                     A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out_tiled
                 )
-            torch.cuda.synchronize()
 
-            start.record()
-            for _ in range(args.iters):
-                torch.ops.bitsandbytes.kbit_scalar_gemv_tiled_(
-                    A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out_tiled
-                )
-            end.record()
-            torch.cuda.synchronize()
-            tiled_us = start.elapsed_time(end) * 1000 / args.iters
+            if args.graph:
+                import statistics
+
+                # CUDA graph replay — measures kernel-only time
+                for fn in (call_flat, call_tiled):
+                    for _ in range(3):
+                        fn()
+                torch.cuda.synchronize()
+
+                def bench_graph(fn, trials, iters):
+                    s = torch.cuda.Stream()
+                    s.wait_stream(torch.cuda.current_stream())
+                    with torch.cuda.stream(s):
+                        g = torch.cuda.CUDAGraph()
+                        with torch.cuda.graph(g, stream=s):
+                            fn()
+                    torch.cuda.synchronize()
+                    times = []
+                    for _ in range(trials):
+                        start.record()
+                        for _ in range(iters):
+                            g.replay()
+                        end.record()
+                        torch.cuda.synchronize()
+                        times.append(start.elapsed_time(end) * 1000 / iters)
+                    return statistics.mean(times), statistics.stdev(times) if len(times) > 1 else 0.0
+
+                flat_us, flat_std = bench_graph(call_flat, args.trials, args.iters)
+                tiled_us, tiled_std = bench_graph(call_tiled, args.trials, args.iters)
+            else:
+                # CUDA events timing (includes Python dispatch overhead)
+                for _ in range(args.warmup):
+                    call_flat()
+                torch.cuda.synchronize()
+                start.record()
+                for _ in range(args.iters):
+                    call_flat()
+                end.record()
+                torch.cuda.synchronize()
+                flat_us = start.elapsed_time(end) * 1000 / args.iters
+
+                for _ in range(args.warmup):
+                    call_tiled()
+                torch.cuda.synchronize()
+                start.record()
+                for _ in range(args.iters):
+                    call_tiled()
+                end.record()
+                torch.cuda.synchronize()
+                tiled_us = start.elapsed_time(end) * 1000 / args.iters
 
             diff_pct = (tiled_us - flat_us) / flat_us * 100
-            print(f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}  {flat_us:>8.1f} {tiled_us:>8.1f} {diff_pct:>+7.1f}%")
+            if args.graph:
+                print(
+                    f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}"
+                    f"  {flat_us:>8.1f} {flat_std:>5.1f}σ {tiled_us:>8.1f} {tiled_std:>5.1f}σ {diff_pct:>+7.1f}%"
+                )
+            else:
+                print(f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}  {flat_us:>8.1f} {tiled_us:>8.1f} {diff_pct:>+7.1f}%")
 
         # Correctness check (once per shape/k)
         assert torch.equal(out_flat, out_tiled) or torch.allclose(out_flat, out_tiled, rtol=0.05, atol=0.1), (