Apple GPU Evolution for AI Compute

SoC Evolution

Apple appears to follow a tick-tock cadence between the CPU + ANE subsystem¹ and the GPU subsystem.

GPU Evolution

The A20 Pro point uses a public Geekbench 6.7.2 result from September 21, 2026.⁴

Each entry reads total / per GPU core for the highest-core base A- or M-series configuration. For Apple7–9, this table treats peak FP16 throughput as equal to FP32. For Apple10 and the reported Apple11 result, SIMD-group FP16 is 2× FP32. Values marked measured come from sustained peak microbenchmarks; ~ values are estimates. A20 remains blank because no reliable public measurement is available. Except for Apple’s published 2.6-TFLOPS M1 figure, earlier values come from reported GPU width and clock rate and are rounded to two decimals. Pro, Max, and Ultra variants are not included unless named.⁹

Basic architecture

GPU subsystem:

GPU subsystem → L2 cache → GPU core → 4× Shader Core → 32 SIMD lanes → ALU pipes

Shader Core (Streaming Multiprocessor / SM Core):

Preliminary — Occupancy

Higher occupancy leads to higher throughput: more of the available TFLOPS can be achieved when more SIMDgroups are ready to use the execution pipelines.⁵

Dynamic Caching

1. Dynamic register allocation
2. Dynamic on-chip memory
3. Occupancy Manage (Not covered in this post): decides when to schedule more SIMDgroups and when to schedule fewer to prevent register spilling⁵

Neural Accelerator (NAX)

The Tensor Core era of the Apple GPU

M5 and A19 place dedicated matrix-multiplication hardware inside every shader core.⁶

Speedup from NAX

NAX changes prompt processing far more than token generation.

The cleanest complete comparison in the llama.cpp reports uses the same 40-core M5 Max and Llama 2 7B workload: pp512, tg128, and full GPU offload. The legacy result is the mean of two independent 8e672ef submissions, which agree within 1.4%; the NAX result is the accepted c1d0e7a Tensor API submission.¹⁰

NAX makes prefill 3.08× faster. Decode is effectively unchanged because single-token generation remains dominated by moving model weights through memory.

GPU AI Software Stack

TensorOps

• MatMul + Conv
• Same API across Apple GPU generations
• Uses NAX on supported GPUs
• Can be mixed with other programmable Metal shader code⁶

Important features:

MLX Metal GEMM: from dispatch to NAX

Source snapshot: 3a62199. The following blocks are source excerpts, not standalone programs; omitted arguments use ....¹¹

1. Layer to dispatch

Linear.__call__ chooses addmm when bias is present and matmul otherwise:¹²

if "bias" in self:
    x = mx.addmm(self["bias"], x, self["weight"].T)
else:
    x = x @ self["weight"].T

In Matmul::eval_gpu, layout preparation is followed by shape-based routing:

if (M == 1 && N == 1 && batch_size_out == 1 &&
    a.flags().row_contiguous && b.flags().row_contiguous &&
    a.dtype() != complex64) {
  return dot_product(...);
}

if (!a_transposed && b_transposed && gemv_wide(...)) {
  return;
}

if (std::min(M, N) == 1) {
  return gemv(...);
}

return steel_matmul(...);

This separates dot products, narrow/wide GEMV cases, and Steel GEMM.

Within Steel, a long reduction dimension K can select a split-K kernel. The regular NAX path below is one branch, not the implementation for every matrix multiplication.

2. Select a Metal kernel

NAX eligibility checks both the platform and dtype:

bool use_nax = metal::is_nax_available() &&
    !issubdtype(a.dtype(), complexfloating) &&
    (env::enable_tf32() || a.dtype() != float32);

is_nax_available() checks build support, OS version (26.2 or newer), and GPU architecture. FP16 passes the dtype filter, but still needs a supported platform.¹³

For regular NAX GEMM, MLX chooses a threadgroup tile bm by bn, reduction block bk, and SIMD-group arrangement wm by wn:

int bm = 128, bn = 128, bk = 512;
int wm = 4, wn = 4;

char devc = d.get_architecture().back();
if (devc == 's' || devc == 'c' || devc == 'd') {
  bk = (K >= 8192 && K > (M + N)) ? 64 : 256;
  bm = 64;
  wm = 2;
}

kname << "steel_gemm_fused_nax_"
      << (transpose_a ? 't' : 'n')
      << (transpose_b ? 't' : 'n')
      << "_" << type_to_name(a)
      << "_" << type_to_name(out)
      << "_bm" << bm << "_bn" << bn << "_bk" << bk
      << "_wm" << wm << "_wn" << wn;

The base kernel name encodes transpose flags, dtypes, and tile sizes. Batch, alignment, and epilogue flags further specialize the cached pipeline.

Without NAX, regular Steel GEMM uses steel_gemm_fused_.... Its initial tile is bm = 64, bn = 64, bk = 16, with wm = wn = 2; GEMM_TPARAM_MACRO then applies device/shape tuning.¹⁴

3. Map tiles to threadgroups

The output tile counts are rounded up, then swizzled for launch. Shown here is regular NAX GEMM; the non-NAX variant uses get_steel_gemm_fused_kernel for pipeline lookup:

int tn = (N + bn - 1) / bn;
int tm = (M + bm - 1) / bm;

// GEMMParams retains the logical tile counts and matrix strides.
int tile = 1 << swizzle_log;
tm = (tm + tile - 1) / tile;
tn = tn * tile;

MTL::Size group_dims = MTL::Size(32, wn, wm);
MTL::Size grid_dims = MTL::Size(tn, tm, batch_size_out);

auto& compute_encoder = metal::get_command_encoder(s);
auto kernel = get_steel_gemm_fused_nax_kernel(...);
compute_encoder.set_compute_pipeline_state(kernel);
compute_encoder.set_input_array(a, 0);
compute_encoder.set_input_array(b, 1);
compute_encoder.set_output_array(out, 3);
compute_encoder.set_bytes(params, 4);
// Optional bias/epilogue and batch bindings omitted.
compute_encoder.dispatch_threadgroups(grid_dims, group_dims);

swizzle_log comes from a device/shape heuristic. Each threadgroup has 32 * wn * wm threads; GEMMParams carries sizes, strides, and the swizzle needed to map the launch back to output tiles. Alignment flags select handling for partial tiles.

The binding indices match the Metal kernel’s [[buffer(0)]] input A, [[buffer(1)]] input B, [[buffer(3)]] output D, and [[buffer(4)]] parameters. MLX’s dispatch wrapper inserts dependency barriers and records the native Metal dispatch:¹⁵

void CommandEncoder::dispatch_threadgroups(
    MTL::Size grid_dims, MTL::Size group_dims) {
  maybeInsertBarrier();
  buffer_ops_++;
  get_command_encoder()->dispatchThreadgroups(grid_dims, group_dims);
}

This encodes work into a command buffer. MLX submits it with buffer_->commit() when it flushes the stream; dispatch itself does not wait for GPU completion.

Inside the NAX kernel, wm/wn partition the output tile, not the reduction dimension. The JIT passes host values bm, bn, bk, wm, wn as template arguments BM, BN, BK, WM, WN. Metal supplies simd_group_id via [[simdgroup_index_in_threadgroup]]; the kernel uses it to select a sub-tile:¹⁶

// steel_gemm_fused_nax.h: after batch and threadgroup tile offsets.
constexpr short SM = BM / WM;
constexpr short SN = BN / WN;
constexpr short TM = SM / 16;
constexpr short TN = SN / 16;

const short tm = SM * (simd_group_id / WN);
const short tn = SN * (simd_group_id % WN);

A += transpose_a ? tm : (tm * params->lda);
B += transpose_b ? (tn * params->ldb) : tn;
D += tm * params->ldd + tn;

NAXTile<AccumType, TM, TN> Dtile;
// gemm_loop<T, SM, SN, ...> accumulates this tile across K.

For BM = BN = 128 and WM = WN = 4, one threadgroup has 16 SIMD-groups (512 threads). Each group computes a 32 x 32 output sub-tile: group 5 starts at (tm, tn) = (32, 32) within the threadgroup tile. With TM = TN = 2, its accumulator contains a 2 x 2 grid of 16 x 16 fragments. The inner (16, 32, 16) operation below is a compute step, not the whole group’s tile. Kernel-local tm/tn are element offsets, unlike the host’s launch tile counts.

4. Load, compute, write

Non-NAX: shared staging, SIMD-group MMA, global store. The kernel stages input tiles in threadgroup memory before computing:¹⁴

// steel_gemm_fused.h: aligned main loop; offsets and epilogue omitted.
threadgroup T As[gemm_kernel::tgp_mem_size_a];
threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
thread mma_t mma_op(simd_group_id, simd_lane_id);
thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);

for (int k = 0; k < gemm_k_iterations; k++) {
  threadgroup_barrier(mem_flags::mem_threadgroup);
  loader_a.load_unsafe();
  loader_b.load_unsafe();
  threadgroup_barrier(mem_flags::mem_threadgroup);
  mma_op.mma(As, Bs);
  loader_a.next();
  loader_b.next();
}
threadgroup_barrier(mem_flags::mem_none);
return mma_op.store_result(D, params->ldd);

mma_op.mma loads fragments from As/Bs and reaches this primitive in mma.h:

METAL_FUNC static constexpr void mma(
    thread mat_type& D, thread mat_type& A,
    thread mat_type& B, thread mat_type& C) {
  simdgroup_multiply_accumulate(D, A, B, C);
}

Here mat_type is metal::simdgroup_matrix<T, 8, 8>: each of 32 threads holds two elements per fragment. This is the non-NAX matrix primitive, distinct from the NAX cooperative-tensor operation below.

NAX: private tiles, cooperative-tensor MMA, global store.

Inside steel/gemm/nax.h, the inner operation uses a (16, 32, 16) descriptor at SIMD-group scope. Cooperative tensors hold operand and accumulator fragments distributed across the group:¹⁷¹⁸

A fragment is 16x16, so 32 threads will hold kElemsPerFrag = 16x16/32 = 8

However this kernel is 16x32x16 (M, N, K), so it requires 1 A fragment(tile), 2 B fragments, 2 C fragments

constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor(
    16, 32, 16, transpose_a, transpose_b, true,
    mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);

mpp::tensor_ops::matmul2d<desc, metal::execution_simdgroup> gemm_op;

auto ct_a = gemm_op
    .template get_left_input_cooperative_tensor<AType, BType, CType>();
auto ct_b = gemm_op
    .template get_right_input_cooperative_tensor<AType, BType, CType>();
auto ct_c = gemm_op.template get_destination_cooperative_tensor<
    metal::remove_addrspace_t<decltype(ct_a)>,
    metal::remove_addrspace_t<decltype(ct_b)>,
    CType>();

for (short i = 0; i < kElemsPerFrag; i++) {
  ct_a[i] = A[i];
  ct_b[i] = Bn0[i];
  ct_b[kElemsPerFrag + i] = Bn1[i];
  ct_c[i] = Cn0[i];
  ct_c[kElemsPerFrag + i] = Cn1[i];
}

gemm_op.run(ct_a, ct_b, ct_c);

for (short i = 0; i < kElemsPerFrag; i++) {
  Cn0[i] = ct_c[i];
  Cn1[i] = ct_c[kElemsPerFrag + i];
}

run performs the multiply-accumulate using initialized fragments; the final loop copies the accumulator back to MLX’s fragment layout.

Cn0[i] = ct_c[i] copies a private fragment, not an output buffer. In this helper, Cn0/Cn1 are declared thread dtype_frag_t<CType>&; A and Bn0/Bn1 are const thread references. dtype_frag_t<U> is a vector of eight elements per thread.¹⁷

The inner fragment named A is not the outer kernel’s device pointer named A. thread does not guarantee physical register allocation: the compiler can spill. Likewise, a device access can hit cache rather than DRAM.¹⁹

The regular NAX path loads from device buffers directly into private tiles, then stores the completed tile through the output device pointer. These excerpts show the missing input/output boundary:²⁰

// gemm_nax.h: aligned input loads inside the reduction loop.
// A and B are const device T*.
Atile.load(A + A_offset, lda);
Btile.load(B + B_offset, ldb);
tile_matmad_nax(
    Dtile, Atile, metal::bool_constant<transpose_a>{},
    Btile, metal::bool_constant<transpose_b>{});

// steel_gemm_fused_nax.h: after gemm_loop and any epilogue.
// D is device T*; partial output tiles use store_safe instead.
Dtile.store(D, int(params->ldd));

So the regular non-NAX path explicitly stages through threadgroup memory, while this NAX path does not use that staging step. In both, fragment assignment and global-buffer storage are separate operations.

This is the GPU control surface visible in MLX: dispatch, tiling, launch geometry, and matrix primitives. The Metal compiler still owns instruction lowering.

Profiling & Debugging

Model level: Instruments

Similar to Nsight Systems

Instruments shows CPU, GPU, memory, display, command-buffer execution, and performance counters on a system timeline.⁶

Kernel level: Xcode Metal debugger

Similar to Nsight Compute

GPU trace replay isolates work for cost graphs, counters, register use, divergence, and instruction analysis.⁶

References