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

flash_attention_ragged

def flash_attention_ragged[mask_t: MHAMask, type: DType, q_layout: Layout, //, config: MHAConfig[type] = MHAConfig(SIMD[IntTuple](q_layout.shape[Int((add q_layout.rank(), -2))]), SIMD[IntTuple](q_layout.shape[Int((add q_layout.rank(), -1))]), Optional(None), Optional(None), Optional(None), Optional(None), Optional(None), Int(4), Int(1), FlashAttentionAlgorithm(Int(-1)), TensorMapSwizzle.SWIZZLE_128B), decoding_warp_split_k: Bool = False, naive_kernel: Bool = False](output: LayoutTensor[element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], q: LayoutTensor[type, q_layout, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment], v: LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment], input_row_offsets: LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin], max_prompt_len: LayoutTensor[DType.uint32, element_layout=max_prompt_len.element_layout, layout_int_type=max_prompt_len.layout_int_type, linear_idx_type=max_prompt_len.linear_idx_type, masked=max_prompt_len.masked, alignment=max_prompt_len.alignment], mask_functor: mask_t, scale: Float32, ctx: DeviceContext, num_partitions: Optional[Int] = None)

Run flash attention on ragged (variable-length) batch inputs.

Accepts Q/K/V as flat rank-3 tensors with shape [total_seq_len, num_heads, head_dim] and a CSR-style input_row_offsets tensor of length batch + 1 that encodes per-sequence boundaries. Dispatches the same kernel paths as the dense overload but wraps K/V in RaggedMHAOperand adapters.

Parameters:

mask_t (MHAMask): Attention mask type implementing MHAMask.
type (DType): Element data type for Q/K/V and the output.
q_layout (Layout): Compile-time layout of the query tensor.
config (MHAConfig[type]): Tile/pipeline configuration; defaults from query shape.
decoding_warp_split_k (Bool): Enable warp-level split-K for decode.
naive_kernel (Bool): Force the fallback naive attention kernel.

Args:

output (LayoutTensor[element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment]): Mutable output tensor, same shape as Q.
q (LayoutTensor[type, q_layout, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment]): Query tensor [total_seq_len, num_heads, head_dim].
k (LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment]): Key tensor [total_seq_len, kv_heads, head_dim].
v (LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment]): Value tensor [total_seq_len, kv_heads, head_dim].
input_row_offsets (LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]): CSR row offsets [batch + 1].
max_prompt_len (LayoutTensor[DType.uint32, element_layout=max_prompt_len.element_layout, layout_int_type=max_prompt_len.layout_int_type, linear_idx_type=max_prompt_len.linear_idx_type, masked=max_prompt_len.masked, alignment=max_prompt_len.alignment]): Scalar tensor holding the maximum sequence length.
mask_functor (mask_t): Mask instance.
scale (Float32): Softmax temperature scale.
ctx (DeviceContext): GPU device context.
num_partitions (Optional[Int]): Override split-K partition count; None for auto.