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

mla_prefill_v2_ragged

def mla_prefill_v2_ragged[k_nope_t: MHAOperand, k_rope_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, qkv_dtype: DType, output_dtype: DType, //, config: MlaConfigV2, compile_options: StringSlice[ImmStaticOrigin] = CompilationTarget.default_compile_options()](q_ptr: Pointer[Scalar[qkv_dtype], ImmutAnyOrigin, _safe=False], k_nope: k_nope_t, k_rope: k_rope_t, v: v_t, output_ptr: Pointer[Scalar[output_dtype], MutAnyOrigin, _safe=False], mask_functor: mask_t, scale: Float32, input_row_offsets_ptr: Pointer[UInt32, ImmutAnyOrigin, _safe=False], max_prompt_len: Int, batch_size: Int, ctx: DeviceContext)

Host launcher for ragged MLA prefill.

Standard ragged-prefill signature/boilerplate (grid/block derivation, three operands, mask functor, scale, input row offsets). Grid: (NUM_HEADS, ceildiv(max_prompt_len, BM), batch_size).

Parameters:

k_nope_t (MHAOperand): Compile-time operand type for the K nope segment (inferred).
k_rope_t (MHAOperand): Compile-time operand type for the K rope segment (inferred).
v_t (MHAOperand): Compile-time operand type for the V segment (inferred).
mask_t (MHAMask): Compile-time mask predicate type (causal, null, etc.) (inferred).
qkv_dtype (DType): Element type of the Q, K, and V tensors; must equal config.dtype (inferred).
output_dtype (DType): Element type of the output tensor; must equal config.output_dtype (inferred).
config (MlaConfigV2): Shape configuration (MlaConfigV2).
compile_options (StringSlice[ImmStaticOrigin]): Compilation options for the kernel (defaults to the device's default compile options).

Args:

q_ptr (Pointer[Scalar[qkv_dtype], ImmutAnyOrigin, _safe=False]): Pointer to the ragged-batch Q tensor data.
k_nope (k_nope_t): K nope-segment operand; supplies cache_length for start_pos.
k_rope (k_rope_t): K rope-segment operand (unused by the kernel; kept for contract parity).
v (v_t): V operand (the nope segment of the latent cache).
output_ptr (Pointer[Scalar[output_dtype], MutAnyOrigin, _safe=False]): Pointer to the ragged-batch output tensor data.
mask_functor (mask_t): Per-tile mask predicate (causal / null / ...).
scale (Float32): Softmax scale (1/sqrt(d_qk)).
input_row_offsets_ptr (Pointer[UInt32, ImmutAnyOrigin, _safe=False]): Pointer to the per-sequence row offset prefix-sum [batch_size+1] uint32 array.
max_prompt_len (Int): Maximum prompt length in the batch; bounds the grid's BM-tile dimension.
batch_size (Int): Number of sequences in the ragged batch.
ctx (DeviceContext): Device context used to compile and enqueue the kernel.