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

kv_cache_row_offsets_ragged_paged_kernel

def kv_cache_row_offsets_ragged_paged_kernel[CacheRowOffsetsLayoutType: TensorLayout, InputRowOffsetsLayoutType: TensorLayout, CacheLengthsLayoutType: TensorLayout](cache_row_offsets: TileTensor[DType.uint32, CacheRowOffsetsLayoutType, MutUntrackedOrigin], input_row_offsets: TileTensor[DType.uint32, InputRowOffsetsLayoutType, ImmUntrackedOrigin], cache_lengths: TileTensor[DType.uint32, CacheLengthsLayoutType, ImmUntrackedOrigin])

Computes cumulative valid-cache row offsets for one batch index in a ragged prefill batch.

Each thread accumulates the running sum of valid cache lengths plus the ragged sequence deltas for all batches before its output index and writes the result to cache_row_offsets.

Parameters:

CacheRowOffsetsLayoutType (TensorLayout): Memory layout of the cache_row_offsets output tensor.
InputRowOffsetsLayoutType (TensorLayout): Memory layout of the input_row_offsets tensor.
CacheLengthsLayoutType (TensorLayout): Memory layout of the cache_lengths tensor.

Args:

cache_row_offsets (TileTensor[DType.uint32, CacheRowOffsetsLayoutType, MutUntrackedOrigin]): Output tensor receiving the cumulative valid-cache row offsets, with shape (batch_size + 1,).
input_row_offsets (TileTensor[DType.uint32, InputRowOffsetsLayoutType, ImmUntrackedOrigin]): Tensor with shape (batch_size + 1,) denoting the start of each sequence along the ragged sequence dimension.
cache_lengths (TileTensor[DType.uint32, CacheLengthsLayoutType, ImmUntrackedOrigin]): Tensor with shape (batch_size,) giving the number of valid cached tokens per batch.