Skip to main content

Mojo struct

GlobalMemoryManager

struct GlobalMemoryManager[dtype: DType, BM: UInt32, BN: UInt32, BK: UInt32, depth: UInt32, num_heads: UInt32, group: UInt32, token_gen: Bool]

Fields

  • q_offset (UInt32):
  • q_runtime_layout (RuntimeLayout[Layout.__init__(IntTuple[__origin_of()](Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), IntTuple[__origin_of()](Int.__init__[UInt32]((num_heads * depth)), 1)) if (token_gen ^ True) else Layout.row_major(Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), element_type=DType.int32, linear_idx_type=DType.int32]):

Implemented traits

AnyType, UnknownDestructibility

Aliases

__del__is_trivial

alias __del__is_trivial = True

kv_gmem_layout

alias kv_gmem_layout = Layout.__init__(IntTuple[__origin_of()](Int.__init__[UInt32](BN), Int.__init__[UInt32](depth)), IntTuple[__origin_of()](Int.__init__[UInt32](((num_heads // group) * depth)), 1))

kv_num_heads

alias kv_num_heads = (num_heads // group)

q_gmem_layout

alias q_gmem_layout = Layout.__init__(IntTuple[__origin_of()](Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), IntTuple[__origin_of()](Int.__init__[UInt32]((num_heads * depth)), 1)) if (token_gen ^ True) else Layout.row_major(Int.__init__[UInt32](BM), Int.__init__[UInt32](depth))

Methods

__init__

__init__(out self, q_tile_idx: UInt32, kv_head_idx: UInt32, seq_len: Int)

get_q_tensor

get_q_tensor[qtype: DType](self, ptr: UnsafePointer[Scalar[qtype]]) -> LayoutTensor[qtype, Layout.__init__(IntTuple[__origin_of()](Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), IntTuple[__origin_of()](Int.__init__[UInt32]((num_heads * depth)), 1)) if (token_gen ^ True) else Layout.row_major(Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), MutableAnyOrigin, layout_int_type=DType.int32, linear_idx_type=DType.int32, masked=True]

Returns:

LayoutTensor

get_output_tensor

get_output_tensor[out_type: DType](self, ptr: UnsafePointer[Scalar[out_type]]) -> LayoutTensor[out_type, Layout.__init__(IntTuple[__origin_of()](Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), IntTuple[__origin_of()](Int.__init__[UInt32]((num_heads * depth)), 1)) if (token_gen ^ True) else Layout.row_major(Int.__init__[UInt32](BM), Int.__init__[UInt32](depth)), MutableAnyOrigin, layout_int_type=DType.int32, linear_idx_type=DType.int32, masked=True]

Returns:

LayoutTensor

get_kv_tensor

get_kv_tensor[kvtype: DType, //](self, ptr: UnsafePointer[Scalar[kvtype], address_space=address_space, alignment=alignment, mut=mut, origin=origin], kv_tile_num_rows: UInt32) -> LayoutTensor[kvtype, Layout.__init__(IntTuple[__origin_of()](Int.__init__[UInt32](BN), Int.__init__[UInt32](depth)), IntTuple[__origin_of()](Int.__init__[UInt32](((num_heads // group) * depth)), 1)), origin, address_space=address_space, masked=True, alignment=alignment]

Returns:

LayoutTensor

Was this page helpful?