metalworkingGitHub ↗
/techniques/flash-attention

Flash Attention

Flash attention computes softmax(Q·Kᵀ/√d)·V without ever materializing the L×L score matrix: streaming K and V in tiles, holding each tile's scores in registers, folding contributions into a running output via online softmax.

CUDA equivalent: FlashAttention itself; the algorithm is the same. The platform inflection is why it matters even more here. Writing an L×L intermediate to memory is a bandwidth catastrophe on any GPU, but on a bandwidth-bound machine with laptop DRAM it's disqualifying: at 32K context, the naive intermediate is a billion floats of round-trip traffic per layer.

The kernel anatomy, common to every implementation (walked line-by-line in the MLX case study):

  1. Q tile loads once and stays; K and V stream past via cooperative loads.
  2. S = Q·Kᵀ lands in a register tile. The score matrix exists only as one simdgroup's simdgroup_matrix fragments, is masked and softmaxed in place, multiplies V, and dies. (This transform-in-place requirement is why steel forked its attention mma.h from the GEMM one.)
  3. Online softmax makes the streaming legal; the base-2 exp path makes it cheap.
  4. Causal masking is a skip, not a mask, where possible: whole tiles above the diagonal are never visited.
  5. Specialization (function constants, head-dim enumeration, or codegen) strips masking/alignment code the shape doesn't need.

the L×L score matrix never materializes: at 32K context it would be a billion floats of round-trip traffic per layer only this K / V STREAM tiles pass through once, cooperative loads K tile V tile ONE SIMDGROUP'S REGISTERS Q tile loads once, stays resident S tile = Q·Kᵀ · masked and softmaxed in place born and dies here, per K tile O accumulator += P · V tile rescaled by the online-softmax correction c each step m, l running max and sum O written to memory once, after the last tile (O /= l) causal masking is a skip, not a mask: whole tiles above the diagonal are never visited

The whole trick in one picture: everything green lives in one simdgroup's registers, everything amber streams through once, and the dashed matrix on the right is the thing the algorithm exists to avoid writing.

What to take from the three-implementation comparison (MLX · MFA · llama.cpp): the algorithm is settled; the engineering disagreements are where the insight lives. Codegen vs templates vs enumeration; spill-tolerance vs spill-avoidance; branch-guarded vs unconditional correction. Same hardware, same math, three defensible kernels.

Boundaries of the technique on this platform, both load-bearing for practice. Decode is a different problem: one query row can't fill an 8×8 tile, so every implementation ships a separate vector kernel. The backward pass is unfinished business. MLX's fused attention has no Metal backward; the entire GPU implementation, verbatim:

bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) {
  return true;
}

void ScaledDotProductAttentionVJP::eval_gpu(
    const std::vector<array>& inputs,
    std::vector<array>& outputs) {
  throw std::runtime_error("NYI");

— MLX scaled_dot_product_attention.cpp:796-803↗ (still true on main; training falls back to the unfused graph). And the one open-source backward (MFA's split dQ / dK-dV design, forced by emulated float atomics) ships in Draw Things, not in a framework. If you're looking for the ecosystem's most valuable unwritten kernel, it's this one.

Edit this page on GitHub ↗