Use fast (#38)

* update to latest mlx swift and use fast norms
* gpu usage -> memory usage

Libraries/LLM/Gemma.swift

@@ -9,7 +9,6 @@ import MLXNN
 
 // specialized norm for gemma
 private class RMSNorm: Module, UnaryLayer {
-
     let weight: MLXArray
     let eps: Float
 
@@ -19,13 +18,8 @@ private class RMSNorm: Module, UnaryLayer {
         super.init()
     }
 
-    func norm(_ x: MLXArray) -> MLXArray {
-        x * rsqrt(x.square().mean(axis: -1, keepDims: true) + eps)
-    }
-
     public func callAsFunction(_ x: MLXArray) -> MLXArray {
-        let output = norm(x.asType(Float.self)).asType(x.dtype)
-        return (1 + weight) * output
+        return MLXFast.rmsNorm(x, weight: 1.0 + self.weight, eps: self.eps)
     }
 }
 

Libraries/LLM/Phi.swift

@@ -7,12 +7,6 @@ import MLXNN
 
 // https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/phi.py
 
-private class LayerNorm: MLXNN.LayerNorm {
-    override func callAsFunction(_ x: MLXArray) -> MLXArray {
-        super.callAsFunction(x.asType(Float.self)).asType(x.dtype)
-    }
-}
-
 private class PhiAttention: Module {
 
     let args: PhiConfiguration