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_cross_attention_kv_cache

def generic_cross_attention_kv_cache[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], q_input_row_offsets: LayoutTensor[DType.uint32, element_layout=q_input_row_offsets.element_layout, layout_int_type=q_input_row_offsets.layout_int_type, linear_idx_type=q_input_row_offsets.linear_idx_type, masked=q_input_row_offsets.masked, alignment=q_input_row_offsets.alignment], q_max_seq_len: LayoutTensor[DType.uint32, element_layout=q_max_seq_len.element_layout, layout_int_type=q_max_seq_len.layout_int_type, linear_idx_type=q_max_seq_len.linear_idx_type, masked=q_max_seq_len.masked, alignment=q_max_seq_len.alignment], kv_input_row_offsets: LayoutTensor[DType.uint32, element_layout=kv_input_row_offsets.element_layout, layout_int_type=kv_input_row_offsets.layout_int_type, linear_idx_type=kv_input_row_offsets.linear_idx_type, masked=kv_input_row_offsets.masked, alignment=kv_input_row_offsets.alignment], 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: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Dispatches cross-attention flash attention over a ragged batch against a paged KV cache.

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, such as "cpu" or "gpu".
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(q_seq_lens), num_heads, head_size).
q_input_row_offsets (LayoutTensor[DType.uint32, element_layout=q_input_row_offsets.element_layout, layout_int_type=q_input_row_offsets.layout_int_type, linear_idx_type=q_input_row_offsets.linear_idx_type, masked=q_input_row_offsets.masked, alignment=q_input_row_offsets.alignment]): Tensor with shape (batch_size + 1,) denoting the start of each query sequence along the ragged sequence dimension.
q_max_seq_len (LayoutTensor[DType.uint32, element_layout=q_max_seq_len.element_layout, layout_int_type=q_max_seq_len.layout_int_type, linear_idx_type=q_max_seq_len.linear_idx_type, masked=q_max_seq_len.masked, alignment=q_max_seq_len.alignment]): Scalar tensor holding the maximum query sequence length.
kv_input_row_offsets (LayoutTensor[DType.uint32, element_layout=kv_input_row_offsets.element_layout, layout_int_type=kv_input_row_offsets.layout_int_type, linear_idx_type=kv_input_row_offsets.linear_idx_type, masked=kv_input_row_offsets.masked, alignment=kv_input_row_offsets.alignment]): Tensor with shape (batch_size + 1,) denoting the start of each KV 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(q_seq_lens), num_heads, head_size).
context (DeviceContext): The call context pointer, passed by the graph compiler.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-batch attention sink weights applied to the leading cache slots.