Tutorial 16: Writing Custom Metal Kernels

MLX provides high-level operations that compile to Metal automatically. But for performance-critical paths — like quantized matmul, fused RMSNorm, or paged attention — a hand-written Metal kernel can be 2-10× faster. This tutorial explains how Metal kernels work and how to integrate them into the engine.

---

1. When You Need Custom Kernels

MLX operations are general-purpose. They handle any tensor shape, any dtype, and any combination of operations. But general-purpose comes with overhead:

• Kernel launch overhead: Each MLX operation launches a separate GPU kernel. Fusing square + mean + rsqrt into one kernel avoids two extra launches.
• Memory bandwidth: Each intermediate tensor is written to and read from global memory. Fused kernels keep intermediate values in registers.
• Specialized algorithms: Quantized matmul can use SIMD-group reductions that the general-purpose matmul kernel does not.

When NOT to Write Custom Kernels

• The operation is already fast enough (most MLX ops are well-optimized).
• The kernel would only work for specific tensor shapes.
• The code complexity is not justified by the performance gain.

Our engine follows the principle: start with pure MLX, profile, then optimize only the bottlenecks.

---

2. Metal Kernel Architecture

A Metal kernel is a function that runs on the GPU. Each thread processes a portion of the data:

GPU Hardware
├── Command Queue      ← host submits work
├── Compute Pipeline   ← compiled shader code
├── Threadgroups       ← groups of threads that share local memory
│   ├── Thread 0       ← each thread has a unique ID
│   ├── Thread 1
│   ├── ...
│   └── Thread N-1
└── Output buffers     ← results written here

A minimal Metal kernel looks like this (in MSL — Metal Shading Language):

#include <metal_stdlib>
using namespace metal;

// Each thread computes one element of the output
kernel void add_arrays(
    device const float* a [[buffer(0)]],     // input array A
    device const float* b [[buffer(1)]],     // input array B
    device float* result [[buffer(2)]],      // output array
    uint id [[thread_position_in_grid]]       // unique thread ID
) {
    result[id] = a[id] + b[id];             // one addition per thread
}

---

3. The Extension System

Our engine loads Metal kernels through a nanobind/Metal extension:

# ext.py — lazy proxy for the native extension
class _NativeKernelProxy:
    def _ensure_loaded(self):
        # Try to import the compiled extension
        from anllm import anllm_ext as module
        return module

    def __getattr__(self, item):
        module = self._ensure_loaded()
        return getattr(module, item)  # e.g., native.rope, native.rms_norm

native = _NativeKernelProxy()

When the extension is not compiled, every attribute access raises RuntimeError. The engine catches this and falls back to the pure-Python path:

norm_cls = FastRMSNorm if features.kernels else RMSNorm

---

4. Example: Fused RMSNorm Kernel

The pure-Python RMSNorm does three operations:

# Pure Python: 3 MLX operations → 3 kernel launches, 2 intermediate tensors
x = x * mx.rsqrt(mx.mean(mx.square(x), axis=-1, keepdims=True) + eps)

The Metal kernel fuses them into one:

# Metal kernel: 1 kernel launch, 0 intermediate tensors
return native.rms_norm(mx.contiguous(x), mx.contiguous(self.weight), self.eps)

What the Metal Kernel Does Internally

For each row of the input:
  1. Load head_dim elements into threadgroup memory
  2. Square each element (in parallel across threads)
  3. Parallel reduction to compute the sum of squares
  4. Divide by head_dim, add epsilon, take rsqrt
  5. Multiply each element by (rsqrt × weight)
  6. Store the result

The parallel reduction is the key optimization — it computes the mean of squares in O(log n) steps instead of O(n).

---

5. Example: Fused RoPE Kernel

RoPE requires splitting vectors into pairs and applying 2D rotations. The Metal kernel does this efficiently:

class FastRoPE:
    def __init__(self, dims, seq_len, base=10000, traditional=False):
        self.dims = dims
        self.seq_len = seq_len
        self.base = base
        self.traditional = traditional

    def __call__(self, x, offset=0):
        """
        Apply RoPE using a Metal kernel.

        The kernel handles:
        - Splitting vectors into pairs
        - Loading precomputed cos/sin from device memory
        - Applying the 2D rotation for each pair
        - Handling batch offsets correctly
        """
        batch_size = x.shape[0]

        # Convert offset to a batch-sized array
        if isinstance(offset, int):
            offset = mx.full((batch_size,), offset, dtype=mx.int32)
        elif isinstance(offset, list):
            offset = mx.array(offset, dtype=mx.int32)

        return native.rope(
            mx.contiguous(x),
            mx.contiguous(offset.astype(mx.int32)),
            self.dims,
            self.base,
            self.traditional,
        )

The Metal kernel avoids the reshape/split/concat overhead of the Python version. It reads each vector, applies the rotation in registers, and writes the result — all in a single pass.

---

6. Example: SwiGLU Kernel

SwiGLU computes silu(gate) * up — a fused multiply-activate:

def swiglu(gate, up):
    """Fused SiLU gating — one Metal kernel instead of two ops + one multiply."""
    return native.swiglu(mx.contiguous(gate), mx.contiguous(up))

The kernel reads both gate and up from memory once, computes silu(gate) in registers, multiplies by up, and writes the result. The pure-Python path would allocate an intermediate tensor for silu(gate) and read it back for the multiply.

---

7. Performance Impact

Operation	Pure Python	Metal Kernel	Speedup
RMSNorm	3 kernel launches	1 kernel launch	~2×
RoPE	5+ ops (split, mul, concat)	1 kernel launch	~3×
SwiGLU	2 ops + 1 multiply	1 kernel launch	~2×
Quantized matmul	Dequant + matmul	Fused dequant-matmul	~2-4×
Decode attention	Full SDPA	Specialized decode kernel	~3-5×
Paged attention	Gather + SDPA	Direct page reads	~2-3×

These speedups compound across all layers. For a 32-layer model, the total speedup from Metal kernels is significant.

---

8. The mx.contiguous() Contract

Before passing tensors to Metal kernels, we call mx.contiguous():

return native.rope(mx.contiguous(x), ...)

This ensures the tensor is laid out in memory in the expected order (no transposes, no strides). Metal kernels access GPU memory directly — they need to know the exact memory layout. mx.contiguous() returns the tensor as-is if it is already contiguous, or makes a copy if it is not.

---

9. Building the Extension

The Metal extension is compiled separately from the Python code:

# From the project root
cd ext/
# Build with nanobind + Metal shader compilation
pip install .

The extension contains: 1. Metal shader files (.metal) — the GPU kernel code 2. C++ bindings (nanobind) — bridges Metal to Python 3. Python module (anllm_ext) — importable from Python

When the extension is not built, the engine works entirely through MLX's built-in operations. The extension is an optimization, not a requirement.

---

10. Summary

• Custom Metal kernels fuse multiple operations into a single GPU dispatch.
• When to use: Quantized matmul, normalization, RoPE, attention — the hot paths that run billions of times per inference.
• When not to use: Everything else — MLX's built-in ops are fast enough.
• Graceful degradation: The extension is optional; pure Python always works.
• Performance: 2-5× speedup per operation, compounding across layers.

Next: Tutorial 17 — Building an OpenAI-Compatible API Server