the path · 0/15
start the path

the path · chapter 15 of 15 · part ii, the track

stage 05
title Capstone: shipped upstream
schedule weeks 11-14
gate YOURS

Prove the track against production standards: one kernel, end to end, in front of the community.

mastery work · this chapter0/5
  1. go →
  2. go →
  3. go →
  4. auto
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
the chapter

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.

the chapter

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 chapter

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.

deliverables

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
sources

Readings

  • JAXBench · the benchmark and the gaps its baselines leave open; capstone target candidates live here
  • Tokamax · the natural upstream venue; its contribution surface is the autotuning option
the plan

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
the self-test, before the gate

Can you

0/7
checked state lives in your browser only · the gate below is the public half
the finale is yours

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.

01

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 ↗
02

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 ↗
03

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, annotated

A 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 beat
where it ships

Tokamax ↗ · 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