IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

MmaOpApple

struct MmaOpApple[out_type: DType, in_type: DType, num_m_mmas: Int, num_n_mmas: Int, *, b_type: DType = in_type, transpose_a: Bool = False, transpose_b: Bool = False]

Apple Silicon simdgroup-level MMA abstraction for tile-based matrix multiply.

Each simdgroup (32 threads) owns one instance and accumulates a (num_m_mmas * 16) x (num_n_mmas * 16) output tile using Apple simdgroup matrix multiply-accumulate instructions. Supports both dense and fused online-im2col A-operand paths.

Parameters​

  • ​out_type (DType): Element type for the output accumulator (float32 or float16).
  • ​in_type (DType): Element type for the A operand.
  • ​num_m_mmas (Int): Number of 16x16 MMA tiles along the M dimension.
  • ​num_n_mmas (Int): Number of 16x16 MMA tiles along the N dimension.
  • ​b_type (DType): Element type for the B operand; defaults to in_type.
  • ​transpose_a (Bool): When True, the A tile is stored in transposed (K, M) layout.
  • ​transpose_b (Bool): When True, the B tile is stored in transposed (N, K) layout.

Fields​

  • ​rb (Int):
  • ​cb (Int):

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

AccumType​

comptime AccumType = Array[SIMD[out_type, SIMDLength(8)], MmaOpApple[out_type, in_type, num_m_mmas, num_n_mmas, b_type=b_type, transpose_a=transpose_a, transpose_b=transpose_b].num_accum]

FRAG_SIZE​

comptime FRAG_SIZE = 8

MMA_K​

comptime MMA_K = 16

MMA_M​

comptime MMA_M = 16

MMA_N​

comptime MMA_N = 16

num_accum​

comptime num_accum = (num_m_mmas * num_n_mmas)

Methods​

__init__​

def __init__(out self)

zero_accum​

static def zero_accum() -> Self.AccumType

Returns:

Self.AccumType

load_fragment​

def load_fragment[dtype: DType, bounded: Bool = False](self, tile: TileTensor[dtype, Storage=tile.Storage, address_space=tile.address_space, linear_idx_type=tile.linear_idx_type], valid_rows: Int = Int(16)) -> SIMD[dtype, SIMDLength(8)]

Loads one 16x16 simdgroup fragment from a TileTensor sub-tile.

The _apple_frag_layout bit-scatter is owned here at the MmaOpApple layer, not via a TileTensor distribute (KB exceptions/apple-mma-fragment-is-not-distribute-expressible). Used for the V operand of a P.V MMA, where A (P) is a register-resident score fragment and so cannot go through mma().

bounded=True zero-fills rows >= valid_rows instead of reading them -- needed for the V load on the last KV sub-tile when num_keys % 16 != 0, where reading an OOB row makes 0 * V_oob == NaN and poisons the fp32 PV accumulator (KB kernels/apple-m5-fa-prefill). Depth columns are always in-bounds, so only rows are predicated.

Parameters:

  • ​dtype (DType): The element dtype of the source tile.
  • ​bounded (Bool): When True, zero-fill rows >= valid_rows.

Args:

Returns:

SIMD[dtype, SIMDLength(8)]: This lane's 8-element fragment.

mma​

def mma[bounded: Bool = False](self, mut accum: Array[SIMD[out_type, SIMDLength(8)], Self.num_accum], a_tile: TileTensor[in_type, Storage=a_tile.Storage, address_space=a_tile.address_space, linear_idx_type=a_tile.linear_idx_type], b_tile: TileTensor[b_type, Storage=b_tile.Storage, address_space=b_tile.address_space, linear_idx_type=b_tile.linear_idx_type], a_valid_rows: Int = (num_m_mmas * Int(16)), b_valid_cols: Int = (num_n_mmas * Int(16)), k_valid: Int = a_tile.LayoutType.static_shape[Int(1)])

Process K elements across all M/N tile positions.

The K depth is inferred from the A tile's column dimension and must be a multiple of 16. For K=16 this is one MMA step; for K=32 this is two steps, etc. The struct iterates K internally.

Tiles may be row-major or col-major. The stride layout is detected from static_stride and the hardware transpose flag is derived via XOR with the transpose parameter: hw_flag = is_col_major XOR transpose_param.

Use mma() (bounded=False) for interior tiles where all memory is in-bounds. Use mmabounded=True for edge tiles -- zero-fills OOB elements. The kernel should check once per simdgroup, not per load.

Parameters:

  • ​bounded (Bool): When True, zero-fill out-of-bounds A/B elements instead of reading them (defaults to False).

Args:

mma_dense_x2​

def mma_dense_x2[bounded: Bool = False](self, mut accum: Array[SIMD[out_type, SIMDLength(8)], Self.num_accum], a_tile: TileTensor[in_type, Storage=a_tile.Storage, address_space=a_tile.address_space, linear_idx_type=a_tile.linear_idx_type], b_tile: TileTensor[b_type, Storage=b_tile.Storage, address_space=b_tile.address_space, linear_idx_type=b_tile.linear_idx_type], a_valid_rows: Int = (num_m_mmas * Int(16)), b_valid_cols: Int = (num_n_mmas * Int(16)), k_valid: Int = Int(32))

BK=32 "double-strip" dense NT MMA: process 32 K per call.

A (SG_M, 32) row-major and B (32, SG_N) col-major are both K contiguous, so one call reads two BK=16 strips' worth of K -- halving the K-loop's iteration count -- via native load[width=4]s (see _load_frag_x2).

Correctness: the 2*cb split permutes K identically on both operands (f0 = K-cols {2cb..2cb+3}, f1 = {2cb+4..2cb+7}), and matmul contraction is K-permutation-invariant, so f0+f1 sum to the full 32-K contraction -- the same trick mma_im2col's width-8 gather uses.

Args:

mma_im2col​

def mma_im2col[input_origin: ImmOrigin, bounded: Bool = True, c_aligned: Bool = False](self, mut accum: Array[SIMD[out_type, SIMDLength(8)], Self.num_accum], input_ptr: Pointer[Scalar[in_type], input_origin, _safe=False], conv: ConvIm2colParams, b_tile: TileTensor[b_type, Storage=b_tile.Storage, address_space=b_tile.address_space, linear_idx_type=b_tile.linear_idx_type], h_base: Array[Int32, (num_m_mmas * Int(2))], w_base: Array[Int32, (num_m_mmas * Int(2))], batch_base: Array[Int32, (num_m_mmas * Int(2))], c0: Int32, r: Int32, s: Int32, k_base: Int32, m_valid: Int32, k_total: Int32, b_valid_cols: Int, k_valid: Int)

Fused-conv MMA step: A from online im2col gather, B as in mma.

The A row anchors (h_base, w_base, batch_base) are prebaked by the loader (K-invariant m -> pixel decomposition, hoisted out of the K-loop); indexed per (mi, half) as ri = mi*2 + half. The strip's K-decomposition (c0, r, s) (= k0base -> tap/channel) is carried by the loader and passed in, so the gather does no ///% on the K axis. Apple M5 has no LDS stage, so the A operand is gathered per fragment rather than staged. Design: KB kernels/apple-conv2d-im2col.

Parameters:

  • ​input_origin (ImmOrigin): Memory origin of the NHWC input pointer (inferred).
  • ​bounded (Bool): When True, zero-fill out-of-bounds A/B elements instead of reading them (defaults to True).
  • ​c_aligned (Bool): When True, assume conv.C is a multiple of 8 so the channel run is contiguous and a single width-8 load suffices (defaults to False).

Args:

store​

def store(self, accum: Array[SIMD[out_type, SIMDLength(8)], Self.num_accum], d_tile: TileTensor[out_type, Storage=d_tile.Storage, address_space=d_tile.address_space, linear_idx_type=d_tile.linear_idx_type])

Store all accumulators to output tile (unconditional).

Caller guarantees all elements are in-bounds.

Args:

store_bounded​

def store_bounded(self, accum: Array[SIMD[out_type, SIMDLength(8)], Self.num_accum], d_tile: TileTensor[out_type, Storage=d_tile.Storage, address_space=d_tile.address_space, linear_idx_type=d_tile.linear_idx_type], valid_rows: Int, valid_cols: Int)

Stores accumulators where (row < valid_rows) and (col < valid_cols).

Assumes row-major d_tile; for col-major, mirror _do_load's swap.

Args:

load_accum​

def load_accum(self, d_tile: TileTensor[out_type, Storage=d_tile.Storage, address_space=d_tile.address_space, linear_idx_type=d_tile.linear_idx_type]) -> Self.AccumType

Inverse of store: seed an accumulator from a previously-written output tile (chained-accumulate split-K, e.g. clamp_v2's second pass, instead of a separate partials buffer + reduce kernel).

Caller guarantees all elements are in-bounds.

Returns:

Self.AccumType

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact