Before the first training step, we timed every kernel the 256M model will run. Two of the planned picks lost, FP8 cleared its speed bar, and the GPU had to slow down to stay up.

01How it was measured.

A microbatch is 16,384 tokens of real documents drawn from the token shards. Their lengths matter: the median document is 457 tokens, but 45% of all tokens sit in documents of 4,096 or more. No document attends across its own boundary, and each row is cut at the context length, so a kernel sees between 18 and 46 separate segments per microbatch.

Before anything was timed, every kernel's outputs and gradients were compared with a 32-bit reference. All of them pass, with relative errors between 0.2% and 0.5%. Each time on this page is the median of 20 runs after warmup.

Every number here was measured with the GPU's graphics clock capped at 2,400 MHz (section 07). The cap makes these kernels about 13% slower and changes none of the rankings.

02FlashAttention-2 wins the training step.

The plan picked cuDNN's attention because a public benchmark ranked its forward pass fastest. Here too it wins the forward pass: 141 TFLOPS against FlashAttention-2's 109 on global layers at 4,096 context. Training runs the backward pass as well, though, and there cuDNN falls behind. Over a full step, FlashAttention-2 (Dao) takes 19 to 31% less time on sliding-window layers (Fig. 3) and 9 to 17% less on global ones (Fig. 4).

cuDNN has two more problems. It builds a new execution graph whenever a batch holds a document count it hasn't seen, which costs 36 to 56 ms each time, so batches would need padding to fixed sizes. It also fails under CUDA graphs, which the training loop needs (section 06). PyTorch's own cuDNN path refuses our GPU outright. It only ran through NVIDIA's frontend library, whose kernel for our GPU is still marked experimental.

FlexAttention (Dong et al.) came last, and building its mask costs another 6.5 to 10 ms per microbatch. FlashAttention-2 ships inside PyTorch with sliding windows and grouped-query attention, so nothing has to be compiled from source.

The training loop now uses FlashAttention-2.

03What linear attention costs.

The plan's local-layer contest puts sliding-window attention up against Gated DeltaNet (Yang et al.) and Kimi Delta Attention (Kimi Linear), scored on loss per hour of training. Day 0 measures the hours. Both linear kernels, from the flash-linear-attention library, match a token-by-token 32-bit reference, and their cost barely moves with context length.

Neither is cheap. A Gated DeltaNet layer takes 1.65 to 1.75 ms per step and a Kimi Delta layer 3.15 to 3.32 ms, against 0.69 to 0.95 ms for a sliding-window layer. Counting the five global layers too, a microbatch spends 33 ms in these kernels with Gated DeltaNet and 57 ms with Kimi Delta, against 21.5 ms with the window (Fig. 1).

The whole model is about 25 × 10¹² operations per microbatch, roughly 140 ms at the measured matrix-multiply rate. Measured against that, Gated DeltaNet adds about 8% to every step and Kimi Delta about 25%, before their extra projections and convolutions. To stay in the design, each has to lower the loss by more than it slows the clock.

04FP8 feed-forward layers: 1.5 times faster.

The plan lets the feed-forward layers run in FP8 if that's more than 10% faster with no loss gap. All three matrix multiplies in a layer, one forward and two backward, run in 8-bit floating point with one scale per row.

A forward and backward pass through one squared-ReLU feed-forward layer drops from 4.95 to 3.28 ms, 1.51 times faster. Bare matrix multiplies run at 180 TFLOPS in bf16 and 334 to 355 in FP8. The full layer reaches only 167 and 251, because scaling and casting around each multiply take time.

Outputs and gradients differ from bf16 by about 6%. Whether that shows in the loss was a question for a proxy run, so FP8 had passed half of its test. The proxy run passed the other half: 1.3116 bits per byte against 1.3130 in bf16, at 11.7% more tokens per second, so the main run uses FP8.

05The loss head: plain PyTorch wins.

The last layer scores each token against all 32,768 vocabulary entries and computes the loss. Cut Cross-Entropy (Wijmans et al.) does this without ever storing the scores, and the plan picked it to save memory.

At 16,384 tokens, plain PyTorch under torch.compile takes 19.4 ms and Cut Cross-Entropy 27.6 ms. The head's matrix multiplies alone come to about 3.3 × 10¹² operations, so the compiled version runs close to the multiply rate. Cut Cross-Entropy saves 1.1 GB, which a model this size doesn't need. The training loop now uses the compiled PyTorch version.

06CUDA graphs work, with padding.

Launching thousands of small kernels per step wastes CPU time, so the training loop will record each step once as a CUDA graph and replay it. Replay needs fixed tensor shapes, but document boundaries change with every batch. The fix is to pad the list of boundaries with empty documents up to a fixed size.

Empty documents change nothing in any kernel: outputs match bit for bit. Replayed over three batches with different boundaries, FlashAttention-2, FlexAttention, Gated DeltaNet and Kimi Delta all matched an unrecorded run exactly, and each step took 9 to 20% less time. cuDNN failed to record.

07The GPU had to slow down.

Timing the linear-attention kernels reset the whole machine twice, with nothing in the kernel log beforehand. On the next boot, the platform reported that an uncorrected hardware error had forced the reset. Both resets came during autotuning, when the compiler runs dozens of kernel variants back to back and power draw swings hard.

Capping the graphics clock at 2,400 MHz, down from about 2,900 under load at stock settings, stopped it. The full linear-attention suite then ran clean, peaking at 483 W in 100 ms samples. The cap is now applied at every boot and costs about 13% on these kernels.

The resets also left hundreds of empty files in the compiler's cache, which broke the next run until they were deleted. The root cause isn't known yet. Power transients and the GPU's link to the board are the main suspects.

08What changes.

Both switches are made: FlashAttention-2 for attention and compiled PyTorch for the loss head. FP8 feed-forward layers move on to a proxy run that checks the loss. Gated DeltaNet and Kimi Delta now have a measured cost that their loss has to beat. Kimi Delta went on to win the bake-off and is now the 256M's local layer.

The learning-rate sweep, run 2 in the proxy queue, and the rest of the queue have run since. Their results are on the pre-run page, and the main run is live on the dashboard.