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

generic_flare_mla_decode_kv_cache_ragged

def generic_flare_mla_decode_kv_cache_ragged[collection_t: KVCollectionT, q_dtype: DType, //, mask_str: StringSlice[ImmStaticOrigin], target: StringSlice[ImmStaticOrigin], local_window_size: Int = Int(-1), per_token_scale_rope_aware: Bool = False, sparse_mla: Bool = False, fold_shared_index: Bool = False](q: TileTensor[q_dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], scalar_args_buf: TileTensor[DType.int64, Storage=scalar_args_buf.Storage, linear_idx_type=scalar_args_buf.linear_idx_type], context: DeviceContext, 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[collection_t.CacheType] = 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 MLA decode attention over a ragged batch against a paged KV cache.

Parameters:

collection_t (KVCollectionT): The KV cache collection type storing the K cache for this layer (inferred).
q_dtype (DType): Data type of the query tensor (inferred).
mask_str (StringSlice[ImmStaticOrigin]): Attention mask name selecting the masking strategy, such as "causal", "null", or "sliding_window_causal".
target (StringSlice[ImmStaticOrigin]): Target device string for kernel dispatch; must be a GPU target.
local_window_size (Int): Sliding-window size in tokens for windowed masks; -1 for masks that ignore it (defaults to -1).
per_token_scale_rope_aware (Bool): Whether q and the KV cache use the interleaved FP8+BF16 rope-aware layout (defaults to False).
sparse_mla (Bool): Whether to use sparse attention with pre-computed physical KV row indices via gather4 TMA (defaults to False).
fold_shared_index (Bool): Whether to use the read-once shared-index MTP fold threaded to flare_mla_decoding (defaults to False).

Args:

q (TileTensor[q_dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query tile tensor with shape (batch_size, num_heads, q_head_size).
input_row_offsets (TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type]): Tile tensor with shape (batch_size + 1,) denoting the start of each Q entry in the batch.
kv_collection (collection_t): The collection storing the KVCache entries for this layer, retrieved via layer_idx.
layer_idx (UInt32): The index of the layer being executed, used to retrieve the KVCache objects from kv_collection.
scale (Float32): The scaling factor in scaled dot-product attention, usually rsqrt(head_size).
output (TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type]): The pre-allocated output tile tensor to write results to.
scalar_args_buf (TileTensor[DType.int64, Storage=scalar_args_buf.Storage, linear_idx_type=scalar_args_buf.linear_idx_type]): Packed MLA dispatch metadata buffer.
context (DeviceContext): The call context pointer, passed by the graph compiler.
q_scale_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Per-token Q scale pointer (float32 array, one per Q token); defaults to null (sigma_Q = 1.0).
d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional device pointer to packed int32 physical KV row indices for sparse decode.
indices_stride (Int): Stride between batch rows in d_indices (e.g. max top-k).
topk_lengths (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional per-batch valid top-k counts.
attn_sink_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional per-batch attention sink weights.
extra_k (OptionalReg[collection_t.CacheType]): Optional second KV cache operand for the extra stream.
extra_d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional extra KV stream sparse indices.
extra_indices_stride (Int): Stride for extra_d_indices.
extra_topk_lengths (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional per-batch lengths for the extra stream.
extra_scales_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional extra stream scales.
num_partitions_in (Optional[Int]): Capturable-graph num_partitions override forwarded from the MoGG op so SM100 grid sizing matches the kernel.