Add global scale support to quantized layers (#426)

* Add global scale support to quantized layers

Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations.

Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there.

Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface.

* Keep nvfp4 global scale on the quantized path for Metal

Since MLX encodes nvfp4 group scales pre-multiplied by (448 * 6) / globalScale,
the global scale factors out of dequantization linearly. Exploit this by always
running quantizedMM (supported on Metal without globalScale) and applying
globalScale / (448 * 6) to the output, instead of dequantizing the full weight
and falling back to a plain matmul.

This keeps weights quantized end-to-end on all backends, removes the
dequantizedWeight fallback from QuantizedLinear and QuantizedEmbedding, and
applies the scale before the bias add so the bias is unaffected.

Add Metal tests that validate the GPU quantized path against CPU
dequantized(globalScale:) references for QuantizedLinear (matmul + bias) and
QuantizedEmbedding (asLinear and index lookup).

Author: aleroot

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"])
+    }
 }