robertlangdonn/mlx-metal-kernels

mlx-metal-kernels is a personal learning project documenting an attempt to write low-level GPU kernels for running language models, from scratch, using Apple's Metal graphics technology on a 16 GB M3 Mac. It uses MLX, Apple's machine learning framework, and its metal_kernel feature to write custom code that runs directly on the GPU. The author is upfront that the kernels themselves are ordinary and not novel research, the real point of the project is the honest comparison against Apple's own expert-written kernels and the lessons learned about when hand-writing a kernel is actually worth it. The project walks through five stages, called rungs, each with its own write-up file. The first rung fuses a common neural network operation and finds it only matches the speed of MLX's built-in compiler, showing that simple elementwise operations are usually already optimized well. The second rung writes a normalization kernel that beats the built-in compiler by roughly two times, because it involves a cross-thread reduction the compiler cannot generate on its own. The third rung builds honest timing tooling so results can be fairly ranked. The fourth rung refines the normalization kernel further to match Apple's own expert kernel in speed. The fifth rung implements an attention technique called online softmax, which turns out correct but thirteen to thirty three times slower than Apple's expert version, showing that the real difficulty lies in tiling and matrix multiplication hardware use rather than the algorithm itself. Each rung can be run as a standalone Python script. The project requires MLX and Apple Silicon hardware to run, and results are measured on wall clock time rather than hardware performance counters, with no end-to-end speed impact demonstrated yet. It is aimed at developers curious about GPU kernel programming and Apple's MLX ecosystem.