metalworkingGitHub ↗
/mlx/steel

Steel

Steel is MLX's kernel library, the CUTLASS of Apple Silicon: GEMM and attention decomposed into composable C++ templates (BlockLoader, BlockMMA, epilogues) that one wiring kernel assembles per shape.

CUDA equivalent: CUTLASS, faithfully. Same decomposition philosophy, same "the kernel is just wiring" endgame, adapted to simdgroup_matrix fragments and 32 KB threadgroup memory. If you've read CUTLASS's Gemm hierarchy, steel reads like a haiku version of it.

Root: mlx/backend/metal/kernels/steel/↗. The components, each with its own case study:

  • BlockLoader (gemm/loader.h): the cooperative load as a template. It derives each thread's slice of the tile at compile time from the threadgroup size, with load_unsafe (vectorized, unchecked) and load_safe (bounds-checked, zero-filled) variants that function constants select between.
  • BlockMMA (gemm/mma.h): the register-blocked multiply. A TM × TN grid of 8×8 fragments per simdgroup, marched along K.
  • Epilogues (gemm/transforms.h): the ending as a plug-in type (TransformAdd, TransformAxpby), applied while results are still in registers.
  • The wiring kernels (gemm/kernels/): steel_gemm_fused.h plus siblings that reuse the same components for other shapes: steel_gemm_splitk.h (huge-K reductions), steel_gemm_masked.h (block-sparse), steel_gemm_gather.h (MoE gather-GEMM), steel_gemm_segmented.h.
  • The attention fork (attn/): flash attention needs its middle matrix transformed in place between two matmuls, so attn/ carries its own mma.h/loader.h with a different fragment layout. Don't mix the two subsystems up when reading.
  • The NAX fork (gemm/nax.h, attn/nax.h, quantized_nax.metal, fp_quantized_nax.metal, gemm_nax.h): a parallel steel targeting the M5 neural accelerators through Metal 4 tensor ops, not a variant file. Diffing a _nax kernel against its plain sibling shows exactly what the new hardware changes: the MMA layer.

THE COMPONENTS (steel/gemm/) BlockLoader loader.h · cooperative load as a template load_unsafe / load_safe / next() BlockMMA mma.h · TM×TN fragment grid, K-march mma() / store_result() Epilogue transforms.h · the ending as a plug-in type TransformAdd / TransformAxpby steel_gemm_fused.h · the wiring kernel ~350 lines, mostly function-constant dispatch: the least interesting file in the library, which is the achievement steel_gemm_splitk huge-K: split, then combine steel_gemm_masked block-sparse steel_gemm_gather MoE gather-GEMM gemm_nax twins Metal-4/M5 tensor ops: only the MMA layer moves same components, different wiring: the whole argument for the decomposition THE ATTENTION FORK (steel/attn/) its own loader.h + mma.h different fragment layout steel_attention.h flash attention from the same vocabulary why the fork: a GEMM accumulator only accumulates; the score tile must be read and transformed in place between two matmuls. same philosophy, new geometry CUDA equivalent: CUTLASS · unlike CUTLASS, steel is not consumable standalone: its ideas travel by reading

Three components, many wirings. Every sibling kernel at the bottom reuses the same green boxes; the attention fork on the right is the one place the geometry itself had to change.

Which instantiation runs is decided host-side: the tile-selection macro maps chip class and problem shape to BM/BN/BK/WM/WN.

Two facts frame steel's place in the ecosystem. First, it's the thing to beat, and it usually wins: multiple independent projects report hand-written replacement kernels coming back 0.5-0.8× stock steel (the failures), so the profitable moves are unlocking its fast paths, not replacing it. Second, it's not consumable standalone: unlike CUTLASS there's no packaged way to use steel outside MLX, so its ideas travel by reading, which is what this glossary's case studies are for.

Edit this page on GitHub ↗