Reduce model lookup and MLX decision request overhead (#18806)

* Reduce model lookup and MLX decision request overhead

Avoid decoding unrelated manifests when resolving existing model names.
Reuse small Metal scratch buffers between decision requests, releasing
larger working sets and clearing scratch after scoring errors.

* address comments

manifest/manifest.go

@@ -1564,6 +1564,26 @@ func checkBlobDigestReader(r io.Reader, digest string) error {
 	return nil
 }
 
+// FindName returns the stored spelling of a case-insensitive name match, or
+// os.ErrNotExist. Like Manifests, it ignores invalid manifests and prefers v2
+// entries, but only reads manifest contents for matching names.
+func FindName(n model.Name) (model.Name, error) {
+	refs, err := namedManifestRefs(true)
+	if err != nil {
+		return model.Name{}, err
+	}
+
+	for name, ref := range refs {
+		if !name.EqualFold(n) {
+			continue
+		}
+		if _, err := parseManifestFile(name, ref.path, ref.root, runnerPreferences()); err == nil {
+			return name, nil
+		}
+	}
+	return model.Name{}, os.ErrNotExist
+}
+
 func Manifests(continueOnError bool) (map[model.Name]*Manifest, error) {
 	refs, err := namedManifestRefs(continueOnError)
 	if err != nil {

manifest/manifest_test.go

@@ -146,6 +146,9 @@ func TestV2ManifestPathEscapesHostPort(t *testing.T) {
 	if err := WriteManifestData(name, data); err != nil {
 		t.Fatal(err)
 	}
+	if got, err := FindName(name); err != nil || got.String() != name.String() {
+		t.Fatalf("FindName = %v, %v; want %v, nil", got, err, name)
+	}
 
 	manifestPath, err := V2PathForName(name)
 	if err != nil {
@@ -730,6 +733,9 @@ func TestPartialManifestListTracksPresentAndMissingChildren(t *testing.T) {
 	if err := WriteManifestData(name, parentData); err != nil {
 		t.Fatal(err)
 	}
+	if got, err := FindName(name); err != nil || got.String() != name.String() {
+		t.Fatalf("FindName = %v, %v; want %v, nil", got, err, name)
+	}
 	parentSum := sha256.Sum256(parentData)
 	parentDigest := fmt.Sprintf("sha256:%x", parentSum)
 
@@ -1051,6 +1057,9 @@ func TestParseNamedManifestRejectsUnsafeSymlinks(t *testing.T) {
 		if err == nil || !strings.Contains(err.Error(), "not a sha256 blob") {
 			t.Fatalf("err = %v, want not a sha256 blob", err)
 		}
+		if _, err := FindName(name); !errors.Is(err, os.ErrNotExist) {
+			t.Fatalf("FindName error = %v, want os.ErrNotExist", err)
+		}
 	})
 
 	t.Run("blob basename outside blob store", func(t *testing.T) {
@@ -1071,6 +1080,9 @@ func TestParseNamedManifestRejectsUnsafeSymlinks(t *testing.T) {
 		if err == nil || !strings.Contains(err.Error(), "does not match blob") {
 			t.Fatalf("err = %v, want does not match blob", err)
 		}
+		if _, err := FindName(name); !errors.Is(err, os.ErrNotExist) {
+			t.Fatalf("FindName error = %v, want os.ErrNotExist", err)
+		}
 	})
 }
 
@@ -1138,6 +1150,62 @@ func TestManifestsV2ShadowsLegacy(t *testing.T) {
 	}
 }
 
+func TestFindName(t *testing.T) {
+	for _, tc := range []struct {
+		name   string
+		legacy string
+		v2     string
+		found  bool
+	}{
+		{name: "missing"},
+		{name: "legacy", legacy: `{}`, found: true},
+		{name: "v2", v2: `{}`, found: true},
+		{name: "v2 shadows corrupt legacy", legacy: `{`, v2: `{}`, found: true},
+		{name: "corrupt legacy", legacy: `{`},
+		{name: "corrupt v2 shadows legacy", legacy: `{}`, v2: `{`},
+		{
+			name: "manifest list", found: true,
+			v2: `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp","config":{"digest":"sha256:abc"}}]}`,
+		},
+		{
+			name: "manifest list with missing child",
+			v2:   `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp","digest":"sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}]}`,
+		},
+		{
+			name: "unsupported runner",
+			v2:   `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"unsupported","config":{"digest":"sha256:abc"}}]}`,
+		},
+		{
+			name: "nested manifest list",
+			v2:   `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp","mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp"}]}]}`,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv("OLLAMA_MODELS", t.TempDir())
+			name := model.ParseName("MyOrg/MyModel:Q8")
+			if tc.legacy != "" {
+				if err := WriteLegacyManifestData(name, []byte(tc.legacy)); err != nil {
+					t.Fatal(err)
+				}
+			}
+			if tc.v2 != "" {
+				digest := writeManifestBlobForTest(t, []byte(tc.v2))
+				if err := linkManifest(name, digest); err != nil {
+					t.Fatal(err)
+				}
+			}
+			got, err := FindName(model.ParseName("myorg/mymodel:q8"))
+			if tc.found {
+				if err != nil || got.String() != name.String() {
+					t.Fatalf("FindName = %v, %v; want %v, nil", got, err, name)
+				}
+			} else if !errors.Is(err, os.ErrNotExist) {
+				t.Fatalf("FindName error = %v, want os.ErrNotExist", err)
+			}
+		})
+	}
+}
+
 func TestManifests(t *testing.T) {
 	cases := map[string]struct {
 		ps               []string

mlx/stream.go

@@ -60,6 +60,11 @@ type Stream struct {
 	ctx C.mlx_stream
 }
 
+// Synchronize waits for submitted work and its completion handlers on the stream.
+func (s Stream) Synchronize() {
+	mlxCheck(C.mlx_synchronize(s.ctx))
+}
+
 func (s Stream) LogValue() slog.Value {
 	str := mlxCheck(C.mlx_string_new())
 	mlxCheck(C.mlx_stream_tostring(&str, s.ctx))

mlxrunner/score.go

@@ -22,6 +22,10 @@ type scoreRow struct {
 	candidates []int32
 }
 
+// Retain enough scratch space for short decision requests without keeping the
+// temporary working set of a large prompt resident between requests.
+const scoreScratchLimit = 512 << 20
+
 // Models may project just the requested vocabulary rows at higher precision.
 type candidateUnembedder interface {
 	UnembedCandidates(hidden, candidates *mlx.Array) *mlx.Array
@@ -34,8 +38,22 @@ func (r *Runner) scoreHandler(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 	result, err := mlxthread.Call(req.Context(), r.mlxThread, func() (llm.ScoreResponse, error) {
-		defer mlx.ClearCache()
-		return r.score(req.Context(), input)
+		if !mlx.MetalIsAvailable() {
+			defer mlx.ClearCache()
+			return r.score(req.Context(), input)
+		}
+		keepScratch := false
+		defer func() {
+			// Keep small reusable buffers, but release large working sets while idle.
+			if !keepScratch || mlx.CacheMemory() > scoreScratchLimit {
+				mlx.ClearCache()
+			}
+		}()
+		result, err := r.score(req.Context(), input)
+		// Eval waits for results, but completion handlers can still free buffers.
+		mlx.DefaultStream().Synchronize()
+		keepScratch = err == nil
+		return result, err
 	})
 	if err != nil {
 		status := http.StatusInternalServerError

mlxrunner/score_test.go

@@ -200,7 +200,7 @@ func TestScoreValidation(t *testing.T) {
 	}
 }
 
-func TestScoreHandlerClearsScratchCache(t *testing.T) {
+func TestScoreHandlerScratchCache(t *testing.T) {
 	worker := mlxtest.Worker(t)
 	if err := worker.Do(context.Background(), func() error {
 		if !mlx.GPUIsAvailable() {
@@ -211,11 +211,19 @@ func TestScoreHandlerClearsScratchCache(t *testing.T) {
 	}); err != nil {
 		t.Skipf("MLX GPU not available: %v", err)
 	}
+	t.Cleanup(func() {
+		if err := worker.Do(context.Background(), func() error {
+			mlx.ClearCache()
+			return nil
+		}); err != nil {
+			t.Error(err)
+		}
+	})
 	r := &Runner{Tokenizer: newTestTokenizer(t, []int32{7}), contextLength: 512, mlxThread: worker}
 	defer worker.Do(context.Background(), func() error { r.Close(); return nil })
 	for _, n := range []int{16, 128, 512, 16} {
 		for _, cancelled := range []bool{false, true} {
-			// Model prior requests leaving differently sized scratch buffers.
+			// Start with reusable scratch from an earlier request.
 			// These CPU-filled arrays are freed synchronously into MLX's cache.
 			const scratchBytes = 4 << 20
 			if err := worker.Do(context.Background(), func() error {
@@ -250,14 +258,17 @@ func TestScoreHandlerClearsScratchCache(t *testing.T) {
 				t.Fatalf("tokens=%d cancelled=%t: status %d: %s", n, cancelled, response.Code, response.Body)
 			}
 			memory, err := mlxthread.Call(context.Background(), worker, func() (int, error) {
-				return mlx.CacheMemory(), nil
+				memory := mlx.CacheMemory()
+				if !cancelled && mlx.MetalIsAvailable() && memory == 0 {
+					return 0, errors.New("successful scoring discarded all reusable scratch buffers")
+				}
+				return memory, nil
 			})
 			if err != nil {
 				t.Fatal(err)
 			}
-			// The injected scratch buffer must be cleared. Active allocations
-			// belong to live model caches or other tests on the shared worker.
-			if memory >= scratchBytes {
+			// Successful requests may reuse scratch; cancelled requests release it.
+			if memory > scoreScratchLimit || (cancelled && memory >= scratchBytes) {
 				t.Fatalf("tokens=%d cancelled=%t: retained %d scratch bytes", n, cancelled, memory)
 			}
 		}

server/routes.go

@@ -1445,22 +1445,19 @@ func getExistingName(n model.Name) (model.Name, error) {
 		return n, nil
 	}
 
+	// The usual request names an existing model. Avoid decoding every installed
+	// manifest just to resolve its spelling; partial matches below are for new names.
+	if name, err := manifest.FindName(n); err == nil {
+		return name, nil
+	} else if !errors.Is(err, os.ErrNotExist) {
+		return zero, err
+	}
+
 	existing, err := manifest.Manifests(true)
 	if err != nil {
 		return zero, err
 	}
-	// First pass: look for a full case-insensitive match across all four
-	// parts. If found, return the on-disk canonical name directly.
-	for e := range existing {
-		if strings.EqualFold(e.Host, n.Host) &&
-			strings.EqualFold(e.Namespace, n.Namespace) &&
-			strings.EqualFold(e.Model, n.Model) &&
-			strings.EqualFold(e.Tag, n.Tag) {
-			return e, nil
-		}
-	}
-
-	// Second pass: find the single manifest with the longest consecutive
+	// Find the single manifest with the longest consecutive
 	// case-insensitive prefix match (host -> namespace -> model) and copy
 	// only the matching prefix parts from that manifest. The tag is left
 	// as-is so that an unrelated manifest with a matching tag cannot

server/routes_test.go

@@ -34,6 +34,72 @@ import (
 	"github.com/ollama/ollama/version"
 )
 
+func TestGetExistingName(t *testing.T) {
+	t.Setenv("OLLAMA_MODELS", t.TempDir())
+	for _, name := range []string{"MyOrg/MyModel:Q4", "OtherOrg/OtherModel:Q8", "DefaultModel:latest", "localhost:12345/MyOrg/MyModel:Q4"} {
+		if err := manifest.WriteManifestData(model.ParseName(name), []byte(`{}`)); err != nil {
+			t.Fatal(err)
+		}
+	}
+	for _, tc := range []struct{ input, want string }{
+		{"MyOrg/MyModel:Q4", "MyOrg/MyModel:Q4"},
+		{"MYORG/MYMODEL:q4", "MyOrg/MyModel:Q4"},
+		{"myorg/mymodel:q8", "MyOrg/MyModel:q8"},
+		{"myorg/newmodel:q8", "MyOrg/newmodel:q8"},
+		{"neworg/newmodel:q8", "neworg/newmodel:q8"},
+		{"defaultmodel", "DefaultModel:latest"},
+		{"registry.ollama.ai/myorg/mymodel:q4", "MyOrg/MyModel:Q4"},
+		{"ollama.com/myorg/mymodel:q4", "ollama.com/myorg/mymodel:q4"},
+		{"registry.ollama.com/myorg/mymodel:q4", "registry.ollama.com/myorg/mymodel:q4"},
+		{"LOCALHOST:12345/myorg/mymodel:q4", "localhost:12345/MyOrg/MyModel:Q4"},
+		{"sha256-" + strings.Repeat("a", 64), "sha256-" + strings.Repeat("a", 64)},
+	} {
+		t.Run(tc.input, func(t *testing.T) {
+			got, err := getExistingName(model.ParseName(tc.input))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if want := model.ParseName(tc.want); got.String() != want.String() {
+				t.Errorf("getExistingName = %v, want %v", got, want)
+			}
+		})
+	}
+}
+
+func BenchmarkGetExistingName(b *testing.B) {
+	// Tensor manifests are large; lookup should only decode the requested model.
+	layers := make([]manifest.Layer, 512)
+	for i := range layers {
+		layers[i] = manifest.Layer{
+			MediaType: manifest.MediaTypeImageTensor,
+			Digest:    fmt.Sprintf("sha256:%064x", i),
+			Size:      1024,
+		}
+	}
+	data, err := json.Marshal(manifest.Manifest{SchemaVersion: 2, Layers: layers})
+	if err != nil {
+		b.Fatal(err)
+	}
+	for _, models := range []int{1, 10, 100} {
+		b.Run(fmt.Sprintf("models=%d", models), func(b *testing.B) {
+			b.Setenv("OLLAMA_MODELS", b.TempDir())
+			for i := range models {
+				name := model.ParseName(fmt.Sprintf("Test/Model%d:Latest", i))
+				if err := manifest.WriteManifestData(name, data); err != nil {
+					b.Fatal(err)
+				}
+			}
+			name := model.ParseName("test/model0:latest")
+			b.ReportAllocs()
+			for b.Loop() {
+				if _, err := getExistingName(name); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
 func TestPsHandlerUsesRunningManifestAndRunner(t *testing.T) {
 	gin.SetMode(gin.TestMode)