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

def flare_mla_decoding[rank: Int, cache_t: KVCacheT, mask_t: MHAMask, dtype: DType, //, config: MHAConfig[dtype], ragged: Bool = False, 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: cache_t, mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], 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], q_max_seq_len: OptionalReg[Int] = None, 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[cache_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)

MLA decoding kernel that would only be called in the optimized compute graph.

The Q input has a shape of [seq_len, num_heads, depth]. The K input has a shape of [seq_len, 1, depth]. The V tensor is derived by reusing K, where V = K[:, :, :depth_v].

Specifically, for DeepSeek V2/3, depth = 576 and depth_v = 512.

When per_token_scale_rope_aware is True, Q and KV cache have an interleaved FP8+BF16 layout: FP8 content (512 bytes) + BF16 rope (128 bytes) = 640 bytes/row. Q's last dimension is 640 (FP8 elements) but represents 576 logical dimensions (512 nope + 64 rope).

This kernel computes attention without needing to load V twice. This kernel only handles decoding requests. In this case q_max_seq_len = 1.

This kernel handles batches with different valid lengths (i.e., before the padding). Such lengths are passed in valid_length argument.

Parameters:

rank (Int): Tensor rank of q and output (inferred). Must be 4 for padded inputs or 3 for ragged inputs.
cache_t (KVCacheT): Paged KV cache type backing k (inferred). Carries the KV layout, dtype, and head geometry.
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.
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).
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 (cache_t): Paged 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.
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.
q_max_seq_len (OptionalReg[Int]): Maximum query sequence length (query tokens per batch). Defaults to the cache's max_prompt_length when None.
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[cache_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.

def flare_mla_decoding[mask_t: MHAMask, dtype: DType, //, config: MHAConfig[dtype], decoding_warp_split_k: 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: TileTensor[Storage=k.Storage, linear_idx_type=k.linear_idx_type], mask_functor: mask_t, 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], num_partitions: Optional[Int] = None)