One kernel’s journey
Eight steps from a slow kernel to a cheaper one, recorded end to end.
Everything below happened on 9 October 2026 on kerneln.com, to kernels/attention.py in our demo repository. It ran in Kerneln’s CPU sandbox, so these are CPU numbers. The loop is the same on a GPU runner.
01 / 08 · The kernel
A slow attention kernel.
The hot path of a serving stack: multi-head attention, written for clarity. Python loops walk every batch and every head, and the softmax is done by hand.
It is correct. It is also 48 ms per call at batch 32.
- File
- kernels/attention.py
- Shape
- 8 heads · seq 128 · dim 64
- Runtime
- PyTorch eager, float32
- Where
- CPU sandbox
def run(q, k, v): batch, heads, seq, dim = q.shape scale = 1.0 / math.sqrt(dim) out = torch.empty_like(q) for b in range(batch): for h in range(heads): scores = torch.matmul(q[b, h], k[b, h].transpose(0, 1)) * scale row_max = scores.max(dim=-1, keepdim=True).values weights = torch.exp(scores - row_max) weights = weights / weights.sum(dim=-1, keepdim=True) out[b, h] = torch.matmul(weights, v[b, h]) return out| batch 1 | 1.24 ms | |
|---|---|---|
| batch 8 | 15.25 ms | |
| batch 32 | 48.26 ms |
02 / 08 · Profile
Profiled. The bottleneck has a name.
Kerneln links the file from GitHub and profiles every new version. The op profile ranks where the time goes: aten::mm owns 37.6%, launched 1,280 times by that loop, with aten::exp close behind.
An Amdahl what-if turns the ranking into a plan: what each op is worth fixing, end to end.
- Top op
- aten::mm · 37.62%
- Calls
- 1,280
- Biggest lever
- 5× mm → 1.43× overall
- Profiler
- torch.profiler
Self-time hotspots · hottest first
| Op | Share | Percent |
|---|---|---|
| aten::mm | 37.62% | |
| aten::exp | 26.47% | |
| aten::max | 10.3% | |
| aten::div | 5.19% | |
| aten::copy_ | 4.26% | |
| aten::mul | 4.09% | |
| aten::sub | 3.11% | |
| aten::select | 2.97% |
aten::mm
- Type
- MATMUL
- Share
- 37.62%
- Total
- 94.25 ms
- Calls
- 1,280
- Likely
- COMPUTE-bound
Biggest lever
Making aten::mm 5× faster speeds the whole run up 1.43×.
Amdahl what-if, end to end.
03 / 08 · Rewrite
AI rewrites it.
Optimize hands Claude the kernel, its profile and its benchmark. Three model calls, three candidate rewrites.
The first one replaces the loops with batched matmuls and an in-place softmax that keeps the same sequence of operations. Plausible is not the bar, though.
- Model
- Claude Opus
- Rewrites
- 3 (default budget)
- Run
- #73
- Lines
- −10 +9
@@ def run(q, k, v): @@ batch, heads, seq, dim = q.shape out = torch.empty_like(q) for b in range(batch): for h in range(heads): scores = torch.matmul(q[b, h], k[b, h].transpose(0, 1)) * scale row_max = scores.max(dim=-1, keepdim=True).values weights = torch.exp(scores - row_max) weights = weights / weights.sum(dim=-1, keepdim=True) out[b, h] = torch.matmul(weights, v[b, h]) return out dim = q.shape[-1] scores = torch.matmul(q, k.transpose(-1, -2)) scores.mul_(scale) row_max = scores.amax(dim=-1, keepdim=True) scores.sub_(row_max) scores.exp_() denom = scores.sum(dim=-1, keepdim=True) scores.div_(denom) return torch.matmul(scores, v)04 / 08 · Verify
The verifier proves it, or it doesn’t ship.
Every rewrite must match the reference (within per-dtype tolerance) through five correctness stages, survive an anti-cheat audit for skipped work and ignored inputs, and win a Bayesian speed verdict outside a ±2% noise band.
All three rewrites were correct. Two were slower and were reverted. One was kept: 7.64× faster. Since then the verifier has grown stricter still: hidden input distributions and a peak-memory gate.
- Correct
- 3 of 3
- Faster
- 1 of 3
- Kept
- rewrite #1 · 7.64×
- Best median
- 3.058 ms
Correctness gate
- Smoke test runs, returns finite output pass
- Shape sweep batch, sequence and head shapes pass
- Adversarial stability scaled and transformed inputs pass
- Determinism repeat runs agree pass
- Edge cases non-power-of-two, boundary shapes pass
Anti-cheat and verdict
- Fresh inputs no memorized outputs pass
- Ignored-input check output responds to every input pass
- Poisoned buffers no reading stale memory pass
- Speed verdict P(faster) ≥ 95% outside ±2% pass
- Hidden input distributions inputs the model never sees added since
- Peak-memory gate no trading memory for speed added since
Every rewrite
1 of 3 kept
| # | Decision | Correct | Anti-cheat | Speedup |
|---|---|---|---|---|
| 1 | kept | pass | clean | 7.64× |
| 2 | reverted | pass | clean | 0.37× |
| 3 | reverted | pass | clean | 0.40× |
Kept · verified faster
best median 3.058 ms per call
7.64×
05 / 08 · Pull request
A pull request goes up, and CI checks it again.
The kept rewrite arrives as pull request #3, with the full audit in its description. The server checks the PR source byte for byte against the logged model output.
Then the repository’s own CI runs the Kerneln benchmark gate on the change, independently. It passed a minute later.
- Opened
- 18:24 UTC
- Gate passed
- 18:25 UTC
- Action
- rnx-benchmark
- Branch
- optimize-73
Optimize attention.py: 7.64× faster (Kerneln run #73) #3
optimize-73 → main · 1 commit · 6e4c9a0
Kerneln rewrote this kernel and verified the result before proposing it.
Speedup: 7.64× on the CPU sandbox (best median 3.058 ms).
| Rewrite | Decision | Correct | Anti-cheat | Speedup |
|---|---|---|---|---|
| #1 | KEEP | yes | clean | 7.64× |
| #2 | REVERT | yes | clean | 0.369× |
| #3 | REVERT | yes | clean | 0.4× |
Your CI's Kerneln benchmark gate checks this pull request independently.
06 / 08 · Merge
Merged.
A person on the team merged it into main four minutes after it opened. Nothing about how the team ships changed: a branch, a PR, a green check, a merge.
Kerneln’s auto-sync noticed the new version of the linked file and started its profile and benchmark jobs on its own.
- Merge commit
- e4429b6
- Merged
- 18:28 UTC
- Auto-sync
- profile + benchmark
- Checks
- 1 passed
- e4429b6 Merge pull request #3 from optimize-73 main 18:28 UTC
- 6e4c9a0 Optimize attention.py: 7.64× faster (Kerneln run #73) optimize-73 18:24 UTC
- 2108517 Reset demo: original attention kernel earlier
Auto-sync · kernels/attention.py
New version on main. Started 2 jobs:
- Op profile CPU · batch 8 · float32
- Benchmark eager + torch.compile · batch 1, 8, 32
07 / 08 · Re-measure
Re-measured after merge, by itself.
A separate benchmark compared main against the run from before the merge, in every configuration: 3.2–4.2× faster in PyTorch eager and 1.8–2.4× under torch.compile, each with 0% probability of being slower.
That is lower than 7.64× because it is a different measurement, over more configurations and a second runtime. We publish both.
- Eager
- 3.38× · 4.22× · 3.24×
- torch.compile
- 2.28× · 2.43× · 1.81×
- Batches
- 1 · 8 · 32
- P(slower)
- 0.0%
| Configuration | Verdict | Speedup |
|---|---|---|
| pytorch_eager / float32 / b1 | 3.38× | |
| pytorch_eager / float32 / b8 | 4.22× | |
| pytorch_eager / float32 / b32 | 3.24× | |
| torch_compile / float32 / b1 | 2.28× | |
| torch_compile / float32 / b8 | 2.43× | |
| torch_compile / float32 / b32 | 1.81× |
PyTorch eager
3.2–4.2×
torch.compile
1.8–2.4×
08 / 08 · Cost
The same work, for less compute.
At a fixed price per hour, cost per call moves with time per call. In eager mode the attention kernel now costs about a quarter to a third of what it did per call.
What that is worth to you depends on how much of your GPU time the kernel takes. Put in your own numbers; it is arithmetic, not a promise.
- Per call, eager
- 23.7–30.9% of before
- Formula
- spend × share × (1 − 1/speedup)