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_decompress_k_cache_ragged_paged

def generic_flare_mla_decompress_k_cache_ragged_paged[target: StringSlice[ImmStaticOrigin], dtype: DType](buffer_row_offsets_1d: LayoutTensor[DType.uint32, element_layout=buffer_row_offsets_1d.element_layout, layout_int_type=buffer_row_offsets_1d.layout_int_type, linear_idx_type=buffer_row_offsets_1d.linear_idx_type, masked=buffer_row_offsets_1d.masked, alignment=buffer_row_offsets_1d.alignment], cache_offsets_1d: LayoutTensor[DType.uint32, element_layout=cache_offsets_1d.element_layout, layout_int_type=cache_offsets_1d.layout_int_type, linear_idx_type=cache_offsets_1d.linear_idx_type, masked=cache_offsets_1d.masked, alignment=cache_offsets_1d.alignment], buffer_length: Int32, weight: LayoutTensor[dtype, element_layout=weight.element_layout, layout_int_type=weight.layout_int_type, linear_idx_type=weight.linear_idx_type, masked=weight.masked, alignment=weight.alignment], kv_collection: PagedKVCacheCollection[scale_dtype_=kv_collection.scale_dtype_, quantization_granularity_=kv_collection.quantization_granularity_], layer_idx: UInt32, k_latent_buffer: LayoutTensor[dtype, element_layout=k_latent_buffer.element_layout, layout_int_type=k_latent_buffer.layout_int_type, linear_idx_type=k_latent_buffer.linear_idx_type, masked=k_latent_buffer.masked, alignment=k_latent_buffer.alignment], k_buffer: LayoutTensor[dtype, element_layout=k_buffer.element_layout, layout_int_type=k_buffer.layout_int_type, linear_idx_type=k_buffer.linear_idx_type, masked=k_buffer.masked, alignment=k_buffer.alignment], context: DeviceContext)

Decompresses MLA latent K cache rows into a full K buffer via a matmul.

Gathers compressed latent K vectors from the paged KV cache into k_latent_buffer using buffer_row_offsets and cache_offsets, then multiplies by the down-projection weight to produce the decompressed K buffer.

Parameters:

target (StringSlice[ImmStaticOrigin]): Target device string for kernel dispatch; must be a GPU target.
dtype (DType): Data type of the weight and K buffers (inferred).

Args:

buffer_row_offsets_1d (LayoutTensor[DType.uint32, element_layout=buffer_row_offsets_1d.element_layout, layout_int_type=buffer_row_offsets_1d.layout_int_type, linear_idx_type=buffer_row_offsets_1d.linear_idx_type, masked=buffer_row_offsets_1d.masked, alignment=buffer_row_offsets_1d.alignment]): Tensor denoting the start and end position of each K entry in the ragged buffer.
cache_offsets_1d (LayoutTensor[DType.uint32, element_layout=cache_offsets_1d.element_layout, layout_int_type=cache_offsets_1d.layout_int_type, linear_idx_type=cache_offsets_1d.linear_idx_type, masked=cache_offsets_1d.masked, alignment=cache_offsets_1d.alignment]): Tensor denoting the start position of each K entry in the PagedKVCacheCollection.
buffer_length (Int32): Number of K rows to decompress.
weight (LayoutTensor[dtype, element_layout=weight.element_layout, layout_int_type=weight.layout_int_type, linear_idx_type=weight.linear_idx_type, masked=weight.masked, alignment=weight.alignment]): Down-projection weight tensor applied to the latent K vectors.
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.
k_latent_buffer (LayoutTensor[dtype, element_layout=k_latent_buffer.element_layout, layout_int_type=k_latent_buffer.layout_int_type, linear_idx_type=k_latent_buffer.linear_idx_type, masked=k_latent_buffer.masked, alignment=k_latent_buffer.alignment]): Mutable output buffer receiving the gathered latent K vectors.
k_buffer (LayoutTensor[dtype, element_layout=k_buffer.element_layout, layout_int_type=k_buffer.layout_int_type, linear_idx_type=k_buffer.linear_idx_type, masked=k_buffer.masked, alignment=k_buffer.alignment]): Mutable output buffer receiving the decompressed K vectors.
context (DeviceContext): The call context pointer, passed by the graph compiler.