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
mla
GPU kernels for Multi-head Latent Attention (MLA) decoding and prefill.
Provides the flare_mla_decoding and flare_mla_prefill entrypoints plus
their platform-specific dispatch and kernel implementations targeting NVIDIA
(SM80/SM100) and AMD (gfx950) GPUs, including split-K reduction, multi-token
prediction (MTP) query folding, per-token scale, and sparse-attention support.
comptime values
AMD_MLA_DECODE_FOLD_M_MAX
comptime AMD_MLA_DECODE_FOLD_M_MAX = 128
AMD_MLA_DECODE_FOLD_MAX_NUM_HEADS
comptime AMD_MLA_DECODE_FOLD_MAX_NUM_HEADS = 16
MLA_DECODE_MAX_SEQ_LEN
comptime MLA_DECODE_MAX_SEQ_LEN = 8
Functions
-
flare_mla_decoding: MLA decoding kernel that would only be called in the optimized compute graph. -
flare_mla_decoding_dispatch: Dispatches an MLA decoding request to the platform-specific kernel. -
flare_mla_prefill: MLA prefill kernel that would only be called in the optimized compute graph. Only supports ragged Q/K/V inputs. -
flare_mla_prefill_dispatch: Dispatches an MLA prefill request to the platform-specific kernel. -
mla_decode_max_seq_len: Max query tokens (S) the MLA decode branch can fold for this config. -
mla_decoding: MLA decoding kernel that computes single-batch attention per CTA. -
mla_decoding_single_batch: Flash attention v2 algorithm. -
mla_prefill: MLA prefill kernel that computes attention for sequences longer than one token. -
mla_prefill_plan: This calls a GPU kernel that plans how to process a batch of sequences with varying lengths using a fixed-size buffer. -
mla_prefill_plan_kernel: Plans how to process a batch of varying-length sequences through a fixed-size buffer. -
mla_prefill_single_batch: MLA for encoding where seqlen > 1. -
mla_splitk_reduce: Combines per-partition MLA decode partial outputs into the final result. -
set_buffer_lengths_to_zero: Zeroes out every element of a 1D buffer-lengths tensor.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!