Modern GEMM Architecture

The three-tier tile pyramid — same shape, two ecosystems

State-of-the-art GEMM on both CUDA and Metal converges on the same three-tier tiling structure. Outer tile fits in shared / threadgroup memory. Middle tile fits in registers. Inner tile is the matrix unit's native shape. The differences are spelling and exact dimensions, not concept.

CUDA (Ada Lovelace / Hopper)

┌── 128×128  thread-block tile (CTA) ───┐
│   ┌─ 64×64  warp tile ─┐               │
│   │ 16×8×16 tensor-core │  ·         │  ← registers + tensor cores
│   │ 16×8×16 tensor-core │  ·         │
│   └────────────────────┘  ·         │
└──────────────────────────────────────┘ ← shared memory (double-buffered)

Metal (Apple M-series)

┌── 128×128  threadgroup tile ───┐
│   ┌─ 64×64  SIMD-group tile ─┐  │
│   │ 8×8 matrix unit (FP16)   │··│  ← registers + matrix coprocessor
│   └─────────────────────────┘··│
└────────────────────────────────┘ ← threadgroup memory (double-buffered)

Three takeaways:

Same pyramid, different micro-tile size. NVIDIA Tensor Cores work on 16×8×16 (m×n×k) FP16 tiles per instruction. Apple's matrix unit works on 8×8 FP16. Different vendor IP, identical role.
Double-buffering at the smem layer is mandatory. While the warp computes on tile i, the next-tile load for i+1 happens in parallel. Without it, the matrix unit stalls.
Tile sizes are auto-tuned per shape. 128×128 is typical, but for tall-skinny GEMMs (m=4, n=4096, k=4096 — single-token decode) the library may pick 32×64 with split-K instead.

Code

Tensor Core micro-kernel (CUDA WMMA API) — what each instruction does·cuda

#include <mma.h>
using namespace nvcuda::wmma;

// One tensor-core instruction: 16×16 = 16×16 + 16×16 (FP16 in, FP32 accum)
fragment<matrix_a, 16, 16, 16, half, row_major> a_frag;
fragment<matrix_b, 16, 16, 16, half, col_major> b_frag;
fragment<accumulator, 16, 16, 16, float> c_frag;

fill_fragment(c_frag, 0.0f);

// Inside the kernel's K loop:
load_matrix_sync(a_frag, smemA + offset, lda);
load_matrix_sync(b_frag, smemB + offset, ldb);
mma_sync(c_frag, a_frag, b_frag, c_frag);   // 16*16*16*2 = 8192 FLOPs / instruction

// CUTLASS wraps all of this; you rarely write WMMA by hand.

Apple matrix coprocessor via SIMD-group ops (Metal)·metal

#include <metal_stdlib>
#include <metal_simdgroup_matrix>
using namespace metal;

// Apple's equivalent of Tensor Cores: simdgroup_matrix.
// 8×8 half-precision tile per multiply.
kernel void gemm_simd(
    const device half *A [[buffer(0)]],
    const device half *B [[buffer(1)]],
    device       half *C [[buffer(2)]],
    uint2 tg_id [[threadgroup_position_in_grid]],
    uint  simd_lane [[thread_index_in_simdgroup]])
{
    simdgroup_matrix<half, 8, 8> A_frag, B_frag, C_frag(0);

    // Load 8×8 tiles, multiply-accumulate; library code wraps this.
    simdgroup_load(A_frag, A + offset, lda);
    simdgroup_load(B_frag, B + offset, ldb);
    simdgroup_multiply_accumulate(C_frag, A_frag, B_frag, C_frag);

    simdgroup_store(C_frag, C + offset, ldc);
}

Exercise

Read 50–100 lines of one of these: (a) CUTLASS's cutlass/gemm/kernel/gemm.h, (b) MLX's mlx/backend/metal/kernels/gemm.metal. Don't try to understand every line — just spot the three-tier tile structure (block tile / warp tile / micro-tile) and the double-buffered shared-memory loads. Once you see them, that's the architecture lesson made concrete in shipping code.