metal: tensor-op GEMMs on the M5's neural accelerators (Metal 4 matmul2d)

THIRD_PARTY_NOTICES.md

@@ -55,7 +55,16 @@ applies to that file.
 weights, reduced to the one instantiation the `metal_gemm_time` probe launches as its reference arm. It is taken from
 [ml-explore/mlx](https://github.com/ml-explore/mlx) v0.32.1 (`mlx/backend/metal/kernels/quantized.h`,
 `steel/gemm/mma.h`, `steel/gemm/loader.h`); the file's header names what was taken out. It is compiled by the probe
-only and never enters the engine's kernels. MLX is under the MIT licence, Copyright (c) 2023 Apple Inc.:
+only and never enters the engine's kernels.
+
+`crates/metal/kernels/gemm_tensor.metal` (the prompt pass's GEMMs on Metal 4's tensor ops) follows the design of MLX's
+NAX quantized matmul, v0.32.1 (`mlx/backend/metal/kernels/quantized_nax.h`, `steel/gemm/nax.h`): the weights
+dequantized into threadgroup memory a K step, the 64-by-64 tile over a 2-by-2 grid of SIMD groups, the `matmul2d`
+descriptor with the weight operand transposed. It is written anew over this crate's own staging and quant owners and
+copies no MLX code; the rule that picks it, `MatrixUnits` in `crates/metal/src/library.rs`, is MLX's
+`is_nax_available` (`mlx/backend/metal/device.cpp`). Both files' comments name the MLX lines they follow.
+
+MLX is under the MIT licence, Copyright (c) 2023 Apple Inc.:
 
 > Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 > documentation files (the "Software"), to deal in the Software without restriction, including without limitation the

crates/metal/build.rs

@@ -1,7 +1,7 @@
 //! The build of the Metal backend: the kernels' one source, and the
 //! metallib when the Metal Toolchain is installed.
 //!
-//! Two outputs land in `OUT_DIR`:
+//! Three outputs land in `OUT_DIR`:
 //!
 //! - `bloomery.metal`: the kernels' source with the fault site's table
 //!   written in front of it, from `bloomery-device`'s table — the one owner
@@ -10,6 +10,9 @@
 //! - `bloomery.metallib`: the same source compiled by `xcrun metal`, when
 //!   the Metal Toolchain is installed. Its presence is the `metal_lib`
 //!   cfg, and the runtime loads it instead of the source.
+//! - `bloomery_tensor.metal`: the tensor-op GEMMs' source, the same table in
+//!   front of the files it reads, which the runtime compiles at Metal 4 on a
+//!   GPU that takes them (`library.rs`).
 
 use std::path::{Path, PathBuf};
 use std::process::Command;
@@ -262,6 +265,7 @@ fn main() {
         "probe",
         "norm",
         "quant",
+        "gemm_shared",
         "gemm",
         "matvec_q4k",
         "matvec_q6k",
@@ -289,6 +293,24 @@ fn main() {
     let source_path = out_dir.join("bloomery.metal");
     std::fs::write(&source_path, &combined).expect("the source is written");
 
+    // The tensor-op GEMMs read the fault's raise, the quant owners and the
+    // shared GEMM ABI, the engine source's own files, in its order.
+    let tensor: Vec<(&str, String)> = kernels
+        .iter()
+        .filter(|(name, _)| ["fault", "quant", "gemm_shared"].contains(name))
+        .cloned()
+        .chain(std::iter::once((
+            "gemm_tensor",
+            std::fs::read_to_string(kernels_dir.join("gemm_tensor.metal"))
+                .expect("kernels/gemm_tensor.metal is read"),
+        )))
+        .collect();
+    std::fs::write(
+        out_dir.join("bloomery_tensor.metal"),
+        combined_source(&header, &tensor),
+    )
+    .expect("the tensor source is written");
+
     match build_metallib(&source_path, &out_dir) {
         Some(_) => println!("cargo:rustc-cfg=metal_lib=\"metallib\""),
         None => println!("cargo:rustc-cfg=metal_lib=\"source\""),

crates/metal/kernels/experts.metal

@@ -48,7 +48,8 @@
 // in every row of the threadgroup, because every row reads every slot.
 //
 // This file comes after matvec_q6k.metal and quant.metal (the decode owners),
-// norm.metal (rms_scale, rms_usable) and fault.metal (raise).
+// norm.metal (rms_scale, rms_usable), gemm_shared.metal (pair_swiglu) and
+// fault.metal (raise).
 
 #include <metal_stdlib>
 
@@ -97,11 +98,6 @@ inline device const char * expert_base(bool shared_slot,
     return shared_slot ? shexp : stack;
 }
 
-// silu(gate) * up. No route weight: the fold multiplies by it.
-inline float pair_swiglu(float gate, float up) {
-    return gate / (1.0f + exp(-gate)) * up;
-}
-
 // True when the routed slot's id is past the stack. Slot 8 is the shared
 // expert and leaves `expert` at 0; the caller reads the shared tensors.
 inline bool pair_expert(device const int * ids, uint slot, uint n_experts,

crates/metal/kernels/gemm.metal

@@ -63,21 +63,16 @@
 // no word past the output's last row or column is written; the residual
 // epilogue reads, adds and writes back per element.
 //
-// This file comes after quant.metal: it reads the quant owners and the
-// fault site's table.
+// This file comes after quant.metal and gemm_shared.metal: it reads the
+// quant owners, the fault site's table, the tile's geometry, the argument
+// blocks and the activations' checked cast (gemm_stage_x).
 
 #include <metal_stdlib>
 
 using namespace metal;
 
 #pragma METAL fp contract(off)
 
-// The tile's geometry: the kernel's own constants, the ones the host's
-// launch module (kernels/gemm.rs) mirrors.
-#define GEMM_ROWS 64     // weight rows a threadgroup owns
-#define GEMM_ACTS 64     // activation rows a threadgroup owns
-#define GEMM_THREADS 128  // four SIMD groups
-
 // The 8x8 fragments one SIMD group owns along each side: two SIMD groups
 // share a side, in 8-row blocks sixteen rows apart.
 #define GEMM_FRAGS_A (GEMM_ACTS / 16)
@@ -87,67 +82,6 @@ static_assert(GEMM_THREADS == 4 * 32, "four SIMD groups of 32 threads");
 static_assert(GEMM_ROWS % 16 == 0 && GEMM_ACTS % 16 == 0,
               "the tile splits into 8x8 fragments, two SIMD groups to a side");
 
-// The K values one step covers: 32, a whole number of every owner's
-// 16-value staging chunks and half the smallest group (A4's 64).
-template <typename T>
-constexpr int gemm_bk() {
-    return 32;
-}
-
-// The staged rows' pitch in elements: the K step plus 16 bytes, steel's
-// padding.
-template <typename T>
-constexpr int gemm_lda() {
-    return gemm_bk<T>() + 16 / (int)sizeof(T);
-}
-
-// The threadgroup bytes one tile stages: the f32 arm's 18432, the f16
-// arm's 10240. The allocation itself lives in the entry kernel (a
-// threadgroup variable may be declared only there), and the tile body takes
-// it as a pointer.
-template <typename T>
-constexpr int gemm_shmem_bytes() {
-    return (GEMM_ROWS + GEMM_ACTS) * gemm_lda<T>() * (int)sizeof(T);
-}
-
-// The argument block of a plain GEMM, bound by value at slot 3.
-struct GemmArgs {
-    uint cols;     // the input columns K, a whole number of the type's blocks
-    uint rows;     // the weight rows N, the output's columns
-    uint acts;     // the activation rows R, the output's rows
-    uint x_stride; // the f32 elements between x's rows, at least cols
-    uint y_stride; // the f32 elements between y's rows, at least rows
-    uint layer;    // the layer a raise is recorded in
-};
-
-static_assert(sizeof(GemmArgs) == 24, "GemmArgs is 24 bytes: the host packs six words");
-
-// One segment of a segmented GEMM: the weight rows of one tensor, and where
-// they sit in the weights and in the output. The mat-vec's segment table's
-// own four fields (norm.metal's MatvecSeg), the GEMM's meaning: the output
-// offset is a column offset in the R-row output, and the first group is a
-// weight-tile threadgroup.
-struct GemmSeg {
-    uint rows;         // the segment's weight rows
-    uint wblocks;      // the segment's first row's block offset in the weights
-    uint out;          // the segment's first column in the output
-    uint first_group;  // the segment's first weight-tile threadgroup
-};
-
-// The argument block of a segmented GEMM, bound by value at slot 3.
-struct GemmSegArgs {
-    uint cols;
-    uint acts;
-    uint x_stride;
-    uint y_stride;
-    uint layer;
-    uint nseg;
-    GemmSeg segs[BLOOMERY_MATVEC_MAX_SEGMENTS];
-};
-
-static_assert(sizeof(GemmSegArgs) == 24 + 16 * BLOOMERY_MATVEC_MAX_SEGMENTS,
-              "GemmSegArgs is the header's six words and the segments' tables");
-
 // The float fragment multiply: a `vec<float, 2>` is the 8x8 simdgroup
 // matrix's per-lane storage, so the fragment aliases the matrix through
 // `thread_elements()` (mma.h's BaseMMAFrag::mma, the reinterpret is MLX's).
@@ -275,30 +209,12 @@ inline void gemm_tile(threadgroup char        * shmem,
         }
     }
 
-    // B: whole 16-value runs of the activation rows, loaded as scalars
-    // (no window alignment is assumed) and checked as float4s. The f16
-    // arm's cast is the activation quantizer's: a value that is not
-    // finite or past the half range (one test: not at most the half
-    // maximum in magnitude, which a NaN fails too) raises the
-    // quantizer's site and stages NaN, so the row's outputs are NaN —
-    // never a silent inf.
+    // B: whole 16-value runs of the activation rows, each through the
+    // checked cast (gemm_stage_x).
     auto stage_b = [&](int k0) {
         FOR_UNROLL(short b = 0; b < B_CHUNKS; ++b) {
             device const float * xr = b_src[b] + k0;
-            FOR_UNROLL(short g = 0; g < 4; ++g) {
-                const float4 v = float4(xr[4 * g], xr[4 * g + 1], xr[4 * g + 2], xr[4 * g + 3]);
-                if (sizeof(T) == 2) {
-                    const bool4 bad = !(fabs(v) <= 65504.f);
-                    if (any(bad)) {
-                        raise(faults, layer, BLOOMERY_FAULT_SITE_QUANT_COLUMN);
-                        b_dst[b][g] = vec<T, 4>(select(v, float4(NAN), bad));
-                    } else {
-                        b_dst[b][g] = vec<T, 4>(v);
-                    }
-                } else {
-                    b_dst[b][g] = vec<T, 4>(v);
-                }
-            }
+            gemm_stage_x<T>(xr, b_dst[b], faults, layer);
         }
     };
 

crates/metal/kernels/gemm_routed.metal

@@ -35,8 +35,9 @@
 // K loop's alone; the tile's height and its split over the SIMD groups do
 // not move it.
 //
-// This file comes after gemm.metal (the helpers and the instantiated tile
-// pieces), experts.metal (pair_swiglu) and route.metal (the table's shape).
+// This file comes after gemm_shared.metal (the routed tile's geometry, its
+// argument block, pair_swiglu and gemm_stage_x), gemm.metal (the helpers
+// and the instantiated tile pieces) and route.metal (the table's shape).
 
 #include <metal_stdlib>
 
@@ -69,8 +70,6 @@ struct F32W {
 
 GEMM_ENTRY(f32w_store_f32, F32W, float, false, gemm_kernel)
 
-#define ROUTED_ROWS 64   // weight rows of a tile
-#define ROUTED_ACTS 64   // activation rows of a tile
 #define ROUTED_FRAGS_A (ROUTED_ACTS / 8)
 #define ROUTED_FRAGS_W 2  // a SIMD group's 8-row weight fragments: its 16 weight rows
 
@@ -86,20 +85,6 @@ constexpr int routed_shmem_bytes() {
     return (ROUTED_ROWS + ROUTED_ACTS) * gemm_lda<T>() * (int)sizeof(T);
 }
 
-// Eight words, bound by value.
-struct RoutedArgs {
-    uint cols;          // the input columns K
-    uint x_stride;      // f32 elements between the activation rows
-    uint y_stride;      // f32 elements between the output rows
-    uint layer;
-    uint n_experts;     // the stack's experts; the table's group n_experts is shared
-    uint stack_blocks;  // blocks of one expert's rows in the stack
-    uint gather;        // 1: an activation row is the slot's pass row; 0: the place's
-    uint rows;          // outputs per expert: the pair's rows of gate, the down's rows
-};
-
-static_assert(sizeof(RoutedArgs) == 32, "RoutedArgs is eight words");
-
 // One tile's K loop and store over its first NA 8-row A fragments: the rows
 // under 8 NA, which hold the run. Every SIMD group multiplies the same NA
 // fragments against its own 16 weight rows, so every fragment the loop holds
@@ -160,20 +145,7 @@ inline void routed_tile(threadgroup T                  * sa,
         if (b_stage) {
             device const float * xr = b_src + k0;
             threadgroup vec<T, 4> * bd = (threadgroup vec<T, 4> *)((threadgroup T *)b_dst + at);
-            FOR_UNROLL(short g = 0; g < 4; ++g) {
-                const float4 v = float4(xr[4 * g], xr[4 * g + 1], xr[4 * g + 2], xr[4 * g + 3]);
-                if (sizeof(T) == 2) {
-                    const bool4 over = !(fabs(v) <= 65504.f);
-                    if (any(over)) {
-                        raise(faults, args.layer, BLOOMERY_FAULT_SITE_QUANT_COLUMN);
-                        bd[g] = vec<T, 4>(select(v, float4(NAN), over));
-                    } else {
-                        bd[g] = vec<T, 4>(v);
-                    }
-                } else {
-                    bd[g] = vec<T, 4>(v);
-                }
-            }
+            gemm_stage_x<T>(xr, bd, faults, args.layer);
         }
         const int pos = k0 + a_koff;
         const int ib = pos / Q::values_per_block();