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 package
gpu
GPU multi-head attention (MHA), cross-attention, and multi-head latent attention (MLA) kernels. Vendor-specific implementations live under amd/ and nvidia/.
Packages
-
amd_rdna: TileTensor-native attention kernels for AMD RDNA3+ (gfx11xx/gfx12xx). -
amd_structured: TileTensor-native attention kernels for AMD gfx950 (MI355X). -
apple: Apple (Metal) GPU attention kernels. -
nvidia: NVIDIA GPU attention kernels and tile-scheduling utilities.
Modules
-
mha: GPU flash-attention kernels and dispatch logic for prefill and decode. -
mha_cross: Implements a naive GPU multihead cross attention kernel supporting ragged batched inputs and a paged KV cache. -
mha_decode_partition_heuristic: Split-K partition-count heuristics for MHA and MLA decode kernels. -
mla: GPU kernels for Multi-head Latent Attention (MLA) decoding and prefill. -
mla_decode_dispatch_scalars: Device-dispatched MLA decode dispatch-metadata scalars. -
mla_graph: Provides manually fused GPU graph kernels for Multi-head Latent Attention (MLA). -
mla_index_fp8: MLA FP8 index kernel for computing attention scores with paged KV cache. -
sparse_index_fp8_sm100: SM100 (B200) tensor-core FP8 MLA lightning-indexer score kernel. -
sparse_index_fp8_sm100_prefill: SM100 (B200) warp-specialized PREFILL variant of the FP8 MLA indexer scorer. -
sparse_indexer_common: Shared device functions for the MiniMax-M3 sparse-attention (MSA) indexer.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!