An open attention backward kernel within 10% of FlashAttention-3 for head dimension 256
Difficulty 33: five matmuls per tile pair, cross-block dQ accumulation and the register budget at head dimension 256 make it clearly harder than the forward parity task (20), and no open kernel publishes a comparison. Application value 38: training cost for large-head models only, narrower than forward attention (49). Leverage 25: a few downstream kernels.
The problem
Can an open program for the learning-phase of attention run within 10% of the best one at large head sizes?
- 1Learning needs a backward attention pass
- 2Five multiplies per block pair
- 3Large heads overflow fast storage
- 4Match best program within 10%
- 5Open, modifiable training kernels
What is asked. Training a neural network needs a backward pass: after the model makes a prediction, the chain rule from calculus is applied layer by layer to find how the error changes with every input and weight. For attention, this means computing three gradients, , and , for the query, key and value matrices. FlashAttention-3 has a fast backward kernel on NVIDIA's H100, including for large 256 (the length of each query and key vector). The problem asks for an open kernel, in any language other than the FlashAttention-3 code itself, that reaches 90% of its backward speed at head dimension 256. The test uses BF16 numbers, causal masking, and sequences of 4096 to 16384 tokens.
Path to understanding
foundations → problemThe research entries this problem rests on, in the order to read them. Each links to its explainer.
- Start with
- Step 2
- Step 3
- FlashAttention: IO-aware exact attentionAlgorithm · host
- Then
- An open attention backward kernel within 10% of FlashAttention-3 for head dimension 256The problem
State of the art
The backward pass has lower arithmetic efficiency than the forward pass: it performs five matrix multiplications per tile pair instead of two, needs gradient accumulation into dQ across blocks (often through atomic adds or a separate pass), and keeps more state in registers and shared memory. Large head dimensions (192 and 256) stress the register and shared-memory budget further. FlashAttention-3's hopper/ directory ships head-dimension-256 backward kernels (BF16 and FP16), and the newer CuTe-DSL FlashAttention-4 code accepts head dimensions up to 256 on Hopper; no ThunderKittens, TileLang, Triton or other open kernel publishes a head-dimension-256 backward comparison against FlashAttention-3.
Realistic applications
Training cost for long-context models depends on the backward pass; large head dimensions appear in some multi-head latent attention and vision models.
What counts as solved
A public kernel meeting the 90% bar on the full sweep in two independent reruns.
How it gets verified
GPU benchmarkA kernel run and timed on stated hardware against a baseline.
Hardware: one H100 SXM 80 GB. Shapes: head dim 256, causal, BF16, sequence 4096, 8192, 16384, total tokens near 16384. Baseline: FlashAttention-3 backward from the hopper/ directory of Dao-AILab/flash-attention at a pinned commit. Metric: backward TFLOP/s (FlashAttention benchmark convention, 2.5x forward FLOPs), median of at least 100 iterations after warmup, three runs. Correctness: dQ, dK, dV max absolute error versus an FP32 PyTorch autograd reference no larger than 2x that of FlashAttention-3, and bitwise-identical results across two runs are not required but run-to-run variation is reported. Harness: public script at a pinned commit.
Claims under verification
+ add a claimNo claimed results on the Verification board for this problem yet.
Where the work is
Kernel engineering dominated by resource budgeting; limited theory.
Bottlenecks
ranked by votes- #1EngineeringRegister and shared-memory budget at head dim 256
Holding Q, K, V, dO tiles and accumulators exceeds on-chip budgets, forcing smaller tiles and lower tensor-core utilisation.
Community grades
Sign in to add your own grade.
Connected problems
- This problem is related to A Triton attention forward kernel within 5% of FlashAttention-3 on H100Same hardware and baseline; the backward pass is harder because of atomics and recomputation.