Skip to main content
Log in

Mojo struct

PagedKVCache

@register_passable(trivial) struct PagedKVCache[type_: DType, kv_params_: KVCacheStaticParams, page_size: Int]

The PagedKVCache is a wrapper around the KVCache blocks for a given layer. It is used to access the KVCache blocks for PagedAttention.

Aliases

  • type = type_:
  • kv_params = kv_params_:
  • blocks_shape = __init__[::Indexer,::Indexer,::Indexer,::Indexer](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_, 4, MutableAnyOrigin, __init__[::Indexer,::Indexer,::Indexer,::Indexer](Dim(-31337), Dim(page_size), Dim(kv_params_.num_heads), Dim(kv_params_.head_size)), _strides_from_shape[::DimList,::Int]()]:

Fields

  • blocks (NDBuffer[type_, 4, MutableAnyOrigin, __init__[::Indexer,::Indexer,::Indexer,::Indexer](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]):

Implemented traits

AnyType, Copyable, ExplicitlyCopyable, KVCacheT, Movable, UnknownDestructibility

Methods

__init__

__init__(blocks: NDBuffer[type_, 4, MutableAnyOrigin, __init__[::Indexer,::Indexer,::Indexer,::Indexer](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]) -> Self

max_tile_size

static max_tile_size() -> Int

Returns the maximum tile size for the KVCache.

cache_lengths_nd

cache_lengths_nd(self) -> NDBuffer[uint32, 1, MutableAnyOrigin]

cache_length

cache_length(self, batch_idx: Int) -> Int

Returns the length of the cache for a given batch index.

load

load[width: Int](self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int) -> SIMD[type_, width]

Loads an element from the given index.

store

store(self, bs: Int, head_idx: Int, tok_idx: Int, head_dim_idx: Int, val: SIMD[type_, size])

Stores an element at the given index.

empty_cache

empty_cache(self) -> Bool

Returns true if the cache_lengths for all requests is 0, false otherwise.

max_prompt_length

max_prompt_length(self) -> SIMD[uint32, 1]

Returns the maximum sequence length across all batches of the current request.

max_context_length

max_context_length(self) -> SIMD[uint32, 1]

Returns the maximum cache length used across all batches of the current request.

block_paged_ptr

block_paged_ptr[tile_size: Int](self, batch_idx: Int, start_tok_idx: Int, head_idx: Int, head_dim_idx: Int = 0) -> UnsafePointer[SIMD[type_, 1]]

Was this page helpful?