Mojo struct
KVCacheIterator
struct KVCacheIterator[cache_t: MHAOperand, tile_size: Int, kv_num_heads: Int, depth: Int]
Fields
- cache (
cache_t): - end (
Int): - tile_start_row (
Int): - batch_idx (
Int): - kv_head_idx (
Int):
Implemented traits
AnyType,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = Int.__del__is_trivial if Int.__del__is_trivial if Int.__del__is_trivial if Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if Int.__del__is_trivial if Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial else Int.__del__is_trivial if cache_t.__del__is_trivial else cache_t.__del__is_trivial
kv_gmem_layout
alias kv_gmem_layout = Layout.__init__(IntTuple[__origin_of()](Int.__init__[Int](tile_size), Int.__init__[Int](depth)), IntTuple[__origin_of()](Int.__init__[Int]((kv_num_heads * depth)), 1))
Methods
__init__
__init__(out self, cache: cache_t, batch_idx: Int, kv_head_idx: Int, end: Int)
next_unsafe
next_unsafe(mut self) -> LayoutTensor[cache_t.dtype, Layout.__init__(IntTuple[__origin_of()](Int.__init__[Int](tile_size), Int.__init__[Int](depth)), IntTuple[__origin_of()](Int.__init__[Int]((kv_num_heads * depth)), 1)), MutableAnyOrigin, masked=True]
Returns:
increment
increment(mut self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!