Optimization: Triton training-kernel speedup

The agent gets the training codebase with the block implemented in plain PyTorch, profiling tools, and a research loop: profile the step, write Triton forward and backward kernels for the block, prove them equal, measure, iterate. Everything about the kernels is open, tiling, fusion, recomputation strategy, memory layout. The gradients are not: the block's outputs and every gradient it produces must match the reference implementation.

Training throughput is usually won or lost in a handful of hot blocks, and the standard way to buy it back is to replace the framework's composed ops with a fused kernel. Writing that kernel by hand is exacting work: the forward pass has to match, and the backward pass has to reproduce every gradient, because a training run that silently drifts is worse than a slow one.

We took the core block of a genomics sequence-to-function model we work with and made that job the task.

The verifier runs the training step with the submitted kernels against the unmodified implementation, checks forward outputs and gradients for agreement, and scores the measured training-step speedup. Kernels that drift numerically score nothing, however fast.

Numerics separated no one: every submission stayed inside the gradient-parity budgets, so the scores rank pure training-step speed. The two runs at the top of the scale earned it by opposite routes. One profiled and benchmarked on a live card until its fused stage ran at the memory-bandwidth floor. It tried to fuse the MLP as well and dropped the idea when the measurements favored the stock path. The other, given no usable GPU, rewrote its kernel algorithm in plain PyTorch. It verified that stand-in against the reference's own gradients across shapes and dtypes up to the full working size. The kernels it shipped had never run.

The weakest run was blind by misdiagnosis rather than necessity. It read two crashed probes as proof the machine had no card; sibling runs found theirs. Then it checked its algebra, compiled its kernels offline, and shipped without ever timing them.

Every graded run cleared the gradient parity gate; the spread is training-step speed alone. Two runs reached full marks by opposite routes, one measuring everything on a live card, one never executing a kernel. The weakest run also worked blind.

A strong run

1. Read the numerics before the kernel language. It read the reference block and the tests that set the parity budgets, then traced how the training loop drives the block under autocast, pinning the dtypes and cast points to reproduce.
2. Choose what to save and where to round. It fused the dilated convolution and the per-example normalization into one differentiable node, leaving the MLP to the framework. The backward recomputes the convolution rather than storing it, and the kernels round the convolution output to the input dtype before the fp32 norm math, matching the reference.
3. No GPU, so prove the math another way. With no card to run them on, it mirrored the kernel algorithm in plain PyTorch, same one-pass variance, same rounding, same tap order, diffed the mirror against autograd of the reference across shapes and dtypes, then reread the kernel source for GPU-only syntax mistakes. Eleven minutes. At grading every parity check passed and the speedup topped the scale.

A failed run

1. Same boundary, four kernels. The weakest run read the same reference and wrote a comparable fused forward with a three-kernel backward. Its plan was fine.
2. Declare the machine blind. Its probes crashed on an unrelated missing dependency, it took the crash as proof there was no GPU, and it never reran the check. Other runs ran the probe, found the card, and measured.
3. Verify everything but speed. It checked its gradient algebra in plain PyTorch and compiled all four kernels offline against a GPU target, catching a real syntax error. It never timed a launch.
4. Correct and slow. Every parity check passed at grading, and the training step came in near 1.3 times baseline, the bottom of the range; the top runs land near 1.7.