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_ragged_sink

def generic_flash_attention_kv_cache_ragged_sink[collection_t: KVCollectionT, dtype: DType, //, *, target: StringSlice[ImmStaticOrigin], mask_str: StringSlice[ImmStaticOrigin], local_window_size: Int = Int(-1), output_dtype: DType = dtype](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], input_row_offsets: LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, output: LayoutTensor[output_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: LayoutTensor[dtype, element_layout=sink_weights.element_layout, layout_int_type=sink_weights.layout_int_type, linear_idx_type=sink_weights.linear_idx_type, masked=sink_weights.masked, alignment=sink_weights.alignment], decode_dispatch_metadata: MHADecodeDispatchMetadata)

Dispatches flash attention over a ragged batch with attention sink weights.

Parameters:

collection_t (KVCollectionT): The KV cache collection type storing the K and V caches for this layer (inferred).
dtype (DType): Data type of the query tensor (inferred).
target (StringSlice[ImmStaticOrigin]): Target device string for kernel dispatch.
mask_str (StringSlice[ImmStaticOrigin]): Attention mask name selecting the masking strategy, such as "causal", "null", or "sliding_window_causal".
local_window_size (Int): Sliding-window size in tokens for windowed masks; -1 for masks that ignore it (defaults to -1).
output_dtype (DType): Data type of the output tensor (defaults to dtype).

Args:

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]): Query tensor with shape (sum(seq_lens), num_heads, head_size).
input_row_offsets (LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]): Tensor with shape (batch_size + 1,) denoting the start of each sequence along the ragged sequence dimension.
kv_collection (collection_t): The collection storing the KVCache entries for this layer, retrieved via layer_idx.
layer_idx (UInt32): The index of the layer being executed, used to retrieve the KVCache objects from kv_collection.
scale (Float32): The scaling factor in scaled dot-product attention, usually rsqrt(head_size).
output (LayoutTensor[output_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]): The pre-allocated output buffer to write results to, with shape (sum(seq_lens), num_heads, head_size).
context (DeviceContext): The call context pointer, passed by the graph compiler.
sink_weights (LayoutTensor[dtype, element_layout=sink_weights.element_layout, layout_int_type=sink_weights.layout_int_type, linear_idx_type=sink_weights.linear_idx_type, masked=sink_weights.masked, alignment=sink_weights.alignment]): Per-batch attention sink weights applied to the leading cache slots.
decode_dispatch_metadata (MHADecodeDispatchMetadata): Precomputed dispatch metadata used to select decode kernels for the GPU target.