Triton training-kernel speedup

We task the agent with making one block of a genomics model train faster on a GPU. The block takes a batch of sequences a few thousand positions long, 128 channels each. It runs a three-tap dilated depthwise convolution along the sequence, normalizes the result over the whole example, then applies a small MLP and a scaled residual. As shipped it is plain PyTorch, slow on a GPU and slowest in the backward pass.

The model is frozen. Its constructor, parameters and outputs stay as they are, and the PyTorch code is the source of truth for the numerics. With gradients on and the input on the GPU, the convolution and the normalization must run through a hand-written Triton kernel, forward and backward, wrapped as an autograd function. The MLP may stay in PyTorch. No torch.compile, no TorchScript, no FX. Every gradient the block produced before, for the input, the convolution taps and both linear weights, must still be right.

The agent has one L4, torch and Triton, and no package index. It submits its committed workspace as often as it likes, and each submission is graded on the same class of GPU within minutes. There is no private split. The run's result is its best submission, graded again on a separate GPU at the end.

Sequence-to-function models in genomics read a stretch of DNA and predict what the cell does with it. Their backbones lean on dilated convolutions, which let a stack of cheap three-tap filters reach across tens of thousands of bases (Kelley et al., 2018). The block here is the ConvNeXt-style unit of a compact model of that kind (Liu et al., 2022).

The forward fusion of convolution and normalization is a short Triton kernel (Tillet et al., 2019); training is where the difficulty lives. The backward pass has to produce four gradients, and the normalization ties every position and channel of an example to one mean and one variance, so its gradient is a reduction over the whole example followed by a second sweep. What to save and what to recompute shows up in both the timing and the numerics. A training run that silently drifts is worse than a slow one, so the gradients are frozen and speed is what is left to win.

Correctness comes first and is all or nothing. The grader runs the block forward and backward on a seeded set of shapes, from odd small widths to production size, against its own fp32 recomputation of the block. The forward must match within 2e-4 in fp32, and each of the four gradient tensors within a relative error of 3e-3, with looser budgets in bf16 and fp16. A missing or wrong gradient on any shape scores the submission 0. With both linear weights zeroed the block must also be the exact identity, gradient included.

Speed is measured on production shapes, in fp32 and bf16, across a range of dilations and batch sizes. Each shape times a full training step, forward and backward with a dense random upstream gradient, median of repeated timings after warm-up. The agent's block, the plain PyTorch block and our reference Triton kernel are timed in the same process on the same GPU. Writing $r$ for the agent's speedup over plain PyTorch on one shape and $g$ for the reference kernel's speedup measured beside it, the shape scores

and the task score is the mean over shapes, with fp32 weighted more heavily than bf16. A plain autograd path runs near 1.0x and scores 0. Beating the reference kernel by ten percent scores 1, and the scale keeps going. Our reference fuses the convolution and the normalization and leaves the MLP in PyTorch.