bedvibe.studio

BedVibe Studios — Engineering

I Wrote the Kernel. Then I Measured It, and at My Model's Shape It Was 1.67× Slower Than Doing Nothing.

A hand-written Triton LayerNorm against PyTorch eager and torch.compile, on one RTX 5080, at six batch sizes. Above roughly 4,096 rows the kernel wins on time and wins harder on memory — 1.36× faster at 65,536 rows using half the peak memory. Below it the kernel loses, and my model trains at 128 rows. The measurement cost less than the kernel would have, and it is the part worth keeping.


I train a 24-layer decoder at d_model 1024 with a batch size of 1 and gradient accumulation of 8. Somewhere in that step there are a lot of LayerNorms, and the standard reflex is that a fused hand-written kernel beats whatever the framework does. So I wrote one in Triton and set out to find how much it bought.

The honest answer turned out to be: less than nothing, at the only shape I actually run.

What counts as a baseline

Comparing a hand-written Triton kernel against PyTorch eager is not a fair fight and not an interesting one. torch.compile lowers to Triton through Inductor — it is generating the same class of kernel I just wrote by hand. So the only version of the question worth asking is can I beat torch.compile, and all three arms are reported throughout.

Everything below is one machine on one day: RTX 5080 (sm_120), driver 610.47, CUDA 12.8, torch 2.11.0+cu128, triton 3.8.0, fp32, N = 1024, 60 repetitions per cell, median with inter-quartile range, L2 flushed between repetitions.

Percentages of peak are a fraction of this machine measured today, not of a spec sheet:

ceilingmeasured
device-to-device copy814.2 GB/s
square matmul fp3239.2 TFLOP/s
square matmul bf16118.6 TFLOP/s

Forward: the kernel does not matter

rowseagercompiletritonbest % of measured BW
1280.0068 ms0.00550.004926.5%
5120.00860.00680.006875.9%
10240.01290.01090.010994.7%
40960.04510.04300.044395.8%
163840.16850.16460.1683100.2%
655360.66040.65630.6601100.5%

At 65,536 rows the three arms are 0.6563, 0.6601 and 0.6604 ms — a spread under one percent. They agree because none of them is the limiting factor. DRAM is. Above about 1,024 rows this kernel is saturated and there is nothing left to win by writing it better.

Two cells read above 100% of the measured ceiling. That is not a kernel beating physics. It means my copy-based ceiling slightly under-estimates what a streaming read-plus-write can achieve, so at the top of the table the instrument is the limiting error term. I have recorded them as measured rather than clipping them to 100%, because clipping would hide exactly the thing that tells you the ceiling is approximate.

Backward: the kernel matters, and at small shapes it hurts

Forward plus backward, same conditions. Peak memory is the allocation attributable to the step.

rowseagercompiletritontriton vs eagereager peaktriton peak
1280.0150 ms0.01290.02511.67× SLOWER1.0 MB2.5 MB
5120.01910.02510.03531.85× slower4.0 MB4.0 MB
10240.02930.04270.05251.79× slower8.0 MB6.0 MB
40960.13180.13370.11761.12× faster32.0 MB18.0 MB
163840.56590.43220.43881.29× faster128.1 MB66.1 MB
655362.23191.71271.63891.36× faster512.5 MB258.5 MB

The crossover is at roughly 4,096 rows. Below it the hand-written kernel is strictly worse: it launches two kernels instead of one and allocates lock and partial-gradient buffers, and at 128 rows that overhead is the entire runtime. Above it, it wins on time and wins harder on memory — at 65,536 rows it is 1.36× faster while using half the peak memory of eager, 258.5 MB against 512.5 MB.

Memory is the axis usually left out of kernel benchmarks, and here it is the larger of the two effects.

Which column my model actually lives in

This is where the benchmark stops being a table and starts being a decision. My training config is d_model 1024, 24 layers, batch size 1, gradient accumulation 8. That puts every LayerNorm in the training step in the far-left column of both tables — the region where my hand-written kernel is 1.67× slower than doing nothing at all, and where every arm sits at 13–27% of the machine's bandwidth.

Nothing at 128 rows is bandwidth-bound. The cost is kernel launch. So the correct optimisation at that shape is not a better kernel, it is fewer launches — CUDA graphs.

That is consistent with, and explains, a result I had already measured separately on the inference side of the same model family: 9.0×, 144.45 against 16.06 tokens/sec, from Inductor plus CUDA-graph compilation. Same machine, same shape regime, and the speedup came from removing launches rather than from any kernel being better.

So: do not hand-write this kernel for this model. I am keeping the measurement and throwing away the kernel, and the measurement is worth more than the kernel would have been.

Defects the instrument found in itself

Two, and both are in the record because a benchmark whose failures are invisible is not a benchmark.

The L2 flush scratch was billed to the first kernel measured. The 256 MB buffer used to cold-start the cache between repetitions was allocated lazily inside the first timed repetition, so whichever arm happened to run first was charged 256 MB of peak memory belonging to the harness. It is visible in the first forward run as eager/fwd 128×1024 → 256.5 MB against 0.5 MB for the two arms that followed. Fixed by allocating the scratch during ceiling measurement, before any kernel is timed.

One cell is unstable and I cannot explain it. compile/fwd+bwd at 1024×1024 returned a median of 0.0427 ms with an IQR of 0.0806 ms — a spread larger than the value — and a peak-memory reading of 0.0 MB. Every other cell in that column has an IQR under 0.01 ms. It is in the table as measured, and it should not be quoted as a number until somebody understands it.

What this is not

The Triton kernels are not novel. Their structure follows the standard Triton fused-layer-norm tutorial, including the lock-based group accumulation for dw and db. That is stated in the source file too. What I am offering here is the measurement, not the kernel.

This experiment was not pre-registered. The prediction — "Triton will not win at batch 1" — was written into the runner's docstring before the first run, and the forward result matched it. But a docstring is editable and a frozen protocol is not, so that is weak evidence of intent and nothing stronger. I have pre-registered other work precisely because this distinction matters; I did not do it here, and I am not going to claim the credit for it retroactively.

One machine, one dtype (fp32), one width (N = 1024), one session. No claim is made about other cards, other dtypes, or other widths. torch.compile timings include no compilation cost, and warmup precedes every measurement.

The generalisable part is not the numbers. It is the shape of the question: before writing a kernel, find out which column of the table your model actually sits in. Mine sat in the one where the answer was no.