Part 5 of the series. An MLA prefill implementation built around an M5-targeted 16x16 matrix path compiled on the M3 host, but Metal rejected that pipeline at runtime. The M3 path uses an 8x8 operation with different fragment ownership, softmax state, and paging behavior. The two low-level implementations share the same algorithm while specializing the hardware-dependent fragment work. The evidence here is a source comparison, an M3 launch-and-finiteness smoke, and an M5-targeted compile. M5 hardware execution remains open.
The setup
You can read this on its own. The only background it assumes: I am building Fornax, a project to build an engine that spreads one large Mixture-of-Experts language model across a fleet of ordinary machines, including Apple Silicon Macs. A 128 GB Mac offers a large physical memory pool shared by the CPU and GPU. macOS, the runtime, KV cache, activations, allocator reserve, and operating headroom all reduce the portion available for model weights.
To use a Mac as a compute stage, the model’s operators have to run on its GPU through Modular’s MAX compiler. The MLA attention path used by DeepSeek-style models was unavailable on my Apple target. I began with a 16x16 specialization for the newer Apple GPU family. Metal rejected that pipeline on the M3, so I added an 8x8 specialization for the older operation. Incompatible fragment layouts make separate low-level implementations the clearest design. Traits or policy types can still share dispatch, loop orchestration, and the common attention algorithm.
The documented M5 boundary
Apple’s chips (M1 in 2020 through M4, then M5 in late 2025) share a design idea that matters here: unified memory. The CPU and GPU address the same physical RAM, so a model’s weights can remain in that shared pool instead of being copied across PCIe into separate device memory. On a 128 GB Mac, both processors can address the physical pool. The usable model budget is lower and has to reserve memory for the operating system, runtime state, KV cache, activations, temporary buffers, and allocator headroom. The large shared pool still makes a Mac attractive as one stage in the Fornax fleet.
Apple publicly introduced dedicated GPU Neural Accelerators with M5. Its MLX write-up describes them as providing dedicated matrix-multiplication operations and reports substantial gains for the large-matrix prefill workloads it tested. Those results apply to the reported models, shapes, and MLX implementation. Apple’s public material lacks a complete execution-unit diagram for the earlier GPUs. Third-party microbenchmarks indicate that the older 8x8 simdgroup-matrix path uses the general arithmetic machinery. That mapping comes from reverse-engineered evidence.
M5 exposes its Neural Accelerators through Tensor APIs in Metal 4. The local MAX source represents M5 as a newer compute capability and contains a 16x16 matrix path that Metal restricts to GPU family 10 and later. That compiler/runtime gate is the boundary this kernel work observed.
For this MAX kernel, “an Apple GPU” maps to two code paths. M1 through M4 use the 8x8 helper. M5 can use the 16x16 specialization when the tensor dimensions satisfy that path’s constraints. Metal rejects a 16x16 pipeline on M3 during pipeline creation.
What a simdgroup and a simdgroup matrix are
Two Metal concepts make the rest of this readable, so here they are plainly.
A simdgroup is a group of GPU threads (32 on Apple hardware) that execute in lockstep, one instruction at a time, across all of them. They are the unit the hardware schedules. Threads in a simdgroup can shuffle values directly between each other’s registers without going through memory, which is fast and is the basis for the reductions attention needs.
A simdgroup matrix is a small fixed-size matrix, such as 8x8 or 16x16, whose elements are spread across the registers of all 32 threads in the simdgroup, so the Metal matrix operation can consume distributed fragments. Each thread holds a few elements of the tile, called its fragment. The programming model resembles an NVIDIA warp-level MMA, although the underlying hardware differs across Apple generations. The relevant constraint is that the tile layout across threads is part of the instruction, and Apple’s 8x8 and 16x16 operations distribute them differently. Any code that reads or writes those fragments has to know which layout it is dealing with.
What the kernel has to compute
The operator I needed is called MLA prefill, and you can follow the argument without prior knowledge of it. Multi-head Latent Attention (MLA) is the attention variant DeepSeek models use. Ordinary attention builds each query-key score from one matrix multiply. MLA builds it from two, because the query and key are each split into a part that carries positional information (the RoPE part, 64 dimensions in DeepSeek-V2-Lite) and a non-positional part (the non-RoPE part, 128 dimensions), and the value head is another 128 dimensions:
score = Q_nope @ K_nope.T # 128-dim contraction
score += Q_rope @ K_rope.T # 64-dim contraction, K_rope shared across heads
output = softmax(score) @ V # accumulate into the 128-dim value head
Prefill runs attention across the prompt. For sufficiently large prompt and head shapes, its matrix contractions can have enough arithmetic intensity for matrix throughput to dominate. Short prompts, ragged batches, cache access, launch costs, and an inefficient tile schedule can move the same operator toward a different point on the roofline. MLA performs two score contractions before the softmax where ordinary attention performs one, so more values can remain live in each thread’s registers during the score phase. Register pressure became one of the main constraints on the Apple kernel shape, together with tile geometry and memory access.
NVIDIA and AMD already had tuned versions of this kernel. Apple needed an implementation against its simdgroup matrix operations, and that is where one kernel became two.
The instruction that split the kernel in two
My first Apple MLA prefill kernel used the M5-targeted 16x16 simdgroup matrix
operation, wrapped in MAX as MmaOpApple. The specialization compiled on the M3
host. Its evidence stops there; numerical validation and M5 execution remain open.
When the M3 tried to create the pipeline, Metal rejected it at runtime with this:
simdgroup_matrix<T,16,16x16> operations are supported by GPUFamily10 and later
The 16x16 operation is unavailable on the M3 and the other M1 through M4 GPUs. They expose the older 8x8 operation. A flag or tuning parameter cannot change that hardware capability.
The M3 path uses _mma_apple_8x8. Its fragment layout differs from the
16x16 layout, so fragment loads, stores, reductions, and state have separate
implementations. The surrounding attention algorithm and orchestration can remain
shared. This boundary is the concrete, code-level effect of the hardware change
from the previous section.
Fragment ownership: the same tile, held differently
A simdgroup matrix operation distributes the tile’s elements across 32 threads, with each thread holding a fragment in registers. The instruction defines that distribution. A single low-level fragment implementation would obscure the different ownership rules of the 16x16 and 8x8 paths.
| Property | M5-targeted path (16x16) | M3 path (8x8) |
|---|---|---|
| Apple GPU primitive | 16x16 MmaOpApple, GPU family 10 |
8x8 _mma_apple_8x8 |
| Hardware interpretation | dedicated matrix hardware documented by Apple | older 8x8 path; general-ALU mapping inferred from third-party microbenchmarks |
| Availability | M5-targeted, compute_capability() >= 5 with compatible dimensions |
fallback path for other compatible Apple configurations; tested on M3 |
| Evidence in this work | source compile only; M5 runtime and parity remain open | launched on M3; finiteness smoke only |
| Fragment elements per thread | 8 | 2 |
| Rows a thread owns | two rows: rb and rb + 8 |
one row |
| Columns per owned row | four | two |
| Reusable MAX abstraction | the MmaOpApple struct |
lower-level intrinsic plus matmul helper |
| Shared structure | attention algorithm and orchestration | attention algorithm and orchestration |
| Specialized structure | 16x16 fragments, reduction state, loads, and stores | 8x8 fragments, reduction state, loads, and stores |
The fragment layout carries the implementation difference. In the 16x16 path each
thread holds 8 elements across rows rb and rb + 8, with four columns each. On
the M3 each thread holds 2 elements, one row with two columns. Everything the
kernel does after the matrix multiply, the softmax
reduction, the rescaling, and the final store to memory is written in terms of
which rows and columns a thread owns. When ownership changes from “two rows, four
columns each” to “one row, two columns,” those low-level steps change with it. The
16x16 softmax code assumes each thread owns two row fragments packed as an 8-wide
vector. The M3 code owns one row as a 2-wide vector, so the row maximum, the row
sum, the output rescale, and the store use the 8x8 specialization.
MmaOpApple hardcodes 16x16 tiling and 8-element fragments. I wrote the M3 path
against the lower-level 8x8 intrinsic with explicit fragment loads and stores,
following Apple’s 8x8 matmul helper. A higher-level trait or policy can share the
loop structure and operation sequence while leaving fragment ownership and
reduction state to each specialization.
Softmax: same topology, different shape
Both paths use the same online-softmax topology. They reduce across the threads
that share a row with register-shuffle masks of 1 and 8. The 16x16 layout reduces
two thread-owned rows at once, while the M3 layout reduces one. Their intermediate
state has a different shape. The current 8x8 path uses its own
AppleSoftmax8 state object. A shared policy can describe the reduction sequence;
each specialization still supplies fragment-shaped state and updates.
Paged KV: alignment follows the tile
The tile size changes how the kernel walks the key-value cache. MAX stores the KV cache in fixed-size pages. The 16x16 path advances through 16-row sub-tiles, while the M3 path advances through 8-row sub-tiles.
The first M3 implementation uses scalar per-row loads from paged pointers for keys, values, and RoPE keys. Independent per-row addressing covers page boundaries without requiring a complete 8-row tile in one contiguous segment. This favors simple page-boundary semantics during bring-up and leaves a more efficient load pattern for later work. Any stronger page-alignment precondition should come from an enforced source assertion and a page-boundary conformance test.
Dispatch has to know the chip
The runtime also has to launch the specialization that matches the hardware and tensor dimensions. The current source selects the paths as follows:
if compute_capability() >= 5 and dimensions_allow_m5_path:
use 16x16 M5-targeted specialization
else:
use 8x8 M3 specialization
An M3 routed to the 16x16 specialization reaches the Metal pipeline-creation failure shown earlier. Capability and dimension checks belong in the dispatch policy. A compile smoke on an M3 can build the M5-targeted specialization and check that it links. Runtime validation of that specialization requires M5 hardware.
Implications for portable GPU kernels
The fragment work yields several broader lessons.
Portability has an intra-vendor dimension. The rest of this series covers the cost of moving between NVIDIA, AMD, and Apple. The same work appears within one vendor’s product line. A change in matrix-instruction shape led to separate low-level fragment implementations, softmax state, and cache walks. The compiler reaches Metal; the kernel author still specifies how each thread owns an 8x8 or 16x16 fragment.
The reported M5 gains depend on workload shape. Apple’s MLX measurements report up to roughly 4x improvement in time to first token for the tested models, while subsequent-token generation improves by about 19 to 27 percent. Apple attributes the smaller decode gain mainly to the M5’s higher memory bandwidth. These are MLX results for named workloads. This MAX kernel has no M5 performance result. A MAX comparison must control prompt shape, batch shape, dtype, tile schedule, cache layout, and the rest of the roofline conditions.
Support follows capability and evidence. “Apple Silicon” spans several GPU capabilities. This source carries an M5-targeted path and an M1 through M4 path, then selects between them using capability and dimension checks. A probe on the exact chip and path establishes the corresponding support claim.
Apple is exposing more of this hardware through documented interfaces. The M5’s tensor APIs in Metal 4 make the fast path on new Macs a first-class interface, and Apple’s MLX work shows the same Neural Accelerators benefiting language-model prefill. Fornax can use that path on supported hardware while retaining the 8x8 path as the compatibility floor for M1 through M4.
Validation status
The M3 kernel has one launch-and-finiteness result on an M3 Max. MAX reports its internal compute capability as 3, and the bounded smoke produced finite outputs. The decode and MoE paths beside it are the bring-up fallbacks described in Part 4. Numerical parity against a CPU or PyTorch reference remains open, along with ragged-prompt tests, page-boundary tests, RoPE key-sharing checks, comparison with the naive attention fallback, and cross-backend parity against NVIDIA and AMD on identical inputs. Fornax records Apple MLA as local smoke evidence.
The M5-targeted path compiled in this work. M5 runtime, finiteness, numerical parity, and performance evidence remain open. A future result must name the exact chip, macOS and Metal toolchain, MAX build, dtype, shape, command, reference comparison, and benchmark method.
The public Fornax repository remains pre-alpha. It contains the planner, the historical FNX1 mechanism, and the candidate FNX2 ragged reference and loopback path at T0/T1. A physical Mojo/MAX stage adapter remains open. For the Apple kernel, the next evidence step is numerical reference comparison on M3, followed by execution through the physical Stage ABI. M5 requires its own hardware run before the M5-targeted path can carry an execution or performance claim.
Previous: Model Bring-Up. Series index.
Sources: Apple unveils M5 (Apple Newsroom); Exploring LLMs with MLX and the M5 Neural Accelerators (Apple Machine Learning Research); Investigating the GPU Neural Accelerators on Apple A19/M5; Apple GPU microarchitecture notes (metal-benchmarks); Fornax Apple MLA prefill port record.