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

rope_split_store_paged_ragged

def rope_split_store_paged_ragged[dtype: DType, freq_dtype: DType, q_out_dtype: DType = dtype, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), interleaved: Bool = True](qkv: TileTensor[dtype, Storage=qkv.Storage, address_space=qkv.address_space, linear_idx_type=qkv.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], freqs_cis: TileTensor[freq_dtype, Storage=freqs_cis.Storage, address_space=freqs_cis.address_space, linear_idx_type=freqs_cis.linear_idx_type], kv_collection: PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], layer_idx: UInt32, q_output: TileTensor[q_out_dtype, Storage=q_output.Storage, address_space=q_output.address_space, linear_idx_type=q_output.linear_idx_type], ctx: DeviceContext)

Rope+split+store with paged KV cache collection.

Parameters:

dtype (DType): Element type of the QKV input buffer.
freq_dtype (DType): Element type of the freqs_cis RoPE frequency table.
q_out_dtype (DType): Element type of the roped Q output buffer (defaults to dtype).
target (StringSlice[ImmStaticOrigin]): Compile target backend, for example "gpu" or "cpu" (defaults to "cpu").
interleaved (Bool): Whether RoPE uses the interleaved real/imag layout (defaults to True).

Args:

qkv (TileTensor[dtype, Storage=qkv.Storage, address_space=qkv.address_space, linear_idx_type=qkv.linear_idx_type]): Flat matmul output [total_seq_len, q_dim + k_dim + v_dim].
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]): [batch_size + 1] ragged offsets.
freqs_cis (TileTensor[freq_dtype, Storage=freqs_cis.Storage, address_space=freqs_cis.address_space, linear_idx_type=freqs_cis.linear_idx_type]): [max_seq_len, head_dim] interleaved rope frequencies.
kv_collection (PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_]): Paged KV cache collection to store K and V into.
layer_idx (UInt32): Index of the layer whose K/V caches in kv_collection to store into.
q_output (TileTensor[q_out_dtype, Storage=q_output.Storage, address_space=q_output.address_space, linear_idx_type=q_output.linear_idx_type]): Output buffer for roped Q [total_seq_len, q_dim].
ctx (DeviceContext): DeviceContext for GPU.