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 function

mha_decoding

def mha_decoding[q_type: DType, k_t: MHAOperand, v_t: MHAOperand, output_type: DType, mask_t: MHAMask, valid_length_layout: Layout, BM: Int, BN: Int, BK: Int, WM: Int, WN: Int, depth: Int, num_heads: Int, num_threads: Int, num_pipeline_stages: Int, group: Int = Int(1), ragged: Bool = False, is_shared_kv: Bool = False, sink: Bool = False, _use_valid_length: Bool = False, _is_cache_length_accurate: Bool = False, decoding_warp_split_k: Bool = False](q_ptr: Pointer[Scalar[q_type], ImmutAnyOrigin, _safe=False], k: k_t, v: v_t, output_ptr: Pointer[Scalar[output_type], MutAnyOrigin, _safe=False], exp_sum_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False], qk_max_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False], scale: Float32, batch_size: Int, num_partitions: Int, valid_length: LayoutTensor[DType.uint32, valid_length_layout, ImmutAnyOrigin], sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]], mask: mask_t)

Flash-attention decode GPU kernel with optional split-K partitioning.

Each CTA processes one split-K partition for one (batch, head) pair. Computes online softmax over its key slice and writes partial exp_sum and qk_max statistics alongside the partial output so the mha_splitk_reduce kernel can merge them. When num_partitions == 1 the output is final and no reduction is needed.

Parameters:

q_type (DType): Element type of the query tensor.
k_t (MHAOperand): Key operand type (dense or KV-cache).
v_t (MHAOperand): Value operand type (dense or KV-cache).
output_type (DType): Element type of the output and partial output buffer.
mask_t (MHAMask): Attention mask type.
valid_length_layout (Layout): Layout of the per-sequence valid-length tensor.
BM (Int): Query tile height (rows per CTA).
BN (Int): Key tile width (columns per CTA).
BK (Int): Tile size along the head-depth dimension.
WM (Int): Warp tile height.
WN (Int): Warp tile width.
depth (Int): Attention head depth.
num_heads (Int): Number of query heads.
num_threads (Int): Total threads per CTA.
num_pipeline_stages (Int): Number of software pipeline stages for KV loads.
group (Int): GQA group size (query heads per KV head).
ragged (Bool): True for ragged-batch inputs.
is_shared_kv (Bool): True when K and V share an SMEM buffer.
sink (Bool): True to enable attention-sink mode.
_use_valid_length (Bool): True to read per-sequence valid lengths.
_is_cache_length_accurate (Bool): True when cache length is exact.
decoding_warp_split_k (Bool): Enable warp-level split-K within a CTA.

Args:

q_ptr (Pointer[Scalar[q_type], ImmutAnyOrigin, _safe=False]): Pointer to query data.
k (k_t): Key operand.
v (v_t): Value operand.
output_ptr (Pointer[Scalar[output_type], MutAnyOrigin, _safe=False]): Pointer to the partial/final output buffer.
exp_sum_ptr (Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False]): Pointer to the partial exponential-sum buffer.
qk_max_ptr (Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin, _safe=False]): Pointer to the partial softmax-maximum buffer.
scale (Float32): Softmax temperature scale.
batch_size (Int): Number of sequences in the batch.
num_partitions (Int): Number of split-K partitions.
valid_length (LayoutTensor[DType.uint32, valid_length_layout, ImmutAnyOrigin]): Per-sequence valid lengths (or row offsets for ragged).
sink_weights (OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Sink-token weights for attention-sink mode.
mask (mask_t): Mask instance.