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
SM90 (Hopper) FlashAttention-3 multi-head attention kernel and dispatch layer.
This module implements the warp-specialized, TMA-based MHA kernel for NVIDIA H100 GPUs, along with the dispatch chain that materializes comptime configuration choices (scheduler, sink, KV row offsets, ragged valid lengths) into concrete kernel instantiations.
Functions
-
mha_sm90_dispatch: Dispatches the SM90 FlashAttention-3 MHA kernel for a single batch.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!