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

mha_prefill_v2_ragged

def mha_prefill_v2_ragged[k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, qkv_dtype: DType, output_dtype: DType, //, config: MhaConfigV2, cross_attention: Bool = False, sink: Bool = False, compile_options: StringSlice[ImmStaticOrigin] = CompilationTarget.default_compile_options()](q_ptr: Pointer[Scalar[qkv_dtype], ImmutAnyOrigin, _safe=False], k: k_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], kv_input_row_offsets_ptr: Pointer[UInt32, ImmutAnyOrigin, _safe=False], max_prompt_len: Int, batch_size: Int, ctx: DeviceContext, sink_weights_ptr: Pointer[Scalar[qkv_dtype], ImmutAnyOrigin, _safe=False] = Pointer.unsafe_dangling())

Host launcher for ragged MhaPrefillV2 prefill.

Mirrors mha_prefill_v2 for non-ragged, but each block does the per-sequence ragged setup (start_of_seq / q_batch_offset / seq_len / num_keys / start_pos) and constructs its rank-4 BSHD Q/O view internally, so the caller doesn't have to pre-slice per sequence.

cross_attention=False (default): self-attention. K/V length per batch derives from start_pos + (end_of_seq - start_of_seq). kv_input_row_offsets_ptr is unused inside the kernel and the caller may pass any well-typed stub (the dispatcher passes input_row_offsets_ptr itself).

cross_attention=True: encoder-decoder. Pass the kv-side input_row_offsets (uint32 cumulative sum, length batch_size + 1).

Grid: (NUM_HEADS, ceildiv(max_prompt_len, BM), batch_size). Blocks where block_idx.y * BM >= seq_len for this sequence early-return. Partial-Q-tile (seq_len % BM != 0) is handled internally via lane-gated _store_o_to_gmem.

Parameters:

k_t (MHAOperand): K operand type (inferred); any MHAOperand whose block_paged_tile returns (KV_BLOCK, DEPTH) tiles.
v_t (MHAOperand): V operand type (inferred); same tile contract as k_t.
mask_t (MHAMask): Mask functor type (inferred); selects the per-tile masking strategy (causal, sliding-window, null, etc.).
qkv_dtype (DType): Element dtype of Q, K, and V (inferred); must equal config.dtype.
output_dtype (DType): Element dtype of the output buffer (inferred); must equal config.output_dtype.
config (MhaConfigV2): Shape configuration (MhaConfigV2).
cross_attention (Bool): Whether K/V length is independent of Q (encoder-decoder style) (defaults to False).
sink (Bool): Whether to seed the online softmax with attention-sink weights (defaults to False).
compile_options (StringSlice[ImmStaticOrigin]): LLVM compile options string forwarded to ctx.compile_function (defaults to the device's default_compile_options).

Args:

q_ptr (Pointer[Scalar[qkv_dtype], ImmutAnyOrigin, _safe=False]): Pointer to the packed Q buffer, pre-offset per sequence by input_row_offsets_ptr.
k (k_t): K operand (MHAOperand).
v (v_t): V operand (MHAOperand).
output_ptr (Pointer[Scalar[output_dtype], MutAnyOrigin, _safe=False]): Pointer to the output buffer, same packing as q_ptr.
mask_functor (mask_t): Per-tile mask predicate (causal, sliding-window, etc.).
scale (Float32): Softmax scale (1/sqrt(DEPTH)).
input_row_offsets_ptr (Pointer[UInt32, ImmutAnyOrigin, _safe=False]): Cumulative uint32 row offsets for Q sequences, length batch_size + 1.
kv_input_row_offsets_ptr (Pointer[UInt32, ImmutAnyOrigin, _safe=False]): Cumulative uint32 row offsets for K/V sequences; read only when cross_attention is True.
max_prompt_len (Int): Maximum sequence length across the batch; sizes the grid's block_idx.y dimension.
batch_size (Int): Number of sequences in the packed batch; sizes the grid's block_idx.z dimension.
ctx (DeviceContext): Device context used to compile and enqueue the kernel.
sink_weights_ptr (Pointer[Scalar[qkv_dtype], ImmutAnyOrigin, _safe=False]): Per-q-head attention-sink scalar weights; read only when sink is True. Callers may pass unsafe_dangling() when sink=False.