IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.

Skip to main content

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

Implements CPU-based multi-head attention kernels, including flash attention and KV-cache-backed attention variants.

Functions

flash_attention: Computes scaled dot-product flash attention on CPU for the given query, key, value, and mask accessors.
flash_attention_kv_cache: Computes flash attention on CPU using a KV cache with an additive LayoutTensor mask.
flash_attention_split_kv: Variant of flash attention that takes the previous KV cache input_{k,v}_cache_fn and the current KV tensors input_k_fn and input_v_fn as separate arguments.

Functions

Edit this page

Edit this page

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!