Shape-specialized Transformer Inference on Commodity GPUs

How this solution addresses the problem statement

The track asks for GPU kernels that implement a fixed Transformer layer, pass the organizer's test cases within rel < 0.02 / abs < 0.002, may specialize per input shape via shape checks, and are documented in a tech report that states the machine and the AI skills used. Point by point:

"Submit one or several GPU kernels that implement the layers." We ship a shape-keyed dispatcher over sixteen registered backends: two hand-written Triton kernels (a fused scale/mask/cast epilogue over the attention scores, and a fused residual join), a packed QKV projection that collapses three GEMMs into one, whole-layer CUDA-graph capture, float16 tensor-core internals behind a float32 accuracy boundary, and torch.compile variants. Every backend is a drop-in UserOptimizedTransformer.forward, so the organizer's script is the only entry point.
"Decide which parts of the layer should be fused into one kernel." Each fusion was chosen from a measurement at the shape it runs at. REPORT.md section 3 profiles where the time goes at each official shape; section 5 prices every fusion we tried, including the four we rejected on evidence (the fully fused softmax, forced contiguity, SDPA and torch.compile at the shapes where they lose, and the residual-join + LayerNorm fusion).
"The diff should be small enough (relative error < 0.02, abs error < 0.002)." That rule runs before the timer does. A backend is admissible at a shape only if its measured elementwise error passes the official rule in every trial, and a shape with no admissible backend falls through to the organizer's own arithmetic. The published numbers are produced by the organizer's own script, with its own graded comparison, run against the flattened artifact in dist/.
"Different shapes ... the participants can choose different implementations for different shapes by adding shape checks." The dispatcher keys on (batch, seq_len, d_model, n_heads, d_ff, n_layers, causal, dtype) and looks the key up in a table derived from a 44-key tune, 704 gate-checked rows of saved evidence. Unknown shapes take the reference path unchanged.
"All the combinations of input shapes will be told to the participants." All fourteen are in official_shapes.json and all fourteen are addressed - including official_14, which no harness can run on any hardware and which carries 99.9% of the benchmark's FLOPs. The streamed runner in bench/frontier.py completes it on a 12 GB card.
"The use of AI tools is encouraged." The implementation, harness, test suite and these documents were written by coding agents - Claude Code and Codex - under human direction, and REPORT.md section 2 discloses that division alongside the machine. What is kept as a repository artifact is the method record rather than a chat transcript: 110 numbered experiments in experiments/log.md, each carrying its hypothesis, method, result and verdict, plus the saved JSON every published figure is read off.
"Provide a clear tech report ... what the environment is (CPU, GPU, DISK, etc), what kind of optimizations you have done, and the final test results." REPORT.md is that report: section 2 is the environment, sections 3 to 6 are the optimizations and the routing, section 8 is the results, section 9 is the evidence for the method, section 10 the limitations.

Inspiration

We read the gate before the kernel zoo. The graded rule is elementwise abs <= 0.002 OR rel <= 0.02 in float32, and internal quantization is expressly allowed - which prices float16 arithmetic in: an fp16 GEMM's rounding is a relative 2^-11, about 5e-4, and a relative rail of 2e-2 absorbs forty times that. The second thing we read was the shape list: official_14 is 99.87% of the list's FLOPs and cannot run at all in reference form on any hardware, because the reference materializes a per-layer attention-score tensor of 20.5 terabytes at S=100000. So the project split in two: win the 13 shapes that fit, and make the one that carries the score run at all.

(An earlier life of this project spent a hundred experiments under a stricter, self-imposed rule - bit-exactness - because the pre-update harness budget came to 1.3 units in the last place of a bfloat16 output. The rules update replaced that budget with the graded rule above, and the same measure-first machinery repriced everything against the new gate. ### What we learned has what survived.)

What it does

UserOptimizedTransformer.forward is a shape-keyed dispatcher. At construction it reads (batch, seq_len, d_model, n_heads, d_ff, n_layers, causal, dtype), looks the key up in a table tuned on this GPU under the official gate, and picks a backend: whole-layer CUDA-graph capture with float16 internals at the launch-bound small shapes, float16-GEMM variants with an fp32 residual stream at the bandwidth-bound ones, torch.compile where fusion wins, and the unchanged reference wherever nothing earned promotion. A backend that hits out-of-memory at run time degrades to the organizer's own baseline instead of crashing. For official_14 there is a dedicated streamed runner: the input lives in host RAM, batch slices stream through SDPA's is_causal path - no mask tensor; the mask alone would be 10 GB - and outputs stream back. Nothing about the arithmetic changes; where it runs does.

Measured through the flattened submission artifact under the organizer's own script:

13 of 13 harness-runnable officials pass, five graded trials each, in every run of a five-run release series. Geometric mean 3.88x as the median of those five runs, whose geomeans span 3.70%; best case 12.01x at official_02, where CUDA-graph replay collapses the whole layer stack into one submission; no case beneath 1.72x; best throughput case 8.84x at official_13.
official_14 streams to completion at 81.5 s per pass - 134.4% MFU against the card's fp32 peak. Its watched outputs pass the from-scratch float64 referee at 0.436 and 0.433 of the graded error budget, and the organizer's own comparison rule passes it 3 of 3 seeds at a proxy length the dense reference can still run.
FLOP-weighted MFU over all fourteen officials: 133.9% of fp32 peak. official_14 carries 99.9% of the weight, so the graded score is, to three digits, the frontier shape's score. A weighted MFU above 100% is not a typo: the denominator is fp32 peak by deliberate, conservative choice, and the biggest shapes run float16 tensor cores behind the float32 accuracy boundary. REPORT.md section 12 carries the formula and the roofline argument.
One of the five release runs failed - official_06 hit the 1800 s subprocess cap while a resident LLM server re-pinned 7 GB of VRAM mid-run - and it is inside every aggregate above anyway, because dropping failed runs is how a population learns to say what you hoped. Section 8.3 of REPORT.md prints it beside its four clean siblings.

How we built it

Float16 internals behind a float32 boundary. The graded rule prices reduced precision per shape, so precision is a dispatch decision made per shape: the mixed backends keep the residual stream in fp32 and run the GEMMs in fp16, the half backends stream fp16 end to end, each in eager, CUDA-graph and packed flavors. The family holds 19 of 39 table entries and is behind every double-digit speedup. Its full-half member refused itself at the deepest shape during its own qualification, which is the gate doing its job.
Whole-layer CUDA-graph capture. A multi-layer forward issues on the order of a hundred kernel launches, and at the small officials none of them hides behind arithmetic. Capture the sequence once, replay it as a single submission. Combined with float16 internals it is what turns official_02 into 12.01x.
The bit-exact era's kernels, requalified. The packed QKV projection (three GEMMs become one), the fused Triton scale/mask/cast epilogue over the attention scores, and the fused Triton residual join all survive from the project's first hundred experiments and still carry table entries on their own measured merits.
A 44-key tune under the official methodology. 704 gate-checked rows, just under a day of GPU time: every reachable dispatch key timed against every admissible backend, promotion only on gate-pass in every trial plus a credible margin, all of it saved as replayable JSON evidence. When a later rule change required re-deriving the table, the re-derivation ran from that evidence at zero GPU cost - and changed exactly one line.
The official_06 rescue. That shape is two-thirds of the measurable FLOPs and its policy incumbent OOMs outright on 12 GB. The promotion rule learned to fall through to the next admissible backend instead of voiding the case, and a runtime OOM now degrades to baseline output rather than a crash.
The streamed frontier runner (bench/frontier.py). Host-resident, slice-seeded input generation so the 13.1 GB tensor never exists on the device, SDPA is_causal streaming so the 10 GB mask never exists at all, and the same mixed-precision forward the gate-checked backends run at cases 1-13.

The dispatch table's 39 entries are 10 cudagraph_mixed, 5 cudagraph_packed, 5 torch_mixed, 5 triton_softmax, 4 cudagraph_half, 3 triton_packed, 3 compiled_max_autotune, 2 compiled, 1 torch_native and 1 triton_exact. The five triton_softmax rows are fossils the dispatcher refuses at run time; unknown shapes fall through to the reference path unchanged, so a shape we never tuned gets the organizer's own arithmetic rather than a guess.

The harness behind all six is what makes the numbers medians of five full runs rather than a best-of: a correctness gate over 82 configurations including adversarial input scales, a timed gate with a noise band, a tuner that promotes only on gate-pass in every trial, a release gate that runs the real artifact in real subprocesses through the organizer's script, and an exclusion lock that makes every other gate - and the test suite - refuse to start while a timing window is open.

Challenges

The card is the antagonist. 12 GB is the whole budget. official_06 peaks at 10.7 GB and its five-run span is 33% - VRAM-pressure timing, reported as a median over its four clean readings with the failed fifth printed beside them. The exact-arithmetic backend wants about twice the card at that shape, which is why the promotion rule had to learn to fall through. And official_14 does not fit in any form: its input alone exceeds the card before the model runs.

The referee problem. Validating S=100000 causal attention with no trustworthy oracle: the graded comparison cannot exist at that length, and the obvious oracle - model.double() on the organizer's model - is not a true float64 forward, because its softmax hard-casts to fp32; our tests pin that. So the referee is built from the weights up in float64 on the CPU, chunked, sharing nothing with the code under test, applied under the official rule to full-length watched samples - and paired with the organizer's own comparison function verbatim at a proxy length, three seeds. Each half covers the other's gap, and passing both is the strongest claim available for a shape whose graded reference does not exist.

A crash that killed three full frontier runs before it was diagnosed. The float64 referee ran without no_grad against weights that require grad, retained ~400 MB of autograd graph per chunk, and exhausted Windows commit - at which point MKL dies with an access violation rather than a Python error, at a chunk index that moved with whatever else the machine held. One decorator fixed it; finding it took commit-instrumented bisection, because the crash site was the commit wall moving.

The release gate refuted the tune, twice, during the release runs themselves. The tuner samples two seeds; the release gate runs five trials per case, and that asymmetry demoted two float16 promotions whose tune-time worst element had relied on the relative rail (official_04 and official_10 - E106 and E107 in the ledger). A promotion the tuner makes is provisional until the release gate has passed it; the reverse is never true. Two officials still carry the same audit signature and pass the gate's five fixed-seed trials deterministically in every run of the series, and REPORT.md section 10 names that class.

Our own instruments were wrong more often than our kernels. Eighty-one defects found in this project's tooling, against exactly one in the optimized code. A traffic model that inflated 2x in the direction that made the plan look good; a compiled backend that silently stopped compiling mid-sweep and kept reporting timings under the label; a release gate that computed a verdict and never let it decide anything; a regression floor set below the reference's own measured noise. Every one of them produced plausible output, which is why they survived - and why most of this repository is the thing that checks the kernels rather than the kernels.

What we learned

Bit-exactness was the right first policy for the old gate and the wrong one for the graded gate, and the difference has a price we can read off the table. Under the pre-update harness the error budget came to 1.3 units in the last place of a bfloat16 output, which cannot absorb a single rounding step, so the project's first hundred experiments shipped only restructurings that move no reduction at all. The rules update replaced that budget with a graded rule that expressly admits internal quantization, and the same measure-first machinery repriced the whole project in days: the float16 backends the old policy excluded now hold 19 of 39 table entries and every double-digit speedup. Build the gate before the kernels, and when the rules change, the gate reprices the kernels.

What's next

A FlashAttention-capable build. The SDPA comparison in REPORT.md section 8.1 runs against mem_efficient; a flash-enabled wheel is the obvious next backend candidate, priced through the same gate as everything else.

A quiet machine for the float16-fragile rows. Two officials pass the release gate deterministically but carry tune-time error near the budget; re-tuning them on an idle card with more seeds either clears the signature or demotes them, and both outcomes beat not knowing.

Overlap in the frontier runner's stream loop. Host-to-device copies and compute run serially today; overlapping them is bounded at about three percent of the 81.5 s pass, which is why the loop was left synchronous and auditable. The saving is real and priced, and we have not claimed it.