Metal GEMM Kernels for LoRA Training on Apple Silicon

KRNL is a LoRA training engine [1] for Apple Silicon: Rust, hand-written Metal kernels, no framework underneath. This is a technical report on the kernels and the measurements that shaped them.

All numbers are from an M1 Pro (16-core GPU, 16 GB unified memory, macOS 26.5.1). Square GEMM figures are N = K = M = 2048, fp16 operands, median of fifteen runs. The reference workload is one LoRA step [2], forward and backward, over a transformer MLP block: RMSNorm, SwiGLU, LoRA adapters on the gate and up projections, base weights frozen, 512 tokens, d_model 2048, FFN 5632, rank 32. That is 72.4 GFLOP per step.

1. Which arithmetic the matrix units accelerate ¶

The simdgroup_matrix intrinsics [3] do not have uniform performance across types. Same kernel structure, same shapes, only the scalar type changed:

Operands	Accumulator	GF/s
half	float	3296
half	half	122
float	float	223
bfloat	float	83

Three consequences.

The accumulator type selects a hardware path rather than a precision. Declaring simdgroup_matrix<half,8,8> as the accumulator costs a factor of 27: the matrix units accumulate in fp32, anything else falls off the fast path, and nothing in the toolchain warns you. In a kernel templated over a scalar type this is easy to write by accident:

typedef simdgroup_matrix<SCALAR, 8, 8> sgmat;
// operands: sgmat. accumulator: always simdgroup_float8x8.
simdgroup_float8x8 acc[4][4];

bf16 is emulated. Both bfloat arithmetic and simdgroup_bfloat8x8 compile without error, and the kernel runs at 83 GF/s. Compilation success is not evidence of a hardware path. Training here means fp16 storage, fp32 accumulation, and loss scaling [5].

fp32 matmul runs slower on the GPU than on the CPU: 223 GF/s against 1588 GF/s from Accelerate on the same machine. Any fp32 matrix work belongs on the CPU side, which matters in section 7.

Peak FMA throughput on this GPU is about 5.2 TFLOP/s, so 3296 GF/s is $3296/5200 \approx 63\%$ of peak, against $3660/5200 \approx 70\%$ for MPS.

2. GEMM structure ¶

One threadgroup owns a 32×32 output tile [6] and contains a single simdgroup (32 threads). Sixteen simdgroup_float8x8 accumulators, 32 fp32 registers per thread. Per K-step of 8: four simdgroup_load for A, four for B, sixteen simdgroup_multiply_accumulate. Arithmetic intensity is

$$\frac{32 \times 32 \times 8 \ \text{MACs}}{(32 \times 8) + (8 \times 32)\ \text{operands}} = \frac{8192}{512} = 16 \ \text{MACs per element read.}$$

Output leaves through threadgroup memory because the accumulators are fp32 and the tensor is fp16: simdgroup_store into a float tile, barrier, then a converted store to device.

Three transpose variants are needed. Weights are stored [out, in] (PyTorch convention, so a checkpoint can be mapped in directly), which makes the forward projections A @ Bᵀ; weight gradients are Aᵀ @ B.

Variant	product	ours	MPS
nn	A[M,K] @ B[K,N]	3296	3660
tn	A[K,M]ᵀ @ B[K,N]	3135	3759
nt	A[M,K] @ B[N,K]ᵀ	2455 → 2925	3678

MPS [4] shows no transpose penalty. Ours did, and in the forward pass every large GEMM is nt.

3. The transposed operand ¶

simdgroup_load takes a transpose flag, so the naive nt reads B with a K-sized row stride. Four hypotheses were tested; the first three are recorded because they are the obvious ones and they are wrong.

Cache-line fragmentation was the first. A transposed 8×8 load touches eight rows, using 16 bytes of each 64-byte line. Staging the tile through threadgroup memory to make device reads contiguous: 2380 GF/s, no better than the naive version.

Cache-set aliasing was the second. K = 2048 halves is a 4096-byte row stride, so 32 concurrent row reads map to the same set. Padding K to 2080 to break the power of two: no change at all, in either variant.

Third, the transposed load instruction itself. Staging in the natural layout and letting simdgroup_load transpose out of threadgroup memory, with no cache lines involved, gives 2754 against 2925 for the transposing write. The transposed load is more expensive than a linear one, but that is not where the 25% went.

The answer was instruction count. The staging loop was 32 scalar loads per thread feeding four sub-steps of matrix multiply. Vectorising the device reads:

// B is [N, K]; each thread stages one column of the tile, 32 values of K
device const vec<T, 4> *src = (device const vec<T, 4> *)(B + (col0 + tid) * K + k);
for (uint v = 0; v < 8; v++) {
    vec<T, 4> q = src[v];
    sB[(v * 4 + 0) * TILE + tid] = q.x;   // transposing scatter into
    sB[(v * 4 + 1) * TILE + tid] = q.y;   // threadgroup memory: consecutive
    sB[(v * 4 + 2) * TILE + tid] = q.z;   // lanes write consecutive addresses
    sB[(v * 4 + 3) * TILE + tid] = q.w;
}

2455 → 2925 GF/s. Alignment holds because K and the K-offset are both multiples of 32. The transposition happens on the write side, so the subsequent simdgroup_load from threadgroup memory is linear.

The staging was costing more in issue slots than it saved in memory behaviour. All three failed hypotheses were about the memory system.

4. Occupancy against tile size ¶

Standard CUDA practice is larger tiles for better data reuse. Measured here:

Configuration	GF/s	Threadgroup memory
32×32, 1 simdgroup (baseline)	3296	4 KB
64×64, 4 simdgroups, A+B staged, 32×32 output scratch each	1256	20 KB
Same, output written 8×8 at a time	2353	5 KB
32-wide staging → 64-wide in K	2872 → 2624	2 → 4 KB
32×64 per simdgroup, 32 accumulators, no shared memory	261	4 KB

The 64×64 tile doubles the intensity to $\tfrac{64 \times 64 \times 16}{(64 \times 16)+(16 \times 64)} = 32$ MACs per element read, half the device traffic per FLOP, and still loses by more than half. Every configuration that increases threadgroup memory loses monotonically with it. The 32×64 register-blocked version spills: 32 accumulators is past the budget, 16 is the ceiling.

On this GPU occupancy dominates data reuse. Three independent experiments in the same direction, all negative.

Two more negatives. Vectorising the output stores changed 3296 to 3306, i.e. nothing: the output is 32×32 per K-long accumulation loop, so it was never the cost. And writing fp32 straight to device, skipping the output staging entirely, gives 3424 but doubles activation traffic for every consumer downstream.

5. Per-dispatch profiling and the skinny-GEMM problem ¶

Wall-clock timing around a full command buffer cannot attribute cost. Giving each dispatch its own command buffer and reading GPUEndTime - GPUStartTime gives GPU-side time per kernel; comparing the sum against the batched step measures the submission overhead that method introduces.

sum of dispatches        24.874 ms
batched step             25.104 ms   (one command buffer)
difference                0.230 ms   (0.9%)

Under one per cent, so the per-dispatch numbers are usable and the host side is not a factor. Extract of the profile:

fwd gate = h@Wg^T         3.900 ms   15.7%
fwd up   = h@Wu^T         3.901 ms   15.7%
fwd down = s@Wd^T         4.072 ms   16.4%
bwd dS   = dY@Wd          3.436 ms   13.8%
bwd dH  += dG@Wg          3.566 ms   14.3%
bwd dH  += dU@Wu          3.563 ms   14.3%
bwd dGa                   0.465 ms    1.9%   <- 0.184 GFLOP => 397 GF/s
bwd dUa                   0.484 ms    1.9%
rmsnorm/swiglu/axpy/...   ~0.5 ms     2.0%

Six large GEMMs are 89% of the step. But bwd dGa runs at 397 GF/s against 3000 for its neighbours. Its shape is [512, 32] = [512, 5632] @ [5632, 32]: the output is 32 columns, exactly one tile, so the grid is (1, 16): sixteen threadgroups on a sixteen-core GPU.

These are the LoRA projections, and they are narrow because the method is low-rank. Collectively 1.9 ms, 7.5% of the step, invisible in a FLOP budget.

Split-K [7] fixes it. Partition the reduction across grid.z, accumulate partials in fp32 in device memory, reduce:

const uint span = K / chunks;
const uint k0 = tg.z * span, k1 = k0 + span;
...
device float *out = P + tg.z * M * N;   // fp32 partials, no conversion,
                                        // so no output staging either

Eight chunks turns 16 threadgroups into 128.

bwd dGa      0.465 ms -> 0.095 ms   (4.9x)
fwd h@Ag^T   0.156 ms -> 0.051 ms   (3.0x)

Routing is by output-tile count: below 32 tiles, and when K divides evenly, the split-K path is taken.

6. Where the step stands ¶

KRNL      24.89 ms   2909 GFLOP/s   20.6k tok/s
MLX       18.51 ms   3911 GFLOP/s   27.7k tok/s
                                    ratio 1.34x

Ninety-two per cent of the step is GEMM, so the residual is the per-kernel deficit against MPS: −10% on nn, −17% on tn, −20% on nt. There is no remaining structural bottleneck in the schedule; there is a GEMM that schedules worse than Apple’s.

One structural option was costed and rejected: storing activations transposed, [features, tokens], would make every forward GEMM nn and every backward one tn, worth about +3.5% on the GEMM mix. RMSNorm would then reduce along a strided axis and lose more than that.

Note also that the MLX comparison must ask for the gradient with respect to the input, not only the LoRA parameters. In an isolated block the autograd prunes the dX path entirely; in a real stack dX is what reaches the layers below. With that included MLX goes from 12.94 ms to 18.51 ms, and the comparison is like-for-like.

7. CPU matrix units alongside the GPU ¶

Accelerate [8] does fp32 GEMM at 1588 GF/s. Feeding it from Metal shared-storage buffers costs nothing measurable: 1755 GF/s from heap memory, 1849 from an MTLBuffer. Unified memory really does remove the copy.

Running both engines at once:

              alone       concurrent
GPU (fp16)   3097 GF/s    2965 GF/s
CPU (fp32)   1415 GF/s    1711 GF/s
total        3097         4675        (+51%)

The GPU gives up 4%. That is aggregate throughput, however, and a training step is a dependency chain. Assigning one whole GEMM per pass to the CPU (up in the forward pass, one dH contribution in the backward) measured 0.97x, slower than the GPU on its own. Per-phase:

fwd  gpu share 4.91 ms      bwd  gpu share 4.46 ms
fwd  cpu job   6.73 ms      bwd  cpu job   9.70 ms

The CPU becomes the critical path: the GPU does that GEMM in 4.0 ms, the CPU needs 6.7 (and 9.7 in the backward pass, where B is not transposed and the layout suits it worse, 1216 against 1753 GF/s for the same FLOP count).

The correct split is by rows, sized so both engines finish together. With $p$ the CPU’s share of the rows, $T_\text{cpu}$ and $T_\text{gpu}$ the single-engine times for that GEMM, and $O$ the GPU’s other work in the same window, the two sides of the window are

$$p\,T_\text{cpu} \;=\; O + (1-p)\,T_\text{gpu} \qquad\Longrightarrow\qquad p = \frac{O + T_\text{gpu}}{T_\text{cpu} + T_\text{gpu}}$$

Measured: $T_\text{cpu} = 6.73$, $T_\text{gpu} = 4.0$, $O = 4.3$ ms forward, giving $p = 0.77$; and $T_\text{cpu} = 9.70$, $T_\text{gpu} = 3.5$, $O = 4.46$ backward, giving $p = 0.60$. Rounded to the 32-row tile height: 384 and 320 of 512 rows. The slower engine takes the larger share, because the faster one has other work queued.

Splitting by rows requires no kernel change: A and C are entered at a row offset through setBuffer:offset:atIndex:, and the row offset in bytes is a multiple of the tile height times the row pitch.

The bridge has a cost. Accelerate is fp32-only, so tensors crossing over need an fp32 twin, and the CPU-side weight is stored fp32. RMSNorm emits both copies of its output in one pass rather than paying for a separate conversion kernel.

The balanced version is implemented and pinned against the pure-GPU path by test, but has no trustworthy timing yet; see section 9.

8. Attention ¶

The layer is RMSNorm [13] into SwiGLU [14]; attention keeps the quadratic intermediate resident rather than tiling it away [9].

Q, K and V are in head-major layout [H, N, head_dim], so each head is a contiguous matrix and a per-head GEMM is a buffer offset. The forward pass, per head, with $d_h$ the head dimension:

$$S = \frac{Q K^{\top}}{\sqrt{d_h}}, \qquad P = \operatorname{softmax}_{\text{causal}}(S), \qquad O = P V$$

and the backward pass, given $dO$:

$$dV = P^{\top} dO, \qquad dP = dO\, V^{\top}, \qquad dS = \frac{P \odot \bigl(dP - \operatorname{rowsum}(dP \odot P)\bigr)}{\sqrt{d_h}}$$ $$dQ = dS\, K, \qquad dK = dS^{\top} Q$$

Six matmuls, all covered by the same three variants. No batched-GEMM kernel is needed, because the head offset does the batching:

product	variant	M, N, K
$QK^{\top}$	nt	$N,\ N,\ d_h$
$PV$	nn	$N,\ d_h,\ N$
$P^{\top}dO$	tn	$N,\ d_h,\ N$
$dO\,V^{\top}$	nt	$N,\ N,\ d_h$
$dS\,K$	nn	$N,\ d_h,\ N$
$dS^{\top}Q$	tn	$N,\ d_h,\ N$

Two details in the softmax kernel. The $1/\sqrt{d_h}$ scale is applied inside it rather than by pre-scaling Q, which saves a pass over the scores. And the masked half of each row is written as a hard zero rather than left alone: P feeds straight into the P V GEMM, where stale values or NaNs would propagate everywhere.

RoPE [10] is one kernel for both directions. It is a rotation, so the backward pass is the same rotation with the opposite angle, which makes it a sign parameter rather than a second kernel.

Correctness tests: gradients against central finite differences; causality (perturb V at the last position, assert no earlier output moves, the failure that trains beautifully and learns nothing); and RoPE round-trip.

9. What is not measured ¶

The attention kernels are correct and unbenchmarked. The balanced heterogeneous path is correct and unbenchmarked. The only performance figure for heterogeneous execution is the 0.97x of the unbalanced version.

The reason is measurement hygiene rather than time. The control benchmark is gemm_nn at N=2048, which had read 3296 GF/s consistently. It later read 1271, 224 and 1069 GF/s on consecutive runs of the identical binary. A five-fold swing between identical runs is contention, not a degraded device: an iOS Simulator, a Virtualization.framework VM at 33% CPU, Xcode and Docker were all running. Thermal state was nominal, on AC power, low-power mode off.

Which means the numbers above were taken under load conditions that were never controlled, only observed to be quiet. The gate should refuse to print a performance figure when the control benchmark is outside 10% of its reference; until it does, the harness is not measuring the software.

10. On building over MLX instead ¶

Evaluated and rejected. The autograd (grad, value_and_grad, vjp) exists in the MLX [11] C++ core and in the official mlx-c C API, so Rust FFI is feasible without writing a C++ shim. But optimisers, the nn layers and all LoRA tooling live only in the Python layer and would have to be rewritten anyway; and mlx-c [12] publishes no releases at all, only a version string in a CMake file, against a project shipping every three to four weeks with no documented API stability policy.

That trades a moving dependency for GEMMs we can already write at 80–90% of MPS. MLX stays as the benchmark baseline.

Next ¶

The gate measures an MLP block, which is roughly two thirds of a real layer’s FLOPs and has a much friendlier memory profile: no softmax, no mask, no quadratic intermediate. Moving it to a full layer with attention is what turns “competitive GEMMs” into “can train”.

---

References ¶

[1] Hu, Edward J.; Shen, Yelong; Wallis, Phillip; Allen-Zhu, Zeyuan; Li, Yuanzhi; Wang, Shean; Wang, Lu; Chen, Weizhu. “LoRA: Low-Rank Adaptation of Large Language Models.” International Conference on Learning Representations (ICLR 2022). Why the adapter matrices are narrow, and therefore why the skinny-GEMM problem in section 5 is intrinsic to the method rather than an artefact of these shapes.

[2] Dettmers, Tim; Pagnoni, Artidoro; Holtzman, Ari; Zettlemoyer, Luke. “QLoRA: Efficient Finetuning of Quantized LLMs.” Advances in Neural Information Processing Systems 36 (NeurIPS 2023). The memory argument that makes LoRA the only realistic training target on a 16 GB machine, and the source of the base-weights-frozen structure the step assumes.

[3] Apple Inc. Metal Shading Language Specification, developer.apple.com/metal. simdgroup_matrix, simdgroup_load/simdgroup_store with their transpose flag, and the address-space qualifiers that make MSL diverge from C++14.

[4] Apple Inc. “Metal Performance Shaders,” Apple Developer Documentation, developer.apple.com/documentation/metalperformanceshaders. MPSMatrixMultiplication, used here as the reachability reference: it takes transpose flags and shows no penalty for them, which is what established that the deficit in section 3 was ours.

[5] Micikevicius, Paulius; Narang, Sharan; Alben, Jonah; Diamos, Gregory; Elsen, Erich; García, David; Ginsburg, Boris; Houston, Michael; Kuchaiev, Oleksii; Venkatesh, Ganesh; Wu, Hao. “Mixed Precision Training.” International Conference on Learning Representations (ICLR 2018). fp16 storage with fp32 accumulation and loss scaling, the regime section 1 forces, given that bf16 is emulated on this GPU.

[6] Goto, Kazushige; van de Geijn, Robert A. “Anatomy of High-Performance Matrix Multiplication.” ACM Transactions on Mathematical Software, vol. 34, no. 3, 2008. The blocking analysis every tiled GEMM descends from, and the reason larger tiles are expected to win, which section 4 measures as false on this GPU.

[7] NVIDIA. CUTLASS: CUDA Templates for Linear Algebra Subroutines, github.com/NVIDIA/cutlass. Split-K as a standard remedy for GEMMs whose output is too small to fill the device, with the same fp32-partials-plus-reduction structure used in section 5.

[8] Apple Inc. “BLAS,” Accelerate framework, Apple Developer Documentation, developer.apple.com/documentation/accelerate/blas. The only supported route to the CPU matrix units; cblas_sgemm is what section 7 measures, and its fp32-only interface is what forces the fp32 twin tensors.

[9] Dao, Tri; Fu, Daniel Y.; Ermon, Stefano; Rudra, Atri; Ré, Christopher. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” Advances in Neural Information Processing Systems 35 (NeurIPS 2022). The tiled attention this implementation does not do: section 8 materialises the full [H, N, N] intermediate, which is the first thing to revisit once the layer gate exists.

[10] Su, Jianlin; Lu, Yu; Pan, Shengfeng; Murtadha, Ahmed; Wen, Bo; Liu, Yunfeng. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” Neurocomputing, vol. 568, 2024. The rotation whose orthogonality is why one kernel with a sign parameter covers both the forward and the backward pass.

[11] Apple Machine Learning Research. MLX, github.com/ml-explore/mlx. The benchmark baseline throughout, and the reference implementation of the same block in bench/mlx_block.py.

[12] Apple Machine Learning Research. mlx-c, github.com/ml-explore/mlx-c. The official C API that would make a Rust binding feasible, versioned only by a string in CMakeLists.txt, the stability argument in section 10.

[13] Zhang, Biao; Sennrich, Rico. “Root Mean Square Layer Normalization.” Advances in Neural Information Processing Systems 32 (NeurIPS 2019). The normalisation in the reference block; its reduction axis is what makes the transposed-activation layout in section 6 unprofitable.

[14] Shazeer, Noam. “GLU Variants Improve Transformer.” arXiv

.05202, 2020. SwiGLU, whose two independent projections are the pair split across engines in section 7.