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

def mla_prefill_plan[cache_t: KVCacheT](buffer_row_offsets: TileTensor[DType.uint32, Storage=buffer_row_offsets.Storage, address_space=buffer_row_offsets.address_space, linear_idx_type=buffer_row_offsets.linear_idx_type], cache_offsets: TileTensor[DType.uint32, Storage=cache_offsets.Storage, address_space=cache_offsets.address_space, linear_idx_type=cache_offsets.linear_idx_type], buffer_lengths: TileTensor[DType.int32, Storage=buffer_lengths.Storage, address_space=buffer_lengths.address_space, linear_idx_type=buffer_lengths.linear_idx_type], input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, address_space=input_row_offsets.address_space, linear_idx_type=input_row_offsets.linear_idx_type], k_cache: cache_t, buffer_token_size: UInt32, ctx: DeviceContext)

This calls a GPU kernel that plans how to process a batch of sequences with varying lengths using a fixed-size buffer.

Each sequence in the batch has some existing cached tokens and new input tokens. The kernel divides the total tokens into chunks of buffer_token_size.

For each chunk (iteration), it calculates: 1. Buffer offsets for each sequence in each chunk 2. Cache offsets for each sequence in each chunk 3. Total buffer lengths for each processing iteration

Parameters:

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, Storage=buffer_row_offsets.Storage, address_space=buffer_row_offsets.address_space, linear_idx_type=buffer_row_offsets.linear_idx_type]): 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, Storage=cache_offsets.Storage, address_space=cache_offsets.address_space, linear_idx_type=cache_offsets.linear_idx_type]): 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, Storage=buffer_lengths.Storage, address_space=buffer_lengths.address_space, linear_idx_type=buffer_lengths.linear_idx_type]): 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, Storage=input_row_offsets.Storage, address_space=input_row_offsets.address_space, linear_idx_type=input_row_offsets.linear_idx_type]): 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.
ctx (DeviceContext): Device context used to enqueue the kernel.