feat: reduced bench script

benchmark.py

@@ -1,99 +1,28 @@
-# /// script
-# dependencies = [
-#   "numpy",
-#   "torch",
-#   "kernels"
-# ]
-# ///
-import argparse
-import json
-import time
-
 import torch
-from kernels import get_kernel
-
-
-def run_benchmark(kernel_fn, out_tensor, input_tensor, iterations: int, warmup: int) -> dict:
-    """Benchmark a kernel function."""
-
-    # Warmup phase (not measured)
-    for _ in range(warmup):
-        kernel_fn(out_tensor, input_tensor)
-        torch.cuda.synchronize()
-
-    # Benchmark phase
-    times_ms = []
-    for _ in range(iterations):
-        start = time.perf_counter()
-        kernel_fn(out_tensor, input_tensor)
-        torch.cuda.synchronize()
-        end = time.perf_counter()
-        times_ms.append((end - start) * 1000)
-
-    # Calculate statistics
-    mean_ms = sum(times_ms) / len(times_ms)
-    variance = sum((t - mean_ms) ** 2 for t in times_ms) / len(times_ms)
-    std_ms = variance ** 0.5
-
-    return {
-        "mean_ms": round(mean_ms, 4),
-        "std_ms": round(std_ms, 4),
-        "min_ms": round(min(times_ms), 4),
-        "max_ms": round(max(times_ms), 4),
-        "iterations": iterations,
-    }
-
-
-# Reference implementation: SwiGLU computes silu(x[:d]) * x[d:]
-def reference_silu_and_mul(x):
-    d = x.shape[-1] // 2
-    return torch.nn.functional.silu(x[..., :d]) * x[..., d:]
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Benchmark activation kernel")
-    parser.add_argument("--iterations", type=int, default=100, help="Number of iterations")
-    parser.add_argument("--warmup", type=int, default=10, help="Number of warmup iterations")
-    args = parser.parse_args()
-
-    # Setup
-    torch.manual_seed(42)
-    activation = get_kernel("kernels-community/activation")
-    device = torch.device("cuda")
-
-    # Create test tensor for SwiGLU (SiLU and Mul)
-    # Input has shape (num_tokens, 2 * hidden_dim)
-    num_tokens, hidden_dim = 128, 512
-    input_tensor = torch.randn(
-        num_tokens, 2 * hidden_dim, device=device, dtype=torch.float16
-    )
-
-    # Reference output for correctness check
-    out_ref = reference_silu_and_mul(input_tensor)
+from kernels.benchmark import Benchmark
 
-    # Custom kernel SwiGLU
-    out_shape = input_tensor.shape[:-1] + (hidden_dim,)
-    out_kernel = torch.empty(out_shape, dtype=input_tensor.dtype, device=device)
-    out_kernel = activation.silu_and_mul(out_kernel, input_tensor)
 
-    # Correctness check
-    print(f"Reference output: {out_ref.shape}")
-    print(f"Kernel output: {out_kernel.shape}")
-    print(f"Outputs close: {torch.allclose(out_kernel, out_ref, atol=1e-3, rtol=1e-3)}")
+class SiluBenchmark(Benchmark):
+    """Benchmark for SwiGLU activation kernel."""
 
-    # Run benchmark
-    timing_results = run_benchmark(
-        activation.silu_and_mul,
-        out_kernel,
-        input_tensor,
-        args.iterations,
-        args.warmup
-    )
+    kernel_id = "kernels-community/activation"
+    seed = 42
 
-    # Output JSON to stdout (required protocol)
-    output = {"timing_results": timing_results}
-    print(json.dumps(output))
+    def setup(self):
+        # Input has shape (num_tokens, 2 * hidden_dim)
+        num_tokens, hidden_dim = 128, 512
+        self.x = torch.randn(
+            num_tokens, 2 * hidden_dim, device="cuda", dtype=torch.float16
+        )
+        self.out = torch.empty(
+            num_tokens, hidden_dim, device="cuda", dtype=torch.float16
+        )
 
+    def benchmark_silu_and_mul(self):
+        self.kernel.silu_and_mul(self.out, self.x)
 
-if __name__ == "__main__":
-    main()
+    def verify_silu_and_mul(self):
+        # Reference: SwiGLU computes silu(x[:d]) * x[d:]
+        d = self.x.shape[-1] // 2
+        ref = torch.nn.functional.silu(self.x[..., :d]) * self.x[..., d:]
+        return torch.allclose(self.out, ref, atol=1e-3, rtol=1e-3)