Add global scale support to quantized layers

Source/MLXNN/Quantized.swift

@@ -27,6 +27,12 @@ public protocol Quantized: Module {
     var mode: QuantizationMode { get }
 }
 
+private func applyNVFP4GlobalScale(_ value: MLXArray, globalScale: MLXArray?) -> MLXArray {
+    guard let globalScale else { return value }
+    // NVFP4 encodes its E4M3 group scales with (448 * 6) / globalScale.
+    return (value * (globalScale / (448 * 6))).asType(value.dtype)
+}
+
 /// Quantize any ``Quantizable`` layer that is not already quantized.
 public func quantizeSingle(
     layer: Module, groupSize: Int = 64, bits: Int = 4, mode: QuantizationMode = .affine
@@ -159,6 +165,7 @@ open class QuantizedEmbedding: Embedding, Quantized {
     public let mode: QuantizationMode
     public let scales: MLXArray
     public let biases: MLXArray?
+    @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray?
 
     open override var shape: (Int, Int) {
         let (embeddingCount, dimensions) = super.shape
@@ -184,14 +191,16 @@ open class QuantizedEmbedding: Embedding, Quantized {
 
     public init(
         weight: MLXArray, groupSize: Int = 64, bits: Int = 4,
-        mode: QuantizationMode = .affine
+        mode: QuantizationMode = .affine,
+        globalScale: MLXArray? = nil
     ) {
         self.groupSize = groupSize
         self.bits = bits
         self.mode = mode
+        self._globalScale.wrappedValue = globalScale
 
         let (quantizedWeight, scales, biases) = MLX.quantized(
-            weight, groupSize: groupSize, bits: bits, mode: mode)
+            weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale)
 
         self.scales = scales
         self.biases = biases
@@ -201,19 +210,36 @@ open class QuantizedEmbedding: Embedding, Quantized {
         self.freeze()
     }
 
+    /// Initializer meant for subclasses to provide arrays directly.
+    public init(
+        weight: MLXArray, scales: MLXArray, biases: MLXArray?,
+        groupSize: Int, bits: Int,
+        mode: QuantizationMode = .affine,
+        globalScale: MLXArray? = nil
+    ) {
+        self.groupSize = groupSize
+        self.bits = bits
+        self.mode = mode
+        self.scales = scales
+        self.biases = biases
+        self._globalScale.wrappedValue = globalScale
+        super.init(weight: weight)
+    }
+
     open override func callAsFunction(_ x: MLXArray) -> MLXArray {
         let s = x.shape
         let x = x.flattened()
         let out = dequantized(
             weight[x], scales: scales[x], biases: biases == nil ? nil : biases![x],
             groupSize: groupSize, bits: bits, mode: mode)
-        return out.reshaped(s + [-1])
+        return applyNVFP4GlobalScale(out, globalScale: globalScale).reshaped(s + [-1])
     }
 
     open override func asLinear(_ x: MLXArray) -> MLXArray {
-        quantizedMM(
+        let result = quantizedMM(
             x, weight, scales: scales, biases: biases, transpose: true, groupSize: groupSize,
             bits: bits, mode: mode)
+        return applyNVFP4GlobalScale(result, globalScale: globalScale)
     }
 }
 
@@ -243,6 +269,7 @@ open class QuantizedLinear: Linear, Quantized {
     public let mode: QuantizationMode
     public let scales: MLXArray
     public let biases: MLXArray?
+    @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray?
 
     open override var shape: (Int, Int) {
         let shape = weight.shape2
@@ -292,14 +319,16 @@ open class QuantizedLinear: Linear, Quantized {
     /// Initialize a ``QuantizedLinear`` with non-quantized weights and bias.
     public init(
         weight: MLXArray, bias: MLXArray?, groupSize: Int = 64, bits: Int = 4,
-        mode: QuantizationMode = .affine
+        mode: QuantizationMode = .affine,
+        globalScale: MLXArray? = nil
     ) {
         self.groupSize = groupSize
         self.bits = bits
         self.mode = mode
+        self._globalScale.wrappedValue = globalScale
 
         let (quantizedWeight, scales, biases) = MLX.quantized(
-            weight, groupSize: groupSize, bits: bits, mode: mode)
+            weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale)
 
         self.scales = scales
         self.biases = biases
@@ -316,13 +345,15 @@ open class QuantizedLinear: Linear, Quantized {
     public init(
         weight: MLXArray, bias: MLXArray? = nil, scales: MLXArray, biases: MLXArray?,
         groupSize: Int, bits: Int,
-        mode: QuantizationMode = .affine
+        mode: QuantizationMode = .affine,
+        globalScale: MLXArray? = nil
     ) {
         self.groupSize = groupSize
         self.bits = bits
         self.mode = mode
         self.scales = scales
         self.biases = biases
+        self._globalScale.wrappedValue = globalScale
         super.init(weight: weight, bias: bias)
     }
 
@@ -334,7 +365,7 @@ open class QuantizedLinear: Linear, Quantized {
     }
 
     open override func callAsFunction(_ x: MLXArray) -> MLXArray {
-        var x = quantizedMM(
+        var result = quantizedMM(
             x,
             weight,
             scales: scales,
@@ -344,10 +375,11 @@ open class QuantizedLinear: Linear, Quantized {
             bits: bits,
             mode: mode
         )
+        result = applyNVFP4GlobalScale(result, globalScale: globalScale)
         if let bias {
-            x = x + bias
+            result = result + bias
         }
-        return x
+        return result
     }
 
     /// Returns a QuantizedLinear layer that applies the same linear transformation up to the quantization error.

Tests/MLXTests/QuantizationTests.swift

@@ -6,6 +6,74 @@ import MLXNN
 import XCTest
 
 class QuantizationTests: XCTestCase {
+    func testQuantizedLinearAppliesNVFP4GlobalScaleOnMetal() {
+        let (input, quantizedWeight, scales, biases, globalScale, bias, expected) =
+            Device.withDefaultDevice(.cpu) {
+                MLXRandom.seed(7)
+                let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32])
+                let input = MLXRandom.uniform(low: -1, high: 1, [4, 32])
+                let globalScale = MLX.max(abs(weight)).asType(.float32)
+                let bias = MLXRandom.uniform(low: -1, high: 1, [32])
+                let (quantizedWeight, scales, biases) = MLX.quantized(
+                    weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
+                let expected =
+                    matmul(
+                        input,
+                        dequantized(
+                            quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
+                            mode: .nvfp4, globalScale: globalScale, dtype: .float32
+                        ).T)
+                    + bias
+                eval(quantizedWeight, scales, expected)
+                return (input, quantizedWeight, scales, biases, globalScale, bias, expected)
+            }
+
+        let actual = Device.withDefaultDevice(.gpu) {
+            QuantizedLinear(
+                weight: quantizedWeight, bias: bias, scales: scales, biases: biases,
+                groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)(input)
+        }
+
+        assertEqual(actual, expected, rtol: 1e-5, atol: 1e-5)
+    }
+
+    func testQuantizedEmbeddingAppliesNVFP4GlobalScaleOnMetal() {
+        let (
+            input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear,
+            expectedLookup
+        ) = Device.withDefaultDevice(.cpu) {
+            MLXRandom.seed(11)
+            let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32])
+            let input = MLXRandom.uniform(low: -1, high: 1, [4, 32])
+            let indices = MLXArray([0, 3, 7, 31])
+            let globalScale = MLX.max(abs(weight)).asType(.float32)
+            let (quantizedWeight, scales, biases) = MLX.quantized(
+                weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
+            let dequantizedWeight = dequantized(
+                quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
+                mode: .nvfp4, globalScale: globalScale, dtype: .float32)
+            let expectedLinear = matmul(input, dequantizedWeight.T)
+            let expectedLookup = dequantized(
+                quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4,
+                mode: .nvfp4, globalScale: globalScale)[indices].asType(.bfloat16)
+            eval(quantizedWeight, scales, expectedLinear, expectedLookup)
+            return (
+                input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear,
+                expectedLookup
+            )
+        }
+
+        let (actualLinear, actualLookup) = Device.withDefaultDevice(.gpu) {
+            let embedding = QuantizedEmbedding(
+                weight: quantizedWeight, scales: scales, biases: biases,
+                groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)
+            return (embedding.asLinear(input), embedding(indices))
+        }
+
+        assertEqual(actualLinear, expectedLinear, rtol: 1e-5, atol: 1e-5)
+        assertEqual(actualLookup, expectedLookup, rtol: 1e-5, atol: 1e-5)
+    }
+
     func testQuantizedLinearShapeDesc() {
         let linear1 = Linear(512, 1024)
         let quantized1 = linear1.toQuantized(groupSize: 64, bits: 4)
@@ -39,4 +107,62 @@ class QuantizationTests: XCTestCase {
         let quantized = QuantizedLinear(64, 64, groupSize: 32, bits: 4, mode: .mxfp4)
         XCTAssertNil(quantized.biases)
     }
+
+    func testQuantizedLinearStoresGlobalScale() {
+        let globalScale = MLXArray(1.0, dtype: .float32)
+        let quantized = QuantizedLinear(
+            weight: MLXArray.zeros([8, 4], dtype: .uint32),
+            bias: nil,
+            scales: MLXArray.ones([8, 4], dtype: .uint8),
+            biases: nil,
+            groupSize: 16,
+            bits: 4,
+            mode: .nvfp4,
+            globalScale: globalScale)
+
+        XCTAssertNotNil(quantized.globalScale)
+        XCTAssertEqual(quantized.globalScale?.dtype, .float32)
+        XCTAssertNotNil(quantized.parameters()["global_scale"])
+        XCTAssertNil(quantized.parameters()["globalScale"])
+    }
+
+    func testQuantizedEmbeddingStoresGlobalScale() {
+        let globalScale = MLXArray(1.0, dtype: .float32)
+        let quantized = QuantizedEmbedding(
+            weight: MLXArray.zeros([8, 2], dtype: .uint32),
+            scales: MLXArray.ones([8, 2], dtype: .uint8),
+            biases: nil,
+            groupSize: 16,
+            bits: 4,
+            mode: .nvfp4,
+            globalScale: globalScale)
+
+        XCTAssertNotNil(quantized.globalScale)
+        XCTAssertEqual(quantized.globalScale?.dtype, .float32)
+        XCTAssertNotNil(quantized.parameters()["global_scale"])
+        XCTAssertNil(quantized.parameters()["globalScale"])
+    }
+
+    func testQuantizedGlobalScaleIsOptionalParameter() {
+        let linear = QuantizedLinear(
+            weight: MLXArray.zeros([8, 4], dtype: .uint32),
+            bias: nil,
+            scales: MLXArray.ones([8, 4], dtype: .uint8),
+            biases: nil,
+            groupSize: 16,
+            bits: 4,
+            mode: .nvfp4)
+        let embedding = QuantizedEmbedding(
+            weight: MLXArray.zeros([8, 2], dtype: .uint32),
+            scales: MLXArray.ones([8, 2], dtype: .uint8),
+            biases: nil,
+            groupSize: 16,
+            bits: 4,
+            mode: .nvfp4)
+
+        XCTAssertNil(linear.globalScale)
+        XCTAssertNil(linear.parameters()["global_scale"])
+        XCTAssertNil(embedding.globalScale)
+        XCTAssertNil(embedding.parameters()["global_scale"])
+    }
 }