Out for Delivery: Why Our Metal Kernel Stalled at 10% Memory Bandwidth
In my last log (Escape from ABI Hell), I wrote about the productivity leap of deleting hundreds of lines of brittle C++ pybind11 boilerplate and compiling raw Metal shaders at Python runtime via mlx.core.metal.compile(). It felt great: zero compilation flags, clean distribution, and rapid iteration.
Then I ran the first full-layer matrix benchmarks of Ternary-Bonsai-2-27B.
The attention projections looked promising: the KV projections (5120 × 512) ran in 161.33 µs, edging past Apple’s native 2-bit quantization kernel by +9.8%.
Then the MLP gate and up-projection (5120 × 17408) executed.
The kernel took 773.70 µs. Apple’s native kernel completed the exact same matmul in 710.56 µs. That was our first major red flag: on the layer accounting for roughly 65% of the model’s total parameters, we were actively regressing by nearly 9%.
Our effective throughput crawled at 28.80 GB/s, utilizing roughly 10% of the memory bus. The GPU was sitting there, completely starved.
A deep dive revealed what actually happened. In my rush to quickly port the kernel into a pure-Python JIT runtime, I had accidentally broken the memory access pattern. It turns out C++ was never what made the original kernel fast. What made it magical was how it fed the silicon.
The Mental Model: How GPUs Actually Drink Memory
Random fetches are not really a thing for GPUs. Hardware memory controllers don’t deal in single bytes; they deal in 128-byte cache lines.
Think of it like a wholesale delivery truck. When a thread requests a single 4-byte float, the memory controller doesn’t send a courier with a tiny package; it drives a 128-byte pallet to the registers. If the neighboring threads aren’t asking for the rest of that pallet, most of the registers go empty and the payload is thrown away.
Thinking about it this way makes the problem—and the solution—obvious.
On Apple Silicon, execution happens in lockstep units of 32 threads called a SIMD warp. This naturally leads to two distinct scenarios: one where you saturate the memory bus, and one where you choke it.
- Saturating the Bus (Coalesced): All 32 threads in the warp request 32 consecutive, adjacent numbers in the exact same cycle (32 × 4 bytes = 128 bytes). The memory controller fulfills the entire warp in a single trip.
- Choking the Bus (Strided): The 32 threads request numbers separated by wide strides across memory. The controller is forced to issue 32 separate 128-byte dispatches, discarding almost the entire line each time. The threads spend their cycles idling while waiting on memory traffic.
The code was calculating the correct answers, but the hardware was waiting on trucks.
The Magic Sauce: Cooperative 32-Lane Warp
The architectural solution was Cooperative Warp Reduction.
Instead of having a single thread slowly grind through an entire output column on its own, an entire 32-thread SIMD warp collaborates on that single column.
Here is how the mechanics come together:
- Lockstep Ingestion: The 32 lanes split the deep K-dimension (5,120 rows). In every inner-loop iteration, Lane 0 through Lane 31 load adjacent 4-byte packed ternary chunks in lockstep. The memory controller sees a contiguous 128-byte window and fulfills it in a single transaction.
- Single-Cycle Hardware Accumulation: Instead of writing partial sums to shared memory and waiting on threadgroup barriers, the warp executes Metal’s hardware primitive:
float thread_sum = dot(act_chunk, weight_chunk); float col_sum = simd_sum(thread_sum);simd_sum()folds all 32 lane registers directly in silicon in exactly 1 clock cycle, with zero lock contention and zero barrier overhead.
The Architecture Trade-Offs
Thomas Sowell famously observed that “there are no solutions, there are only trade-offs.” Designing custom high-performance kernels for edge AI is an exercise in picking which trade-offs you are willing to live with:
| Architectural Strategy | Core Mechanism | Trade-off Upside | Why We Chose or Rejected |
|---|---|---|---|
| Thread-Per-Column | 1 thread computes 1 entire column. | Simple code, zero inter-thread communication. | Rejected: Destroys memory coalescing. Drops bus efficiency to ~10%. |
| Threadgroup SRAM with Barriers | Stage tiles into shared memory (threadgroup), accumulate with barrier locks. |
Flexible tile shapes across arbitrary matrix dimensions. | Modified: Barrier locks (threadgroup_barrier) stall execution. We use SRAM strictly to cache input activation tiles, never for partial sum reductions. |
| 2D Thread Tiling (2 × 4 Elements/Thread) | Each thread computes multiple rows and columns concurrently. | Amortizes activation loads across more arithmetic operations. | Deferred: High register pressure. On large matrices (5120 × 17408), excessive register allocation limits warp occupancy. |
| Cooperative SIMD Warp (Chosen) | 32 lanes split K, single-cycle simd_sum() register reduction. |
100% coalesced reads, zero barrier stalls, 1-cycle reduction. | Selected: Optimal balance of maximum memory throughput and minimal register footprint. |
The Result
By anchoring the kernel in cooperative 32-lane warps and staging activations into on-chip SRAM, MLP execution latency dropped below Apple’s native baseline, clawing back our bandwidth efficiency from a starved 10% to hardware saturation.
Language runtimes and high-level abstractions don’t dictate kernel speed—respecting cache lines and hardware execution topologies does.
In the next log, we look at the other half of the runtime: how we engineered speculative decoding at 0 MB extra VRAM using Prompt Lookup Decoding and 16-layer self-speculative pruning.