Mojo struct

PagedKVCacheCollection

struct PagedKVCacheCollection[type_: DType, kv_params_: KVCacheStaticParams, page_size: Int, assert_write_mode: Int = 0]

Fields

blocks (NDBuffer[type_, 6, MutableAnyOrigin, DimList(Dim(-31337), Dim(-31337), Dim(-31337), Dim(page_size), Dim(kv_params_.num_heads), Dim(kv_params_.head_size)), _strides_from_shape[::DimList,::Int]()]):
cache_lengths (NDBuffer[uint32, 1, MutableAnyOrigin]):
lookup_table (NDBuffer[uint32, 2, MutableAnyOrigin]):
max_seq_length (SIMD[uint32, 1]):
max_cache_length (SIMD[uint32, 1]):
kv_cache_dynamic_shape (IndexList[4]):
kv_cache_dynamic_strides (IndexList[4]):

Implemented traits

AnyType, Copyable, KVCollectionT, Movable, UnknownDestructibility

Aliases

`blocks_shape`

alias blocks_shape = DimList(Dim(-31337), Dim(-31337), Dim(-31337), Dim(page_size), Dim(kv_params_.num_heads), Dim(kv_params_.head_size))

`blocks_stride`

alias blocks_stride = _strides_from_shape[::DimList,::Int]()

`blocks_type`

alias blocks_type = NDBuffer[type_, 6, MutableAnyOrigin, DimList(Dim(-31337), Dim(-31337), Dim(-31337), Dim(page_size), Dim(kv_params_.num_heads), Dim(kv_params_.head_size)), _strides_from_shape[::DimList,::Int]()]

`CacheType`

alias CacheType = PagedKVCache[type_, kv_params_, page_size, assert_write_mode]

`kv_params`

alias kv_params = kv_params_

`name_str`

alias name_str = "paged"

`type`

alias type = type_

Methods

`init`

__init__(out self, blocks: NDBuffer[type_, 6, MutableAnyOrigin], cache_lengths: NDBuffer[uint32, 1, MutableAnyOrigin], lookup_table: NDBuffer[uint32, 2, MutableAnyOrigin], max_seq_length: SIMD[uint32, 1], max_cache_length: SIMD[uint32, 1])

`copyinit`

__copyinit__(out self, other: Self)

`moveinit`

__moveinit__(out self, owned other: Self)

`copy`

copy(self) -> Self

Explicitly construct a copy of self.

Returns:

A copy of this value.

`get_key_cache`

get_key_cache(self, layer_idx: Int) -> PagedKVCache[type_, kv_params_, page_size, assert_write_mode]

`get_value_cache`

get_value_cache(self, layer_idx: Int) -> PagedKVCache[type_, kv_params_, page_size, assert_write_mode]

`cache_length`

cache_length(self, bs_idx: Int) -> Int

Fields​

Implemented traits​

Aliases​

blocks_shape​

blocks_stride​

blocks_type​

CacheType​

kv_params​

name_str​

type​

Methods​

__init__​

__copyinit__​

__moveinit__​

copy​

get_key_cache​

get_value_cache​

cache_length​