fix: Gemma 4 MoE Metal fusions and Qwen 3.6 GDN prefill pipelining

Libraries/MLXLLM/LLMModel.swift

@@ -40,17 +40,28 @@ extension LLMModel {
         var processed = 0
 
         // Prepare the prompt in chunks if larger than the prefill size.
-        // After each chunk, call the progress hook so the server can emit
-        // llama-server-style slot_update SSE events with real n_past.
+        // asyncEval lets the CPU build chunk N+1's graph while the GPU evaluates
+        // chunk N. Python mlx-lm gets this pipelining for free because its bindings
+        // defer eval until a value is read. The previous `eval(cache)` call was a
+        // blocking sync that drained the GPU pipeline between chunks.
         while y.tokens.size > prefillStepSize {
             let input = y[.newAxis, ..<prefillStepSize]
             _ = self(input, cache: cache.isEmpty ? nil : cache, state: nil)
-            eval(cache)
+
+            var cacheArrays: [MLXArray] = []
+            for c in cache { cacheArrays.append(contentsOf: c.innerState()) }
+            if !cacheArrays.isEmpty {
+                asyncEval(cacheArrays)
+            }
+
             y = y[prefillStepSize...]
             processed += prefillStepSize
             activePrefillProgressHook?(processed, totalTokens)
         }
 
+        // Single sync after the loop — flush any remaining async work.
+        eval(cache)
+
         return .tokens(y)
     }
 

Libraries/MLXLLM/Models/GatedDelta.swift

@@ -86,7 +86,7 @@ private func makeGatedDeltaKernel(hasMask: Bool) -> MLXFast.MLXFastKernel? {
             }
             for (int i = 0; i < n_per_t; ++i) {
               auto s_idx = n_per_t * dk_idx + i;
-              o_state[s_idx] = static_cast<InT>(state[i]);
+              o_state[s_idx] = static_cast<StT>(state[i]);
             }
         """
 
@@ -135,6 +135,7 @@ func gatedDeltaKernel(
     let Hv = v.dim(2)
     let Dv = v.dim(3)
     let inputType = q.dtype
+    let stateType = state.dtype
 
     let selectedKernel: MLXFast.MLXFastKernel?
     var inputs: [MLXArray] = [q, k, v, g, beta, state, MLXArray(T)]
@@ -153,6 +154,7 @@ func gatedDeltaKernel(
         inputs,
         template: [
             ("InT", inputType),
+            ("StT", stateType),
             ("Dk", Dk),
             ("Dv", Dv),
             ("Hk", Hk),
@@ -161,7 +163,7 @@ func gatedDeltaKernel(
         grid: (32, Dv, B * Hv),
         threadGroup: (32, 4, 1),
         outputShapes: [[B, T, Hv, Dv], state.shape],
-        outputDTypes: [inputType, inputType]
+        outputDTypes: [inputType, stateType]
     )
 
     return (outputs[0], outputs[1])
@@ -287,7 +289,12 @@ func gatedDeltaUpdate(
     let Hv = v.dim(2)
     let Dv = v.dim(3)
 
-    let state = state ?? MLXArray.zeros([B, Hv, Dv, Dk], dtype: q.dtype)
+    // State kept in fp32 to match Python mlx-lm. Using q.dtype (bf16) loses
+    // precision across T-step recurrence, compounding rounding error.
+    var state = state ?? MLXArray.zeros([B, Hv, Dv, Dk], dtype: .float32)
+    if state.dtype != .float32 {
+        state = state.asType(.float32)
+    }
 
     let isCPU = Device.defaultDevice().deviceType == .cpu
     if !isCPU, GatedDeltaKernelManager.shared.kernel != nil {

Libraries/MLXLLM/Models/Gemma4Text.swift

@@ -390,7 +390,7 @@ private class Gemma4MLP: Module {
     }
 
     func callAsFunction(_ x: MLXArray) -> MLXArray {
-        downProj(geluApproximate(gateProj(x)) * upProj(x))
+        downProj(safeGeluApproximate(gateProj(x)) * upProj(x))
     }
 }
 
@@ -454,7 +454,7 @@ private class Gemma4TextExperts: Module {
             inputDims: config.hiddenSize,
             hiddenDims: moeIntermediateSize,
             numExperts: numExperts,
-            activation: geluApproximate,
+            activation: safeGeluApproximate,
             bias: false
         )
         super.init()
@@ -596,7 +596,7 @@ private class Gemma4DecoderLayer: Module {
         {
             let residual3 = out
             var g = gate(out)
-            g = geluApproximate(g)
+            g = safeGeluApproximate(g)
             g = g * perLayerInput
             g = proj(g)
             g = norm(g)
@@ -680,7 +680,7 @@ private class Gemma4TextModelInner: Module {
         cache: [KVCache]? = nil
     ) -> MLXArray {
         let inputEmbeddings = embedTokens(inputs)
-        var h = inputEmbeddings * embedScale
+        var h = inputEmbeddings * MLXArray(embedScale, dtype: inputEmbeddings.dtype)
 
         // Compute per-layer inputs (PLE)
         var perLayerInputs: [MLXArray?]
@@ -692,7 +692,7 @@ private class Gemma4TextModelInner: Module {
             // Token-based PLE
             let tokenPLE =
                 embedPerLayer(inputs)
-                * Float(config.hiddenSizePerLayerInput).squareRoot()
+                * MLXArray(Float(config.hiddenSizePerLayerInput).squareRoot(), dtype: h.dtype)
 
             // [B, L, numLayers * hiddenSizePerLayerInput] -> [B, L, numLayers, hiddenSizePerLayerInput]
             let reshapedTokenPLE = tokenPLE.reshaped(
@@ -706,7 +706,7 @@ private class Gemma4TextModelInner: Module {
             let normedModelPLE = projNorm(modelPLE)
 
             // Combine: (model_proj + token_embed) * 2^{-0.5}
-            let perLayerInputScale = pow(Float(2.0), -0.5)
+            let perLayerInputScale = MLXArray(pow(Float(2.0), -0.5), dtype: h.dtype)
             let combined = (normedModelPLE + reshapedTokenPLE) * perLayerInputScale
 
             perLayerInputs = (0 ..< config.numHiddenLayers).map { i in
@@ -796,7 +796,8 @@ public class Gemma4TextModel: Module, LLMModel, KVCacheDimensionProvider {
         } else {
             out = model.embedTokens.asLinear(out)
         }
-        out = tanh(out / config.finalLogitSoftcapping) * config.finalLogitSoftcapping
+        let cap = MLXArray(config.finalLogitSoftcapping, dtype: out.dtype)
+        out = compiledSoftcap(x: out, cap: cap)
         return out
     }
 

Libraries/MLXLLM/Models/Qwen3MoE.swift

@@ -106,7 +106,12 @@ class Qwen3MoEMLP: Module, UnaryLayer {
     }
 
     public func callAsFunction(_ x: MLXArray) -> MLXArray {
-        down(silu(gate(x)) * up(x))
+        let g = silu(gate(x))
+        let u = up(x)
+        // Prevent Metal from auto-promoting float16 element-wise ops to float32.
+        // bfloat16 is natively supported and avoids the extra memory bandwidth cost.
+        let product = g.dtype == .float16 ? g.asType(.bfloat16) * u.asType(.bfloat16) : g * u
+        return down(product)
     }
 }
 

Libraries/MLXLLM/Models/Qwen3Next.swift

@@ -16,6 +16,14 @@ func sigmoidMultiply(_ x: MLXArray, _ gate: MLXArray) -> MLXArray {
     x * sigmoid(gate)
 }
 
+private let _compiledSwiGLUFn = MLX.compile(shapeless: true) { (args: [MLXArray]) -> [MLXArray] in
+    [silu(args[0]) * args[1]]
+}
+
+private func compiledSwiGLU(_ gate: MLXArray, _ x: MLXArray) -> MLXArray {
+    _compiledSwiGLUFn([gate, x])[0]
+}
+
 // MARK: - Model Components
 
 final class Qwen3NextRMSNormGated: Module {
@@ -136,7 +144,7 @@ final class Qwen3NextMLP: Module, UnaryLayer {
     }
 
     func callAsFunction(_ x: MLXArray) -> MLXArray {
-        downProj(silu(gateProj(x)) * upProj(x))
+        downProj(compiledSwiGLU(gateProj(x), upProj(x)))
     }
 }
 

Libraries/MLXLMCommon/Optimizations.swift

@@ -0,0 +1,39 @@
+import Foundation
+import MLX
+import MLXNN
+
+/// Fixes an MLXNN Power primitive compilation crash on M1/M2/M3 GPUs.
+/// The original `geluApproximate` uses `x ** 3`, which maps to the Power primitive.
+/// Power lacks `output_shapes` under `compile(shapeless: true)` and returns 0 outputs on compilation,
+/// causing a crash when `compileState.call([a])[0]` is executed.
+/// This safe replacement uses `x * x * x` which maps to Multiply (supported natively).
+///
+/// Declared as `compile(shapeless: true)` so MLX caches and reuses the compiled Metal kernel
+/// across all call sites instead of re-tracing the graph on every invocation.
+public let safeGeluApproximate: @Sendable (MLXArray) -> MLXArray =
+    compile(shapeless: true) { (x: MLXArray) -> MLXArray in
+        let xFloat = x.asType(.float32)
+        let half = MLXArray(0.5, dtype: .float32)
+        let one = MLXArray(1.0, dtype: .float32)
+        let c1 = MLXArray(Float(sqrt(2 / Float.pi)), dtype: .float32)
+        let c2 = MLXArray(0.044715, dtype: .float32)
+        let out = half * xFloat * (one + tanh(c1 * (xFloat + c2 * xFloat * xFloat * xFloat)))
+        return out.asType(x.dtype)
+    }
+
+/// Module wrapper for `safeGeluApproximate` — drop-in for `GELU(approximation: .precise)`.
+public class SafeGELU: Module, UnaryLayer {
+    public override init() { super.init() }
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        safeGeluApproximate(x)
+    }
+}
+
+/// Fuses logit softcap operations (`tanh(x / cap) * cap`) into a single Metal dispatch.
+private let _compiledSoftcapFn = MLX.compile(shapeless: true) { (args: [MLXArray]) -> [MLXArray] in
+    [tanh(args[0] / args[1]) * args[1]]
+}
+
+public func compiledSoftcap(x: MLXArray, cap: MLXArray) -> MLXArray {
+    return _compiledSoftcapFn([x, cap])[0]
+}

Libraries/MLXLMCommon/SwitchLayers.swift

@@ -4,6 +4,40 @@ import MLXNN
 
 // Port of https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/switch_layers.py
 
+// Compiled activation kernels — fuses gate activation + element-wise multiply into
+// a single Metal dispatch per expert batch. Matches Python's @partial(mx.compile, shapeless=True).
+// This eliminates the separate Metal kernel launch for the activation and the multiply,
+// halving the number of dispatches in every MoE expert invocation.
+private let _compiledSwiGLUFn = MLX.compile(shapeless: true) { (args: [MLXArray]) -> [MLXArray] in
+    [silu(args[0]) * args[1]]
+}
+
+private let _compiledGeGLUFn = MLX.compile(shapeless: true) { (args: [MLXArray]) -> [MLXArray] in
+    let gate = args[0].asType(.float32)
+    let x = args[1].asType(.float32)
+    let half = MLXArray(0.5, dtype: .float32)
+    let one = MLXArray(1.0, dtype: .float32)
+    let c1 = MLXArray(Float(sqrt(2 / Float.pi)), dtype: .float32)
+    let c2 = MLXArray(0.044715, dtype: .float32)
+    let out = (half * gate * (one + tanh(c1 * (gate + c2 * gate * gate * gate)))) * x
+    return [out.asType(args[0].dtype)]
+}
+
+@inline(__always)
+public func compiledSwiGLU(_ gate: MLXArray, _ x: MLXArray) -> MLXArray {
+    _compiledSwiGLUFn([gate, x])[0]
+}
+
+@inline(__always)
+public func compiledGeGLU(_ gate: MLXArray, _ x: MLXArray) -> MLXArray {
+    _compiledGeGLUFn([gate, x])[0]
+}
+
+@inline(__always)
+private func preservePrecisionCast(_ x: MLXArray) -> MLXArray {
+    x.dtype == .float16 ? x.asType(.bfloat16) : x
+}
+
 public func gatherSort(x: MLXArray, indices: MLXArray) -> (MLXArray, MLXArray, MLXArray) {
     let m = indices.dim(-1)
     let indices = indices.flattened()
@@ -44,6 +78,11 @@ public class SwitchGLU: Module, @unchecked Sendable {
     let hiddenDims: Int
     let numExperts: Int
     let activation: (MLXArray) -> MLXArray
+    // Compiled fast-path flags — detected once at init via test inference.
+    // When true the corresponding file-scope compiled closure is used instead
+    // of the generic activation call, saving one extra Metal kernel dispatch.
+    let isSiluActivation: Bool
+    let isGeluActivation: Bool
 
     // ── Async pipeline state (SSD streaming optimization) ──
     // Persistent buffers: allocated once per layer, reused across tokens.
@@ -66,6 +105,18 @@ public class SwitchGLU: Module, @unchecked Sendable {
         self.numExperts = numExperts
         self.activation = activation
 
+        // Detect common activation types for compiled fast path.
+        // We run a test forward pass to compare outputs numerically rather than
+        // using pointer equality (closures cannot be compared directly in Swift).
+        // Use a scalar to keep this essentially free at init time.
+        let testInput = MLXArray([Float(1.0)])
+        let testOutput = activation(testInput)
+        let siluOut = silu(testInput)
+        // safeGeluApproximate value at x=1: 0.5 * 1 * (1 + tanh(sqrt(2/π) * (1 + 0.044715)))
+        let geluOut = (0.5 as Float) * testInput * (1 + tanh(sqrt(2 / Float.pi) * (testInput + 0.044715 * testInput * testInput * testInput)))
+        self.isSiluActivation = (testOutput .== siluOut).all().item(Bool.self)
+        self.isGeluActivation = !self.isSiluActivation && (testOutput .== geluOut).all().item(Bool.self)
+
         self._gateProj.wrappedValue = SwitchLinear(
             inputDims: inputDims, outputDims: hiddenDims, numExperts: numExperts, bias: bias)
         self._upProj.wrappedValue = SwitchLinear(
@@ -202,7 +253,7 @@ public class SwitchGLU: Module, @unchecked Sendable {
                     let usedDown = Array(_persistentDown![0..<ranges.count])
                     let xGate = qGate.computeExperts(x, buffers: usedGate, ranges: ranges)
                     let xUp = qUp.computeExperts(x, buffers: usedUp, ranges: ranges)
-                    let intermediate = activation(xGate) * xUp
+                    let intermediate = activation(preservePrecisionCast(xGate)) * preservePrecisionCast(xUp)
                     x = qDown.computeExperts(intermediate, buffers: usedDown, ranges: ranges)
 
                 } else {
@@ -354,7 +405,8 @@ public class SwitchGLU: Module, @unchecked Sendable {
                     // Lazy compute (no eval — next layer forces it)
                     let xGate = qGate.computeExperts(x, buffers: usedGate, ranges: ranges)
                     let xUp = qUp.computeExperts(x, buffers: usedUp, ranges: ranges)
-                    let intermediate = activation(xGate) * xUp
+                    var intermediate = activation(preservePrecisionCast(xGate)) * preservePrecisionCast(xUp)
+                    if intermediate.dtype != xGate.dtype { intermediate = intermediate.asType(xGate.dtype) }
                     x = qDown.computeExperts(intermediate, buffers: usedDown, ranges: ranges)
                 }
 
@@ -414,7 +466,8 @@ public class SwitchGLU: Module, @unchecked Sendable {
                 // Lazy compute (no eval — next layer forces it)
                 let xGate = qGate.computeExperts(x, buffers: gateBuffers, ranges: ranges)
                 let xUp = qUp.computeExperts(x, buffers: upBuffers, ranges: ranges)
-                let intermediate = activation(xGate) * xUp
+                var intermediate = activation(preservePrecisionCast(xGate)) * preservePrecisionCast(xUp)
+                if intermediate.dtype != xGate.dtype { intermediate = intermediate.asType(xGate.dtype) }
                 x = qDown.computeExperts(intermediate, buffers: downBuffers, ranges: ranges)
             }
 
@@ -424,11 +477,24 @@ public class SwitchGLU: Module, @unchecked Sendable {
             return MLX.squeezed(x, axis: -2)
         }
 
-        // ── Fallback: original sequential path (non-SSD or non-quantized) ──
+        // ── Standard path: non-SSD or non-quantized ─────────────────────────
+        // Use compiled fused activation when possible: saves one Metal kernel
+        // dispatch per expert batch vs calling activation(gate) then multiplying.
+        // Note: bfloat16 promotion is handled inside the compiled closures since
+        // both inputs flow through gatherMM which preserves model weight dtype.
         let xUp = upProj(x, idx, sortedIndices: doSort)
         let xGate = gateProj(x, idx, sortedIndices: doSort)
+        var intermediate: MLXArray
+        if isSiluActivation {
+            intermediate = compiledSwiGLU(preservePrecisionCast(xGate), preservePrecisionCast(xUp))
+        } else if isGeluActivation {
+            intermediate = compiledGeGLU(preservePrecisionCast(xGate), preservePrecisionCast(xUp))
+        } else {
+            intermediate = activation(preservePrecisionCast(xGate)) * preservePrecisionCast(xUp)
+        }
+        if intermediate.dtype != xGate.dtype { intermediate = intermediate.asType(xGate.dtype) }
         x = downProj(
-            activation(xGate) * xUp,
+            intermediate,
             idx,
             sortedIndices: doSort)
 

Libraries/MLXVLM/Models/Gemma4.swift

@@ -495,7 +495,7 @@ private final class Gemma4TextMLP: Module, UnaryLayer {
     }
 
     func callAsFunction(_ x: MLXArray) -> MLXArray {
-        downProj(geluApproximate(gateProj(x)) * upProj(x))
+        downProj(safeGeluApproximate(gateProj(x)) * upProj(x))
     }
 }
 
@@ -555,7 +555,7 @@ private final class Gemma4TextExperts: Module {
             inputDims: config.hiddenSize,
             hiddenDims: moeIntermediateSize,
             numExperts: numExperts,
-            activation: geluApproximate,
+            activation: safeGeluApproximate,
             bias: false
         )
         super.init()
@@ -850,7 +850,7 @@ private final class Gemma4TextDecoderLayer: Module {
         {
             residual = h
             var gated = perLayerInputGate(h)
-            gated = geluApproximate(gated)
+            gated = safeGeluApproximate(gated)
             gated = gated * perLayerInput
             gated = perLayerProjection(gated)
             gated = postPerLayerInputNorm(gated)
@@ -1130,8 +1130,8 @@ private final class Gemma4TextLanguageModel: Module, KVCacheDimensionProvider {
             logits = model.embedTokens.asLinear(output)
         }
         if let finalLogitSoftcapping, finalLogitSoftcapping > 0 {
-            let scale = MLXArray(finalLogitSoftcapping)
-            return LMOutput(logits: tanh(logits / scale) * scale)
+            let scale = MLXArray(finalLogitSoftcapping, dtype: logits.dtype)
+            return LMOutput(logits: compiledSoftcap(x: logits, cap: scale))
         }
         return LMOutput(logits: logits)
     }
@@ -1384,7 +1384,7 @@ private final class Gemma4VisionMLP: Module, UnaryLayer {
     }
 
     func callAsFunction(_ x: MLXArray) -> MLXArray {
-        downProj(geluApproximate(gateProj(x)) * upProj(x))
+        downProj(safeGeluApproximate(gateProj(x)) * upProj(x))
     }
 }