Mojo struct
PagedKVCacheCollection
struct PagedKVCacheCollection[type_: DType, kv_params_: KVCacheStaticParams, page_size: Int]
Aliases
type = type_:kv_params = kv_params_:CacheType = PagedKVCache[type_, kv_params_, page_size]:blocks_shape = DimList(Dim(-31337), Dim(-31337), Dim(-31337), Dim(page_size), Dim(kv_params_.num_heads), Dim(kv_params_.head_size)):blocks_stride = _strides_from_shape[::DimList,::Int]():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]()]:
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 (
Index[4]): - kv_cache_dynamic_strides (
Index[4]):
Implemented traits
AnyType,
Copyable,
KVCollectionT,
Movable,
UnknownDestructibility
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]
get_value_cache
get_value_cache(self, layer_idx: Int) -> PagedKVCache[type_, kv_params_, page_size]
cache_length
cache_length(self, bs_idx: Int) -> Int
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!