Decoder-only Transformer optimization 101: 164K to 3.2M tokens/s on one GPU

Decoder-Only Transformer Optimization 101 is a report I worked on, published in September 2026. It measures how fast a single RTX 6000 Pro can train a decoder-only Transformer to play Super Smash Bros. Melee well enough to beat a level-9 CPU Fox. The naive implementation ran at 163,645 tokens per second. The final one runs at 3,217,492, and after ten minutes of training it still wins 26 games out of 30. This is a walk through the five steps in between and what each one bought.

The question and the model

The model is Eric Gu's HAL-20M: a six-layer decoder-only Transformer with width 512, eight heads of 64, and a 2048-wide MLP, about 20 million parameters. It reads 256 frames of game state and predicts the controller autoregressively, c-stick first, then main stick, then buttons. It is a good subject because it is the language-model pretraining workload in miniature. Causal attention, matrix multiplies, MLPs, backpropagation, and optimizer updates are all there; the tokens happen to be Melee frames. A batch is 512 windows of 256 frames, so 131,072 tokens per step.

Two things made this a different kind of speedrun. The target is gameplay, and next-token error is only a proxy for it, so the PyTorch baselines and the two final configurations are judged on 30 games against the level-9 CPU on Final Destination, with stocks taken, stocks lost, and damage dealt and taken tracked over the first twenty minutes of training. And the run is split in two stages. In the first, the architecture is frozen and every gain comes from running the same computation better on the GPU. In the second, the architecture may change as long as the model still beats the CPU.

For reference, PyTorch eager in BF16 does 401,766 tokens per second on this model, and torch.compile in BF16 does 918,193, although some width-512 BF16 runs in PyTorch developed non-finite gradients or diverged and the root cause was never isolated. The report's own kernels are written in Rust, first through cuTile-rs and later through CUDA-Oxide, and the naive cuTile-rs version starts well below PyTorch.

m1: faster matrix multiplies

Two changes to the GEMMs take m0 from 163,645 to 207,146 tokens per second. The first is tiling. Both versions run on Tensor Cores with BF16 inputs and FP32 accumulation, but the old kernel gave each thread block a 16 by 64 output tile and the new one gives it 128 by 128. Every element of A and B a block loads is reused across far more output elements, so the kernel spends less of its time waiting on memory.

The second is split-K for the gradient reductions. A weight gradient is a long sum over batch and sequence, and a long reduction with few output tiles cannot fill the GPU. Split-K cuts the reduction into k segments, computes their partial sums in parallel, and adds them in a second kernel. The profile shows where the two changes land: the QKV projection drops from 4.77 to 2.57 ms per layer in the forward pass, the two MLP GEMMs from 5.61 and 4.64 to 2.67 and 1.66, and the relative-position gradient in the backward pass from 6.80 to 1.91. Forward kernel time per layer goes from 40.48 to 31.66 ms, backward from 50.35 to 39.23.

m2: fusion and FlashAttention

Fusion is the general trick. If y = g(f(x)) and both functions are elementwise, compute both inside one kernel and keep the intermediate in a register. That removes the read of the intermediate tensor, and when the backward pass does not need it, the write as well. m2 applies it to the pointwise ops and reductions around every layer.

# old: two kernels, U round-trips through global memory
kernel_1:  parallel for i in 0..N:  x = load(X[i]); u = round_U(f(x)); store(U[i], u)
kernel_2:  parallel for i in 0..N:  u = load(U[i]); y = g(u); store(Y[i], y)

# new: one kernel, u stays in a register
fused_kernel:
    parallel for i in 0..N:
        x = load(X[i])
        u = round_U(f(x))
        if backward_needs(U): store(U[i], u)
        y = g(u)
        store(Y[i], y)

The bigger win is attention. The old path was three kernels: one wrote the full T by T score matrix S to global memory, softmax read it and wrote the full probability matrix P, and a third kernel read P to multiply by V. The fused kernel makes two sweeps over the key tiles for each query tile. The first computes the scores tile by tile and accumulates the row normalizers. The second recomputes the scores, normalizes and rounds each probability tile to BF16, and multiplies it into the output accumulator right away. S and P never touch global memory. Backward rebuilds the probability tiles from the saved row maximum and denominator. The FP32 reduction order changes, so the results are close to the old kernels and not bitwise identical.

fused_attention: parallel for query tile I:
    ell = 0                                   # one normalizer per row
    for J in causal_key_tiles(I):             # sweep 1: softmax statistics
        s = mask((Q[I] @ K[J].T + R[I,J]) / sqrt(d_h), J <= I)
        ell += rowsum(exp(s))
    o = 0
    for J in causal_key_tiles(I):             # sweep 2: recompute and consume
        s = mask((Q[I] @ K[J].T + R[I,J]) / sqrt(d_h), J <= I)
        p = bf16(exp(s) / ell)
        o += p @ V[J]                         # no global S or P
    store(O[I], bf16(o))
    store(norm_stats[I], state(ell))          # max and denominator for backward

This is the largest single step in the first stage. The attention group in the forward pass goes from 18.27 ms per layer to 2.32, and the attention group in the backward pass, the softmax derivative together with the dQ, dK, dV, and dR products, from 20.97 to 4.85. Forward kernel time per layer drops to 10.68 ms and backward to 19.84, and throughput more than doubles to 445,372 tokens per second.

m3: stop allocating, zeroing, and freeing

This step changes no arithmetic and adds 48%. For x = a + b, the old code allocated a buffer for x, launched a kernel to zero it, launched the add, which overwrote every element, and freed the buffer afterwards. With hundreds of temporaries per step the fills alone cost 15.74 ms of the 58 ms one forward layer took in the m0 profile. The fix is to allocate reusable storage once at initialization and write results straight into it. A buffer can be reused once its previous value's last consumer has finished, and accumulators that add into an existing value still need initializing. Nothing else does.

Buffer reuse and zero-fill removal together account for a 47.84% gain in tokens per second; the two were not measured separately. Capturing the whole step in a CUDA graph, so one replay replaces hundreds of host submissions, added 0.36%. Preparing and uploading the next batch during the current step measured at -0.04% in the fixed-batch benchmark, which is noise. m3 runs at 660,540 tokens per second.

r16: CUDA-Oxide, one sweep, two streams

cuTile-rs and CUDA-Oxide are both Rust interfaces to the GPU. The difference is that CUDA-Oxide gives control over individual threads, warps, and blocks. The model, the batch, and the sequence length stay exactly what they were in m3. Three things changed, and together they take throughput from 660,540 to 1,125,109 tokens per second, 116.5 ms per batch.

Pipeline control. Blocks load the next operand tile while the Tensor Cores compute the current one, so the compute units are not idle during loads. The GEMM epilogues fuse the bias, the dropout, and the residual add, so the output projection produces R1 = X + dropout(C Wo^T + bo) in one kernel and the intermediate never lands in memory. Other projections store directly in the layout the next kernel wants.

One-sweep attention. The two-sweep kernel from m2 recomputed the scores. The CUDA-Oxide kernel computes each score tile once, relative-position term included, and immediately accumulates the weighted values and the row normalizer, dividing at the end. It keeps a running maximum and rescales both accumulators as the maximum changes. This moves the BF16 rounding from the normalized probabilities to the unnormalized weights, so the floating-point results shift again even though the math is unchanged. Forward avoids the global R, S, and P matrices entirely; backward consumes probability and derivative tiles locally and keeps only dR for the relative-embedding gradient.

Two streams. The resource mix shifts during a step: GEMMs live on the Tensor Cores, LayerNorm and dropout live on memory bandwidth. Splitting the 512 sequences into two independent lanes of 256, each with its own workspace, lets one lane's Tensor Core work overlap the other's memory work. After both lanes finish, their gradients are summed, clipped, and applied in a single AdamW update, so the update is the same as the single-batch version up to floating point. In the Nsight capture, kernel overlap covers 36% of the trace.

Per layer, forward kernel time falls from 10.57 to 6.40 ms. Backward barely moves, 19.91 to 19.01, which says where the remaining time lives. The kernel sums count both lanes separately, so they can exceed the elapsed step time; part of the gain is overlap that no per-kernel number shows.

x2: a smaller model on every other frame

With the first stage done, the second stage changes the model. Two changes. The CPU keeps frames 0, 2, ..., 254 of each 256-frame window, so 128 frames reach the GPU. And the width halves, 512 to 256, with heads down from eight to four and the MLP from 2048 to 1024. Six layers, 64-wide heads, batch 512, and the two lanes stay. The result is a 5.3M-parameter model at 20.37 ms per batch. Tokens per second here count processed frames, 128 per sequence, so the fair comparison is windows per second: 5.72 times r16. Per-layer kernel time falls from 6.40 to 1.12 ms forward and from 19.01 to 3.13 backward. x2 also recomputes the GELU input during backward and never stores it, and that work is included in those numbers.

The cost is the ceiling. After ten minutes of training, x2 won 26 of 30 games against the level-9 CPU; r16 won 28. The report's last page says why the trade works at all. A standard FP32 matrix multiply of an m by k matrix with a k by n one costs about 2mnk FLOPs and at least 4(mk + kn + mn) bytes of traffic, and both are floors that no amount of kernel work removes. Halving all three dimensions cuts the arithmetic by 8 and the minimum traffic by 4. A model that is too small plateaus before it learns enough. One that is merely smaller learns faster and lands lower, and for a fixed target like beating the CPU, faster is what counts.

What carries over

Five lessons, in the order they paid off. Use Tensor Cores and tile so that loaded operands are reused. Fuse so that intermediates stay in registers, attention included. Do not allocate, zero, or free inside the training step for values that will be overwritten anyway. Design kernels around the workload: fused epilogues, loads overlapped with compute, and concurrent streams for work that does not depend on each other. And when the goal is a capable model in the least time, shrink the matrices.

The first four transfer to large Transformer training as they are, and at that scale the percentages are money: a ten percent reduction in training cost on a $100 million run is $10 million. The fifth depends on what you are measuring. Time to a target loss and time to a model that does the job are different races, and this report was run on the second one.