IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /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_seq_len: Int = Int(1)](q_ptr: Pointer[Scalar[q_type], ImmutAnyOrigin], k: k_t, v: v_t, output_ptr: Pointer[Scalar[output_type], MutAnyOrigin], exp_sum_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin], qk_max_ptr: Pointer[Scalar[get_accum_type[q_type]()], MutAnyOrigin], scale: Float32, batch_size: Int32, num_partitions: Int32, 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.
  • ​q_seq_len (Int): Query tokens per sequence folded into the MMA M dimension; 1 is plain decode. A property of the TENSOR, not the geometry — the fold may build a taller tile (mha_decode_fold_tile_q_seq_len). The non-ragged arms take every Q/output/split-K stride from it; ragged recovers the true length from input_row_offsets.

Args:

Models

Cookbook

Puzzles

Blog

Community

Contact