perf(engine/metal): gelu-fused down-projection kernels (#341 phase 1, capability)

lthn_gelu_qmv + lthn_gather_qmv_gelu fold gemma's MLP gate into the down
projection's x-load: each element computes T(gelu_approx(gate)*up) AT
LOAD — the exact expression, type flow (bfloat-typed partial sums before
widening, per-width x_thread scalings) and single bf16 rounding of the
lthn_gelu_gate_mul dispatch it replaces — so the values entering qdot
are byte-identical to the chain's gated buffer and the gelu dispatch
plus its barrier hop disappear from both MoE chains (local down: qmv
shape; routed expert down: the fc-specialised gather with gate/up riding
the same lhs indexing).

lthn_gelu_qmv_impl.h is a verbatim structural copy of quantized.h's
load_vector/load_vector_safe/qmv_impl with only the element reads
swapped; both kernels share it, so the two dispatch shapes cannot drift.
Instantiated gs {32,64,128} x bits {4,8} (the local/expert down widths).

Kernels-only slice: compiled into lthn_kernels.metallib (28 kernels),
unreferenced until the host wiring lands with its byte-parity oracle and
real-model receipts (#341 task carries the emit map).

Co-Authored-By: Virgil <virgil@lethean.io>

Author: Snider

go/engine/metal/kernels/lthn_gather_qmv.metal

@@ -93,6 +93,60 @@ template <typename T, int group_size, int bits>
       w, scales, biases, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
+#include "lthn_gelu_qmv_impl.h"
+
+// lthn_gather_qmv_gelu — the routed experts' DOWN gather with the MLP gate
+// fused into its x-load (#341 phase 1): each route reads its expert's gate/up
+// activation rows and computes gelu(gate)·up at load, byte-identical to the
+// chain's gated buffer (lthn_gelu_qmv_impl.h), so the expert gelu dispatch and
+// its barrier disappear. Same fc prologue as lthn_gather_qmv; gate and up ride
+// the same lhs indexing.
+template <typename T, int group_size, int bits>
+[[kernel]] void lthn_gather_qmv_gelu(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* gate [[buffer(3)]],
+    const device T* up [[buffer(4)]],
+    const device uint32_t* lhs_indices [[buffer(5)]],
+    const device uint32_t* rhs_indices [[buffer(6)]],
+    device T* y [[buffer(7)]],
+    const constant int& in_vec_size [[buffer(8)]],
+    const constant int& out_vec_size [[buffer(9)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  const int64_t w_idx = rhs_indices[tid.z];
+  w += w_idx * lthn_gather_expert_rows * (int64_t(in_vec_size) * bits / 32);
+  const int64_t sb_stride = int64_t(lthn_gather_expert_rows) * (in_vec_size / group_size);
+  scales += w_idx * sb_stride;
+  biases += w_idx * sb_stride;
+  if (lthn_gather_batched_x) {
+    const int64_t xo = int64_t(lhs_indices[tid.z]) * in_vec_size;
+    gate += xo;
+    up += xo;
+  }
+  y += tid.z * out_vec_size;
+  qmv_gelu_impl<T, group_size, bits>(
+      w, scales, biases, gate, up, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+#define instantiate_lthn_gather_qmv_gelu(group_size, bits)                          \
+  template [[host_name("lthn_gather_qmv_gelu_bfloat16_t_gs_" #group_size           \
+                       "_b_" #bits)]] [[kernel]] void                              \
+  lthn_gather_qmv_gelu<bfloat16_t, group_size, bits>(                              \
+      const device uint32_t*, const device bfloat16_t*, const device bfloat16_t*,  \
+      const device bfloat16_t*, const device bfloat16_t*, const device uint32_t*,  \
+      const device uint32_t*, device bfloat16_t*,                                  \
+      const constant int&, const constant int&, uint3, uint, uint);
+
+instantiate_lthn_gather_qmv_gelu(32, 4)
+instantiate_lthn_gather_qmv_gelu(32, 8)
+instantiate_lthn_gather_qmv_gelu(64, 4)
+instantiate_lthn_gather_qmv_gelu(64, 8)
+instantiate_lthn_gather_qmv_gelu(128, 4)
+instantiate_lthn_gather_qmv_gelu(128, 8)
+
 #define instantiate_lthn_gather_qmv(group_size, bits)                              \
   template [[host_name("lthn_gather_qmv_fast_bfloat16_t_gs_" #group_size           \
                        "_b_" #bits)]] [[kernel]] void                              \

go/engine/metal/kernels/lthn_gelu_qmv.metal

@@ -0,0 +1,49 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+// lthn_gelu_qmv — the MoE local expert's DOWN projection with the MLP gate
+// fused into its x-load (#341 phase 1): y = dequant(W_down) · gelu(gate)·up,
+// one dispatch where the chain ran lthn_gelu_gate_mul + affine_qmv with a
+// dependency hop between them. Values are byte-identical to the chain (see
+// lthn_gelu_qmv_impl.h); the gelu dispatch and its barrier disappear.
+//
+// Same include chain as MLX's quantized.metal (built with -I external/mlx).
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/quantized_utils.h"
+#include "mlx/backend/metal/kernels/quantized.h"
+// clang-format on
+
+#include "lthn_gelu_qmv_impl.h"
+
+template <typename T, int group_size, int bits>
+[[kernel]] void lthn_gelu_qmv(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* gate [[buffer(3)]],
+    const device T* up [[buffer(4)]],
+    device T* y [[buffer(5)]],
+    const constant int& in_vec_size [[buffer(6)]],
+    const constant int& out_vec_size [[buffer(7)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  qmv_gelu_impl<T, group_size, bits>(
+      w, scales, biases, gate, up, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+#define instantiate_lthn_gelu_qmv(group_size, bits)                              \
+  template [[host_name("lthn_gelu_qmv_bfloat16_t_gs_" #group_size               \
+                       "_b_" #bits)]] [[kernel]] void                           \
+  lthn_gelu_qmv<bfloat16_t, group_size, bits>(                                  \
+      const device uint32_t*, const device bfloat16_t*, const device bfloat16_t*, \
+      const device bfloat16_t*, const device bfloat16_t*, device bfloat16_t*,   \
+      const constant int&, const constant int&, uint3, uint, uint);
+
+instantiate_lthn_gelu_qmv(32, 4)
+instantiate_lthn_gelu_qmv(32, 8)
+instantiate_lthn_gelu_qmv(64, 4)
+instantiate_lthn_gelu_qmv(64, 8)
+instantiate_lthn_gelu_qmv(128, 4)
+instantiate_lthn_gelu_qmv(128, 8)