Stage 5: Capstone: shipped upstream

mastery work · this chapter0/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 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

·State the gap your kernel closes, and name the two baselines (the XLA floor and the hand-tuned ceiling) you benchmark against.
·Produce a differential test suite covering values and grads, corner cases enumerated, passing at 1e-3 forward and 1e-2 grads in bf16.
·Show a benchmark record where every number states chip, dtype, shapes, and method, for both the floor and the ceiling comparison.
·Point to the review comments that changed your kernel, and explain what each one caught that your own testing didn't.
·Defend every schedule decision (block sizes, grid order, what got remat-ed) with a measured consequence, not a guess.
·Write the derivation section so a reader with no prior context could rebuild your kernel from it alone.
·Confirm the upstream PR is public with the review discussion linked, not merely "ready to send."

checked state lives in your browser only · the gate below is the public half

pass or fail, in public

The gate

GATE 05 QUEUED

criterion	measured
Correct under the differential suite	awaiting bench data
Beats XLA decisively at target shapes, within striking distance of the hand-tuned ceiling	awaiting bench data
The upstream submission exists in public	awaiting bench data

gates pass only on real hardware, published pass or fail · est./computed rows are predictions holding the slot until a run replaces them · records in bench/

Public artifact for this stage: The capstone write-up: derivation, schedule decisions, full benchmark record, what review changed.