Mojo struct
KVCacheIterator
struct KVCacheIterator[cache_t: MHAOperand, tile_size: Int, kv_num_heads: Int, depth: Int]
TileTensor-based DRAM tile iterator.
Returns a TileTensor with RuntimeInt for the row dimension (valid-row count) and ComptimeInt for depth and strides. No RuntimeLayout storage.
Fields
- cache (
cache_t): - end (
Int): - tile_start_row (
Int): - batch_idx (
Int): - kv_head_idx (
Int):
Implemented traits
AnyType,
ImplicitlyDestructible
comptime members
GmemTileLayout
comptime GmemTileLayout = Layout[RuntimeInt[DType.int64], ComptimeInt[depth], ComptimeInt[(kv_num_heads * depth)], ComptimeInt[1]]
GmemTileType
comptime GmemTileType = TileTensor[cache_t.dtype, KVCacheIterator[cache_t, tile_size, kv_num_heads, depth].GmemTileLayout, ImmutAnyOrigin]
Methods
__init__
__init__(out self, cache: cache_t, batch_idx: Int, kv_head_idx: Int, end: Int)
next_tile
next_tile(mut self) -> KVCacheIterator[cache_t, tile_size, kv_num_heads, depth].GmemTileType
Returns a TileTensor for the next DRAM tile.
Returns:
KVCacheIterator
increment
increment(mut self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!