vendor/ollama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mlxrunner: harden linear/embedding layers against over-promotion (#16682)

Browse Source 
Adding/Multiplying a tensor by a scalar w/ a different data type
can cause the tensor to be promoted and cause performance issues.

This change adds several guards against over-promotion.
This commit is contained in:
Patrick Devine
2026-06-11 13:56:25 -07:00
committed by GitHub
parent 1abd56b6e6
commit 82e0ddb6fe
1 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -61,6 +61,9 @@ type Linear struct {
}
func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear {
if bias != nil && bias.Valid() && bias.DType() != weight.DType() {
bias = bias.AsType(weight.DType())
}
return &Linear{Weight: weight, Bias: bias}
}
@@ -95,6 +98,9 @@ func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int,
} else {
mlx.Eval(qw, scales)
}
if bias != nil && bias.Valid() && bias.DType() != weight.DType() {
bias = bias.AsType(weight.DType())
}
return &QuantizedLinear{
Weight: qw,
Scales: scales,
@@ -107,20 +113,22 @@ func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int,
}
func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array {
var out *mlx.Array
out := mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode)
if ql.GlobalScale != nil {
// Double-scale nvfp4 (e.g., NVIDIA ModelOpt): standard quantized_matmul
// followed by global_scale multiply. The global_scale is a per-tensor
// F32 scalar (weight_scale_2 in NVIDIA's format).
// TODO: switch to a fused double-scale matmul once MLX has kernel
// coverage for this path.
out = mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode)
out = mlx.Mul(out, ql.GlobalScale)
} else {
out = mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode)
outDType := out.DType()
out = mlx.Mul(out, ql.GlobalScale).AsType(outDType)
}
if ql.Bias != nil && ql.Bias.Valid() {
out = out.Add(ql.Bias)
bias := ql.Bias
if bias.DType() != out.DType() {
bias = bias.AsType(out.DType())
}
out = out.Add(bias)
}
return out
}
@@ -184,7 +192,8 @@ func (qe *QuantizedEmbedding) Forward(indices *mlx.Array) *mlx.Array {
}
out := mlx.Dequantize(weight, scales, qbiases, qe.GroupSize, qe.Bits, qe.Mode)
if qe.GlobalScale != nil {
out = mlx.Mul(out, qe.GlobalScale)
outDType := out.DType()
out = mlx.Mul(out, qe.GlobalScale).AsType(outDType)
}
return out
}