fix: update bench to produce timing

benchmark.py

@@ -5,20 +5,43 @@
 #   "kernels"
 # ]
 # ///
+import argparse
+import json
+import time
+
 import torch
 from kernels import get_kernel
 
-# 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
-)
+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:]
@@ -27,16 +50,50 @@ def reference_silu_and_mul(x):
     return torch.nn.functional.silu(x[..., :d]) * x[..., d:]
 
 
-out_ref = reference_silu_and_mul(input_tensor)
+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)
+
+    # 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)}")
+
+    # Run benchmark
+    timing_results = run_benchmark(
+        activation.silu_and_mul,
+        out_kernel,
+        input_tensor,
+        args.iterations,
+        args.warmup
+    )
+
+    # Output JSON to stdout (required protocol)
+    output = {"timing_results": timing_results}
+    print(json.dumps(output))
 
-# 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)
 
-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)}")
-# Reference output: torch.Size([128, 512])
-# Kernel output: torch.Size([128, 512])
-# Outputs close: True
+if __name__ == "__main__":
+    main()