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

flare_mla_decoding_dispatch

def flare_mla_decoding_dispatch[k_t: MHAOperand, mask_t: MHAMask, dtype: DType, //, kv_num_heads: Int, config: MHAConfig[dtype], ragged: Bool = False, _is_cache_length_accurate: Bool = False, _use_valid_length: Bool = True, decoding_warp_split_k: Bool = False, per_token_scale_rope_aware: Bool = False, sparse: Bool = False, rope_aware_kv_sparse: Bool = False, fold_shared_index: Bool = False](output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: k_t, mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], max_prompt_len: Int, max_cache_valid_length: Int, scale: Float32, ctx: DeviceContext, scalar_args_buf: NullableTileTensor[DType.int64, linear_idx_type=scalar_args_buf.linear_idx_type, element_size=scalar_args_buf.element_size], kv_input_row_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, num_partitions: Optional[Int] = None, q_scale_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, indices_stride: Int = Int(0), topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, extra_k: OptionalReg[k_t] = None, extra_d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, extra_indices_stride: Int = Int(0), extra_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, extra_scales_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, num_partitions_in: Optional[Int] = None)

Dispatches an MLA decoding request to the platform-specific kernel.

Routes to the SM100 (mla_decode_sm100_dispatch), AMD, or legacy NVIDIA decode path based on the target GPU, selects the block-M tile geometry (including the AMD MTP token-fold and BM=32/64 heuristic), and launches the mla_decoding kernel with optional split-K reduction via mla_splitk_reduce.

Parameters:

k_t (MHAOperand): KV cache operand type backing k (inferred). Either KVCacheMHAOperand (paged cache) or LayoutTensorMHAOperand (TileTensor input).
mask_t (MHAMask): Mask functor type applied to attention scores (inferred).
dtype (DType): Element type of q (inferred). Must be a half-float or float8_e4m3fn.
kv_num_heads (Int): Number of KV attention heads. Must be 1 for MLA.
config (MHAConfig[dtype]): MLA attention config carrying num_heads and depth (576 for DeepSeek V2/3).
ragged (Bool): Whether q is a ragged rank-3 tensor with batch offsets in valid_length instead of a padded rank-4 tensor (defaults to False).
_is_cache_length_accurate (Bool): Workaround unifying KVCache and TileTensor inputs: True when max_cache_valid_length is the accurate latest length (TileTensor case); False when it precedes the latest tokens, for example zero for continuous execution (KV cache case) (defaults to False).
_use_valid_length (Bool): Whether valid_length is needed for masking; False skips it for TileTensor inputs to avoid benchmark overhead (defaults to True).
decoding_warp_split_k (Bool): Whether to use warp-level split-K reduction within the decode kernel (defaults to False).
per_token_scale_rope_aware (Bool): Whether q and the KV cache use the interleaved FP8+BF16 rope-aware layout (640 bytes/row, 576 logical dims) (defaults to False).
sparse (Bool): Whether to use sparse attention with pre-computed physical KV row indices via gather4 TMA (defaults to False).
rope_aware_kv_sparse (Bool): Sparse-only routing flag: True selects the BF16-rope sparse kernel (split FP8 nope + BF16 rope); False selects the all-FP8 sparse kernel (defaults to False).
fold_shared_index (Bool): Whether to use the read-once shared-index fold that packs folded output/LSE slots into one CTA (defaults to False).

Args:

output (TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Output tensor with shape [batch, num_heads, depth_v] (or ragged equivalent). Dtype is bfloat16 when q is float8_e4m3fn, else matches q.
q (TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query tensor. Padded rank-4 shape [batch, q_seq_len, num_heads, depth] (or rank-3 ragged). For per_token_scale_rope_aware, the last dim is 640 FP8 elements representing 576 logical dims.
k (k_t): KV cache operand. V is derived as K[:, :, :depth_v], so V is not loaded separately.
mask_functor (mask_t): Mask functor instance applied to attention scores.
valid_length (TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type]): Per-batch uint32 tensor of valid (pre-padding) sequence lengths. For ragged inputs, carries cumulative row offsets.
max_prompt_len (Int): Maximum query sequence length (query tokens per batch); the MTP token-fold ceiling.
max_cache_valid_length (Int): Total number of cached KV entries (the cache context length).
scale (Float32): Softmax scale factor applied to QK^T.
ctx (DeviceContext): Device context used to enqueue the kernel.
scalar_args_buf (NullableTileTensor[DType.int64, linear_idx_type=scalar_args_buf.linear_idx_type, element_size=scalar_args_buf.element_size]): Optional GPU buffer of pre-computed scalar dispatch args for capturable graph launches; null for the legacy host-computed path.
kv_input_row_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-batch row offsets into a ragged KV layout; None for non-ragged inputs.
num_partitions (Optional[Int]): Optional explicit split-K partition count; None selects via heuristic.
q_scale_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional per-token Q scale array (float32); folded into the softmax as sigma_Q[token_idx]. None means scale 1.0.
d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional sparse indices: d_indices[batch * indices_stride + token] gives the physical KV row index. None for non-sparse.
indices_stride (Int): Allocation stride (max topk across all batches) for d_indices (defaults to 0).
topk_lengths (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional per-batch array of actual valid sparse index counts; None for non-sparse.
attn_sink_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional attention-sink scale pointer (float32); None to disable.
extra_k (OptionalReg[k_t]): Optional separate always-attend KV cache operand, appended after the topk tokens in the attention loop; None to disable.
extra_d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional sparse indices for extra_k, same layout as d_indices; None for non-sparse extra KV.
extra_indices_stride (Int): Allocation stride for extra_d_indices (defaults to 0).
extra_topk_lengths (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional per-batch valid index counts for extra_k; None for non-sparse extra KV.
extra_scales_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional per-token scale pointer for extra_k; None to disable.
num_partitions_in (Optional[Int]): Optional capturable-graph scalar from the Python resolver; when set, the SM100 dispatcher uses it instead of recomputing num_partitions at grid time.