FlashAttention Algorithm › Key Parameters

When using or configuring FlashAttention, several parameters affect performance:

• block_size ($B_r$, $B_c$): The tile sizes for queries and keys/values. Larger blocks increase SRAM usage but reduce iteration overhead. Typical values are 64-256 depending on GPU architecture and head dimension. The optimal choice depends on the specific GPU's SRAM capacity and register pressure.

• d_head: The dimension per attention head. FlashAttention is most efficient when $d_{\text{head}}$ is small (64-128). Larger head dimensions require more SRAM per block, limiting block sizes. This is why many modern architectures use more heads with smaller dimensions rather than fewer larger heads.

• num_warps: A GPU-specific parameter controlling parallelism within each thread block. More warps can hide memory latency but increase register pressure. Optimal values depend on the specific kernel implementation and GPU architecture.

• causal: Whether to use causal masking for autoregressive models. Causal FlashAttention skips computation for future positions. This provides additional speedup beyond the memory benefits.

• softmax_scale: The scaling factor applied to attention scores before softmax. Default is $1/\sqrt{d_k}$, where $d_k$ is the key dimension. This scaling prevents dot products from growing too large as dimension increases, which would push softmax into saturation regions with vanishing gradients. Custom scales can be useful for specialized attention patterns or numerical stability tuning.

• dropout: FlashAttention supports fused dropout during training. Applying dropout within the kernel avoids materializing intermediate results. This provides additional memory savings.