mlx-metal-kernels

MLX Metal Kernels for Apple Silicon

Routing Guide (Pick the Right Template)

• Elementwise / pointwise ops → scripts/batch_elementwise_kernel.py
• Row-wise reductions (RMSNorm / LayerNorm / Softmax) → scripts/rmsnorm_kernel.py, scripts/layernorm_kernel.py, scripts/softmax_kernel.py
• Attention → Prefer MLX built-ins first (mx.fast.scaled_dot_product_attention) and use references/attention-variants-guide.md only when you need a custom layout/mask.
• Long-context partitioned attention → scripts/paged_attention_partitioned_kernel.py + references/paged-attention-patterns.md for two-phase partition + reduce softmax.
• Quantized matvec / dequant patterns → scripts/dequant_matvec_kernel.py
• M3+ matrix ops (simdgroup_matrix) → scripts/simdgroup_matmul_kernel.py + references/simdgroup-matrix-guide.md

Pre-Benchmark Checklist (Avoid “bench lies”)

1. Force evaluation: time with mx.eval(out); mx.synchronize() (MLX is lazy).
2. Contiguity: if your kernel does x[row * D + i], require contiguous inputs (or call mx.ascontiguousarray).
3. Bounds checks: if you round grid up, guard if (tid >= ...) return;.
4. Dtype expectations: float16 I/O is common; accumulate in float32 for stability.
5. Threadgroup limits: threadgroup.x must be ≤ 1024 and a multiple of 32 (one simdgroup).
6. First-call compile: ignore the first run when benchmarking (compile + cache effects).

This skill provides patterns and guidance for developing custom Metal compute kernels using MLX's mx.fast.metal_kernel() API, targeting Apple Silicon GPUs (M1, M2, M3, M4).

Quick Start

import mlx.core as mx

kernel = mx.fast.metal_kernel(
    name="my_relu",
    input_names=["x"],
    output_names=["out"],
    source="uint i = thread_position_in_grid.x; out[i] = max(x[i], T(0));",
)
x = mx.random.normal((1024,))
out = kernel(
    inputs=[x],
    template=[("T", mx.float32)],
    grid=(1024, 1, 1),
    threadgroup=(256, 1, 1),
    output_shapes=[(1024,)],
    output_dtypes=[mx.float32],
)[0]
mx.eval(out)