Split-Q, Not Split-K

FlashAttention-2 changes how the 4 warps inside a thread block partition work versus FlashAttention-1, which split K and V across warps (each warp writing a partial output to shared memory, then synchronizing to sum). What did FA-2 change this to, and why does it matter?

FA-2 keeps FA-1's split-K scheme but overlaps the shared-memory reduction with cp.async prefetches of the next K/V tile to hide the sync cost

FA-2 splits the batch dimension across warps while keeping the full Q, K, and V tiles resident in every warp's registers

FA-2 splits Q across warps instead, keeping K and V visible to all warps, so each warp computes a full output slice via $QK^\top$ then $PV$ with no inter-warp shared-memory reduction

FA-2 assigns whole attention heads to separate warps and drops online softmax rescaling entirely, since each warp now owns an independent softmax

FA-2 moves warp partitioning up to the thread-block level only, relying on L2 cache residency of K/V to avoid the cross-warp reduction instead of changing which operand is split