simdgroup_matrix

simdgroup_matrix<T, 8, 8> is Metal's matrix-multiply primitive: an 8×8 matrix held collectively in one simdgroup's registers, with a single-call multiply-accumulate.

CUDA equivalent: warp-level MMA (wmma fragments / mma.sync) as an API. As hardware, the analogy is looser: pre-M5 Apple GPUs have no separate tensor-core unit; simdgroup_matrix ops execute on the regular FP32 pipes, arranged well (the reverse-engineered Rigel paper↗ confirms even Metal 4.1's newer matmul2d tensor API has no dedicated matrix datapath on M4-class hardware). The performance story is therefore different from NVIDIA's: you use simdgroup_matrix for its register layout and issue efficiency, not for a 10× throughput unlock. A well-written simdgroup-matrix GEMM still reaches 13.5 TFLOPS fp32 on an M5 Max, beating MPS.

The API in one breath (load fragments, accumulate, store):

simdgroup_float8x8 A_simd, B_simd;
#pragma clang loop unroll(full)
for (ushort i = 0; i < DIM * 8; i += 8) {
  simdgroup_load(A_simd, A, DIM * 8, ulong2(i, c_pos.y));
  simdgroup_load(B_simd, B, BN,      ulong2(c_pos.x, i));
  simdgroup_multiply_accumulate(acc, A_simd, B_simd, acc);
}

— m5-gemm sync_copy.metal:29-35↗

Rules of use, all matching your wmma instincts:

• Fragments are opaque. You never index into a simdgroup_float8x8; the 64 elements live distributed across the 32 threads in an undocumented layout. simdgroup_load/simdgroup_store move whole tiles between the fragment and threadgroup or device memory.
• They are registers, and they add up fast. A 4×4 grid of accumulators is 1024 floats per simdgroup; doubling that spilled and ran 10× slower in the measured case. The accumulator budget is the tile-size decision.
• Types: simdgroup_float8x8, simdgroup_half8x8, simdgroup_bfloat8x8. F16 fragments halve the register cost; serious kernels take 16-bit inputs and accumulate in fp32.

On M5, this story forks. The M5 generation adds a real per-core matrix unit (the neural accelerators) reached through Metal 4 tensor ops, with Apple reporting multi-x prefill gains. simdgroup_matrix remains the portable primitive: it runs on every M-series chip, and it is still what most shipped kernels use. Read this page as the baseline; the NAX pages cover the fork.

Everything above simdgroup_matrix is composition: MLX's BlockMMA tiles these 8×8 fragments into simdgroup-level register blocks; attention kernels chain two of those through an online softmax. When Q has only one row and no 8×8 tile can be filled, kernels abandon simdgroup_matrix entirely; that's the decode-vs-prefill split.