diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index f9b63510c..2db3d0f8b 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -4,6 +4,7 @@ import ( "context" "errors" "log/slog" + "maps" "net" "net/http" "slices" @@ -73,6 +74,13 @@ func (r *Runner) Load(modelName string) error { return err } + // On Metal, materialize the loaded tensors with CPU reads before any + // weight graph exists, so the weight eval never commits a command buffer + // that waits on file data. CUDA loads read at dispatch and need no pre-pass. + if mlx.MetalIsAvailable() { + mlx.Eval(slices.Collect(maps.Values(tensors))...) + } + // Assign weights to model (model-specific logic). Target and draft weights // must be loaded before sweeping so tensors from a combined manifest are // not discarded before the draft model can retain them. diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go index 6dff3aa51..64f5d271f 100644 --- a/x/models/cohere2_moe/cohere2_moe.go +++ b/x/models/cohere2_moe/cohere2_moe.go @@ -401,10 +401,7 @@ func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1) - cloned := t.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Transpose(w, 0, 2, 1).Clone() } // loadStackedProjection returns expert weights already stacked as a single 3D diff --git a/x/models/gemma3/gemma3.go b/x/models/gemma3/gemma3.go index e3ef290b3..4d3917e48 100644 --- a/x/models/gemma3/gemma3.go +++ b/x/models/gemma3/gemma3.go @@ -207,11 +207,6 @@ func precomputeGemmaScaledWeights(m *Model) { m.NormScaled = mlx.AddScalar(m.Norm.Weight, 1.0) } - var scaled []*mlx.Array - if m.NormScaled != nil { - scaled = append(scaled, m.NormScaled) - } - for _, layer := range m.Layers { if layer == nil || layer.Attention == nil { continue @@ -219,34 +214,24 @@ func precomputeGemmaScaledWeights(m *Model) { if layer.InputNorm != nil { layer.InputNormScaled = mlx.AddScalar(layer.InputNorm.Weight, 1.0) - scaled = append(scaled, layer.InputNormScaled) } if layer.PostAttnNorm != nil { layer.PostAttnNormScaled = mlx.AddScalar(layer.PostAttnNorm.Weight, 1.0) - scaled = append(scaled, layer.PostAttnNormScaled) } if layer.PreFFNorm != nil { layer.PreFFNormScaled = mlx.AddScalar(layer.PreFFNorm.Weight, 1.0) - scaled = append(scaled, layer.PreFFNormScaled) } if layer.PostFFNorm != nil { layer.PostFFNormScaled = mlx.AddScalar(layer.PostFFNorm.Weight, 1.0) - scaled = append(scaled, layer.PostFFNormScaled) } if layer.Attention.QNorm != nil { layer.Attention.QNormScaled = mlx.AddScalar(layer.Attention.QNorm.Weight, 1.0) - scaled = append(scaled, layer.Attention.QNormScaled) } if layer.Attention.KNorm != nil { layer.Attention.KNormScaled = mlx.AddScalar(layer.Attention.KNorm.Weight, 1.0) - scaled = append(scaled, layer.Attention.KNormScaled) } } - - if len(scaled) > 0 { - mlx.Eval(scaled...) - } } func newModel(root *model.Root) (base.Model, error) { diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index c52c3e9bd..f390289b3 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -185,9 +185,7 @@ func transposeForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1).Clone() - mlx.Eval(t) - return t + return mlx.Transpose(w, 0, 2, 1).Clone() } // collectExpertProjection collects per-expert tensors, stacks them, and @@ -242,15 +240,12 @@ func collectExpertProjection(tensors map[string]*mlx.Array, cfg *TextConfig, pre } stacked := mlx.Stack(weights, 0).Clone() - mlx.Eval(stacked) out := &stackedExpertResult{Weight: stacked, Bits: bits, GroupSize: groupSize, Mode: mode} if len(scales) == len(weights) { out.Scales = mlx.Stack(scales, 0).Clone() - mlx.Eval(out.Scales) } if len(biases) == len(weights) { out.Biases = mlx.Stack(biases, 0).Clone() - mlx.Eval(out.Biases) } return out } diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index 8940e93b9..07751c70e 100644 --- a/x/models/laguna/laguna.go +++ b/x/models/laguna/laguna.go @@ -590,18 +590,14 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { if len(parts) == 0 { return nil } - stacked := mlx.Stack(parts, 0).Clone() - mlx.Eval(stacked) - return stacked + return mlx.Stack(parts, 0).Clone() } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1).Clone() - mlx.Eval(t) - return t + return mlx.Transpose(w, 0, 2, 1).Clone() } func transposeExpertWeightViewForGatherMM(w *mlx.Array) *mlx.Array { @@ -762,9 +758,7 @@ func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { if a == nil || !a.Valid() || b == nil || !b.Valid() { return nil } - out := mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() - mlx.Eval(out) - return out + return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() } func applyExpertGlobalScale(x, globalScale, idx *mlx.Array) *mlx.Array { @@ -1069,9 +1063,7 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { layerPrefix+".mlp.switch_mlp.e_score_correction_bias", ) if moe.EScoreCorrectionBias != nil && moe.EScoreCorrectionBias.DType() != mlx.DTypeFloat32 { - bias := moe.EScoreCorrectionBias.AsType(mlx.DTypeFloat32).Clone() - mlx.Eval(bias) - moe.EScoreCorrectionBias = bias + moe.EScoreCorrectionBias = moe.EScoreCorrectionBias.AsType(mlx.DTypeFloat32).Clone() } gateW := loadStackedProjection(tensors, cfg, useQuantizedExperts, diff --git a/x/models/nemotron_h/nemotron_h.go b/x/models/nemotron_h/nemotron_h.go index c7c86afdf..2a8bdabf7 100644 --- a/x/models/nemotron_h/nemotron_h.go +++ b/x/models/nemotron_h/nemotron_h.go @@ -354,10 +354,7 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { if len(parts) == 0 { return nil } - stacked := mlx.Stack(parts, 0) - cloned := stacked.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Stack(parts, 0).Clone() } func supportsGatherQMM(mode string, bits int) bool { @@ -428,10 +425,7 @@ func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1) - cloned := t.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Transpose(w, 0, 2, 1).Clone() } func sliceAxis(a *mlx.Array, axis int, start, stop int32) *mlx.Array { @@ -447,18 +441,14 @@ func sliceAxis(a *mlx.Array, axis int, start, stop int32) *mlx.Array { } func appendAndClone(dst, src *mlx.Array) *mlx.Array { - out := mlx.Concatenate([]*mlx.Array{dst, src}, 0).Clone() - mlx.Eval(out) - return out + return mlx.Concatenate([]*mlx.Array{dst, src}, 0).Clone() } func stackSlicesAndClone(slices []*mlx.Array) *mlx.Array { if len(slices) == 0 { return nil } - out := mlx.Stack(slices, 0).Clone() - mlx.Eval(out) - return out + return mlx.Stack(slices, 0).Clone() } func foldSharedExperts(m *SparseMoE, cfg *Config) bool { diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index f67dff614..64ca4dbbe 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/x/models/qwen3_5/qwen3_5.go @@ -479,29 +479,21 @@ func stackAndClone(parts []*mlx.Array) *mlx.Array { if len(parts) == 0 { return nil } - stacked := mlx.Stack(parts, 0) - cloned := stacked.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Stack(parts, 0).Clone() } func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { if w == nil || !w.Valid() || w.NumDims() != 3 { return w } - t := mlx.Transpose(w, 0, 2, 1) - cloned := t.Clone() - mlx.Eval(cloned) - return cloned + return mlx.Transpose(w, 0, 2, 1).Clone() } func fuseExpertStacks(a, b *mlx.Array, axis int) *mlx.Array { if a == nil || !a.Valid() || b == nil || !b.Valid() { return nil } - out := mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() - mlx.Eval(out) - return out + return mlx.Concatenate([]*mlx.Array{a, b}, axis).Clone() } // fuseGateUpProjections joins gate and up stacks along the output dimension,