For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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: Flash-attention prefill GPU kernel (FA2/FA3 algorithm).
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.

Structs​

Functions​

Structs

Functions