Skip to main content

/

DocsCode

Mojo struct

PagedKVCacheCollection

struct PagedKVCacheCollection[dtype_: DType, kv_params_: KVCacheStaticParams, page_size: Int, scale_dtype_: DType = DType.invalid, quantization_granularity_: Int = 1]

Fields​

  • ​scales (OptionalReg[PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].scales_type]):
  • ​kv_cache_scales_dynamic_shape (IndexList[4]):
  • ​kv_cache_scales_dynamic_strides (IndexList[4]):
  • ​blocks (PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].blocks_type):
  • ​cache_lengths (PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].cache_lengths_type):
  • ​lookup_table (PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].lookup_table_type):
  • ​max_seq_length (UInt32):
  • ​max_cache_length (UInt32):
  • ​kv_cache_dynamic_shape (IndexList[4]):
  • ​kv_cache_dynamic_strides (IndexList[4]):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, KVCollectionT, Movable

comptime members​

blocks_layout​

comptime blocks_layout = Layout.row_major(PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].blocks_shape)

blocks_shape​

comptime blocks_shape = IntTuple(VariadicList(-1, 2 if not PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params.is_mla.__bool__() else 1, -1, page_size, Int.__init__[UInt](PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params.num_heads), Int.__init__[UInt](PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params.head_size)))

blocks_type​

comptime blocks_type = LayoutTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].dtype, PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].blocks_layout, MutAnyOrigin]

cache_lengths_type​

comptime cache_lengths_type = LayoutTensor[DType.uint32, Layout(IntTuple(-1)), ImmutAnyOrigin]

CacheType​

comptime CacheType = PagedKVCache[PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].dtype, PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params, page_size, PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].scale_dtype, quantization_granularity_]

dtype​

comptime dtype = dtype_

head_dim_granularity​

comptime head_dim_granularity = ceildiv(Int.__init__[UInt](PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params.head_size), PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].CacheType.quantization_granularity)

kv_params​

comptime kv_params = kv_params_

lookup_table_type​

comptime lookup_table_type = LayoutTensor[DType.uint32, Layout.row_major[2](), ImmutAnyOrigin]

name_str​

comptime name_str = "paged"

scale_dtype​

comptime scale_dtype = scale_dtype_

scales_layout​

comptime scales_layout = Layout.row_major(PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].scales_shape)

scales_shape​

comptime scales_shape = IntTuple(VariadicList(-1, 2 if not PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params.is_mla.__bool__() else 1, -1, page_size, Int.__init__[UInt](PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].kv_params.num_heads), PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].head_dim_granularity))

scales_type​

comptime scales_type = LayoutTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].scale_dtype, PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].scales_layout, MutAnyOrigin]

Methods​

__init__​

__init__(out self, blocks: LayoutTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].dtype, Layout.row_major[6](), MutAnyOrigin], cache_lengths: LayoutTensor[DType.uint32, Layout(IntTuple(-1)), ImmutAnyOrigin], lookup_table: LayoutTensor[DType.uint32, Layout.row_major[2](), ImmutAnyOrigin], max_seq_length: UInt32, max_cache_length: UInt32, scales: OptionalReg[LayoutTensor[PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].scale_dtype, Layout.row_major[6](), MutAnyOrigin]] = None)

get_key_cache​

get_key_cache(self, layer_idx: Int) -> PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].CacheType

Returns:

PagedKVCacheCollection

get_value_cache​

get_value_cache(self, layer_idx: Int) -> PagedKVCacheCollection[dtype_, kv_params_, page_size, scale_dtype_, quantization_granularity_].CacheType

Returns:

PagedKVCacheCollection

cache_length​

cache_length(self, bs_idx: Int) -> Int

Returns:

Int

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact