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
MMAConfig
struct MMAConfig[InType: DType, OutType: DType, mma_shape: IndexList[Int(3)], transpose_b: Bool = True]
Configures matrix multiply-accumulate parameters for AMD tensor cores.
Computes the register and K-group sizing derived from the MMA shape and input type so that callers can query the adjusted K dimension used per tensor core load for operands A and B.
Parameters
- InType (
DType): Data type of the input operands. - OutType (
DType): Data type of the accumulator output. - mma_shape (
IndexList[Int(3)]): The [M, N, K] shape of a single MMA operation. - transpose_b (
Bool): Whether matrix B is loaded transposed.
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
k_group_size_a
comptime k_group_size_a = (simd_width_of[InType]() // num_matrix_reg[mma_shape[Int(0)], mma_shape[Int(2)]]())
k_group_size_b
comptime k_group_size_b = (simd_width_of[InType]() // num_matrix_reg[mma_shape[Int(1)], mma_shape[Int(2)]]())
mma
comptime mma = TensorCore()
registers_per_thread_a
comptime registers_per_thread_a = num_matrix_reg[mma_shape[Int(0)], mma_shape[Int(2)]]()
registers_per_thread_b
comptime registers_per_thread_b = num_matrix_reg[mma_shape[Int(1)], mma_shape[Int(2)]]()
simd_width
comptime simd_width = simd_width_of[InType]()
Methods
adjusted_mma_k_shape_a
adjusted_mma_k_shape_b
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!