Motivation and scope
A language model run twice on the same prompt can give two different words, because the GPU adds numbers in a different order when the server is busier. The environment asks an AI agent to write the fast attention kernel that never does that.
Design choices
Fast kernels set the cost of training and serving, and few people can write them. The reward is hard to fool: bytes match or they don't, and speed is measured against a physical limit. The best run so far reaches about half that limit.
Experiment: Thinking Machines, Defeating Nondeterminism in LLM Inference.
Decode-only was a tuning task: one idea solved invariance, then an hour of tuning a constant. With prefill, decode wants to split and merge while prefill wants big ordered tiles. One addition order has to serve both.
Submission interface
Only /app/kernel.cu crosses over. The grader compiles it, loads the binary, and does everything else itself.
Reward hacks
Design choices
Reading the private seed
#include runs at compile time. So the seed is deleted before that.The cubin must be under 512 KiB, so answers can't be baked in. Each compile step times out after 5 minutes. The PTX must have exactly one kernel, attention, plus the four launch settings.
No host code runs: the grader launches one function, attention. So the whole program is the PTX, and scanning it covers everything the agent controls. Data, launching and timing all belong to the grader.
Static analysis of the PTX
The PTX is plain text. A few searches rule out the cheapest ways to be non-deterministic or to cheat the clock.
%clock %globaltimer %pm: a kernel that reads time can change its behaviour with timing.
atom / red on floats: threads add in arrival order, so the last bits change from run to run. Integer atomics are allowed.
cp.reduce.async.bulk, multimem.red: the same arrival-order float sum, done by the copy engine.
The scan only removes the cheap cases. The byte comparison in stage 4 is what enforces invariance.
Reward hacks
Accuracy check
The grader computes the same token three ways and measures how far the kernel is from the truth, using plain bf16 as the yardstick.
Reward hacks
Design choices
Dropping the sink
Reading past the end of a sequence
Accuracy versus invariance
Taken from FlashAttention-3's kvcache test. The yardstick is how wrong plain bf16 already is: 2× at the worst output, 1.5× on average. Anything sloppier than bf16 fails. The output is fp32, so rounding can't hide a difference.
Batch-invariance check
Alone, inside a chunk, inside a big batch, on any pages: the output bytes may not move. This is the check that makes the task hard.
Reward hacks
Design choices
Split chosen by batch size
Order chosen by chunk length
Merging in arrival order
Adding pages in memory order
Why 4095, 4096 and 4097
Why the sweep varies chunks, fillers and batch sizes
Decoded alone is the plainest run: one query, no neighbours. Inputs are seeded by (sequence, block), so a token has the same input bytes in every setting. Each batch size gets a new process, so no GPU state carries over.
The checks sample: 96 targets, up to 16 tokens each, 22 batch sizes, plus rows from the timed launches. That is strong evidence, not a proof.
Timing methodology
The kernel is timed on 24 fixed problems. Each one has its own floor: the least time an H100 needs just to move the bytes or do the math. Your score is how close you get to it.
Reward hacks
Design choices
Replaying a stored answer
__device__ array is certain to live long enough. Fresh inputs make what it holds wrong.Keeping state across rounds
The floor is the H100's spec peak. Any problem faster than its floor fails the whole grade, because it means work was skipped.
Why fresh data every launch
Why the median
Why many launches per timing
The 24 shapes are fixed and public, so scores are comparable, and agents may tune to them. The contents are new random data every launch, from a private seed, and invariance is checked far beyond those shapes. So tuning buys no shortcut.
Verification of timed launches
Speed is measured on 24 fixed shapes, so the grader pulls a few tokens out of those very launches and re-checks them.
Reward hacks
Design choices
Writing into the inputs
Writing into rows nobody owns
Careful on tests, fast on timing
Guessing which rows are checked
Picks come from random.SystemRandom, which uses the operating system. Seeded random numbers can be replayed by anyone with the seed. OS randomness can't, so the kernel can't know which rows are checked.
Why half by row, half by sequence
Why only 16 rows a round, not all
Decoding a token alone costs a launch, and grading must fit in an hour. A shape-dependent path affects every row of that shape, so a few rows expose it. Total: 2,688 rows per grade.
Scoring
Each problem gets an efficiency, the 24 are combined with a geometric mean, and a fixed curve turns that into a score from 0 to 1.
| Astra 1 | Astra 2 | Opus 1 | Opus 2 | |
|---|---|---|---|---|
| score | 0.274 | 0.276 | 0.028 | 0.030 |
| E | 0.49 | 0.49 | 0.05 | 0.06 |
| decode / prefill / mixed | 0.67 / 0.42 / 0.40 | 0.62 / 0.45 / 0.42 | 0.30 / 0.02 / 0.02 | 0.02 / 0.11 / 0.08 |
| first kernel written | minute 2 | minute 1 | minute 72 | minute 71 |
| tensor cores | wgmma + TMA | wgmma | none | wmma |
| cost | $37 | $33 | $13 | $12 |
All four passed every check, invariance included, in their first draft. What separated 0.03 from 0.27 was the Hopper tensor-core pipeline. Four runs describe these attempts, not a ranking of the models. Every turn and command is in the run transcripts.
Design choices
Geometric mean, not arithmetic
Fixed anchors, not a moving band
Where the 0.35 reference point comes from
Why the starter's time is frozen
The starter is slow on purpose, 0.3% of the limit: a fast one would hand over the hard part. Its efficiency, E0 = 0.00285, is fixed, because timing it in every grade would not fit the hour. The cost is about 2% noise in E.
Scaling
The environment is one layer of one model on one GPU. Each of those is a parameter.
Proposal, not built yet: vary layer × precision × cache layout × hardware × workload to get a family of tasks. A memorized kernel fails on a new combination, and held-out combinations show whether the agent learned the method.
Models and architectures
Other parameters
