ml-explore · pcuenca · Mar 18, 2026 · Mar 19, 2026 · Mar 19, 2026 · Mar 19, 2026
diff --git a/mlx_lm/models/activations.py b/mlx_lm/models/activations.py
@@ -11,6 +11,13 @@ def swiglu(gate, x):
     return nn.silu(gate) * x
 
 
+@partial(mx.compile, shapeless=True)
+def precise_swiglu(h, gate, x):
+    gate = nn.silu(gate.astype(mx.float32))
+    x = x.astype(mx.float32)
+    return (gate * x).astype(h.dtype)
+
+
 @partial(mx.compile, shapeless=True)
 def xielu(x, alpha_p, alpha_n, beta, eps):
     alpha_p = nn.softplus(alpha_p)