Why does a flash attention kernel need three warpgroups?

Because the two roles want opposite amounts of the register file. The kernel is launched with num_threads set to _COMPUTE_WGS + 1, which is three warpgroups along an axis named wg, and a warpgroup is four warps or 128 threads acting as one issuing unit. A single lax.cond on wg < 2 sends warpgroups 0 and 1 into the compute role and warpgroup 2 into the memory role. The compute role calls set_max_registers(232, action="increase") because the output accumulator is block_q by head_dim_out floats and lives in registers for the whole key loop. The memory role calls set_max_registers(40, action="decrease"), which is what makes 232 affordable for the other two, and it does no arithmetic at all: one warp of it issues every TMA copy of keys, values, bias and mask into the staged shared-memory slots. The two compute warpgroups split one 2 * block_q query super-tile between them and take turns issuing matrix instructions using a pair of schedule barriers, so they overlap rather than contend.

What does the persistent flag do in the sm90 attention config?

It is the only field the sm90 Config adds to the shared base, and it changes three things at once. A persistent launch goes through static_scheduling_persistent_kernel, which launches one block per streaming multiprocessor and loops over the logical grid inside the kernel rather than letting the hardware scheduler hand out grid points. The loop carries prev_iters, the number of key blocks every earlier tile consumed, so the stage ring keeps its phase across grid steps instead of restarting at zero and colliding with a transfer already in flight. And a second warp of the memory warpgroup prefetches the next tile queries while the previous tile output is still draining, which is exactly why the persistent path keeps q, k and o as separate refs instead of overlapping them in a plgpu.RefUnion. The heuristic sets persistent to the negation of mask.is_causal, because a causal mask makes the per-tile work vary by a large factor and a static schedule handles that badly, and dropping persistence is also its last resort when the shared-memory plan will not fit.

How does the kernel choose block_q, block_kv and num_stages?

By one predicate: does the shared-memory plan fit in 227 KiB. get_heuristics_config starts at block_q 64, block_kv 128, num_stages 2, with persistent set from the causal flag, and returns that if _estimate_shared_mem_usage_bytes comes back under 227 * 1024. If not, block_kv halves to 64 and the same check runs again. If that still fails it turns persistence off, which lets the queries and keys share physical bytes with the output through a ref union and shrinks the plan again. get_autotuning_configs enumerates block_q in 64 and 128, with 128 only when the next power of two of the query sequence length is above 128; block_kv in 64 plus whichever of 128 and 256 are at most the next power of two of the key sequence length; persistent both ways; and num_stages in 2, 3 and 4, skipping non-persistent candidates whose stage count exceeds the number of key blocks. Every survivor is filtered through the same estimate. Nothing is compiled or measured to decide whether a config is legal.

Why does the kernel use exp2 rather than exp?

Because exp2 is a single hardware instruction on the GPU and exp is not. The kernel multiplies its running logits scale by math.log2(math.e) once, right before the softmax, and then uses jnp.exp2 everywhere, so the change of base costs nothing at run time. That decision leaks into the residual contract with the backward pass. The running maximum m_i is in base-2 log space inside the kernel, so when return_residuals is set the epilogue multiplies it by 1 / math.log2(math.e) and stores it in natural-log space, which is what the backward kernel expects to read. Skip that division and the forward output is still exactly right while the gradients are quietly wrong, which is the kind of bug that survives a numerical test on the forward pass alone.

What is rescale_threshold and when does it help?

It is how the kernel skips most of the online-softmax rescales. Textbook online softmax multiplies the whole accumulator by alpha = exp2(m_scale - m_i) every time the running maximum moves. Here that multiply is guarded: the accumulator is only scaled when alpha falls below rescale_threshold, and otherwise the previous scale m_scale is kept and the exponent absorbs the difference on the next block. The accumulator is block_q by head_dim_out float32 values living in registers, so skipping the multiply is a real saving whenever the maximum is barely changing, which is the common case once a few key blocks have been seen. The cost is bookkeeping: the denominator l_i and the accumulator can be left carrying a pending scale, so the epilogue applies alpha to l before storing it as a residual and to the accumulator when normalize_output is off, and a threshold of exactly 1.0 takes its own branch throughout so the old behaviour is still reachable.

Attention for sm90, Line by Line

Q: What does the persistent flag do in the sm90 attention config?

It is the only field the sm90 Config adds to the shared base, and it changes three things at once. A persistent launch goes through static_scheduling_persistent_kernel, which launches one block per streaming multiprocessor and loops over the logical grid inside the kernel rather than letting the hardware scheduler hand out grid points. The loop carries prev_iters, the number of key blocks every earlier tile consumed, so the stage ring keeps its phase across grid steps instead of restarting at zero and colliding with a transfer already in flight. And a second warp of the memory warpgroup prefetches the next tile queries while the previous tile output is still draining, which is exactly why the persistent path keeps q, k and o as separate refs instead of overlapping them in a plgpu.RefUnion. The heuristic sets persistent to the negation of mask.is_causal, because a causal mask makes the per-tile work vary by a large factor and a static schedule handles that badly, and dropping persistence is also its last resort when the shared-memory plan will not fit.

Q: How does the kernel choose block_q, block_kv and num_stages?

By one predicate: does the shared-memory plan fit in 227 KiB. get_heuristics_config starts at block_q 64, block_kv 128, num_stages 2, with persistent set from the causal flag, and returns that if _estimate_shared_mem_usage_bytes comes back under 227 * 1024. If not, block_kv halves to 64 and the same check runs again. If that still fails it turns persistence off, which lets the queries and keys share physical bytes with the output through a ref union and shrinks the plan again. get_autotuning_configs enumerates block_q in 64 and 128, with 128 only when the next power of two of the query sequence length is above 128; block_kv in 64 plus whichever of 128 and 256 are at most the next power of two of the key sequence length; persistent both ways; and num_stages in 2, 3 and 4, skipping non-persistent candidates whose stage count exceeds the number of key blocks. Every survivor is filtered through the same estimate. Nothing is compiled or measured to decide whether a config is legal.

Q: Why does the kernel use exp2 rather than exp?

Because exp2 is a single hardware instruction on the GPU and exp is not. The kernel multiplies its running logits scale by math.log2(math.e) once, right before the softmax, and then uses jnp.exp2 everywhere, so the change of base costs nothing at run time. That decision leaks into the residual contract with the backward pass. The running maximum m_i is in base-2 log space inside the kernel, so when return_residuals is set the epilogue multiplies it by 1 / math.log2(math.e) and stores it in natural-log space, which is what the backward kernel expects to read. Skip that division and the forward output is still exactly right while the gradients are quietly wrong, which is the kind of bug that survives a numerical test on the forward pass alone.

Q: What is rescale_threshold and when does it help?

It is how the kernel skips most of the online-softmax rescales. Textbook online softmax multiplies the whole accumulator by alpha = exp2(m_scale - m_i) every time the running maximum moves. Here that multiply is guarded: the accumulator is only scaled when alpha falls below rescale_threshold, and otherwise the previous scale m_scale is kept and the exponent absorbs the difference on the next block. The accumulator is block_q by head_dim_out float32 values living in registers, so skipping the multiply is a real saving whenever the maximum is barely changing, which is the common case once a few key blocks have been seen. The cost is bookkeeping: the denominator l_i and the accumulator can be left carrying a pending scale, so the epilogue applies alpha to l before storing it as a residual and to the accumulator when normalize_output is off, and a threshold of exactly 1.0 takes its own branch throughout so the old behaviour is still reachable.

A production flash attention kernel for the H100, read whole at a pinned commit of google/tokamax. The algorithm is the one everybody knows; what a real kernel adds on top of it is six hundred lines of scheduling. One thread block launches three warpgroups and each branches immediately into a different role: two do the arithmetic on 232 registers a thread, and the third drops to 40 and does nothing but issue transfers, running a stage ahead of the maths. Automatic synchronisation is switched off at launch, so every handshake in the file is one of sixteen numbered hardware barriers placed by hand, and running out of them is a real failure mode. The page walks the constants, the shared-memory plan, both configuration heuristics, both warpgroup roles and the launch, then steps through one query tile: which slot each warpgroup is holding, how far ahead the loader has run, and where the causal mask starts biting. A second exhibit runs the file own shared-memory estimate so you can watch a tiling stop fitting in 227 KiB.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Attention for sm90, Line by Line — built to be understood, not skimmed.

Questions

Why does a flash attention kernel need three warpgroups?: Because the two roles want opposite amounts of the register file. The kernel is launched with num_threads set to _COMPUTE_WGS + 1, which is three warpgroups along an axis named wg, and a warpgroup is four warps or 128 threads acting as one issuing unit. A single lax.cond on wg < 2 sends warpgroups 0 and 1 into the compute role and warpgroup 2 into the memory role. The compute role calls set_max_registers(232, action="increase") because the output accumulator is block_q by head_dim_out floats and lives in registers for the whole key loop. The memory role calls set_max_registers(40, action="decrease"), which is what makes 232 affordable for the other two, and it does no arithmetic at all: one warp of it issues every TMA copy of keys, values, bias and mask into the staged shared-memory slots. The two compute warpgroups split one 2 * block_q query super-tile between them and take turns issuing matrix instructions using a pair of schedule barriers, so they overlap rather than contend.
What does the persistent flag do in the sm90 attention config?: It is the only field the sm90 Config adds to the shared base, and it changes three things at once. A persistent launch goes through static_scheduling_persistent_kernel, which launches one block per streaming multiprocessor and loops over the logical grid inside the kernel rather than letting the hardware scheduler hand out grid points. The loop carries prev_iters, the number of key blocks every earlier tile consumed, so the stage ring keeps its phase across grid steps instead of restarting at zero and colliding with a transfer already in flight. And a second warp of the memory warpgroup prefetches the next tile queries while the previous tile output is still draining, which is exactly why the persistent path keeps q, k and o as separate refs instead of overlapping them in a plgpu.RefUnion. The heuristic sets persistent to the negation of mask.is_causal, because a causal mask makes the per-tile work vary by a large factor and a static schedule handles that badly, and dropping persistence is also its last resort when the shared-memory plan will not fit.
How does the kernel choose block_q, block_kv and num_stages?: By one predicate: does the shared-memory plan fit in 227 KiB. get_heuristics_config starts at block_q 64, block_kv 128, num_stages 2, with persistent set from the causal flag, and returns that if _estimate_shared_mem_usage_bytes comes back under 227 * 1024. If not, block_kv halves to 64 and the same check runs again. If that still fails it turns persistence off, which lets the queries and keys share physical bytes with the output through a ref union and shrinks the plan again. get_autotuning_configs enumerates block_q in 64 and 128, with 128 only when the next power of two of the query sequence length is above 128; block_kv in 64 plus whichever of 128 and 256 are at most the next power of two of the key sequence length; persistent both ways; and num_stages in 2, 3 and 4, skipping non-persistent candidates whose stage count exceeds the number of key blocks. Every survivor is filtered through the same estimate. Nothing is compiled or measured to decide whether a config is legal.
Why does the kernel use exp2 rather than exp?: Because exp2 is a single hardware instruction on the GPU and exp is not. The kernel multiplies its running logits scale by math.log2(math.e) once, right before the softmax, and then uses jnp.exp2 everywhere, so the change of base costs nothing at run time. That decision leaks into the residual contract with the backward pass. The running maximum m_i is in base-2 log space inside the kernel, so when return_residuals is set the epilogue multiplies it by 1 / math.log2(math.e) and stores it in natural-log space, which is what the backward kernel expects to read. Skip that division and the forward output is still exactly right while the gradients are quietly wrong, which is the kind of bug that survives a numerical test on the forward pass alone.
What is rescale_threshold and when does it help?: It is how the kernel skips most of the online-softmax rescales. Textbook online softmax multiplies the whole accumulator by alpha = exp2(m_scale - m_i) every time the running maximum moves. Here that multiply is guarded: the accumulator is only scaled when alpha falls below rescale_threshold, and otherwise the previous scale m_scale is kept and the exponent absorbs the difference on the next block. The accumulator is block_q by head_dim_out float32 values living in registers, so skipping the multiply is a real saving whenever the maximum is barely changing, which is the common case once a few key blocks have been seen. The cost is bookkeeping: the denominator l_i and the accumulator can be left carrying a pending scale, so the epilogue applies alpha to l before storing it as a residual and to the accumulator when normalize_output is off, and a threshold of exactly 1.0 takes its own branch throughout so the old behaviour is still reachable.

Attention for sm90, Line by Line

Questions

Related explainers