Kernel Benchmark

kernel-benchmark

Use this skill to benchmark a custom GPU kernel against a PyTorch reference implementation.

What It Does

scripts/benchmark.py compares a solution against selectable baselines:

The benchmark validates correctness against the PyTorch eager reference before timing. Selected baselines are also checked against the eager reference. The report is written to <output-dir>/benchmark.md.

Default timing uses KernelBench-style CUDA events with explicit warmup/trial counts and cold-cache L2 thrashing. Use --timing-method=host_time only when diagnosing Python/runtime overhead.

Required Interfaces

Reference And Baselines

The reference file must define:

def reference(**kwargs):
    ...

The function receives the same tensor/scalar kwargs as the solution. It may update output tensors in place or return a tensor, tuple/list of tensors, or dict keyed by output name. Optional scalar and dimension parameters are passed through kwargs.

When --baselines=flashinfer is selected, the runner calls the same function with baseline="flashinfer" in kwargs. When no baseline kwarg is provided, reference(**kwargs) is the PyTorch eager reference.

CUDA-C++ / CUTLASS Solution

CUDA-C++ and CUTLASS solutions are .cu files with a compiled shared library beside them. The script loads the shared library; it does not compile it.

The .cu file must expose:

extern "C" void solve(...);

Compile first:

nvcc -shared -std=c++17 -arch=sm_90 -O3 -Xcompiler -fPIC -o kernel.so kernel.cu

CUTLASS code must be wrapped by solve(...); compile with the required CUTLASS include/library flags.

CuTe DSL / Triton Solution

CuTe DSL and Triton solutions are .py files and must define:

def setup(**kwargs):
    return {"inputs": {...}, "outputs": ["out"]}

def run_kernel(**kwargs):
    ...

setup()["inputs"] must contain every tensor/scalar passed to the kernel and reference. setup()["outputs"] names the tensors compared for correctness.

Supported Implementations

--backend is kept as a compatibility alias for --implementation. Prefer --implementation in new commands.

Usage

CUDA-C++:

python scripts/benchmark.py kernel.cu \
  --implementation=cuda-cpp \
  --baselines=pytorch-eager \
  --ref=ref.py \
  --output-dir=bench_out \
  --M=1024 --N=1024

CUTLASS:

python scripts/benchmark.py cutlass_kernel.cu \
  --implementation=cutlass \
  --baselines=pytorch-eager,torch-compile \
  --ref=ref.py \
  --output-dir=bench_out \
  --M=1024 --N=1024

Triton:

python scripts/benchmark.py kernel.py \
  --implementation=triton \
  --baselines=pytorch-eager,flashinfer \
  --ref=ref.py \
  --output-dir=bench_out \
  --M=1024 --N=1024

CuTe DSL:

python scripts/benchmark.py cute_kernel.py \
  --implementation=cute-dsl \
  --ref=ref.py \
  --output-dir=bench_out \
  --M=1024 --N=1024

--implementation=auto selects Triton for .py files and CUDA-C++ otherwise.

Options

Keep dimensions, seed, pointer sizing, GPU, and tolerances fixed when comparing kernel versions.

Output

If correctness fails, inspect the output tensors before using the timing result for optimization decisions.

For GPU-kernel latency, use the default cuda_event. For end-to-end Python/runtime overhead diagnostics, use host_time.