Context-manager wall-clock timeout. Works on both Unix (SIGALRM) and Windows (thread-based fallback).

A single KernelBench problem.

Describes a single kernel replacement: what to replace and with what.

Context manager that patches a model's submodules to use optimized Triton kernels. Usage: with OptimizedModelContext(model, replacem

SwiGLU MLP (LLaMA-style).

Aggregated stats for one kernel (name + shape combo).

Full verification result.

Wraps nn.LayerNorm to use an optimized kernel_fn.

Wraps nn.Linear to use an optimized matmul kernel_fn.

Wraps RMSNorm-like modules to use an optimized kernel_fn.

Context manager that raises TimeoutError after `seconds`.

BERT-base: hidden_size=768, num_layers=12, num_heads=12, intermediate=3072 (110M params).

GPT-2 124M (small) by default. Configurable via constructor args. Sizes: - GPT-2 Small: n_layer=12, n_head=12, n_embd=768 (124M par

Compact LLaMA (160M params) -- fits on any GPU, good for testing AutoKernel. Config: dim=768, n_layers=12, n_heads=12, n_kv_heads=4, hidden_

LLaMA-7B scale -- requires ~14GB VRAM in fp16. Config: dim=4096, n_layers=32, n_heads=32, n_kv_heads=8, hidden_dim=11008

Example custom model -- a simple CNN + MLP. Replace this with your own architecture.