Mojo struct
VBuffer
struct VBuffer[v_t: MHAOperand, //, mma_shape: IndexList[3], k_group_size: Int, BN: Int, BK: Int, depth: Int, num_threads: Int, kv_num_heads: Int]
Fields
- mma_tile (
LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]): - smem_iter (
LayoutTensorIter[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]): - kv_cache_iter (
KVCacheIterator[v_t, BN, kv_num_heads, depth]): - buffer_idx (
Int):
Implemented traits
AnyType,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = Int.__del__is_trivial if KVCacheIterator[v_t, BN, kv_num_heads, depth].__del__is_trivial if LayoutTensorIter[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensorIter[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else KVCacheIterator[v_t, BN, kv_num_heads, depth].__del__is_trivial if LayoutTensorIter[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensorIter[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial
base_layout
alias base_layout = Layout.row_major(32, 32)
MMA_K
alias MMA_K = mma_shape.__getitem__[3, DType.int64, Int](2)
MMA_M
alias MMA_M = mma_shape.__getitem__[3, DType.int64, Int](0)
MMATileType
alias MMATileType = LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]
num_k_tiles
alias num_k_tiles = ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))
SharedIterType
alias SharedIterType = LayoutTensorIter[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]
SharedTileType
alias SharedTileType = LayoutTensor[v_t.dtype, blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False), MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_index_type(AddressSpace(3)), linear_idx_type=_get_index_type(AddressSpace(3))]
simd_width
alias simd_width = simd_width_of[v_t.dtype]()
smem_layout
alias smem_layout = blocked_product(Layout.row_major(32, 32), Layout.row_major((BN // 32), (depth // 32)), False)
tiler_layout
alias tiler_layout = Layout.row_major((BN // 32), (depth // 32))
Methods
__init__
__init__(out self, v_cache: v_t, batch_idx: UInt, head_idx: UInt, shared_ptr: UnsafePointer[Scalar[v_t.dtype], address_space=AddressSpace(3), mut=mut, origin=origin], end: UInt)
load_from_dram
get_mma_tile
get_mma_tile(self) -> LayoutTensor[v_t.dtype, Layout.row_major(((depth // mma_shape.__getitem__[3, DType.int64, Int](0)) * ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))), simd_width_of[v_t.dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]
Returns:
copy_to_shared
copy_to_shared(self)
load_from_shared
load_from_shared(self, buffer: UInt, bk_tile: UInt)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!