Inspiration
We started by profiling the reference implementation instead of reaching for a list of known optimizations, and the profile said something we didn't expect: there is no single bottleneck to fix.
A Transformer layer computes
$$\text{Attention}(Q,K,V)=\text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$
and why that is slow depends entirely on the input shape. At batch size 1 the GPU is idle, starved waiting on the CPU to submit ~60 kernel launches for 0.13 GFLOP of work. At batch size 10,000 it is bound by memory bandwidth. At \( d_{\text{model}} = 1024 \) it is bound by GEMM throughput. An optimization that is decisive for one of these is irrelevant — sometimes actively harmful — for another.
The problem statement invites you to handle this by "adding shape checks." That is where we disagreed with the obvious approach, and that disagreement became the project.
What it does
A drop-in optimized forward pass for the benchmark's Transformer, satisfying the task's per-element correctness contract:
\( |u - r| \le 0.002 \) or \( |u - r| \le 0.02\,|r| \)
Results on an RTX 4080 Laptop GPU: ~10x geometric-mean speedup across the 13 shapes with a runnable reference, from 2.7x to 35.1x, all 13 numerically correct. The 14th shape has no runnable reference at all and we execute it anyway.
How we built it
We measure instead of hardcoding. A threshold like "use fp16 when batch exceeds 32" is a constant fitted to one machine's cache hierarchy and clocks; it is unfalsifiable in the source and silently wrong elsewhere. Instead, the first time our implementation meets a shape it builds the candidate implementations, runs them on the real input during the benchmark's own untimed warmup, times them, and caches the winner. Numeric precision, execution path, slice size and kernel fusion are all decided that way.
The optimization itself is a mixin over the baseline module tree. It overrides only forward(), adds no parameters and registers no buffers, so state_dict() stays byte-identical and strict weight copying keeps working.
Underneath, nine specific defects in the reference, each mapped to one countermeasure:
| Defect | Countermeasure |
|---|---|
| \( [B,H,S,S] \) score tensor materialized, then again in fp32 for softmax | Fused SDPA — never forms the matrix |
| Full score matrix computed, half then masked to \( -\infty \) | is_causal=True — the masked half is never computed |
| Padding mask applied where it provably cannot change the result | Mask elision, which unlocks SDPA's fastest backends |
| Q, K, V as three separate \( [d,d] \) GEMMs | One fused \( [3d,d] \) GEMM |
| Every sublayer boundary a separate full pass; the dtype cast does no arithmetic | One Triton kernel fusing residual add + LayerNorm + cast |
| FFN intermediate written, re-read for GELU, written again | cuBLASLt GEMM+GELU epilogue — written once |
| ~60 launches for 0.13 GFLOP on small shapes | CUDA-graph capture — one replay |
| Each layer re-reads the whole activation from HBM | Slice-through-the-stack — a slice stays L2-resident |
| Everything in fp32 | Measured per-shape mixed precision, fp32 residual |
Only the last is a precision trade. The other eight remove work the reference performs and discards.
Challenges we ran into
A correctness gate that silently disabled our own kernel. Our Triton kernel is verified against F.layer_norm before use and rejected if they differ by more than \( 5\% \) of the absolute tolerance — a budget of \( 10^{-4} \). But the comparison happens on tensors already stored in fp16, which quantizes in steps of \( 2^{-10} \approx 9.77\times10^{-4} \). Two results agreeing perfectly in fp32 land either bit-identical or one full step apart, with nothing between. The gate could only ever pass an exactly-equal result. It reported "fusion off" across the entire suite — indistinguishable from the kernel never running. Instrumenting it showed two of three cases at exactly 0.000e+00 and the third at 9.766e-4: precisely one fp16 ULP.
A shape that cannot be benchmarked. Shape #14 (\( B=32, S=100{,}000 \)) would need an explicit \( 32 \times 16 \times 100{,}000^2 \) score tensor — about 20 TB — so no reference exists to compare against. Even our path needs care: at fp32 the input plus a separate output is 24.41 GiB. Aliasing the output onto the input halves it, but one fp32 buffer is still 12.21 GiB against 11.99 GiB of VRAM. Only fp16 with aliasing fits, at 6.10 GiB. It runs, and the output statistics are correct (mean \( \approx 0 \), std \( \approx 1.0 \)).
Triton doesn't ship for Windows. Official Triton publishes no win32 wheel, so our fused kernel was dead code until we found the community build. We lost real time to a "why is this never running" that turned out to be an unavailable dependency.
What we learned
Our own confident fixes were wrong twice, and only measurement caught it.
We shipped a rule requiring the fusions to prove themselves >2% faster before adoption. It seemed obviously correct. An ablation showed it was losing throughput — discarding 41–57% of the available speedup on four shapes. Our first fix inverted the gate's direction and changed nothing (15.503 vs 15.456 ms), which is what forced us to stop adjusting the threshold and look at the measurement instead. The gate timed the eager path, while half the suite runs under CUDA-graph capture. On shape #9 the eager comparison calls the fused path 18% slower while the captured end-to-end result is 34% faster. Both numbers are correct; they answer different questions, and the gate was asking the wrong one. Gating on a non-predictive proxy is worse than not gating, so we removed it.
Later, a single sample put shape #4 at 0.78x — fusion slower. Five runs per mode reversed the sign. One sample of a borderline measurement is not a result, and we nearly published one.
The measured policy is not uniform, and we would not have guessed it. Ten of thirteen shapes adopt CUDA-graph capture, three run eagerly. Eleven adopt mixed fp16; two measure fp32 as faster. That partition is not monotonic in any shape parameter — \( B=1 \) picks fp16, \( B=4 \) and \( B=16 \) do not, \( B=64 \) and above do again. No threshold we would have written by hand reproduces it.
What's next
The clearest remaining headroom is a hand-written attention kernel. We rely on PyTorch's SDPA, whose epilogue is fixed — so the output projection and the residual add must each be a separate pass over the \( [B,S,d] \) activation. Fusing them into the FlashAttention epilogue would remove an HBM round-trip per layer. Our own profile says where to look: on the GEMM-bound shape, attention itself is only 2% of the FLOPs.
We would also make calibration time the captured path rather than the eager one, which is the root cause behind the retired speed gate.
Log in or sign up for Devpost to join the conversation.