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

mla_prefill_plan_kernel

def mla_prefill_plan_kernel[BufferRowOffsetsLayoutType: TensorLayout, CacheOffsetsLayoutType: TensorLayout, BufferLengthsLayoutType: TensorLayout, InputRowOffsetsLayoutType: TensorLayout, cache_t: KVCacheT](buffer_row_offsets: TileTensor[DType.uint32, BufferRowOffsetsLayoutType, MutUntrackedOrigin], cache_offsets: TileTensor[DType.uint32, CacheOffsetsLayoutType, MutUntrackedOrigin], buffer_lengths: TileTensor[DType.int32, BufferLengthsLayoutType, MutUntrackedOrigin], input_row_offsets: TileTensor[DType.uint32, InputRowOffsetsLayoutType, ImmUntrackedOrigin], k_cache: cache_t, buffer_token_size: UInt32)

Plans how to process a batch of varying-length sequences through a fixed-size buffer.

For each sequence, computes the per-chunk buffer row offsets, cache offsets, and total buffer lengths needed to divide the cached plus new tokens into fixed-size chunks aligned to the page size, enabling the MLA prefill kernel to process sequences that exceed the buffer in multiple iterations.

Parameters:

BufferRowOffsetsLayoutType (TensorLayout): Compile-time layout of buffer_row_offsets (inferred).
CacheOffsetsLayoutType (TensorLayout): Compile-time layout of cache_offsets (inferred).
BufferLengthsLayoutType (TensorLayout): Compile-time layout of buffer_lengths (inferred).
InputRowOffsetsLayoutType (TensorLayout): Compile-time layout of input_row_offsets (inferred).
cache_t (KVCacheT): Paged KV cache type backing k_cache (inferred). Carries the KV layout, dtype, and page size.

Args:

buffer_row_offsets (TileTensor[DType.uint32, BufferRowOffsetsLayoutType, MutUntrackedOrigin]): Output [MAX_CHUNKS, batch_size + 1] tensor of per-chunk buffer row offsets. buffer_row_offsets[chunk_idx, seq_idx] is the starting row offset within the buffer for sequence seq_idx in chunk chunk_idx.
cache_offsets (TileTensor[DType.uint32, CacheOffsetsLayoutType, MutUntrackedOrigin]): Output [MAX_CHUNKS, batch_size] tensor of per-chunk cache offsets. cache_offsets[chunk_idx, seq_idx] is the starting offset into the paged KV cache for sequence seq_idx in chunk chunk_idx.
buffer_lengths (TileTensor[DType.int32, BufferLengthsLayoutType, MutUntrackedOrigin]): Output [MAX_CHUNKS] tensor of per-chunk total buffer lengths. buffer_lengths[chunk_idx] is the total number of valid rows across all sequences in chunk chunk_idx; -1 marks unused chunks.
input_row_offsets (TileTensor[DType.uint32, InputRowOffsetsLayoutType, ImmUntrackedOrigin]): Input [batch_size + 1] tensor of cumulative row offsets; offsets[b] to offsets[b+1] gives batch b's new token range.
k_cache (cache_t): Paged KV cache providing per-sequence cache lengths.
buffer_token_size (UInt32): Fixed buffer size in tokens; each chunk processes at most this many tokens per sequence.