KernelForge

Inspiration

The track asks for kernels that stay inside rel < 0.02, abs < 0.002 and run fast across fourteen published shapes, and it says outright that "participants can choose different implementations for different shapes by adding shape checks."

That sentence is the whole problem. Fourteen shapes spanning batch 1 to 10,000 and sequence 32 to 100,000 do not have one best kernel, so hand-tuning one is optimising the average of things that want opposite decisions. We built the loop that produces and selects them instead.

What it does

Five stages, each of which can reject the previous one:

Profile the workload and compute a per-stage precision error budget.
Propose a configuration from a hand-written heuristic, or from an LLM given the same spec sheet, profile and budget.
Prove it against the organizer's own tolerance rule, over multiple seeds, before it is ever timed.
Time it honestly, interleaved round-robin, rotating order, medians.
Freeze the winner into that card's dispatch table, then re-verify and demote anything that drifts.

At run time submission.py does none of that. It looks up the GPU, dtype and shape in a frozen table and runs the plan it names: no LLM call, no autotuning stall, no nondeterminism.

What we built

Triton FlashAttention handling causal and key-padding masking in-register. PyTorch's SDPA takes is_causal or an attn_mask, not both cheaply, and the benchmark's generator produces exactly that combination.
Fused add + mask + LayerNorm, fused QKV projection, and CUDA-graph capture removing ~105 kernel launches per forward.
An accuracy gate built from the organizer's own comparison code, collapsed to one number (envelope utilization) so "how close to failing" is measurable rather than a yes/no.

Results

Nine GPU configurations, four architectures, official shapes only, measured through the organizer's unmodified torch_transformer_benchmark.py:

GPU	arch	passed	median vs reference
H100 NVL	sm_90	13/13	7.14x
A100-PCIE-40GB	sm_80	13/13	6.61x
A100 80GB PCIe	sm_80	13/13	5.25x
H100 NVL MIG 3g.47gb	sm_90	13/13	4.45x
A100 80GB MIG 3g.40gb	sm_80	13/13	4.38x
H200 NVL	sm_90	13/13	4.34x
TITAN V	sm_70	12/12	3.55x
TITAN RTX	sm_75	13/13	3.37x
Tesla T4	sm_75	13/13	3.12x

116 shape-runs, 116 passed. Median margin over torch.compile(max-autotune) runs 1.66x–1.67x on Hopper, 1.35x–1.59x on Ampere, and 1.02x–1.23x on Volta and Turing, where there is no TF32 headroom to spend.

Official shape 14 is the one the reference cannot run at all, it would allocate an 18.6 TB attention matrix. Rather than quote no number we measured against PyTorch's own scaled_dot_product_attention substituted into their model: 4.32x on an A100-80, 5.95x on an H100. Correctness checked against a streamed exact reference at full length and full batch, 0 of 3,276,800,000 elements outside tolerance.

How the LLM is used

Two distinct roles, wired in differently because they carry different risk.

As a kernel author. The model writes complete Triton source against a contract. Every candidate compiles, is gated against an exact reference, and is timed by the same harness as everything else: layernorm 10/10 correct at 2.46x over torch, gelu 7/10 at 3.91x. The failures are the interesting part, one kernel substituted an Abramowitz-Stegun polynomial for erf, quantified its own error correctly in a comment, and was right about the number while being 22.8x outside a tolerance measured against exact erf. Only the gate caught it.

As a plan proposer, competing for real slots. cli agent runs after the search, so every shape already holds the best plan the heuristic found. A proposal is promoted only if it is measurably faster than what is frozen:

-> promoted into the table (4.60x vs 3.34x held)
-> not promoted: 1.20x does not beat the 2.31x already frozen

Per card the model won roughly 8 or 9 slots of 13–14 proposals, and all nine GPUs measured faster after the agent round.

Generated source is never shipped, that is a deliberate line. A configuration the gate accepts is as safe as any other configuration; Triton source no person has read is not, however well it measures.

Why it matters

The official shape list is a recommendation ranking workload: 2–4 layers, d_model 32–1024, batch 1 to 10,000 on one model, causal masking with a padding mask throughout, and a 100,000-token behaviour log. That is the user-behaviour sequence model that scores candidate videos against a viewing history, the highest-volume transformer inference in a short-video feed.

Measured end to end on a laptop (RTX 3070 Ti, 8 GB) over a ranking traffic mix, 42% short-history realtime, 18% long-history, 25% coalesced peak, 15% re-rank:

	traffic-weighted	GPUs for 100k QPS
architecture defaults	1.65x	52 → 32 (39% freed)
tuned, then LLM agent	2.66x	52 → 20 (61% freed)

Latency matters as much as capacity: the realtime segments get 2.5x and 2.9x faster, which is headroom spent on a longer history or a bigger candidate set, a better feed, not just a cheaper one.

How we built it

Python, PyTorch, Triton. The LLM proposer speaks plain urllib against any OpenAI-compatible /v1/chat/completions endpoint, OpenAI, Ollama, vLLM, LM Studio, or your own gateway, with Anthropic supported directly. No SDK, no vendor lock-in, and everything works with no credentials at all: a deterministic heuristic proposer runs instead.

Measurements come from a Slurm cluster driven over SSH, but nothing requires it. The whole loop runs on one GPU, which is how the laptop numbers above exist.

Challenges

Measuring honestly was harder than going fast. Two examples that changed what we ship:

The compiler looked better than it was. Our sweep times candidates interleaved, so a torch.compile model is already warm with its CUDA graph captured. The organizer's script uses a fresh process. Across 116 paired shapes, plans delegating to torch.compile came in at a mean 0.891x of our own kernels under their harness while ours averaged 1.017x. A tie in our harness is a loss in theirs, so the selector now makes a compile candidate win by 15% before taking a slot. All nine GPUs got faster.

We were wrong about causal attention three times. Causal masking halves the FLOPs but the kernel cost 0.62–0.72 of non-causal. We blamed load imbalance and built the standard fix, a persistent-tile kernel, slower on every shape. So we decomposed it: the work reduction was already realised, and applying the mask cost 1.16x–1.34x, because the kernel built the causal predicate on every key block though only the diagonal one can be affected. Splitting the loop is worth 1.08x–1.21x and is bit-identical. Then the target itself turned out to be wrong: the floor for a tiled kernel is (M+1)/2M, not 0.50.

What we learned

A gate that runs before the timer changes what you can safely automate. It is what lets an LLM propose freely, it can suggest bfloat16 compute and an fp16 residual, both already ruled out by the error budget, and the gate simply rejects them. A proposer that is never rejected is one that never explores.

What's next

A split-K attention that partitions the key range and combines partial softmax states, which balances by construction. Hopper-specific work, TMA and wgmma, for the two compute-bound shapes where we reach a lower fraction of the H100's ceiling than the A100's. And re-timing a winning proposal in the incumbent's harness before promoting it, which is the one known weakness in how the agent decides.