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)