A different referee
Everything before this was practice under your own referee. The capstone changes the referee: a kernel is done when it is correct under a differential suite, competitive against both the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → floor and the hand-tuned ceiling, and reviewed by people who maintain kernels for a living. Upstream review is the only test that cannot be gamed, and what it changes about your code is part of the artifact.
Pick a real gap
Pick a gap that is real: an attention variant no fast public TPU kernel covers, a block-size autotuning pass for an existing Tokamax kernel with measured wins across a shape sweep, or a JAXBench entry on an operator where published baselines trail the expert bound. Sized right, this is four weeks: one derivation, one forward, one backward, one review cycle.
The write-up is the artifact
The write-up matters as much as the merge. Derivation, schedule decisions with their measured consequences, the full benchmark record with provenance, and what review changed: that document is the proof that the fourteen weeks produced an engineer, and it closes the track in public, where it started.
What gets built
- One kernel closing a real gap: an uncovered attention variant, an autotuning pass, or a benchmark entry
- Forward and backward, differential-tested, corner cases enumerated
- Benchmarked against the XLA floor and the hand-tuned ceiling, recorded in bench/
- Sent upstream, review discussion linked from this site
Readings
The work, in order
weeks 11-14
- Week 11: pick the gap, write the derivation and the plan, open the conversation upstream early
- Week 12: forward, differential-tested; first benchmark sweep against floor and ceiling
- Week 13: backward via custom_vjp; corner cases enumerated; bench records land in the repo
- Week 14: the upstream PR, the review cycle, and the capstone write-up as the track finale
Can you
Your gate
Every gate above was measured by this site and published. This one is yours: the site cannot run it, and that is the point. Three criteria, each with the exact method this course already proved out, and three starting points you can open and run today.
Correct under the differential suite
The gate-close methodology, exactly: f32 pipelines with matmul precision pinned against tight bars, bf16 reported as relative error against an f32 reference, edge cases enumerated. Copy the harness from the gate-close lab.
open the differential harness in Colab ↗Beats XLA decisively at your target shapes
The power-day bar: the compiler is the baseline, not naive code; every number carries chip, dtype, shapes, and method; the hand-tuned ceiling (Tokamax) is named and measured against.
open the benchmark harness in Colab ↗The submission exists in public
A PR to Tokamax or the JAX Pallas ops directory. Upstream review is the one referee that cannot be gamed, and what review changes about your kernel is part of the learning.
Tokamax, the natural venue ↗Three starting points, runnable today
Ragged attention, hardened
The power bench already shows the toy version beating padded XLA by 1.98x, and the profiler proved the compiler pattern-matches only vanilla attention: data-dependent variants are exactly where recognition stops.
Start from the ragged kernel in the power notebook; harden it with splash-style block masks and a backward pass.
the ragged kernel, runnable now ↗Block-mask attention (splash-lite)
Windowed attention ran 4.87x past XLA masked at seq 16k on this bench. General block masks are the production version, and the Splash walk annotates the real machinery line by line.
Start from the windowed kernel and the Splash walk; generalize the loop bounds to an arbitrary block mask fed by scalar prefetch.
the Splash walk, annotatedA fused quantized decode kernel that wins
The int8 fight ended 0.95x: the compiler kept that one. Finding the shape and quantization scheme where the fused kernel wins, and proving it with provenance, is a capstone with a built-in adversary.
Start from the dequant kernel in the power notebook; sweep shapes and schemes until the win is decisive or the negative result is airtight.
the current record to beatTokamax ↗ · custom accelerator kernels for TPU and GPU, built on Pallas; read its contribution surface before you write a line
JAX Pallas ops ↗ · the in-tree kernel directory; smaller surface, same referee