FlashAttention

FlashAttention is integrated into several major machine learning frameworks and libraries, allowing for optimized Transformer operations without manual implementation. Key integrations include:

• PyTorch: Integrated into the core nn.Transformer module.
• Huggingface Transformers: Supported within the library.
• Microsoft DeepSpeed: Integrated into the DeepSpeed inference engine.
• NVIDIA Megatron-LM: Used for training large transformer language models at scale.
• MosaicML Composer: Used for efficient neural network training.
• EleutherAI GPT-NeoX: A research library for large language models.
• PaddlePaddle: Accessible via the paddle.nn.functional.flash_attention API.

This CUDA extension provides a fused implementation of dropout, residual connections, and LayerNorm (building on Apex's FastLayerNorm). Key capabilities include:

• Fused Operations: Combines dropout + residual + LayerNorm.
• Architecture Support: Works for both pre-norm and post-norm architectures.
• RMSNorm: Supports RMSNorm as an option.
• Parallel Residuals: Supports layer norm with parallel residuals (e.g., used in GPT-J, GPT-NeoX, PaLM).
• Dimension Support: Supports hidden dimensions up to 8192 (must be divisible by 8). For dimensions larger than 8k, contact the maintainers via an issue.

The number of Warp Groups (num_wg) determines the tile_m size and the available register budget per thread. Increasing num_wg improves M-direction parallelism but increases shared memory (smem) usage and reduces the register budget.

Hardware Constraints (H100):

• SMEM: ~224 KB available for tensor buffers (after reserving ~3 KB for LSE, dPsum, and mbarriers).
• GMMA atom: Always M=64. The effective M dimension (after swap) must be divisible by 64. The N dimension must be divisible by atom_layout_n * 8.

Rank your evidence to understand the risk of corruption or hallucination:

• Tier 1 (Trustworthy): Deterministic source-level facts, reproducible pass/fail results, or divergence against a reference implementation at a specific tensor index.
• Tier 2 (Usable, needs corroboration): PTX (via CUTE_DSL_KEEP_PTX=1), dumped SASS (CUTE_CUBIN_PATH), shared-memory layout offsets, or cute.printf traces.
• Tier 3 (Contaminated by definition): Anything captured at the moment of a trap (register values, faulting addresses, block/thread IDs, CUTE_DSL_LINEINFO attribution, or cuda-gdb backtraces).

Warning: Tier 3 data is highly dangerous for building hypotheses because the corruption that caused the fault also produced the details that make the fault look like it fits your theory. Use Tier 3 only to confirm a prediction made using Tier 1 or 2 evidence.

The second operand (Operand A) in an HGMMA instruction determines the memory access mode:

• SS Mode (Shared Memory): Indicated by a shared memory descriptor like gdesc[UR..]. In this mode, both operands A and B are read from shared memory.
• RS Mode (Register Source): Indicated by a plain register (e.g., R<N>). In this mode, operand A is read directly from registers, while operand B is read from shared memory.

When to use RS: RS is used to reduce shared memory traffic when data is already present in registers from a previous computation (e.g., feeding dS from a softmax backward directly into a GEMM).

# RS Mode: A=R232 (dS in regs), B=gdesc (K in smem)
HGMMA.64x192x16.F32.BF16 R24, R232, gdesc[UR4], ...

# SS Mode: both from smem
HGMMA.64x192x16.F32.BF16 R24, gdesc[UR16], gdesc[UR4], ...

The mma_dkv_is_rs (Register-Source for dKV) optimization can significantly reduce shared memory traffic and latency.

Conditions for activation: When AtomLayoutMSdP == 1 AND AtomLayoutNdKV == num_wg AND SdP_swapAB is true AND dKV_swapAB is false.

Benefits:

• Eliminates sP from smem: Saves tile_m * tile_n * 2 bytes.
• Eliminates P R2S store: Reduces shared memory traffic.
• Eliminates A operand reads: Removes the need to read the A operand for dK and dV GEMMs by feeding P and dS matrices directly from registers.

In variable-length (varlen) operations, buffers such as dq_accum are laid out with tile-aligned gaps between sequences. The calculation of the padded_offset_q depends strictly on the tile_m value. If the tile size used for preprocessing differs from the tile_m used by the kernel, the preprocessing will zero out the wrong memory locations, leading to corrupted data in subsequent batches.

padded_offset_q = ((offset_q + batch_idx * tile_m) // tile_m) * tile_m

The atom_layout distributes Warp Groups (WGs) across the M and N dimensions of an MMA output.

Given num_wg MMA WGs and atom_layout_m = A:

• M direction: A warp groups, each handling M/A rows.
• N direction: num_wg/A warp groups, each handling N/(num_wg/A) columns.

Performance Trade-off:

• Larger wg_n (more WGs in N direction): Each instruction reads a smaller B slice, but there are more total instructions reading overlapping A slices.
• Smaller wg_n (fewer WGs in N direction): Fewer instructions total, but each instruction reads a larger B slice.
• Rule of thumb: Smaller wg_n typically results in less total smem traffic.

The swap_AB option allows an MMA to swap its A and B operands, transposing the output tile. This is used to exchange which dimension maps to M (which must be divisible by 64) and which maps to N.

When to use swap_AB:

• If the natural M dimension is not divisible by 64, but N is (e.g., tile_m=80).
• To change which operand resides in registers versus shared memory.

Application in Backward Pass (5 MMAs):

• SdP (S=Q@K^T, dP=dO@V^T): Swap if tile_m % 64 != 0.
• dKV (dK=dS^T@Q, dV=P^T@dO): Swap if tile_n % 64 != 0 but hdim % 64 == 0.
• dQ (dQ=dS@K): Swap if tile_m % 64 != 0 but hdim % 64 == 0.