IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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.

Padded pages and page_stride​

page_stride is the distance in elements from one page to the next; -1 means the pages are packed and the product of the dimensions inside a page is already right. Padding lands entirely in that one stride.

The tensor passed in still describes packed pages, so a padded page_stride addresses beyond the extent it declares. The allocator owes total_num_blocks * page_stride elements behind it.

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[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[(kv_params_ * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)], ComptimeInt[Int(1)]]()], 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[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[kv_params_.head_size]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * kv_params_)], ComptimeInt[(kv_params_ * kv_params_)], ComptimeInt[kv_params_.head_size], ComptimeInt[Int(1)]]()]

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[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[kv_params_.head_size]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * kv_params_)], ComptimeInt[(kv_params_ * kv_params_)], ComptimeInt[kv_params_.head_size], ComptimeInt[Int(1)]]()], 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, DefaultEngine, cache_lengths_origin, DefaultEngine, lookup_table_origin, DefaultEngine, 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[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[(kv_params_ * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)], ComptimeInt[Int(1)]]()]

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[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[(kv_params_ * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)], ComptimeInt[Int(1)]]()], 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[.uint32, Layout(IntTuple(Int(-1))), cache_lengths_origin], lookup_table: LayoutTensor[.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[.uint32, Layout.row_major[Int(2)](), lookup_table_origin]] = None, page_stride: Int = Int(-1), scales_page_stride: Int = Int(-1))

Construct from LayoutTensor params (MOGG boundary).

def __init__(out self, blocks: TileTensor[Self.dtype, Layout[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[kv_params_.head_size]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * kv_params_)], ComptimeInt[(kv_params_ * kv_params_)], ComptimeInt[kv_params_.head_size], ComptimeInt[Int(1)]]()], blocks_origin], cache_lengths: TileTensor[.uint32, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], cache_lengths_origin], lookup_table: TileTensor[.uint32, Layout[TypeList[Int64, Int64](), TypeList[Int64, ComptimeInt[Int(1)]]()], lookup_table_origin], max_seq_length: UInt32, max_cache_length: UInt32, scales: OptionalReg[TileTensor[Self.scale_dtype, Layout[TypeList[Int64, ComptimeInt[Int(2) if (xor kv_params_.is_mla, True) else Int(1)], Int64, ComptimeInt[page_size], ComptimeInt[kv_params_.num_heads], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)]](), TypeList[Int64, Int64, ComptimeInt[(Int((mul page_size, kv_params_.num_heads)) * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[(kv_params_ * ceildiv(kv_params_.head_size, quantization_granularity_))], ComptimeInt[ceildiv(kv_params_.head_size, quantization_granularity_)], ComptimeInt[Int(1)]]()], scales_origin]] = None, scales_lookup_table: OptionalReg[TileTensor[.uint32, Layout[TypeList[Int64, Int64](), TypeList[Int64, ComptimeInt[Int(1)]]()], lookup_table_origin]] = None, page_stride: Int = Int(-1), scales_page_stride: Int = Int(-1))

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

cache_length​

def cache_length(self, bs_idx: Int) -> Int

Returns:

Int

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community