upstream-archive/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-07-23 09:10:53 -05:00
Code Issues Packages Projects Releases Wiki Activity

mlx: allow the embedding layer to use the nvfp4 global scale (#16527)

Browse Source
This commit is contained in:
Patrick Devine
2026-06-04 17:40:01 -07:00
committed by GitHub
parent 1a7786be14
commit 3ef69ef784
4 changed files with 121 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
x/mlxrunner/model/embedding.go
View File

@@ -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)

48
x/mlxrunner/model/embedding_test.go
View File

@@ -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)
}
}

43
x/models/nn/nn.go
View File

@@ -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,
}
}

37
x/models/nn/nn_test.go
View File

@@ -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")
}
}