For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

RegisterAccumulatorLayout

struct RegisterAccumulatorLayout[MMA_M: Int, MMA_N: Int, num_m_mmas: Int, num_n_mmas: Int, consumer_group_size: Int, *, frag_simdwidth: Int = Int(2)]

Describes how UMMA accumulator fragments are distributed across the threads of a consumer warp group.

Parameters

MMA_M (Int): The M dimension of a single UMMA instruction tile.
MMA_N (Int): The N dimension of a single UMMA instruction tile.
num_m_mmas (Int): Number of UMMA tiles along the M dimension.
num_n_mmas (Int): Number of UMMA tiles along the N dimension.
consumer_group_size (Int): Number of threads in the consumer warp group.
frag_simdwidth (Int): SIMD width of each accumulator fragment (defaults to 2).

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

`comptime` members

`element_layout`

comptime element_layout = Layout.row_major(Int(1), frag_simdwidth)

`frag_size`

comptime frag_size = ((MMA_M * MMA_N) // consumer_group_size)

`num_row_blocks_per_mma`

comptime num_row_blocks_per_mma = 2

`rows_of_frags_layout`

comptime rows_of_frags_layout = Layout.row_major((num_m_mmas * num_n_mmas), (Int((mul MMA_M, MMA_N)) // consumer_group_size))

`vec_output_layout`

comptime vec_output_layout = Layout(IntTuple(IntTuple(Int(2), num_m_mmas), IntTuple(((Int((mul MMA_M, MMA_N)) // consumer_group_size) // Int((mul frag_simdwidth, 2))), num_n_mmas), __list_literal__=NoneType(None)), IntTuple(IntTuple(frag_simdwidth, (Int((mul MMA_M, MMA_N)) // consumer_group_size)), IntTuple((Int(2) * frag_simdwidth), Int((mul (Int((mul MMA_M, MMA_N)) // consumer_group_size), num_m_mmas))), __list_literal__=NoneType(None)))

Methods

`description`

static def description() -> RegisterAccumulatorDescription

Returns:

RegisterAccumulatorDescription

Parameters​

Implemented traits​

comptime members​

element_layout​

frag_size​

num_row_blocks_per_mma​

rows_of_frags_layout​

vec_output_layout​

Methods​

description​