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

generic_flash_attention_kv_cache_padded

def generic_flash_attention_kv_cache_padded[collection_t: KVCollectionT, dtype: DType, //, *, target: StringSlice[StaticConstantOrigin], mask_str: StringSlice[StaticConstantOrigin], local_window_size: Int = Int(-1), num_heads: Int = Int(-1)](q: LayoutTensor[dtype, 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], kv_collection: collection_t, layer_idx: UInt32, valid_lengths: LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin], scale: Float32, output: LayoutTensor[dtype, 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], context: DeviceContext, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)