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 struct
PagedKVCacheCollection
struct PagedKVCacheCollection[dtype_: DType, kv_params_: KVCacheStaticParams, page_size: Int, blocks_origin: MutOrigin, cache_lengths_origin: ImmOrigin, lookup_table_origin: ImmOrigin, scales_origin: MutOrigin, *, scale_dtype_: Optional[DType] = None, quantization_granularity_: Int = Int(1)]
Paged pair of key and value caches backed by a block-allocated tensor.
Stores both the K and V caches in a single 6D block tensor of shape
[total_num_blocks, 2, num_layers, page_size, num_heads, head_size]
(the 2 collapses to 1 under Multi-head Latent Attention), along with
per-request cache lengths and a lookup table mapping logical batches to
physical blocks. Supports optional quantization scales stored in a parallel
tensor with head_dim_granularity as the inner dimension.
Fields
- scales (
OptionalReg[TileTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].scale_dtype, Layout[*(), *()], scales_origin]]): - kv_cache_scales_dynamic_shape (
Coord[Int64, Int64, Int64, Int64]): - kv_cache_scales_dynamic_strides (
Coord[Int64, Int64, Int64, Int64]): - blocks (
PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].blocks_tt_type): - cache_lengths (
PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].CacheType.cache_lengths_tt_type): - lookup_table (
PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].CacheType.lookup_table_tt_type): - scales_lookup_table (
PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].CacheType.lookup_table_tt_type): - max_seq_length (
UInt32): - max_cache_length (
UInt32): - kv_cache_dynamic_shape (
Coord[Int64, Int64, Int64, Int64]): - kv_cache_dynamic_strides (
Coord[Int64, Int64, Int64, Int64]):
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
KVCollectionT,
Movable
comptime members
blocks_layout
comptime blocks_layout = Layout.row_major(PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].blocks_shape)
blocks_shape
comptime blocks_shape = IntTuple(Int(-1), Int(2) if (xor kv_params_.is_mla, True) else Int(1), Int(-1), page_size, kv_params_, kv_params_)
blocks_tt_layout
comptime blocks_tt_layout = Layout[*(), *()]
blocks_tt_type
comptime blocks_tt_type = TileTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].dtype, Layout[*(), *()], blocks_origin]
CacheType
comptime CacheType = PagedKVCache[PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].dtype, PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].kv_params, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_]
dtype
comptime dtype = dtype_
head_dim_granularity
comptime head_dim_granularity = ceildiv(kv_params_.head_size, PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].CacheType.quantization_granularity)
kv_params
comptime kv_params = kv_params_
name_str
comptime name_str = "paged"
scale_dtype
comptime scale_dtype = scale_dtype_.or_else(dtype_)
scales_layout
comptime scales_layout = Layout.row_major(PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].scales_shape)
scales_shape
comptime scales_shape = IntTuple(Int(-1), Int(2) if (xor kv_params_.is_mla, True) else Int(1), Int(-1), page_size, kv_params_, ceildiv(kv_params_.head_size, quantization_granularity_))
scales_tt_layout
comptime scales_tt_layout = Layout[*(), *()]
scales_tt_type
comptime scales_tt_type = TileTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].scale_dtype, Layout[*(), *()], scales_origin]
Methods
__init__
def __init__[scales_dtype: DType = PagedKVCacheCollection[dtype_, kv_params_, page_size, blocks_origin, cache_lengths_origin, lookup_table_origin, scales_origin, scale_dtype_=scale_dtype_, quantization_granularity_=quantization_granularity_].scale_dtype](out self, blocks: LayoutTensor[Self.dtype, Layout.row_major[Int(6)](), blocks_origin], cache_lengths: LayoutTensor[DType.uint32, Layout(IntTuple(Int(-1))), cache_lengths_origin], lookup_table: LayoutTensor[DType.uint32, Layout.row_major[Int(2)](), lookup_table_origin], max_seq_length: UInt32, max_cache_length: UInt32, scales: OptionalReg[LayoutTensor[scales_dtype, Layout.row_major[Int(6)](), scales_origin]] = OptionalReg(), scales_lookup_table: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major[Int(2)](), lookup_table_origin]] = None)
Construct from LayoutTensor params (MOGG boundary).
def __init__(out self, blocks: TileTensor[Self.dtype, Layout[*(), *()], blocks_origin], cache_lengths: TileTensor[DType.uint32, Layout[*(), *()], cache_lengths_origin], lookup_table: TileTensor[DType.uint32, Layout[*(), *()], lookup_table_origin], max_seq_length: UInt32, max_cache_length: UInt32, scales: OptionalReg[TileTensor[Self.scale_dtype, Layout[*(), *()], scales_origin]] = None, scales_lookup_table: OptionalReg[TileTensor[DType.uint32, Layout[*(), *()], lookup_table_origin]] = None)
Construct from TileTensor fields directly.
get_key_cache
def get_key_cache(self, layer_idx: Int) -> Self.CacheType
Returns:
Self.CacheType
get_value_cache
def get_value_cache(self, layer_idx: Int) -> Self.CacheType
Returns:
Self.CacheType