From 3ef69ef7849e8723dbf29bdf4d251b22f545588e Mon Sep 17 00:00:00 2001 From: Patrick Devine Date: Thu, 4 Jun 2026 17:40:01 -0700 Subject: [PATCH] mlx: allow the embedding layer to use the nvfp4 global scale (#16527) --- x/mlxrunner/model/embedding.go | 18 ++++++++++- x/mlxrunner/model/embedding_test.go | 48 +++++++++++++++++++++++++++++ x/models/nn/nn.go | 43 ++++++++++++-------------- x/models/nn/nn_test.go | 37 ++++++++++++++++++++++ 4 files changed, 121 insertions(+), 25 deletions(-) diff --git a/x/mlxrunner/model/embedding.go b/x/mlxrunner/model/embedding.go index d45d48e2..1b3c14ef 100644 --- a/x/mlxrunner/model/embedding.go +++ b/x/mlxrunner/model/embedding.go @@ -35,7 +35,23 @@ func MakeEmbeddingLayer( scales, ) - return nn.NewQuantizedEmbedding(w, scales, qbiases, groupSize, bits, mode) + // Check for per-tensor global scale (NVIDIA double-scale nvfp4). + // NVIDIA ModelOpt stores this as "weight_scale_2"; our import + // pipeline maps it to "weight.global_scale". + globalScale := tensors[path+".weight.global_scale"] + if globalScale == nil { + globalScale = tensors[path+".weight_scale_2"] + } + + return &nn.QuantizedEmbedding{ + Weight: w, + Scales: scales, + QBiases: qbiases, + GlobalScale: globalScale, + GroupSize: groupSize, + Bits: bits, + Mode: mode, + } } return nn.NewEmbedding(w) diff --git a/x/mlxrunner/model/embedding_test.go b/x/mlxrunner/model/embedding_test.go index 2935bdc3..975e4f83 100644 --- a/x/mlxrunner/model/embedding_test.go +++ b/x/mlxrunner/model/embedding_test.go @@ -76,3 +76,51 @@ func TestMakeEmbeddingLayerQuantized(t *testing.T) { t.Fatalf("AsLinear type = %T, want *nn.QuantizedLinear", emb.AsLinear()) } } + +func TestMakeEmbeddingLayerQuantizedGlobalScale(t *testing.T) { + weight := &mlx.Array{} + scales := &mlx.Array{} + globalScale := &mlx.Array{} + + emb := MakeEmbeddingLayer(map[string]*mlx.Array{ + "model.embed_tokens.weight": weight, + "model.embed_tokens.weight_scale": scales, + "model.embed_tokens.weight.global_scale": globalScale, + }, "model.embed_tokens", 16, 4, "nvfp4", nil) + + qemb, ok := emb.(*nn.QuantizedEmbedding) + if !ok { + t.Fatalf("embedding type = %T, want *nn.QuantizedEmbedding", emb) + } + if qemb.GlobalScale != globalScale { + t.Fatalf("GlobalScale = %p, want %p", qemb.GlobalScale, globalScale) + } + + linear, ok := qemb.AsLinear().(*nn.QuantizedLinear) + if !ok { + t.Fatalf("AsLinear type = %T, want *nn.QuantizedLinear", qemb.AsLinear()) + } + if linear.GlobalScale != globalScale { + t.Fatalf("AsLinear GlobalScale = %p, want %p", linear.GlobalScale, globalScale) + } +} + +func TestMakeEmbeddingLayerQuantizedGlobalScaleFallback(t *testing.T) { + weight := &mlx.Array{} + scales := &mlx.Array{} + globalScale := &mlx.Array{} + + emb := MakeEmbeddingLayer(map[string]*mlx.Array{ + "model.embed_tokens.weight": weight, + "model.embed_tokens.weight_scale": scales, + "model.embed_tokens.weight_scale_2": globalScale, + }, "model.embed_tokens", 16, 4, "nvfp4", nil) + + qemb, ok := emb.(*nn.QuantizedEmbedding) + if !ok { + t.Fatalf("embedding type = %T, want *nn.QuantizedEmbedding", emb) + } + if qemb.GlobalScale != globalScale { + t.Fatalf("GlobalScale = %p, want %p", qemb.GlobalScale, globalScale) + } +} diff --git a/x/models/nn/nn.go b/x/models/nn/nn.go index 4410848b..9aa333bc 100644 --- a/x/models/nn/nn.go +++ b/x/models/nn/nn.go @@ -166,23 +166,13 @@ func (e *Embedding) AsLinear() LinearLayer { // QuantizedEmbedding performs row-wise embedding lookup from affine/nvfp4/etc. // packed weights and dequantizes only the selected rows. type QuantizedEmbedding struct { - Weight *mlx.Array - Scales *mlx.Array - QBiases *mlx.Array - GroupSize int - Bits int - Mode string -} - -func NewQuantizedEmbedding(weight, scales, qbiases *mlx.Array, groupSize, bits int, mode string) *QuantizedEmbedding { - return &QuantizedEmbedding{ - Weight: weight, - Scales: scales, - QBiases: qbiases, - GroupSize: groupSize, - Bits: bits, - Mode: mode, - } + Weight *mlx.Array + Scales *mlx.Array + QBiases *mlx.Array + GlobalScale *mlx.Array // Per-tensor global scale for double-scale nvfp4 (nil for standard) + GroupSize int + Bits int + Mode string } func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array { @@ -192,17 +182,22 @@ func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array { if qe.QBiases != nil && qe.QBiases.Valid() { qbiases = qe.QBiases.TakeAxis(indices, 0) } - return mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode) + out := mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode) + if qe.GlobalScale != nil { + out = mlx.Mul(out, qe.GlobalScale) + } + return out } func (qe *QuantizedEmbedding) AsLinear() LinearLayer { return &QuantizedLinear{ - Weight: qe.Weight, - Scales: qe.Scales, - QBiases: qe.QBiases, - GroupSize: qe.GroupSize, - Bits: qe.Bits, - Mode: qe.Mode, + Weight: qe.Weight, + Scales: qe.Scales, + QBiases: qe.QBiases, + GlobalScale: qe.GlobalScale, + GroupSize: qe.GroupSize, + Bits: qe.Bits, + Mode: qe.Mode, } } diff --git a/x/models/nn/nn_test.go b/x/models/nn/nn_test.go index 8ec4b757..8c475600 100644 --- a/x/models/nn/nn_test.go +++ b/x/models/nn/nn_test.go @@ -185,3 +185,40 @@ func TestQuantizedLinearMXFP4MatchesDequantizedWeight(t *testing.T) { } } } + +func TestQuantizedEmbeddingAsLinearPreservesGlobalScale(t *testing.T) { + weight := &mlx.Array{} + scales := &mlx.Array{} + qbiases := &mlx.Array{} + globalScale := &mlx.Array{} + + embedding := &QuantizedEmbedding{ + Weight: weight, + Scales: scales, + QBiases: qbiases, + GlobalScale: globalScale, + GroupSize: 16, + Bits: 4, + Mode: "nvfp4", + } + + linear, ok := embedding.AsLinear().(*QuantizedLinear) + if !ok { + t.Fatalf("AsLinear type = %T, want *QuantizedLinear", embedding.AsLinear()) + } + if linear.Weight != weight { + t.Fatalf("AsLinear Weight = %p, want %p", linear.Weight, weight) + } + if linear.Scales != scales { + t.Fatalf("AsLinear Scales = %p, want %p", linear.Scales, scales) + } + if linear.QBiases != qbiases { + t.Fatalf("AsLinear QBiases = %p, want %p", linear.QBiases, qbiases) + } + if linear.GlobalScale != globalScale { + t.Fatalf("AsLinear GlobalScale = %p, want %p", linear.GlobalScale, globalScale) + } + if linear.GroupSize != 16 || linear.Bits != 4 || linear.Mode != "nvfp4" { + t.Fatalf("AsLinear quant params = (%d, %d, %q), want (16, 4, %q)", linear.GroupSize, linear.Bits, linear.Mode, "nvfp4") + } +}