Mojo struct
SharedMemoryManager
struct SharedMemoryManager[dtype: DType, BM: Int, BN: Int, BK: Int, depth: Int, num_rowwise_warps: Int, token_gen: Bool]
Fields
- p_smem (
UnsafePointer[SIMD[dtype, 1], address_space=AddressSpace(3)]): - k_v_smem (
UnsafePointer[SIMD[dtype, 1], address_space=AddressSpace(3)]):
Implemented traits
AnyType,
Defaultable,
UnknownDestructibility
Aliases
accum_type
alias accum_type = get_accum_type[::DType,::DType]()
alignment
alias alignment = alignof[::AnyType,__mlir_type.!kgen.target]()
k_v_smem_size
alias k_v_smem_size = (pad[::DType,::Int,::Int]() * BK)
p_smem_size
alias p_smem_size = (BM * BN) if token_gen else 0
simd_width
alias simd_width = simdwidthof[::DType,__mlir_type.!kgen.target]()
Methods
__init__
__init__(out self)
get_kv_ptr
get_kv_ptr[dtype: DType](self) -> UnsafePointer[SIMD[dtype, 1], address_space=AddressSpace(3), alignment=alignof[::AnyType,__mlir_type.!kgen.target]()]
Returns:
get_p_ptr
get_p_ptr(self) -> UnsafePointer[SIMD[dtype, 1], address_space=AddressSpace(3), alignment=alignof[::AnyType,__mlir_type.!kgen.target]()]
Returns:
get_k_iter
get_k_iter(self) -> LayoutTensorIter[dtype, row_major(BN, BK), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]
Returns:
get_v_iter
get_v_iter(self) -> LayoutTensorIter[dtype, row_major(BK, BN), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]
Returns:
get_p_iter
get_p_iter(self) -> LayoutTensorIter[dtype, row_major(BM, BK), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]
Returns:
get_warp_scratch_tensor
get_warp_scratch_tensor(self) -> LayoutTensor[get_accum_type[::DType,::DType](), row_major((num_rowwise_warps * 2), BM), MutableAnyOrigin, address_space=AddressSpace(3)]
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!