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_ragged_paged_plan

def generic_flare_mla_prefill_ragged_paged_plan[target: StringSlice[ImmStaticOrigin]](input_row_offsets: LayoutTensor[DType.uint32, element_layout=input_row_offsets.element_layout, layout_int_type=input_row_offsets.layout_int_type, linear_idx_type=input_row_offsets.linear_idx_type, masked=input_row_offsets.masked, alignment=input_row_offsets.alignment], kv_collection: PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], layer_idx: UInt32, buffer_token_size: UInt32, buffer_row_offsets: LayoutTensor[DType.uint32, element_layout=buffer_row_offsets.element_layout, layout_int_type=buffer_row_offsets.layout_int_type, linear_idx_type=buffer_row_offsets.linear_idx_type, masked=buffer_row_offsets.masked, alignment=buffer_row_offsets.alignment], cache_offsets: LayoutTensor[DType.uint32, element_layout=cache_offsets.element_layout, layout_int_type=cache_offsets.layout_int_type, linear_idx_type=cache_offsets.linear_idx_type, masked=cache_offsets.masked, alignment=cache_offsets.alignment], buffer_lengths: LayoutTensor[DType.int32, element_layout=buffer_lengths.element_layout, layout_int_type=buffer_lengths.layout_int_type, linear_idx_type=buffer_lengths.linear_idx_type, masked=buffer_lengths.masked, alignment=buffer_lengths.alignment], context: DeviceContext)

Computes the MLA prefill plan for a ragged paged KV cache.

Populates buffer_row_offsets, cache_offsets, and buffer_lengths from the input row offsets and paged key cache so the subsequent MLA prefill kernel can gather K/V entries into a contiguous buffer.

Parameters:

target (StringSlice[ImmStaticOrigin]): Target device string for kernel dispatch; must be a GPU target.

Args:

input_row_offsets (LayoutTensor[DType.uint32, element_layout=input_row_offsets.element_layout, layout_int_type=input_row_offsets.layout_int_type, linear_idx_type=input_row_offsets.linear_idx_type, masked=input_row_offsets.masked, alignment=input_row_offsets.alignment]): Tensor with shape (batch_size + 1,) denoting the start of each sequence along the ragged sequence dimension.
kv_collection (PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): 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.
buffer_token_size (UInt32): Token capacity of each prefill buffer row.
buffer_row_offsets (LayoutTensor[DType.uint32, element_layout=buffer_row_offsets.element_layout, layout_int_type=buffer_row_offsets.layout_int_type, linear_idx_type=buffer_row_offsets.linear_idx_type, masked=buffer_row_offsets.masked, alignment=buffer_row_offsets.alignment]): Mutable output tensor receiving the start and end position of each K entry in the ragged K/V buffer.
cache_offsets (LayoutTensor[DType.uint32, element_layout=cache_offsets.element_layout, layout_int_type=cache_offsets.layout_int_type, linear_idx_type=cache_offsets.linear_idx_type, masked=cache_offsets.masked, alignment=cache_offsets.alignment]): Mutable output tensor receiving the start position of each K entry in the PagedKVCacheCollection.
buffer_lengths (LayoutTensor[DType.int32, element_layout=buffer_lengths.element_layout, layout_int_type=buffer_lengths.layout_int_type, linear_idx_type=buffer_lengths.linear_idx_type, masked=buffer_lengths.masked, alignment=buffer_lengths.alignment]): Mutable output tensor receiving the number of valid tokens per buffer row.
context (DeviceContext): The call context pointer, passed by the graph compiler.