tl;dr: We extended our kernel generation research to CUDA, and benchmarked results on an H100. The generated kernels achieve a ~1.8X speedup on KernelBench problems over the baseline PyTorch (including torch.compile).
Introduction
Autonomous kernel generation by AI models is an active research area, because it offers the chance to speed up AI workloads without low-level kernel expertise. This approach could potentially allow ML researchers to stay in pure PyTorch, while gaining the performance of optimized kernels. It can also allow for more immediate optimized support for new models across different hardware platforms.
In an earlier post, we shared benchmark results for Metal kernels generated for Apple devices. We found in that work that an agentic swarm (where multiple agents work together to find the best implementation) provides significantly better performance than a single agent.
We wanted to extend this work to generating CUDA kernels, since NVIDIA devices are widely used for AI workloads. Going into the investigation, it wasn’t clear whether the agentic swarm would be better at improving performance for Metal or CUDA kernels. On one hand, there are far more examples of CUDA kernels online for the agents to learn from. On the other hand, more time has been spent optimizing PyTorch for NVIDIA hardware than for Apple devices.
It turns out that the results are comparable, with AI-generated CUDA kernels performing slightly better than AI-generated Metal kernels. Let’s dive into the results, but first cover how these results were collected and compared against the baseline.
System architecture for our agentic swarm. It takes as input a PyTorch implementation as the prompt, and multiple work to identify the best solution. Candidate kernels are verified for numerical equivalence and performance against corresponding reference implementations.
Benchmarking methodology
Benchmarking kernel performance has a lot of subtle pitfalls, so we are sharing the benchmark code used to assess kernel performance. It’s important when running PyTorch benchmarks to ensure that warmups are performed, caches are cleared, and synchronization is conducted.
import argparse
import time
from dataclasses import dataclass, field
from pathlib import Path
import polars as pl
import torch
import torch.nn as nn
import torch.utils.benchmark as benchmark
if torch.cuda.is_available():
import triton
TRITON_AVAILABLE = True
else:
TRITON_AVAILABLE = False
def set_seed(seed: int):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.mps.manual_seed(seed)
@dataclass(repr=False)
class Perf:
name: str
ellapsed_time: list[float]
n_trials: int
time_mean: float = field(init=False)
time_median: float = field(init=False)
time_std: float = field(init=False)
time_min: float = field(init=False)
time_max: float = field(init=False)
def __post_init__(self):
assert len(self.ellapsed_time) == self.n_trials
time_th = torch.Tensor(self.ellapsed_time)
self.time_mean = time_th.mean()
self.time_median = time_th.median()
self.time_std = time_th.std()
self.time_min = time_th.min()
self.time_max = time_th.max()
def __repr__(self) -> str:
return (
f"Perf: time in ms\n"
f"{'n_trials':<10} = {self.n_trials}\n"
f"{'mean':<10} = {self.time_mean:.4f}\n"
f"{'median':<10} = {self.time_median:.4f}\n"
f"{'std':<10} = {self.time_std:.4f}\n"
f"{'min':<10} = {self.time_min:.4f}\n"
f"{'max':<10} = {self.time_max:.4f}\n"
)
def stat_dict(self) -> dict:
return {
"name": self.name,
"n_trials": self.n_trials,
"mean": self.time_mean,
"median": self.time_median,
"std": self.time_std,
"min": self.time_min,
"max": self.time_max,
}
@dataclass(repr=False)
class ComparePerf:
perf_results: list[Perf]
def __repr__(self) -> str:
perf_df = pl.from_dicts(perf.stat_dict() for perf in self.perf_results)
return f"{perf_df}"
@dataclass(repr=False)
class CompareMeasurement:
perf_results: list[benchmark.Measurement]
def __repr__(self) -> str:
perf_df = pl.from_dicts(
{
"name": perf.task_spec.summarize().replace("\n", " "),
"n_trials": perf.number_per_run,
"median": perf.median * 1e3,
"mean": perf.mean * 1e3,
}
for perf in self.perf_results
)
return f"{perf_df}"
def get_empty_cache_for_benchmark(device):
cache_size = 256 * 1024 * 1024
return torch.empty(int(cache_size // 4), dtype=torch.int, device=device)
def clear_cache(cache):
cache.zero_()
def compare_perf(perf_results: list[Perf | benchmark.Measurement]):
if isinstance(perf_results[0], Perf):
return ComparePerf(perf_results)
elif isinstance(perf_results[0], benchmark.Measurement):
# proceed with benchmark.Compare
return CompareMeasurement(perf_results)
def load_reference(ref_src: str, ctx):
try:
compile(ref_src, "<string>", "exec")
except SyntaxError as e:
raise e
try:
exec(ref_src, ctx)
except Exception as e:
raise RuntimeError(f"Error in executing src model. {e}") from e
# these should be defined in the original model code and present in the context
get_init_inputs_fn = ctx.get("get_init_inputs")
get_inputs_fn = ctx.get("get_inputs")
Model = ctx.get("Model")
return Model, get_init_inputs_fn, get_inputs_fn
def load_target(tgt_src: str, ctx: dict, build_dir: str = None) -> nn.Module:
if build_dir:
tgt_src = (
f"""
import os
os.environ['TORCH_EXTENSIONS_DIR'] = '{build_dir}'
"""
) + tgt_src
local_ctx = {}
try:
compile(tgt_src, "<string>", "exec")
except SyntaxError as e:
raise e
try:
exec(tgt_src, local_ctx)
except Exception as e:
raise RuntimeError(f"Error in executing tgt model. {e}") from e
ModelNew = local_ctx.get("ModelNew")
# update execution context with ModelNew only
ctx["ModelNew"] = ModelNew
return ModelNew
@torch.inference_mode()
def benchmark_triton(model, inputs, num_trials=100, num_warmup=3) -> Perf:
def func():
model(*inputs)
bench_res = triton.testing.do_bench(
func,
warmup=num_warmup,
rep=num_trials,
return_mode="all",
)
perf = Perf(model.__class__.__name__, bench_res, len(bench_res))
return perf
@torch.inference_mode()
def benchmark_event(model, inputs, num_trials=100, num_warmup=3) -> Perf:
device = inputs[0].device
# warm up
for _ in range(num_warmup):
_ = model(*inputs)
torch.accelerator.synchronize()
cache = get_empty_cache_for_benchmark(device)
start_event = [
torch.Event(device=device, enable_timing=True) for _ in range(num_trials)
]
end_event = [
torch.Event(device=device, enable_timing=True) for _ in range(num_trials)
]
for i in range(num_trials):
clear_cache(cache)
start_event[i].record()
_ = model(*inputs)
end_event[i].record()
torch.accelerator.synchronize()
exec_time = [start_event[i].elapsed_time(end_event[i]) for i in range(num_trials)]
perf = Perf(model.__class__.__name__, exec_time, num_trials)
return perf
@torch.inference_mode()
def benchmark_simple(model, inputs, num_trials=100, num_warmup=3) -> Perf:
# warm up
for _ in range(num_warmup):
_ = model(*inputs)
torch.accelerator.synchronize()
exec_time = []
for _ in range(num_trials):
start_time = time.perf_counter()
_ = model(*inputs)
torch.accelerator.synchronize()
end_time = time.perf_counter()
exec_time.append((end_time - start_time) * 1e3)
perf = Perf(model.__class__.__name__, exec_time, num_trials)
return perf
@torch.inference_mode()
def check_correctness(model_src, model_tgt, inputs):
src_out = model_src(*inputs)
tgt_out = model_tgt(*inputs)
is_close = torch.allclose(src_out, tgt_out, atol=1e-02, rtol=1e-02)
print(f"Results match: {is_close}")
if not is_close:
abs_diff = torch.abs(src_out - tgt_out)
max_diff = torch.max(abs_diff).item()
avg_diff = torch.mean(abs_diff).item()
print(f"max abs diff: {max_diff:.8f}")
print(f"mean abs diff: {avg_diff:.8f}")
return is_close
def benchmark_problem(ref_src, tgt_src, device, num_trials, num_warmup):
set_seed(args.seed)
ctx = {}
Model, get_init_inputs_fn, get_inputs_fn = load_reference(ref_src, ctx)
ModelNew = load_target(tgt_src, ctx)
init_data = [
itm.to(device=args.device) if isinstance(itm, torch.Tensor) else itm
for itm in get_init_inputs_fn()
]
inputs = [
itm.to(device=args.device) if isinstance(itm, torch.Tensor) else itm
for itm in get_inputs_fn()
]
set_seed(args.seed)
model_tgt = ModelNew(*init_data)
set_seed(args.seed)
model_src = Model(*init_data)
model_tgt.to(args.device)
model_src.to(args.device)
torch.accelerator.synchronize()
bench_funcs = [benchmark_simple, benchmark_event]
if TRITON_AVAILABLE:
bench_funcs.append(benchmark_triton)
perfs = {}
for bench_fn in bench_funcs:
perfs[bench_fn] = {}
for model in [model_src, model_tgt]:
perf_result = bench_fn(model, inputs, args.num_trials, args.num_warmup)
perfs[bench_fn][model.__class__.__name__] = perf_result
comparison = compare_perf(list(perfs[bench_fn].values()))
print(bench_fn.__name__, comparison)
check_correctness(model_src, model_tgt, inputs)
return perfs
def main(args):
# benchmark all problems in the target directory
if args.tgt_dir:
from kernel_gen.dataset import construct_kernelbench_dataset
from kernel_gen.utils.data import fetch_ref_arch_from_problem_id
from kernel_gen.utils.file import get_level_problem_id_from_gen_arc_path
args.tgt_dir = Path(args.tgt_dir)
for tgt_problem_path in args.tgt_dir.glob("**/*kernel.py"):
lvl_id, prob_id = get_level_problem_id_from_gen_arc_path(tgt_problem_path)
dataset = construct_kernelbench_dataset(lvl_id)
ref_src = fetch_ref_arch_from_problem_id(dataset, prob_id, "local")
with open(tgt_problem_path, "r") as f:
tgt_src = f.read()
benchmark_problem(
ref_src,
tgt_src,
args.device,
args.num_trials,
args.num_warmup,
)
# benchmark a single problem
else:
with open(args.ref_problem, "r") as f:
ref_src = f.read()
with open(args.tgt_problem, "r") as f:
tgt_src = f.read()
benchmark_problem(
ref_src,
tgt_src,
args.device,
args.num_trials,
args.num_warmup,
)
if __name__ == "__main__":
# torch.set_num_threads(1)
parser = argparse.ArgumentParser()
parser.add_argument("--ref-problem", type=str)
parser.add_argument("--tgt-problem", type=str)
parser.add_argument("--tgt-dir", type=str, default=None)
parser.add_argument("--num-trials", type=int, default=100)
parser.add_argument("--num-warmup", type=int, default=3)
parser.add_argument("--device")
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
args.device = torch.device(args.device)
main(args)For this post, we benchmarked the CUDA kernel generation on an H100, and will follow with results on additional NVIDIA GPUs.
For the input PyTorch models, we used the original KernelBench Levels 1-31. For those who are unfamiliar, KernelBench is a popular benchmark of neural network operators, building blocks and end to end architectures of varying complexity implemented as PyTorch modules. Level 1 problems are the least complex, and Level 3 problems are the most complex. We omitted a few cases from the original set of problems, because the agents found (technically correct) shortcuts compared to the original implementations that provided disproportionate speedups.
In this work, we always use 1X as the baseline speedup (no speedup), because our system will automatically fall back to the baseline implementation when the generated kernels are slower. Additionally, as before, we use geometric mean to compute the average speedup2.
What about torch.compile?
torch.compile is a great PyTorch feature that automatically speeds up PyTorch code using JIT compilation, kernel fusion, and graph analysis. It’s widely adopted, and it provides great support for CUDA (Metal support is in progress).
We want to compare generated kernels against a realistic scenario, not an artificially low bar. Actual practitioners who care about improving model speed would likely at least try torch.compile. As a result, we compare our generated kernels against torch.compile baselines3.
Sometimes eager mode is actually faster in the version of KernelBench we tried. This surprised us, and we welcome any feedback on our benchmarking methodology. This is more common for the simpler modules (e.g. Level 1) - which makes sense because complex modules are more likely to benefit from kernel fusion. Additionally, for smaller runtimes, the overhead of torch.compile may dominate. As a result, we compare our generated kernels to whichever baseline is the higher bar for each problem. (We expect that the win rate of torch.compile increases for the later version of KernelBench, which we will switch to in future posts.)
Chart showing whether torch.compile or eager mode is faster in our measurements on an H100, broken down by KernelBench level. We can see that torch.compile starts gaining an advantage as the problems become more complex.
Show me the results!
We see similar results to the generated Metal kernels: about 1.8X speedup when compared to the better-performing PyTorch mode (eager mode or torch.compile).
Results for the AI-generated CUDA kernels on the H100. We can see that the average speedup is about 1.8X when compared to the better-performing PyTorch mode (eager vs torch.compile).
Focusing on the green bars, which provide the most rigorous comparison: We can see that the results are strongest for the less complex PyTorch modules (2.3X), and remain strong but drop as we get to Level 3 (1.4X). This tells us that the agentic swarm likely has room for improvement, since more complex modules should actually have more opportunities for performance optimization. (Individual ops are already well-optimized in PyTorch).
Kernel fusion (where multiple kernels are rewritten as a single kernel) is a key part of improving performance, and most of the AI generated kernels make use of fusion. However, other techniques beyond fusion show up as well. For example, the agentic swarm generated fused kernels for Level 2 Problem 94 (Gemm_BiasAdd_Hardtanh_Mish_GroupNorm), but it also added a few other optimizations. There is a hardware-specific optimization of enabling tensor cores (which are specific to NVIDIA). There was also an algorithmic optimization, where different kernel implementations are selected based on the group size in the GroupNorm layer.
Gemm_BiasAdd_Hardtanh_Mish_GroupNorm (Level 2 Problem 94)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
A model that performs a GEMM, BiasAdd, Hardtanh, Mish, and GroupNorm operations in sequence.
"""
def __init__(self, in_features, out_features, bias_shape, num_groups):
super(Model, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.hardtanh = nn.Hardtanh()
self.mish = nn.Mish()
self.groupnorm = nn.GroupNorm(num_groups=num_groups, num_channels=out_features)
def forward(self, x):
"""
Args:
x (torch.Tensor): Input tensor of shape (batch_size, in_features).
Returns:
torch.Tensor: Output tensor of shape (batch_size, out_features).
"""
x = self.gemm(x)
x = x + self.bias
x = self.hardtanh(x)
x = self.mish(x)
x = self.groupnorm(x)
return x
batch_size = 128
in_features = 512
out_features = 1024
bias_shape = (out_features,)
num_groups = 32
def get_inputs():
return [torch.randn(batch_size, in_features)]
def get_init_inputs():
return [in_features, out_features, bias_shape, num_groups]import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
// Simple checks
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT32(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be float32")
#define CHECK_INPUT_F32(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x); CHECK_FLOAT32(x)
#ifndef WARP_SIZE
#define WARP_SIZE 32
#endif
__device__ __forceinline__ float hardtanhf_fast(float x, float min_v, float max_v) {
return fminf(fmaxf(x, min_v), max_v);
}
// Numerically-stable softplus: log1p(exp(-|x|)) + max(x, 0)
__device__ __forceinline__ float softplusf_stable_fast(float x) {
float ax = fabsf(x);
return log1pf(__expf(-ax)) + fmaxf(x, 0.0f);
}
__device__ __forceinline__ float mishf_fast(float x) {
float sp = softplusf_stable_fast(x);
return x * tanhf(sp);
}
// Warp reduce sum using shuffles
__device__ __forceinline__ float warp_sum(float v) {
unsigned mask = 0xffffffffu;
v += __shfl_down_sync(mask, v, 16);
v += __shfl_down_sync(mask, v, 8);
v += __shfl_down_sync(mask, v, 4);
v += __shfl_down_sync(mask, v, 2);
v += __shfl_down_sync(mask, v, 1);
return v;
}
static inline int next_pow2_int(int v) {
if (v <= 1) return 1;
v--;
v |= v >> 1;
v |= v >> 2;
v |= v >> 4;
v |= v >> 8;
v |= v >> 16;
return v + 1;
}
// Fast path: One warp handles one (sample, group). No shared memory, keep values in registers.
// Requires: group_size <= 32
__global__ void fused_bias_act_gn_warp_kernel(
const float* __restrict__ z, // [N, C]
const float* __restrict__ b1, // [C] linear bias (may be nullptr)
const float* __restrict__ b2, // [C] extra bias (may be nullptr)
const float* __restrict__ gamma, // [C]
const float* __restrict__ beta, // [C]
float* __restrict__ y, // [N, C]
int N, int C, int G,
float ht_min, float ht_max, float eps
){
const int n = blockIdx.y;
const int g = blockIdx.x;
const int lane = threadIdx.x & 31;
const int group_size = C / G;
if (group_size > WARP_SIZE) return;
const int c0 = g * group_size;
const int base = n * C + c0;
float m = 0.0f;
if (lane < group_size) {
const int c = c0 + lane;
float v = z[base + lane];
if (b1) v += b1[c];
if (b2) v += b2[c];
v = hardtanhf_fast(v, ht_min, ht_max);
m = mishf_fast(v);
}
float sum = warp_sum(m);
float sq = warp_sum(m * m);
unsigned mask = 0xffffffffu;
float sum_all = __shfl_sync(mask, sum, 0);
float sq_all = __shfl_sync(mask, sq, 0);
float mean = sum_all / (float)group_size;
float var = sq_all / (float)group_size - mean * mean;
float rstd = rsqrtf(var + eps);
if (lane < group_size) {
const int c = c0 + lane;
float gmul = gamma ? gamma[c] : 1.0f;
float gbias = beta ? beta[c] : 0.0f;
float gn = (m - mean) * rstd;
y[base + lane] = gn * gmul + gbias;
}
}
// General kernel: shared mem reduction, any group size.
__global__ void fused_bias_act_gn_block_kernel(
const float* __restrict__ z, // [N, C]
const float* __restrict__ b1, // [C]
const float* __restrict__ b2, // [C]
const float* __restrict__ gamma, // [C]
const float* __restrict__ beta, // [C]
float* __restrict__ y, // [N, C]
int N, int C, int G,
float ht_min, float ht_max, float eps
){
const int n = blockIdx.y;
const int g = blockIdx.x;
const int group_size = C / G;
const int c0 = g * group_size;
const int base = n * C + c0;
extern __shared__ float smem[];
float* sm_vals = smem; // group_size floats for values
float* ssum = sm_vals + group_size; // blockDim.x floats for sum
float* ssum2 = ssum + blockDim.x; // blockDim.x floats for sum2
const int tid = threadIdx.x;
const int stride = blockDim.x;
float local_sum = 0.0f;
float local_sum2 = 0.0f;
// Pass 1: compute activation and accumulate sums; keep values in shared memory
for (int i = tid; i < group_size; i += stride) {
const int c = c0 + i;
float v = z[base + i];
if (b1) v += b1[c];
if (b2) v += b2[c];
v = hardtanhf_fast(v, ht_min, ht_max);
float m = mishf_fast(v);
sm_vals[i] = m;
local_sum += m;
local_sum2 += m * m;
}
ssum[tid] = local_sum;
ssum2[tid] = local_sum2;
__syncthreads();
// Block reduction
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
if (tid < offset) {
ssum[tid] += ssum[tid + offset];
ssum2[tid] += ssum2[tid + offset];
}
__syncthreads();
}
float mean = ssum[0] / (float)group_size;
float var = ssum2[0] / (float)group_size - mean * mean;
float rstd = rsqrtf(var + eps);
// Pass 2: normalize and affine
for (int i = tid; i < group_size; i += stride) {
const int c = c0 + i;
float m = sm_vals[i];
float gmul = gamma ? gamma[c] : 1.0f;
float gbias = beta ? beta[c] : 0.0f;
float gn = (m - mean) * rstd;
y[base + i] = gn * gmul + gbias;
}
}
// Persistent group kernel: cache per-channel parameters in shared memory and process a tile of N.
// Greatly reduces parameter bandwidth when N is large.
__global__ void fused_bias_act_gn_group_persistent_kernel(
const float* __restrict__ z, // [N, C]
const float* __restrict__ b1, // [C]
const float* __restrict__ b2, // [C]
const float* __restrict__ gamma, // [C]
const float* __restrict__ beta, // [C]
float* __restrict__ y, // [N, C]
int N, int C, int G,
float ht_min, float ht_max, float eps,
int tileN
){
const int g = blockIdx.x;
const int group_size = C / G;
const int c0 = g * group_size;
// Tile of batch this block will process
const int n_start = blockIdx.y * tileN;
const int n_end = min(N, n_start + tileN);
const int tid = threadIdx.x;
const int stride = blockDim.x;
extern __shared__ float smem[];
float* s_bsum = smem; // [group_size]
float* s_gamma = s_bsum + group_size; // [group_size]
float* s_beta = s_gamma + group_size; // [group_size]
float* sm_vals = s_beta + group_size; // [group_size]
float* ssum = sm_vals + group_size; // [blockDim.x]
float* ssum2 = ssum + blockDim.x; // [blockDim.x]
// Load per-channel params once per block
for (int i = tid; i < group_size; i += stride) {
const int c = c0 + i;
float bsum = 0.0f;
if (b1) bsum += b1[c];
if (b2) bsum += b2[c];
s_bsum[i] = bsum;
s_gamma[i] = gamma ? gamma[c] : 1.0f;
s_beta[i] = beta ? beta[c] : 0.0f;
}
__syncthreads();
for (int n = n_start; n < n_end; ++n) {
const int base = n * C + c0;
float local_sum = 0.0f;
float local_sum2 = 0.0f;
// Pass 1: load, bias, hardtanh, mish; accumulate stats
for (int i = tid; i < group_size; i += stride) {
float v = z[base + i] + s_bsum[i];
v = hardtanhf_fast(v, ht_min, ht_max);
float m = mishf_fast(v);
sm_vals[i] = m;
local_sum += m;
local_sum2 += m * m;
}
ssum[tid] = local_sum;
ssum2[tid] = local_sum2;
__syncthreads();
// Reduce sums across block
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
if (tid < offset) {
ssum[tid] += ssum[tid + offset];
ssum2[tid] += ssum2[tid + offset];
}
__syncthreads();
}
float mean = ssum[0] / (float)group_size;
float var = ssum2[0] / (float)group_size - mean * mean;
float rstd = rsqrtf(var + eps);
// Pass 2: groupnorm affine and store
for (int i = tid; i < group_size; i += stride) {
float gn = (sm_vals[i] - mean) * rstd;
y[base + i] = gn * s_gamma[i] + s_beta[i];
}
__syncthreads();
}
}
torch::Tensor fused_bias_act_groupnorm(
torch::Tensor z, // [N, C], float32, CUDA
torch::Tensor b1, // [C], float32, CUDA
torch::Tensor b2, // [C], float32, CUDA
torch::Tensor gamma, // [C], float32, CUDA
torch::Tensor beta, // [C], float32, CUDA
int64_t G,
double ht_min,
double ht_max,
double eps
) {
CHECK_INPUT_F32(z);
CHECK_INPUT_F32(b1);
CHECK_INPUT_F32(b2);
CHECK_INPUT_F32(gamma);
CHECK_INPUT_F32(beta);
TORCH_CHECK(z.dim() == 2, "z must be 2D (N, C)");
int N = (int)z.size(0);
int C = (int)z.size(1);
TORCH_CHECK((int)gamma.size(0) == C && (int)beta.size(0) == C, "gamma/beta must have size C");
TORCH_CHECK((int)b1.size(0) == C && (int)b2.size(0) == C, "b1/b2 must have size C");
TORCH_CHECK(C % (int)G == 0, "num_channels must be divisible by num_groups");
auto y = torch::empty_like(z);
const int group_size = C / (int)G;
dim3 grid((unsigned int)G, (unsigned int)N, 1);
// Heuristic kernel choice:
// - Small groups => warp kernel.
// - Otherwise, try persistent group kernel (caches params, processes tile of N).
// - Fallback to generic block kernel if shared memory is too large for persistent.
bool use_warp = (group_size <= WARP_SIZE) && (N <= 64);
if (use_warp) {
dim3 block(WARP_SIZE, 1, 1);
fused_bias_act_gn_warp_kernel<<<grid, block, 0>>>(
z.data_ptr<float>(),
b1.data_ptr<float>(),
b2.data_ptr<float>(),
gamma.data_ptr<float>(),
beta.data_ptr<float>(),
y.data_ptr<float>(),
N, C, (int)G,
(float)ht_min, (float)ht_max, (float)eps
);
} else {
// Prefer the persistent kernel to reduce parameter bandwidth when N >= 64.
bool prefer_persistent = (N >= 64);
if (prefer_persistent) {
int threads = next_pow2_int(group_size);
if (threads > 256) threads = 256;
if (threads < 32) threads = 32;
// Choose a batch tile to get enough blocks while still reusing parameters
int tileN;
if (N >= 256) tileN = 16;
else if (N >= 128) tileN = 8;
else if (N >= 64) tileN = 4;
else tileN = 2;
dim3 grid_p((unsigned int)G, (unsigned int)((N + tileN - 1) / tileN), 1);
dim3 block((unsigned int)threads, 1, 1);
// Shared memory: 3*group_size (params) + group_size (vals) + 2*threads (sums)
size_t shmem = sizeof(float) * ( (size_t)group_size * 4 + (size_t)threads * 2 );
// If shared memory too large, fallback to the generic block kernel
if (shmem <= 96 * 1024) {
fused_bias_act_gn_group_persistent_kernel<<<grid_p, block, shmem>>>(
z.data_ptr<float>(),
b1.data_ptr<float>(),
b2.data_ptr<float>(),
gamma.data_ptr<float>(),
beta.data_ptr<float>(),
y.data_ptr<float>(),
N, C, (int)G,
(float)ht_min, (float)ht_max, (float)eps,
tileN
);
} else {
// Generic block kernel fallback
int threads_g = next_pow2_int(group_size);
if (threads_g > 256) threads_g = 256;
if (threads_g < 64) threads_g = 64;
dim3 block_g((unsigned int)threads_g, 1, 1);
size_t shmem_g = sizeof(float) * (group_size + threads_g * 2);
dim3 grid_g((unsigned int)G, (unsigned int)N, 1);
fused_bias_act_gn_block_kernel<<<grid_g, block_g, shmem_g>>>(
z.data_ptr<float>(),
b1.data_ptr<float>(),
b2.data_ptr<float>(),
gamma.data_ptr<float>(),
beta.data_ptr<float>(),
y.data_ptr<float>(),
N, C, (int)G,
(float)ht_min, (float)ht_max, (float)eps
);
}
} else {
// Generic block kernel (works for any shape)
int threads = next_pow2_int(group_size);
if (threads > 256) threads = 256;
if (threads < 64) threads = 64;
dim3 block((unsigned int)threads, 1, 1);
size_t shmem = sizeof(float) * (group_size + threads * 2);
dim3 grid_b((unsigned int)G, (unsigned int)N, 1);
fused_bias_act_gn_block_kernel<<<grid_b, block, shmem>>>(
z.data_ptr<float>(),
b1.data_ptr<float>(),
b2.data_ptr<float>(),
gamma.data_ptr<float>(),
beta.data_ptr<float>(),
y.data_ptr<float>(),
N, C, (int)G,
(float)ht_min, (float)ht_max, (float)eps
);
}
}
cudaError_t err = cudaGetLastError();
TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));
return y;
}
"""
cpp_src = r"""
torch::Tensor fused_bias_act_groupnorm(
torch::Tensor z,
torch::Tensor b1,
torch::Tensor b2,
torch::Tensor gamma,
torch::Tensor beta,
int64_t G,
double ht_min,
double ht_max,
double eps
);
"""
# Build the extension once
kernels = load_inline(
name="fused_epilogue_gn_kernels_v3",
cpp_sources=cpp_src,
cuda_sources=cuda_src,
functions=["fused_bias_act_groupnorm"],
extra_cflags=["-O3"],
extra_cuda_cflags=[
"-O3",
"--use_fast_math",
"-Xptxas=-O3,-dlcm=ca" # prefer caching in L2
],
verbose=False,
)
class ModelNew(nn.Module):
"""
Optimized model:
- cuBLAS-backed GEMM via F.linear (no bias in GEMM).
- Single fused CUDA kernel for: BiasAdd (gemm.bias + extra bias) + Hardtanh + Mish + GroupNorm.
- Heuristic kernel chooser:
* For small groups and small N: warp kernel (no shared memory).
* For larger N: persistent per-group kernel that caches parameters in shared memory and
processes a tile of the batch to reduce parameter bandwidth and improve cache locality.
* Generic block kernel as a fallback for very large groups or limited shared memory.
"""
def __init__(self, in_features, out_features, bias_shape, num_groups):
super().__init__()
# Keep bias=True to maintain parameter parity with original, but don't add it in GEMM
self.gemm = nn.Linear(in_features, out_features, bias=True)
self.bias = nn.Parameter(torch.randn(bias_shape))
self.groupnorm = nn.GroupNorm(num_groups=num_groups, num_channels=out_features)
self.ht_min = -1.0
self.ht_max = 1.0
@torch.no_grad()
def _prepare(self, x):
if self.gemm.weight.device != x.device or self.gemm.weight.dtype != x.dtype:
self.to(x.device, dtype=x.dtype)
# Enable TF32 on tensor cores for GEMM where applicable
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
def forward(self, x):
if not x.is_cuda:
# CPU fallback path: reference ops
z = F.linear(x, self.gemm.weight, bias=None)
z = z + self.gemm.bias + self.bias
z = torch.clamp(z, min=self.ht_min, max=self.ht_max)
z = z * torch.tanh(F.softplus(z))
z = self.groupnorm(z)
return z
self._prepare(x)
# GEMM (no bias here)
z = F.linear(x, self.gemm.weight, bias=None)
# Fused BiasAdd(+gemm.bias + extra bias) + Hardtanh + Mish + GroupNorm
y = kernels.fused_bias_act_groupnorm(
z.contiguous(),
self.gemm.bias.contiguous(),
self.bias.contiguous(),
self.groupnorm.weight.contiguous(),
self.groupnorm.bias.contiguous(),
int(self.groupnorm.num_groups),
float(self.ht_min),
float(self.ht_max),
float(self.groupnorm.eps),
)
return y
# Default shapes
batch_size = 128
in_features = 512
out_features = 1024
bias_shape = (out_features,)
num_groups = 32
def get_inputs():
return [torch.randn(batch_size, in_features, device="cuda", dtype=torch.float32)]
def get_init_inputs():
return [in_features, out_features, bias_shape, num_groups]Going back to evaluation, we wanted to look at the subset of KernelBench problems where torch.compile outperformed eager mode in our benchmarking on the H100. These 124 problems are meatier and more representative of cases that we care about optimizing - optimizing a smaller, already-fast module may not provide as much benefit. (Note that the later version of KernelBench includes larger input sizes, which helps mitigate this problem.) When we compare our generated kernels to the subset that’s better on torch.compile, we see that there is still a consistent speedup.
Results for the AI-generated CUDA kernels on the H100, only considering the 124 problems where torch.compile outperformed eager mode execution. We can see for this meatier subset of problems, there is still an average 1.6X speedup for the generated kernels.
In this work, we have been generating CUDA kernels for eager execution. However, in 2 cases, the swarm deviated from eager mode - it actually rewrote the PyTorch code to use torch.compile instead.
For example, on Problem 3, Level 9 (which is a ResNet model), the swarm rewrote the PyTorch code to create a fused implementation in Python, only dropping down to CUDA for a few operations. Then it wrapped this fused PyTorch implementation in torch.compile so torch.compile could handle the fusion. This approach achieved an 1.66X speedup over torch.compile and a 2.1X speedup over eager mode.
ResNet18 (Level 3 Problem 9)
import torch
import torch.nn as nn
import torch.nn.functional as F
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_channels, out_channels, stride=1, downsample=None):
"""
:param in_channels: Number of input channels
:param out_channels: Number of output channels
:param stride: Stride for the first convolutional layer
:param downsample: Downsample layer for the shortcut connection
"""
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.downsample = downsample
self.stride = stride
def forward(self, x):
"""
:param x: Input tensor, shape (batch_size, in_channels, height, width)
:return: Output tensor, shape (batch_size, out_channels, height, width)
"""
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class Model(nn.Module):
def __init__(self, num_classes=1000):
"""
:param num_classes: Number of output classes
"""
super(Model, self).__init__()
self.in_channels = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(BasicBlock, 64, 2, stride=1)
self.layer2 = self._make_layer(BasicBlock, 128, 2, stride=2)
self.layer3 = self._make_layer(BasicBlock, 256, 2, stride=2)
self.layer4 = self._make_layer(BasicBlock, 512, 2, stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes)
def _make_layer(self, block, out_channels, blocks, stride=1):
downsample = None
if stride != 1 or self.in_channels != out_channels * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(out_channels * block.expansion),
)
layers = []
layers.append(block(self.in_channels, out_channels, stride, downsample))
self.in_channels = out_channels * block.expansion
for _ in range(1, blocks):
layers.append(block(self.in_channels, out_channels))
return nn.Sequential(*layers)
def forward(self, x):
"""
:param x: Input tensor, shape (batch_size, 3, height, width)
:return: Output tensor, shape (batch_size, num_classes)
"""
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# Test code
batch_size = 2
num_classes = 1000
input_shape = (batch_size, 3, 224, 224)
def get_inputs():
return [torch.randn(input_shape)]
def get_init_inputs():
return [num_classes]import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
# Custom CUDA kernel for Add+ReLU (channels-last)
add_relu_cuda_src = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void add_relu_kernel(const float* a, const float* b, float* out, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
float val = a[idx] + b[idx];
out[idx] = val > 0.0f ? val : 0.0f;
}
}
torch::Tensor add_relu_cuda(torch::Tensor a, torch::Tensor b) {
auto size = a.numel();
auto out = torch::empty_like(a);
const int block_size = 256;
const int num_blocks = (size + block_size - 1) / block_size;
add_relu_kernel<<<num_blocks, block_size>>>(
a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size
);
return out;
}
"""
add_relu_cpp_src = "torch::Tensor add_relu_cuda(torch::Tensor a, torch::Tensor b);"
add_relu = load_inline(
name="add_relu",
cpp_sources=add_relu_cpp_src,
cuda_sources=add_relu_cuda_src,
functions=["add_relu_cuda"],
verbose=False,
)
# Helper to fold BN into Conv2d (for inference)
def fuse_conv_bn(conv, bn):
# Only for inference!
assert not conv.training and not bn.training
fused_conv = nn.Conv2d(
conv.in_channels,
conv.out_channels,
conv.kernel_size,
conv.stride,
conv.padding,
conv.dilation,
conv.groups,
bias=True,
device=conv.weight.device,
dtype=conv.weight.dtype,
)
# Prepare parameters
w_conv = conv.weight
if conv.bias is not None:
b_conv = conv.bias
else:
b_conv = torch.zeros(conv.weight.size(0), device=w_conv.device, dtype=w_conv.dtype)
w_bn = bn.weight
b_bn = bn.bias
mean = bn.running_mean
var = bn.running_var
eps = bn.eps
std = torch.sqrt(var + eps)
w_bn = w_bn / std
fused_conv.weight.data = w_conv * w_bn.reshape([-1, 1, 1, 1])
fused_conv.bias.data = (b_conv - mean) / std * bn.weight + bn.bias
# Set channels-last for weights
fused_conv.weight.data = fused_conv.weight.data.contiguous(memory_format=torch.channels_last)
fused_conv.bias.data = fused_conv.bias.data.contiguous()
return fused_conv
# Fused Conv+BN+ReLU module (for inference)
class FusedConvBNReLU(nn.Module):
def __init__(self, conv, bn, relu=True):
super().__init__()
self.relu = relu
self.fused_conv = fuse_conv_bn(conv, bn)
self.fused_conv.weight.requires_grad = False
self.fused_conv.bias.requires_grad = False
def forward(self, x):
x = self.fused_conv(x)
if self.relu:
return F.relu(x)
else:
return x
# Fused Add+ReLU module
class AddReLUFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, a, b):
return add_relu.add_relu_cuda(a, b)
@staticmethod
def backward(ctx, grad_output):
# Not needed for inference
return grad_output, grad_output
class AddReLU(nn.Module):
def forward(self, a, b):
return AddReLUFunction.apply(a, b)
# Fused MaxPool+Conv+BN+ReLU for stem
class FusedStem(nn.Module):
def __init__(self, conv, bn, relu, maxpool):
super().__init__()
self.maxpool = maxpool
self.fused_conv = fuse_conv_bn(conv, bn)
self.fused_conv.weight.requires_grad = False
self.fused_conv.bias.requires_grad = False
def forward(self, x):
x = self.fused_conv(x)
x = F.relu(x)
x = self.maxpool(x)
return x
# Downsample block with fusion
class DownsampleNew(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super().__init__()
self._conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)
self._bn = nn.BatchNorm2d(out_channels)
self.conv = None
def fuse(self):
self.conv = FusedConvBNReLU(self._conv, self._bn, relu=False)
def forward(self, x):
if self.conv is not None:
return self.conv(x)
else:
return self._bn(self._conv(x))
# Optimized BasicBlock with fused ops and channels-last
class BasicBlockNew(nn.Module):
expansion = 1
def __init__(self, in_channels, out_channels, stride=1, downsample=None):
super().__init__()
self.stride = stride
self.downsample = downsample
self.conv1 = None
self.conv2 = None
self._conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
self._bn1 = nn.BatchNorm2d(out_channels)
self._conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
self._bn2 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=False)
self.add_relu = AddReLU()
# Set channels-last for weights
self._conv1.weight.data = self._conv1.weight.data.contiguous(memory_format=torch.channels_last)
self._conv2.weight.data = self._conv2.weight.data.contiguous(memory_format=torch.channels_last)
def fuse(self):
# Call this after loading weights and switching to eval()
self.conv1 = FusedConvBNReLU(self._conv1, self._bn1, relu=True)
self.conv2 = FusedConvBNReLU(self._conv2, self._bn2, relu=False)
if self.downsample is not None and hasattr(self.downsample, "fuse"):
self.downsample.fuse()
def forward(self, x):
identity = x
if self.conv1 is not None:
out = self.conv1(x)
out = self.conv2(out)
else:
out = self._conv1(x)
out = self._bn1(out)
out = self.relu(out)
out = self._conv2(out)
out = self._bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
if self.conv1 is not None:
out = self.add_relu(out, identity)
else:
out = out + identity
out = self.relu(out)
return out
# Optimized Model
class ModelNew(nn.Module):
def __init__(self, num_classes=1000):
super().__init__()
self.in_channels = 64
self._conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self._bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=False)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
# Set channels-last for stem weights
self._conv1.weight.data = self._conv1.weight.data.contiguous(memory_format=torch.channels_last)
self.layer1 = self._make_layer(BasicBlockNew, 64, 2, stride=1)
self.layer2 = self._make_layer(BasicBlockNew, 128, 2, stride=2)
self.layer3 = self._make_layer(BasicBlockNew, 256, 2, stride=2)
self.layer4 = self._make_layer(BasicBlockNew, 512, 2, stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * BasicBlockNew.expansion, num_classes)
self.stem = None
def fuse(self):
# Call this after loading weights and switching to eval()
self.stem = FusedStem(self._conv1, self._bn1, self.relu, self.maxpool)
for m in self.modules():
if hasattr(m, "fuse"):
m.fuse()
def _make_layer(self, block, out_channels, blocks, stride=1):
downsample = None
if stride != 1 or self.in_channels != out_channels * block.expansion:
downsample = DownsampleNew(self.in_channels, out_channels * block.expansion, stride)
layers = []
layers.append(block(self.in_channels, out_channels, stride, downsample))
self.in_channels = out_channels * block.expansion
for _ in range(1, blocks):
layers.append(block(self.in_channels, out_channels))
return nn.Sequential(*layers)
def forward(self, x):
# Use channels-last for input
x = x.to(memory_format=torch.channels_last)
if self.stem is not None:
x = self.stem(x)
else:
x = self._conv1(x)
x = self._bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# Compile the model for further fusion (PyTorch 2.0+)
try:
import torch._dynamo
ModelNew = torch.compile(ModelNew, mode="reduce-overhead")
except Exception:
pass # torch.compile not available, fallback to normal classAs seen from this case, AI kernel generation and optimizations like torch.compile have the potential to be complementary when used together. The goal is to get faster PyTorch performance using any techniques available, not focus exclusively on any one technique. Generating low-level kernels alongside calls to torch.compile will generate the best result in many cases.
What’s next?
One common question with performance optimization - whether it’s done by a human or AI - is the upper limit on the gains you can achieve. In order to answer this question rigorously, we can use the roofline model as a tool. The roofline model helps us understand the theoretical maximum performance of a given workload, based on its arithmetic intensity - the ratio of operations performed to bytes of memory accessed. Compute-bound workloads have high arithmetic intensity and can (in theory) go as fast as the max compute capability of the GPU, whereas memory-bound workloads have low arithmetic intensity and cannot saturate the GPU's compute resources because memory bandwidth becomes the bottleneck. ```
Depiction of a roofline model. We can see that compute-bound workloads can achieve the maximum FLOPs on a given accelerator, but memory-bound workloads cannot. Source: Google DeepMind.
We are incorporating roofline models into this work. They are very useful when characterizing workloads and understanding how far off baseline performance is from theoretical best performance, and which resource the workload is bound by (e.g. compute, bandwidth, etc). These models will guide the agentic swarm to perform better optimization of the underlying kernels.
We’ll continue to expand this research in other directions, too - more use cases, more models, more target hardware platforms, supporting backward passes for training - in order to achieve the goal of autonomously improving model performance. If you’re interested in this topic, give us a shout at hello@gimletlabs.ai.
Footnotes
We will be comparing against the newer version of KernelBench in subsequent posts. ↩
Arithmetic means will artificially inflate results when comparing averages of speedups. For example, if you have two speedups - 0.5X and 2X, the arithmetic mean of them will be 1.25X - even though one doubled the speed and one halved it. To account for this, the geometric mean is the typical approach. In a geometric mean, the average of 0.5 and 2 will be 1. ↩
We compare against the default
torch.compilemode here, but it would be interesting to compare againsttorch.compilewith the different configuration options - e.g.max-autotune. ↩