One kernel, end to end: a slow attention kernel to a merged, verified PR

On 9 October 2026 we pointed Kernalin at one deliberately slow kernel and let it run the whole loop, recorded end to end on kerneln.com. This post walks through what happened, with the numbers as they were recorded, including the ones that are less flattering.

Context first. Everything here ran in Kernalin’s CPU sandbox (x86_64, 2 threads, PyTorch 2.5.1+cpu, no network). These are CPU numbers. The loop is the same on a GPU runner; the measurements would not be.

1. The kernel

kernels/attention.py in JashomTech/kernalin-demo is multi-head attention written for clarity: Python loops over every batch and every head, and a hand-rolled softmax.

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

Benchmarked in PyTorch eager, float32: a median of 1.24 ms at batch 1, 15.25 ms at batch 8 and 48.26 ms at batch 32.

2. The profile

Linking the file from GitHub turns on auto-sync: every new version is profiled and benchmarked. The op profile (batch 8, float32) ranked the time:

Op	Share of self time
aten::mm	37.62%
aten::exp	26.47%
aten::max	10.30%
aten::div	5.19%
aten::copy_	4.26%

aten::mm was called 1,280 times in the profiled window: one small matmul at a time, launched by the loop. Kernalin’s Amdahl what-if put a number on it: making aten::mm 5× faster would make the whole run 1.43× faster.

3. The rewrites

Optimize run #73 gave Claude Opus the kernel, its profile and its benchmark, with the default budget of three rewrites. Rewrite #1 replaced the loops with batched matmuls and an in-place softmax that follows the same op sequence:

def run(q, k, v):
    dim = q.shape[-1]
    scale = 1.0 / math.sqrt(dim)
    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)

4. The verdict

Every rewrite went through the five-stage correctness gate (smoke, shape sweep, adversarial stability, determinism, edge cases), the anti-cheat audit and a Bayesian speed verdict.

Rewrite	Decision	Correct	Anti-cheat	Speedup
#1	kept	pass	clean	7.64×
#2	reverted	pass	clean	0.37×
#3	reverted	pass	clean	0.40×

All three were correct. Two were slower, so they were reverted, and you can see them in the run. One was kept, at a best median of 3.058 ms per call.

5. The pull request and the gate

The kept rewrite became pull request #3, opened at 18:24 UTC with the audit table in its description. The repository’s own CI runs the rnx-benchmark action on every pull request (PyTorch eager, batch 1 and 8, on a GitHub-hosted runner). The gate passed at 18:25 UTC, and the PR was merged at 18:28 UTC.

6. The re-measurement

After the merge, auto-sync benchmarked main again, and Kernalin compared it against the run from before the merge:

Configuration	Speedup vs before
pytorch_eager / float32 / batch 1	3.38×
pytorch_eager / float32 / batch 8	4.22×
pytorch_eager / float32 / batch 32	3.24×
torch_compile / float32 / batch 1	2.28×
torch_compile / float32 / batch 8	2.43×
torch_compile / float32 / batch 32	1.81×

Every configuration came back “faster”, with 0.0% probability of being slower. It is not 7.64×, and it should not be expected to be: it is a separate measurement, over more configurations and a second runtime. We report both, with what each measured.

What this does and doesn’t show

It shows the loop working end to end, on a real repository, with a person merging a normal pull request. It does not show a GPU result; for those, see the H100 dry run. And the original kernel was written for clarity, so most of the win comes from removing Python-level loops, which is exactly the kind of fix an AI is good at and a verifier should check.