FlashAttention Implementation › Key Parameters

When using FlashAttention, several parameters directly affect behavior and performance:

• block_size (B_r, B_c): The tile dimensions for processing attention. Typically 64-128, determined automatically based on head dimension and SRAM availability. Larger blocks amortize memory transfer overhead but require more SRAM.

• head_dim (d): The dimension of each attention head. FlashAttention supports specific values: 8, 16, 32, 64, 128, 256. Performance is best at 64 or 128.

• causal: Whether to use causal (autoregressive) attention masking. Causal attention is optimized in FlashAttention to skip computation for masked positions entirely, not just mask them after computation.

• dropout_p: Dropout probability during training. FlashAttention implements dropout within the kernel, avoiding the need to store dropout masks in HBM.

• softmax_scale: Scaling factor applied before softmax, typically $1/\sqrt{d}$ where $d$ is the head dimension. This prevents dot products from growing too large with higher dimensions, which would push softmax outputs toward one-hot distributions. Can be customized for specialized attention variants.

• dtype: Data type for computation. fp16 and bf16 are fully supported and recommended. fp32 falls back to non-FlashAttention kernels in PyTorch.

• window_size: For FlashAttention-2's sliding window support, the size of the local attention window. Set to (-1, -1) for full attention or specific values like (256, 256) for symmetric windows.