For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

mha

GPU flash-attention kernels and dispatch logic for prefill and decode.

Implements FA2 and FA3 flash-attention for NVIDIA and AMD GPUs, a naive two-BMM reference path, split-K decode partitioning, and the host-side dispatch layer (flash_attention_dispatch) that selects among them based on dtype, head depth, and target architecture.

Structs

• ​MHADecodeDispatchMetadata: Runtime metadata required to dispatch an MHA decode kernel launch.

Functions

• ​depth_supported_by_gpu: Return True if the given head depth is supported for flash attention on this GPU.
• ​flash_attention: Run flash attention with a dense mask tensor on the current device.
• ​flash_attention_dispatch: Dispatch a flash-attention kernel for prefill or decode over dense or KV-cache operands.
• ​flash_attention_hw_supported: Return True if the current GPU supports flash attention for qkv_type.
• ​flash_attention_ragged: Run flash attention on ragged (variable-length) batch inputs.
• ​get_mha_decoding_max_num_partitions: Return the maximum number of split-K partitions for CUDA-graph-stable launches.
• ​get_mha_decoding_num_partitions: Return the recommended number of split-K partitions for MHA decoding.
• ​get_waves_per_eu: Return the recommended rocdl.waves_per_eu hint for an AMD MHA kernel.
• ​mha:
• ​mha_decoding: Flash-attention decode GPU kernel with optional split-K partitioning.
• ​mha_decoding_single_batch: Flash attention v2 algorithm.
• ​mha_decoding_single_batch_pipelined: Flash attention v2 decode kernel for a single batch element with pipelined multistage MMA.
• ​mha_gpu_naive: Launch the naive (two-pass BMM) GPU attention implementation.
• ​mha_single_batch: MHA for token gen where seqlen = 1 and num_keys >= 1.
• ​mha_single_batch_pipelined: MHA for token gen where seqlen = 1 and num_keys >= 1.
• ​mha_splitk_reduce: Single-warp reduction kernel that merges split-K partial attention outputs.
• ​q_num_matrix_view_rows: Return the number of matrix rows when viewing Q as a 2-D tensor for TMA.
• ​scale_and_mask_helper: Apply softmax scaling and masking to one P = Q·Kᵀ MMA result tile in registers.