Optimizing a GPT-2-Class Transformer on a GPU — Martin S. Kristiansen

Bottom line: one GPT-2-small-class transformer (85M params, fp32, seq 128)<br>taken from a naive CUDA port at 78.2ms per forward pass to a tuned hybrid at<br>1.60ms — 49&times; — on a single RTX 3080 Ti , in fifteen explicit steps, ending<br>past torch.compile's 1.72ms on the same workload. Batching then repaid the one<br>factor no kernel could reach, finishing at 136,000 tokens/second on a $1,000<br>consumer card.

One layout bit was worth 5.3&times;. Storing the weights transposed —<br>so that consecutive threads read consecutive addresses — was the single largest rung of<br>the entire ladder. The same change on the CPU was worth 1.24&times;.

Hand-written kernels beat cuBLAS — on the narrow shapes. The library won<br>the wide ones. So the fastest binary of the campaign is neither hand-written nor library:<br>it's a measured routing table between them (2.27ms).

The compiler's entire margin was organizational, not kernel-level.<br>torch.compile's lead decomposed into three named tricks — fused epilogues,<br>flash attention, CUDA graph capture — each rebuilt by hand and measured alone. Their sum<br>landed within 0.01ms of the predicted margin, and the result passed the compiler.

The last 1.75&times; was never in the code; it was in the question. A floor<br>analysis split the remaining gap into nameable craft and a "shape tax" no implementation<br>escapes at 128 rows. Batching repaid the tax at exactly the predicted factor.

The method itself became a program: a discrete search with the campaign's<br>priors installed at five distinct ports rediscovered the human-found dispatch table from<br>measurements alone — then named, by regression, the constraint its authors had mispriced<br>all along.

Results

The CPU campaign ended with a tiny<br>transformer certified against the physics of one core. The obvious next question — the one<br>everybody asks about five minutes into any performance conversation — is what about the<br>GPU? So: same discipline, bigger model (GPT-2-small-class, 85M parameters, 340 MB of<br>fp32 weights, sequence length 128), and a machine actually worth optimizing for — an RTX<br>3080 Ti, which is 80 streaming multiprocessors, ~34 TF/s of fp32, and 912 GB/s of<br>memory bandwidth, sitting in a desktop.

The same model in naive single-core C runs at 1,336ms per forward pass. The naive CUDA port<br>of that exact code is where this ladder starts.

rungchangetimevs naive port

—naive C, one CPU core (i9-10850K)1336 ms—<br>1naive CUDA port78.2 ms1.00&times;<br>2coalescing (transposed weights)14.77 ms5.3&times;<br>3shared-memory tiling9.08 ms8.6&times;<br>4register tiling + fused qkv7.57 ms10.3&times;<br>5smaller tiles (deliberate falsification)8.03 msnull<br>6split-K on the narrow GEMMs4.48 ms17.5&times;<br>—oracle: all GEMMs to cuBLAS (fp32 / TF32)3.37 / 2.98 ms—<br>7parallel attention rebuild2.91 ms26.9&times;<br>8SASS-driven fp32 kernel (beats cuBLAS fp32)2.80 ms27.9&times;<br>9hand-built tensor cores, WMMA (beats cuBLAS TF32)2.47 ms31.7&times;<br>10the CPU campaign's algebra, ported2.29 ms34.1&times;<br>11the dispatch table (hand + library, per shape)2.27 ms34.5&times;<br>12profiler-guided assault on w22.30 msnull<br>13acuBLASLt fused epilogues2.17 ms36.0&times;<br>13bhand-built flash attention2.03 ms38.5&times;<br>13cCUDA graph capture1.60 ms 48.9&times;<br>14a/bautotuning tournament / persistent megakernel1.61 / 1.99 mstwo nulls<br>15batch axis: B sequences per pass0.94 ms/seq136K tok/s

The full ladder, re-measured back-to-back in one five-minute window;<br>every historical rung reproduced within 1%. Each row links to its source — including the<br>failed experiments, which are part of the record. Reference points: torch.compile<br>max-autotune on the identical workload = 1.72ms; PyTorch eager = 5.6ms.

The complete ladder. Blue: the hand-written climb. Green:<br>hand-written beating the library. Gold: the cross-vendor dispatch champion. Grey: the<br>informative null.

Two house rules carried over from the CPU campaign, because they earned it. First, every<br>version fills its weights from the same seeded random-number generator — deterministic, so<br>all nineteen versions compute on bit-identical weights — and prints its first two output<br>logits: 0.095916 0.033000. A timing without that check attached simply<br>does not count as a result. This sounds bureaucratic until you learn it caught six real bugs<br>across nineteen versions, every single one before a wrong conclusion got drawn from a<br>fast-but-broken binary. Second, timings compare only within one measurement window, because<br>the same binary drifts up to 10% with thermal state, and a comparison across windows is a<br>comparison of weather.

One law governs everything below. A pass that must execute F floating-point operations and<br>move B bytes at some level of the memory hierarchy cannot finish faster than<br>max(F/peak, B/bandwidth) — the roofline. The interesting part is never the formula; it's<br>which term binds, because that tells you which currency you're paying in, and an optimization<br>that saves the other currency...