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_flare_mla_prefill_kv_cache_ragged

def generic_flare_mla_prefill_kv_cache_ragged[collection_t: KVCollectionT, input_dtype: DType, dtype: DType, //, mask_str: StringSlice[ImmStaticOrigin], target: StringSlice[ImmStaticOrigin], local_window_size: Int = Int(-1)](q: TileTensor[input_dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: TileTensor[input_dtype, Storage=k.Storage, linear_idx_type=k.linear_idx_type], v: TileTensor[input_dtype, Storage=v.Storage, linear_idx_type=v.linear_idx_type], buffer_row_offsets: TileTensor[DType.uint32, Storage=buffer_row_offsets.Storage, linear_idx_type=buffer_row_offsets.linear_idx_type], cache_offsets: TileTensor[DType.uint32, Storage=cache_offsets.Storage, linear_idx_type=cache_offsets.linear_idx_type], input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, output: TileTensor[dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type], context: DeviceContext)

Dispatches MLA prefill attention over a ragged batch against a paged KV cache.

Parameters:

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

Args:

q (TileTensor[input_dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query tile tensor with shape (total_seq_len, num_heads, q_head_size).
k (TileTensor[input_dtype, Storage=k.Storage, linear_idx_type=k.linear_idx_type]): Key tile tensor with shape (total_seq_len, num_heads, kv_head_size).
v (TileTensor[input_dtype, Storage=v.Storage, linear_idx_type=v.linear_idx_type]): Value tile tensor with shape (total_seq_len, num_heads, kv_head_size).
buffer_row_offsets (TileTensor[DType.uint32, Storage=buffer_row_offsets.Storage, linear_idx_type=buffer_row_offsets.linear_idx_type]): Tile tensor denoting the start and end position of each K entry in the ragged K/V tensor.
cache_offsets (TileTensor[DType.uint32, Storage=cache_offsets.Storage, linear_idx_type=cache_offsets.linear_idx_type]): Mutable tile tensor denoting the start position of each K entry in the PagedKVCacheCollection.
input_row_offsets (TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type]): Tile tensor denoting the start and end position of each Q entry in the batch.
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 (TileTensor[dtype, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): The pre-allocated output tile tensor to write results to, with shape (total_seq_len, num_heads, kv_head_size).
context (DeviceContext): The call context pointer, passed by the graph compiler.