For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
generic_get_paged_cache_with_scales
def generic_get_paged_cache_with_scales[dtype: DType, scale_dtype: DType, kv_params: KVCacheStaticParams, page_size: Int, quantization_granularity: Int](blocks: LayoutTensor[dtype, Layout.row_major[Int(6)]()], cache_lengths: LayoutTensor[DType.uint32, Layout(IntTuple(Int(-1)))], lookup_table: LayoutTensor[DType.uint32, Layout.row_major[Int(2)]()], max_prompt_length: LayoutTensor[DType.uint32, Layout.row_major[Int(1)]()], max_cache_length: LayoutTensor[DType.uint32, Layout.row_major[Int(1)]()], scales: LayoutTensor[scale_dtype, Layout.row_major[Int(6)]()], scales_lookup_table: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major[Int(2)](), lookup_table.origin]] = None, out result: PagedKVCacheCollection[dtype, kv_params, page_size, blocks.origin, cache_lengths.origin, lookup_table.origin, scales.origin, scale_dtype_=scale_dtype, quantization_granularity_=quantization_granularity])
Create a PagedKVCacheCollection with scales for MLA attention.
Args:
- blocks (
LayoutTensor[dtype, Layout.row_major[Int(6)]()]): KV cache blocks tensor [num_blocks, kv_dim, num_layers, page_size, num_heads, head_dim]. - cache_lengths (
LayoutTensor[DType.uint32, Layout(IntTuple(Int(-1)))]): Cache lengths per batch [batch_size]. - lookup_table (
LayoutTensor[DType.uint32, Layout.row_major[Int(2)]()]): Page lookup table [batch_size, max_pages]. - max_prompt_length (
LayoutTensor[DType.uint32, Layout.row_major[Int(1)]()]): Max prompt (query) length scalar tensor [1]. - max_cache_length (
LayoutTensor[DType.uint32, Layout.row_major[Int(1)]()]): Max cache length scalar tensor [1]. - scales (
LayoutTensor[scale_dtype, Layout.row_major[Int(6)]()]): Scales tensor [num_blocks, kv_dim, num_layers, page_size, num_heads, head_dim_granularity]. - scales_lookup_table (
OptionalReg[LayoutTensor[DType.uint32, Layout.row_major[Int(2)](), lookup_table.origin]]): Page lookup table for the scales [batch_size, max_pages]. Pass this when the scales are paged independently of the values, so a request's scale pages carry their own ids. When absent the scales resolve throughlookup_table, which is correct only while the two share one block-id space.
Returns: