Optimizing GEMM on Apple's M5 from scratch, Part 2: Outperforming MPS's GEMM (in big 2026)

Recap of Part 1

This is the second part of my worklog on writing an FP32 GEMM kernel for Apple’s M5 GPU. Part 1 has introduced us to the hardware, the roofline model, and the GPU-counter method, and built kernels 1 to 4. We still use the same setup: row-major, non-transposed FP32 matrices with $M=N=K=4096$, on a MacBook Pro with an M5 (10-core GPU).

Some numbers from Part 1 that we need to remember:

• The FP32 compute ceiling is 4.04 TFLOP/s, the measured memory bandwidth is 118 GB/s, and the ridge point is about 34 FLOP/byte.
• Measured intensity uses the bytes the profiler sees leave the GPU: GPU Read Bandwidth × GPU time (method).
• MPS reaches 3365 GFLOP/s and MLX 3332 GFLOP/s on this problem.
• Kernel 4, 1D thread coarsening, reaches 1222 GFLOP/s, 36% of MPS. Its Instruction Throughput Limiter stays at 78%, and every multiply-add is still a separate scalar instruction.

Kernel 5: SIMD-Group Matrix Multiplication

Thread coarsening let one thread reuse an input across several outputs but we still compute those outputs with “raw” scalar multiply-adds though. To increase arithmetic, Metal also gives us a matrix multiply-accumulate primitve that the threads of a SIMD group execute together.

This time, each SIMD group will compute an $8 \times 8$ tile of $C$Show sidenoteSo we can think each thread will compute $4$ elements but it’s not true, we still don’t know how Apple did the dark magic with simdgroup primitives.. At each step along $K$, the group loads an $8\times8$ tile from $A$ and another from $B$, multiplies them, and adds the product to its output tile. We still walk along $K$; each step now adds a small matrix product.

One SIMD Group per Output Tile

Our earlier kernels already ran in SIMD groups. Each thread held its own accumulator and computed its own outputs. With simdgroup_float8x8, the 32 threads hold one logical matrix together. The loads, multiply-accumulates, and stores all operate on that matrix cooperatively [1].

Let’s reuse our example on $C[35, 93]$ to see how what “cooperatively” means. We first divide the output into $8\times8$ tiles, so it belongs to the tile starting at row 32 and column 88. That tile covers rows 32 through 39 and columns 88 through 95. Its SIMD group holds an accumulator for all 64 outputs, and $C[35,93]$ is entry $(3,5)$ within it.

At k = 0, the group loads rows 32 through 39 and columns 0 through 7 of $A$, together with rows 0 through 7 and columns 88 through 95 of $B$. Their matrix product adds eight terms to each output. For our chosen element, that update is:

At k = 8, it adds the next eight terms, and so on. After $4096/8 = 512$ steps, the accumulator contains the full output tile.

Which thread holds entry $(3,5)$? Metal leaves the mapping from matrix elements to threads unspecified. We can identify the entry in the logical matrix, but the API does not give us a portable lane number for it. We work with the tile through the collective operations insteadShow sidenoteThe Metal Shading Language specification describes the matrix types in section 2.4 and their operations in section 6.8 [1]. All threads in the SIMD group must execute these operations under uniform control flow: we cannot put a matrix load or multiply inside a branch that only some of the threads take..

Implementation

We only use this kernel when $M, N$ and $K$ are multiples of $8$.

1
#include <metal_stdlib>
2
#include <metal_simdgroup_matrix>
3

4
using namespace metal;
5

6
kernel void simdgroup_8x8(device const float* a [[buffer(0)]],
7
                          device const float* b [[buffer(1)]],
8
                          device float* c [[buffer(2)]],
9
                          uint M, uint N, uint K,
10
                          uint2 group [[threadgroup_position_in_grid]]) {
11
  const uint row = group.y * 8;
12
  const uint col = group.x * 8;
13
  simdgroup_float8x8 acc = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
14
  for (uint k = 0; k < K; k += 8) {
15
    simdgroup_float8x8 ta, tb, next;
16
    simdgroup_load(ta, a + row * K + k, K);
17
    simdgroup_load(tb, b + k * N + col, N);
18
    simdgroup_multiply_accumulate(next, ta, tb, acc);
19
    acc = next;
20
  }
21
  simdgroup_store(acc, c + row * N + col, N);
22
}

The main primitive is: simdgroup_multiply_accumulate(next, ta, tb, acc). It computes the matrix update $\text{next} = \text{ta} * \text{tb} + \text{acc}$. We keep that result for the next iteration, then store the completed tile into $C$, whose row stride is also N.

There are no need for threadgroup memory here. The group loads directly from the device buffers, keeps the matrices in registers, and stores its result. We also need no threadgroup barrier: there is no shared-memory data being passed between SIMD groups.

Counting the Loads

How much reuse does this small tile give us? Each step loads $2\times8^2 = 128$ floats, or 512 requested bytes. It performs $8^3 = 512$ multiply-adds, or 1024 FLOPs. Ignoring the final output store, which happens once after all 512 steps, our no-cache estimate of arithmetic intensity is:

Across the grid, the kernel requests about $137.4\ \text{GFLOP}/2 = 68.7$ GB from the device buffers. That is eight times the 1D kernel’s 8.6 GB. We changed the way we express the arithmetic, but also shrank the output tile from $64\times64$ to $8\times8$, giving up much of the reuse between outputs.

As before, these are requested bytes assuming every load reaches device memory. Neighbouring groups load overlapping inputs, so caches can reduce the actual traffic. The estimate tells us how much reuse the kernel explicitly arranges; it does not place the measured kernel on the roofline.

Results

For the same $M=N=K=4096$ benchmark:

Our first SIMD-group matrix kernel reaches about 10% less throughput than 1D coarsening. We now have collective matrix arithmetic, but much less reuse in the kernel.

Can we blame device memory? With the no-cache estimate, its bandwidth roof would be only $2\times118 = 236$ GFLOP/s. The kernel runs well above that, so the 68.7 GB count clearly overstates the traffic reaching device memory. We need this kernel’s GPU counters to measure that traffic and identify what limits it; the timing alone does not settle either question.

So let’s measure it:

The integer pipeline falls from 35% to 12% busy: the matrix calls did remove most of the bookkeeping around each multiply-add. Now look at the memory rows. The kernel reads 119 GB/s from outside the GPU, our device-memory ceiling, and Last Level Cache Limiter is at 83%. Over its 121 ms, that is 14.5 GB, about a fifth of the 68.7 GB it requests, so the caches still help. At $137.4 / 14.5 \approx 9.5$ FLOP/byte, the roof is $9.5 \times 118 \approx 1.12$ TFLOP/s, and the kernel runs right at it.

The count does suggest a next step we can test. In the 1D kernel, one thread reused an input by holding several outputs. Here, a SIMD group can do the same with whole tiles: hold several $8\times8$ accumulators and reuse a loaded input tile across their matrix products. Let’s give each group a $16\times16$ output tile, made of four $8\times8$ matrices, and see whether that reuse helps.

Kernel 6: Larger Output Tiles

Each SIMD group now holds four $8\times8$ accumulators, covering a $16\times16$ output tile. At each step along $K$, it loads two tiles of $A$ (a0, a1) and two of $B$ (b0, b1). Each input tile feeds two matrix products: a0 updates both top accumulators, a1 both bottom ones, and the two $B$ tiles do the same along the columns.

Four SIMD groups cover a $32\times32$ threadgroup tile, which gives the kernel its name, simdgroup_32x32. sg selects one of its four $16\times16$ quadrants. The harness launches 128 threads per threadgroup and requires $M,N$ to be multiples of 32 and $K$ a multiple of 8.

1
kernel void simdgroup_32x32(device const float* a [[buffer(0)]],
2
                            device const float* b [[buffer(1)]],
3
                            device float* c [[buffer(2)]],
4
                            constant GemmShape& p [[buffer(3)]],
5
                            uint2 group [[threadgroup_position_in_grid]],
6
                            uint sg [[simdgroup_index_in_threadgroup]]) {
7
  const uint row = group.y * 32 + (sg / 2) * 16;
8
  const uint col = group.x * 32 + (sg % 2) * 16;
9
  simdgroup_float8x8 c00 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
10
  simdgroup_float8x8 c01 = c00, c10 = c00, c11 = c00;
11
  for (uint k = 0; k < p.K; k += 8) {
12
    simdgroup_float8x8 a0, a1, b0, b1, next;
13
    simdgroup_load(a0, a + row * p.K + k, p.K);
14
    simdgroup_load(a1, a + (row + 8) * p.K + k, p.K);
15
    simdgroup_load(b0, b + k * p.N + col, p.N);
16
    simdgroup_load(b1, b + k * p.N + col + 8, p.N);
17
    simdgroup_multiply_accumulate(next, a0, b0, c00);
18
    c00 = next;
19
    simdgroup_multiply_accumulate(next, a0, b1, c01);
20
    c01 = next;
21
    simdgroup_multiply_accumulate(next, a1, b0, c10);
22
    c10 = next;
23
    simdgroup_multiply_accumulate(next, a1, b1, c11);
24
    c11 = next;
25
  }
26
  simdgroup_store(c00, c + row * p.N + col, p.N);
27
  simdgroup_store(c01, c + row * p.N + col + 8, p.N);
28
  simdgroup_store(c10, c + (row + 8) * p.N + col, p.N);
29
  simdgroup_store(c11, c + (row + 8) * p.N + col + 8, p.N);
30
}

We now do four matrix multiply-accumulates for four input loads; the $8\times8$ kernel did one for two loads. Per SIMD group, that is 4096 FLOPs for 1024 requested bytes, so the no-cache intensity doubles to 4 FLOP/byte. Across the grid, the requested reads fall from 68.7 to 34.4 GB, again ignoring the final output stores.

The benchmark reaches 1687 GFLOP/s, $1.54\times$ the $8\times8$ kernel and $1.38\times$ 1D coarsening. In the separate profiler runs:

The read rate still reaches our device-memory ceiling, but external reads fall to about 8.4 GB per matrix product, down from 14.5 GBShow sidenoteThe counters cover three timed dispatches over 214 ms, so the average is $117.75\ \text{GB/s}\times0.214\ \text{s}/3 \approx 8.4$ GB per product. The throughput above uses the median GPU time.. Its measured intensity rises to about $137.4/8.4 \approx 16$ FLOP/byte, putting its bandwidth roof near 1.9 TFLOP/s, close to the profiled 2 TFLOP/s. Reuse has reduced the traffic, and device memory still sets the pace.

There is more reuse available within the threadgroup. The two SIMD groups side by side load the same $A$ tiles; the two above and below load the same $B$ tiles. Caches can answer some of those repeated requests. Our next kernel will arrange that sharing explicitly: load each input tile once into threadgroup memory, then let the SIMD groups read it from there.

Kernel 7: Sharing Inputs via Threadgroup Memory

We keep the same output layout: four SIMD groups, each holding four accumulators. Now the 128 threads first cooperate to load a $32\times32$ tile of $A$ and another of $B$ into threadgroup memory. The SIMD groups take their $8\times8$ fragments from those shared tiles and reuse them in registers as before.

There are now two loops along $K$. The outer loop stages 32 columns of $A$ and 32 rows of $B$, starting at kb. The inner loop consumes that data in four steps, at local k = 0, 8, 16, 24. The accumulators stay alive across both loops.

1
kernel void simdgroup_32x32_shared(device const float* a [[buffer(0)]],
2
                                   device const float* b [[buffer(1)]],
3
                                   device float* c [[buffer(2)]],
4
                                   constant GemmShape& p [[buffer(3)]],
5
                                   uint2 group [[threadgroup_position_in_grid]],
6
                                   uint sg [[simdgroup_index_in_threadgroup]],
7
                                   uint tid [[thread_index_in_threadgroup]]) {
8
  threadgroup float As[32 * 33];
9
  threadgroup float Bs[32 * 33];
10
  const uint block_row = group.y * 32;
11
  const uint block_col = group.x * 32;
12
  const uint subrow = (sg / 2) * 16;
13
  const uint subcol = (sg % 2) * 16;
14
  simdgroup_float8x8 c00 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
15
  simdgroup_float8x8 c01 = c00, c10 = c00, c11 = c00;
16

17
  for (uint kb = 0; kb < p.K; kb += 32) {
18
    for (uint i = tid; i < 32 * 32; i += 128) {
19
      const uint r = i / 32;
20
      const uint col = i % 32;
21
      As[r * 33 + col] = a[(block_row + r) * p.K + kb + col];
22
      Bs[r * 33 + col] = b[(kb + r) * p.N + block_col + col];
23
    }
24
    threadgroup_barrier(mem_flags::mem_threadgroup);
25
    for (uint k = 0; k < 32; k += 8) {
26
      simdgroup_float8x8 a0, a1, b0, b1, next;
27
      simdgroup_load(a0, As + subrow * 33 + k, 33);
28
      simdgroup_load(a1, As + (subrow + 8) * 33 + k, 33);
29
      simdgroup_load(b0, Bs + k * 33 + subcol, 33);
30
      simdgroup_load(b1, Bs + k * 33 + subcol + 8, 33);
31
      simdgroup_multiply_accumulate(next, a0, b0, c00);
32
      c00 = next;
33
      simdgroup_multiply_accumulate(next, a0, b1, c01);
34
      c01 = next;
35
      simdgroup_multiply_accumulate(next, a1, b0, c10);
36
      c10 = next;
37
      simdgroup_multiply_accumulate(next, a1, b1, c11);
38
      c11 = next;
39
    }
40
    threadgroup_barrier(mem_flags::mem_threadgroup);
41
  }
42
  const uint row = block_row + subrow;
43
  const uint col = block_col + subcol;
44
  simdgroup_store(c00, c + row * p.N + col, p.N);
45
  simdgroup_store(c01, c + row * p.N + col + 8, p.N);
46
  simdgroup_store(c10, c + (row + 8) * p.N + col, p.N);
47
  simdgroup_store(c11, c + (row + 8) * p.N + col + 8, p.N);
48
}

tid runs from 0 to 127. The loading loop assigns each thread eight positions in each input tile, spaced 128 elements apart. Together they cover all 1024 positions exactly once. Within a SIMD group, neighbouring threads read 32 consecutive floats from each device buffer, so the cooperative copy keeps those reads contiguous.

Why 32 * 33? Each shared row has 32 values followed by one unused slot. Its stride is therefore 33 floats, or 132 bytes, and the simdgroup_load calls must use that stride. The two arrays occupy 8.25 KiB togetherShow sidenoteThe padding shifts successive rows’ alignment by one float compared with a stride of 32. It is intended to change bank alignment, as the kernel’s source comment says. We have not benchmarked an unpadded version on the M5, so these results do not establish whether the padding improves performance..

The first barrier waits for all 128 threads to finish the copy before any SIMD group reads the shared tiles. The second waits for all groups to finish reading before the next outer iteration overwrites them. A group that finishes early must wait for the others; its matrix operations only coordinate the threads within that SIMD groupShow sidenoteAs in kernels 2 and 3, mem_flags::mem_threadgroup orders the threadgroup-memory accesses across the barrier. All threads reach both barriers on every outer iteration [1]. This kernel requires $M,N,K$ to be multiples of 32, so every staged tile is complete..

Per outer iteration, the threadgroup loads $2\times32^2\times4 = 8192$ bytes from the device buffers and performs $2\times32^3 = 65{,}536$ FLOPs. The no-cache intensity is now 8 FLOP/byte, and requested device reads fall from 34.4 to 17.2 GB. The SIMD groups still load their fragments separately, but those repeated loads now come from threadgroup memory.

The benchmark reaches 2722 GFLOP/s, $1.61\times$ the previous kernel and 81% of MPS. The separate profiler runs show:

The FP32 units are busier, and the new threadgroup-memory reads are visible. External reads fall again, from 8.4 to about 5.6 GB per matrix productShow sidenoteThe shared kernel’s counters cover three timed dispatches over 148.3 ms: $113.40\ \text{GB/s}\times0.1483\ \text{s}/3 \approx 5.6$ GB per product.. That gives an intensity of about $137.4/5.6 \approx 25$ FLOP/byte and a bandwidth roof of 2.9 TFLOP/s, close to the profiled 2.8 TFLOP/s.

Instruction Throughput Limiter rises to 82%, but it includes stalls as well as work. With external bandwidth still near its ceiling, that number alone does not establish an instruction-issue bottleneck. We have reduced the traffic enough to run much faster while adding the copies and barriers needed to share inputs explicitly.

Kernel 8: TensorOps

Kernel 7 still needs manual work for every $8\times8$ tile:the input copies, two barriers per step and dot product. A larger output tile would need more accumulators and more calls. But fortunately, with new Metal 4 from Apple (introduced with M5 in WWDC25 [4]), we can let Metal Performance Primitives (MPP) to do “dark magic” with the whole threadgroup tile instead.

MPP’s matmul2d TensorOp is new API for our kernel. We choose the output tile and the SIMD groups that work on it, then MPP will supply the loads and the matrix arithmetic [1]. Kernel 8 passes views of the device buffers directly to it, without threadgroup-memory staging. Apple’s programming guide recommends this on Apple silicon, where the caches can provide the reuse [5].

Tensors and Slices

The buffers hold the same row-major FP32 matrices. A tensor adds their dimensions and strides. Metal lists dimensions from the innermost outward, so $A$ has extents $(K,M)$, $B$ has $(N,K)$, and $C$ has $(N,M)$. This describes the existing layout; nothing is transposed or copied.

Here is the full-tile path of matmul_tensorops, with TM = TN = 64 and SG = 4. The file includes <metal_tensor> and <MetalPerformancePrimitives/MetalPerformancePrimitives.h> and uses the mpp::tensor_ops namespace:

1
auto A = tensor<device float, dextents<int, 2>, tensor_inline>(
2
    a, dextents<int, 2>((int)p.K, (int)p.M));
3
auto B = tensor<device float, dextents<int, 2>, tensor_inline>(
4
    b, dextents<int, 2>((int)p.N, (int)p.K));
5
auto C = tensor<device float, dextents<int, 2>, tensor_inline>(
6
    c, dextents<int, 2>((int)p.N, (int)p.M));
7

8
constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>);
9
matmul2d<desc, execution_simdgroups<SG>> op;
10
const int row = (int)group.y * TM;
11
const int col = (int)group.x * TN;
12

13
auto ta = A.slice<dynamic_extent, TM>(0, row);
14
auto tb = B.slice<TN, dynamic_extent>(col, 0);
15
auto tc = C.slice<TN, TM>(col, row);
16
op.run(ta, tb, tc);

The slices select 64 rows of $A$ and 64 columns of $B$, each across all of $K$, and the matching $64\times64$ output tile. Slices are also views: constructing ta loads nothing. op.run performs the loads and writes the product to tc. The descriptor’s default mode is multiply, so tc does not need to be initialized.

dynamic_length_v<int> lets MPP traverse the full $K$ dimension, and execution_simdgroups<4> makes all 128 threads participate in the call. The fixed slice extents give the compiler the full tile sizes; the source has a separate path with dynamic extents for partial edge tilesShow sidenoteThe tensor coordinates put columns first, so the output view is C.slice<TN, TM>(col, row). The descriptor uses the mathematical order (M, N, K). For our $C[35,93]$, the $64\times64$ tile begins at row 0 and column 64..

Kernel 8 reaches 3035 GFLOP/s, up from 2722 for kernel 7 and about 90% of MPS. MPP controls everything inside the call; kernel 9 changes what happens around it.

Kernel 9: Synchronizing SIMD Groups Along K

In kernel 8, the four SIMD groups move through MPP’s internal $K$ loop independently. If one runs ahead, the groups use inputs from distant parts of $A$ and $B$ at the same time, which enlarges the active input set and makes cache reuse harder. Apple recommends periodically synchronizing the groups with a threadgroup barrier [5].

To place those barriers, we need to own the $K$ loop. In matmul_tensorops_sync, each call handles one BK-wide chunk, and a cooperative tensor holds the output between calls: MPP partitions its elements across the participating threads in thread memory and chooses their layout [1].

After constructing the same buffer views, the kernel runs:

1
const int row = (int)group.y * TM;
2
const int col = (int)group.x * TN;
3
constexpr auto desc = matmul2d_descriptor(
4
    TM, TN, BK, false, false, false,
5
    matmul2d_descriptor::mode::multiply_accumulate);
6
matmul2d<desc, execution_simdgroups<4>> op;
7
auto first_a = A.slice<BK, TM>(0, row);
8
auto first_b = B.slice<TN, BK>(col, 0);
9
auto acc = op.template get_destination_cooperative_tensor<
10
    decltype(first_a), decltype(first_b), float>();
11

12
for (uint16_t i = 0; i < acc.get_capacity(); ++i) {
13
  if (acc.is_valid_element(i)) acc[i] = 0.0f;
14
}
15
for (int k = 0; k < (int)p.K; k += BK) {
16
  threadgroup_barrier(mem_flags::mem_none);
17
  auto ta = A.slice<BK, TM>(k, row);
18
  auto tb = B.slice<TN, BK>(col, k);
19
  op.run(ta, tb, acc);
20
}
21
auto tc = C.slice<TN, TM>(col, row);
22
acc.store(tc);

Each thread zeroes the valid elements in its own part of acc: get_capacity() counts its storage slots, and is_valid_element skips slots that hold no tensor element. Every call then adds a partial product into the same accumulator, and we store it to $C$ once, after the last chunk, so the chunks cost no round trips through the output buffer.

The barrier uses mem_none: every thread must reach the chunk boundary before any continues, but no memory fence is added. Unlike kernel 7, there is no shared tile to protect from being overwritten; the barrier only keeps the groups in step while they read the input buffers. The groups can still drift within a chunk, so BK sets how often they synchronize.

Choosing the Tile and Chunk Size

This leaves two parameters: the output tile and BK. For the same $4096^3$ problem, the saved benchmark runs give:

Larger output tiles reuse each input more. If each threadgroup loads every input it needs once per chunk, every load reaches device memory, and we ignore the final output store, the tile model gives:

That is 16 FLOP/byte for $64\times64$, about 21.3 for $64\times128$, and 32 for $128\times128$. The model ranks $128\times128$ first, yet it runs about three times slower than the other two, so reuse alone does not determine the best tile.

BK cancels out of the formula. Splitting $K$ does not change the input bytes the tile model needs, only when the groups request them and how much data they use at once. Smaller chunks synchronize more often; larger chunks give the groups more room to drift. The $64\times64$ results are nearly identical across 256, 512, and 1024, so 256 is a reasonable choice here, not a universal optimum.

The model does not tell us what MPP actually reads, so we profile the kernels as before:

Multiplying each read bandwidth by its kernel time, the $128\times128$ tile reads 18.4 GB per product, 3.4× the 5.3 GB of the $64\times64$ tile, although the model predicts half as many bytes. It does not spill, and its occupancy is the same 16% as $64\times64$, so the loss comes from traffic rather than registers or occupancy. At 7.5 FLOP/byte it sits on the memory roof, with F32 utilization at 29%. The counters show that MPP’s $128\times128$ path reads more, but not why; MPP does not expose how it loads the tile.

The full-K $64\times64$ and $64\times128$ kernels read less than the model predicts: 26 and 30 FLOP/byte against 16 and 21. Part of the re-reads is served before device memory, most likely by the GPU’s caches, which is consistent with Apple’s guide. The synchronized kernels read 3.3 and 2.6 GB, 38% and 43% less than their full-K versions, which places both to the right of the ridge. F32 Utilization reaches 92%: for the first time in this post, the FP32 units are the limiting resourceShow sidenoteBoth synchronized kernels spill a few registers, 144 and 368 bytes, and are still the fastest kernels here. Bytes per product are GPU Read Bandwidth × median GPU time, for example $71.63\ \text{GB/s}\times36.28\ \text{ms}\approx2.6$ GB for the winner..

In the benchmark, the fastest combination is tensorops_sync_64x128_k256: a $64\times128$ output tile, four SIMD groups, and 16 calls covering $K=4096$. It reaches 3.686 TFLOP/s, about 1.095× MPS. On GPU time alone, the synchronized $64\times64$ kernel is equally fast; the two trade places from run to runShow sidenoteTwo re-timings without counters (run_kernel, 15 iterations each, GPU time): 3808 and 3682 GFLOP/s for synchronized $64\times64$, and 3651 and 3810 for synchronized $64\times128$. The full-K pair even swaps order against the table: 3444 and 3285 for $64\times64$, and 3180 and 3168 for $64\times128$. The table’s 2% gap is within this spread.. The synchronized versions also differ in more than the barrier: they use a static $K$ in the descriptor and a cooperative destination, so these timings do not isolate the barrier’s contribution.

Precision, Timing, and the Neural Accelerators

Two caveats apply to this result. First, precision. Both TensorOps versions keep FP32 buffers and leave relaxed_precision false: by default in the first descriptor and explicitly in the second. This disables the option to truncate FP32 mantissas before multiplication [1]. The benchmark checks agreement with MPS using $|C-C_{\text{MPS}}|\leq10^{-3}+10^{-3}|C_{\text{MPS}}|$; passing that check is a comparison with MPS, not with an exact reference.

Second, timing. The table averages the median throughput from the three saved claude_full_0925_run1/2/3.csv runs. Those timings cover complete host calls after warming and checking the kernels, including command encoding, submission, and waiting. MPS and our shaders reuse allocated buffersShow sidenoteThese are the benchmark timings used in the opening ladder, rather than the GPU-only timings from the SIMD-group profiler. The comparison is for packed, non-transposed FP32 GEMM at $M=N=K=4096$ on this M5. The synchronized winner requires $M$ divisible by 64, $N$ by 128, and $K$ by 256..

TensorOps also provides access to the M5’s GPU Neural Accelerators [5], which raises the question of whether our FP32 kernels use them. Neither the API name nor the timings answer it: the winner’s 91% of our 4.04 TFLOP/s estimate is a comparison with the ordinary FP32 peak, not a measurement. I recorded a trace with the accelerator counters and added two controls, MLX’s FP32 matmul and MPS on FP16 matrices, which should use the accelerators:

They do not. Our FP32 winner, MPS, and MLX all show 0% Neural Accelerator utilization; their FLOPs run on the ordinary FP32 ALUs. The FP16 control confirms that the counter works: MPS keeps the accelerators 93% busy and reaches 15.7 TFLOP/s, almost four times our FP32 roof. The winner therefore beats MPS and MLX on the same FP32 units. It keeps them 92% busy against 83%, reads about half the bytes, and reads nothing from threadgroup memory, while both libraries stage through itShow sidenoteMLX doesn’t label its command buffers, so for MLX gpu_counters.py --min-ms 20 --process python selects the 13 command buffers with at least 20 ms of GPU time: 3 warm-up and 10 timed products. The GPU time-slices each buffer with other processes, so one buffer shows up as several intervals in the trace..

We finish at 3.686 TFLOP/s, ahead of MPS on this $4096^3$ benchmark. TensorOps supplies the matrix implementation; we choose the tile and control progress through $K$. The chunked $K$ loop leaves the tile model’s arithmetic intensity unchanged but reduces the bytes that actually reach memory.

Conclusion

Let’s end where we started. Here is the chart from the introduction of part 1 again; this time, you can read every dot on it:

If you take one thing from this post, take the loop we ran for every kernel: count the FLOPs and the bytes, put the kernel on the roofline, ask the profiler which limit it actually hit, change one thing, and measure again. Every big jump came out of that loop, and so did every surprise: tiling that took the pressure off device memory and then ran into an instruction wall, matrix instructions that came out slower than plain coarsening, and a $128\times128$ tile that read 3.4 times the bytes of a smaller one.

There’s plenty we didn’t do. Everything here is one shape, $4096^3$, packed and non-transposed, in FP32; other shapes would need other tiles and, for the synchronized kernel, edge handling. And the Neural Accelerators sat idle in every FP32 kernel, Apple’s included. MPP has a relaxed_precision flag that we left off; whether turning it on lets FP32 work reach the accelerators is an experiment I haven’t run yet. If you try it, I’d love to hear what you find.

Finally, thank you. This worklog turned out really long, and if you’ve read all the way down here, you gave it hours you could have spent on anything else. From the bottom of our hearts (mine, and whatever Claude and Codex have instead), thank you for reading. If you find a mistake or a faster kernel, the code is on GitHub: open an issue, and let’s make it faster together.

References