Mojo struct
GlobalMemoryManager
struct GlobalMemoryManager[dtype: DType, BM: SIMD[uint32, 1], BN: SIMD[uint32, 1], BK: SIMD[uint32, 1], depth: SIMD[uint32, 1], num_heads: SIMD[uint32, 1], group: SIMD[uint32, 1], token_gen: Bool]
Fields
- q_offset (
SIMD[uint32, 1]): - q_runtime_layout (
RuntimeLayout[Layout(IntTuple(Int(BM), Int(depth)), IntTuple(Int((num_heads * depth)), 1)) if (token_gen ^ True) else row_major(Int(BM), Int(depth)), element_type=int32, linear_idx_type=int32]):
Implemented traits
AnyType,
UnknownDestructibility
Methods
__init__
__init__(out self, q_tile_idx: SIMD[uint32, 1], kv_head_idx: SIMD[uint32, 1], seq_len: Int)
get_q_tensor
get_q_tensor[qtype: DType](self, ptr: UnsafePointer[SIMD[qtype, 1]]) -> LayoutTensor[qtype, Layout(IntTuple(Int(BM), Int(depth)), IntTuple(Int((num_heads * depth)), 1)) if (token_gen ^ True) else row_major(Int(BM), Int(depth)), MutableAnyOrigin, layout_int_type=int32, linear_idx_type=int32, masked=True]
get_output_tensor
get_output_tensor[out_type: DType](self, ptr: UnsafePointer[SIMD[out_type, 1]]) -> LayoutTensor[out_type, Layout(IntTuple(Int(BM), Int(depth)), IntTuple(Int((num_heads * depth)), 1)) if (token_gen ^ True) else row_major(Int(BM), Int(depth)), MutableAnyOrigin, layout_int_type=int32, linear_idx_type=int32, masked=True]
get_kv_tensor
get_kv_tensor[kvtype: DType, //](self, ptr: UnsafePointer[SIMD[kvtype, 1], address_space=address_space, alignment=alignment, mut=mut, origin=origin], kv_tile_num_rows: SIMD[uint32, 1]) -> LayoutTensor[kvtype, Layout(IntTuple(Int(BN), Int(depth)), IntTuple(Int(((num_heads // group) * depth)), 1)), origin, address_space=address_space, masked=True, alignment=alignment]
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!