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_prefill_dispatch

def flare_mla_prefill_dispatch[k_t: MHAOperand, v_t: MHAOperand, k_rope_t: MHAOperand, mask_t: MHAMask, dtype: DType, output_type: DType, //, kv_num_heads: Int, config: MHAConfig[dtype], q_depth: Int = Int(192), cache_depth: Int = Int(576), _ndbuffer_mha_operand: Bool = False](output: TileTensor[output_type, 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, v: v_t, k_rope: k_rope_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, scale: Float32, ctx: DeviceContext, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]] = None)

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

Routes to mla_sm100_prefill on SM100 GPUs or enqueues the generic mla_prefill kernel on other NVIDIA and AMD targets, computing the grid and shared-memory layout from the MHA config.

Parameters:

k_t (MHAOperand): Key operand type backing k (inferred). Either RaggedMHAOperand or LayoutTensorMHAOperand.
v_t (MHAOperand): Value operand type backing v (inferred). Either RaggedMHAOperand or LayoutTensorMHAOperand.
k_rope_t (MHAOperand): KV cache operand type backing k_rope (inferred). Either KVCacheMHAOperand or LayoutTensorMHAOperand.
mask_t (MHAMask): Mask functor type applied to attention scores (inferred).
dtype (DType): Element type of q (inferred). Must be bfloat16 or float8_e4m3fn.
output_type (DType): Element type of output (inferred). bfloat16 when q is float8_e4m3fn, else matches q.
kv_num_heads (Int): Number of KV attention heads. Must be 1 for MLA.
config (MHAConfig[dtype]): MLA attention config carrying num_heads, depth, and tile geometry (block_m, block_n, block_k).
q_depth (Int): Q head dimension in elements (defaults to 192). For DeepSeek V2/3, 192 = 128 nope + 64 rope.
cache_depth (Int): KV cache head dimension in elements (defaults to 576). The absorbed-latent width of the paged MLA cache.
_ndbuffer_mha_operand (Bool): Whether the K/V/k_rope operands use ND-buffer layout instead of the default TileTensor layout (defaults to False).

Args:

output (TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Output tensor with ragged rank-3 shape [total_q_tokens, num_heads, v_depth]. 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 with ragged rank-3 shape [total_q_tokens, num_heads, q_depth].
k (k_t): Key operand (nope part). Carries the ragged K tensor and row offsets.
v (v_t): Value operand. Carries the ragged V tensor and row offsets.
k_rope (k_rope_t): KV cache operand providing the rope part of K, with shape [cache_len, 1, q_depth - depth].
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 cumulative row offsets with batch_size + 1 entries; offsets[b] to offsets[b+1] gives batch b's token range.
max_prompt_len (Int): Maximum query sequence length (tokens per batch across all batches); drives the grid's M dimension.
scale (Float32): Softmax scale factor applied to QK^T.
ctx (DeviceContext): Device context used to enqueue the kernel.
cache_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]]): Optional per-batch uint32 tensor of starting offsets into the paged KV cache; None when the cache is contiguous.