mlxrunner: avoid Metal GPU timeouts when loading models from slow storage

Model load code eagerly evaluated every weight fold (expert stacking,
gather transposes, gate/up fusing) as it was built, with the folds
running on the GPU against lazily loaded tensors: Metal committed
command buffers that waited on file reads, and macOS kills command
buffers that stall too long, so loading a large model from a slow
volume aborted with "Command buffer execution failed". The eager evals
also kept every layer's fold sources alive until the post-load sweep,
transiently holding roughly twice the expert weights on MoE models.

Build the folds lazily and let the runner's weight eval run them, and
on Metal materialize the loaded tensors with CPU reads before any
weight graph exists: no command buffer is ever committed waiting on
file data, at any storage speed, and fold sources free as their folds
execute. CUDA loads read at dispatch and skip the pre-pass. Models no
longer evaluate weights at load; on Metal, tensors the model does not
retain are now read before the sweep frees them.

Measured on an M5 Max, warm page cache, greedy outputs bit-identical:

                                    before           after
  nemotron-3.5-lightning:30b-mlx    1.9s  39.7GiB    1.45s    24.7GiB
  qwen3.6:35b-mlx                   1.27s 22.5GiB    1.1-1.2s 22.4GiB
  nemotron, reads at ~60MB/s        aborts in 6s     loads in 346s

Fixes #17902

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.

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

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

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
 }

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,

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 {

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,