extras

CMakeLists.txt

@@ -86,9 +86,9 @@ if(CMAKE_CUDA_COMPILER)
     )
 endif()
 
-set(WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX "^gfx(906|908|90a|1200|1201):xnack[+-]$"
+set(WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX "^gfx(906|908|90a):xnack[+-]$"
     CACHE STRING
-    "Regular expression describing AMDGPU_TARGETS not supported on Windows. Override to force building these targets. Default \"^gfx(906|908|90a|1200|1201):xnack[+-]$\"."
+    "Regular expression describing AMDGPU_TARGETS not supported on Windows. Override to force building these targets. Default \"^gfx(906|908|90a):xnack[+-]$\"."
 )
 
 check_language(HIP)
@@ -97,7 +97,7 @@ if(CMAKE_HIP_COMPILER)
 
     find_package(hip REQUIRED)
     if(NOT AMDGPU_TARGETS)
-        list(FILTER AMDGPU_TARGETS INCLUDE REGEX "^gfx(900|94[012]|101[02]|1030|110[012]|120[01])$")
+        list(FILTER AMDGPU_TARGETS INCLUDE REGEX "^gfx(900|94[012]|101[02]|1030|110[012])$")
     elseif(WIN32 AND WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX)
         list(FILTER AMDGPU_TARGETS EXCLUDE REGEX ${WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX})
     endif()

CMakePresets.json

@@ -56,7 +56,7 @@
       "name": "ROCm 6",
       "inherits": [ "ROCm" ],
       "cacheVariables": {
-        "AMDGPU_TARGETS": "gfx900;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1151;gfx1200;gfx1201;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-"
+        "AMDGPU_TARGETS": "gfx900;gfx940;gfx941;gfx942;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx906:xnack-;gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-"
       }
     }
   ],

README.md

@@ -285,7 +285,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [Bionic GPT](https://github.com/bionic-gpt/bionic-gpt)
 - [HTML UI](https://github.com/rtcfirefly/ollama-ui)
 - [Saddle](https://github.com/jikkuatwork/saddle)
-- [TagSpaces](https://www.tagspaces.org) (A platform for file based apps, [utilizing Ollama](https://docs.tagspaces.org/ai/) for the generation of tags and descriptions)
 - [Chatbot UI](https://github.com/ivanfioravanti/chatbot-ollama)
 - [Chatbot UI v2](https://github.com/mckaywrigley/chatbot-ui)
 - [Typescript UI](https://github.com/ollama-interface/Ollama-Gui?tab=readme-ov-file)
@@ -325,7 +324,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [RWKV-Runner](https://github.com/josStorer/RWKV-Runner) (RWKV offline LLM deployment tool, also usable as a client for ChatGPT and Ollama)
 - [Ollama Grid Search](https://github.com/dezoito/ollama-grid-search) (app to evaluate and compare models)
 - [Olpaka](https://github.com/Otacon/olpaka) (User-friendly Flutter Web App for Ollama)
-- [Casibase](https://casibase.org) (An open source AI knowledge base and dialogue system combining the latest RAG, SSO, ollama support and multiple large language models.)
 - [OllamaSpring](https://github.com/CrazyNeil/OllamaSpring) (Ollama Client for macOS)
 - [LLocal.in](https://github.com/kartikm7/llocal) (Easy to use Electron Desktop Client for Ollama)
 - [Shinkai Desktop](https://github.com/dcSpark/shinkai-apps) (Two click install Local AI using Ollama + Files + RAG)
@@ -394,10 +392,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [1Panel](https://github.com/1Panel-dev/1Panel/) (Web-based Linux Server Management Tool)
 - [AstrBot](https://github.com/Soulter/AstrBot/) (User-friendly LLM-based multi-platform chatbot with a WebUI, supporting RAG, LLM agents, and plugins integration)
 - [Reins](https://github.com/ibrahimcetin/reins) (Easily tweak parameters, customize system prompts per chat, and enhance your AI experiments with reasoning model support.)
-- [Ellama](https://github.com/zeozeozeo/ellama) (Friendly native app to chat with an Ollama instance)
-- [screenpipe](https://github.com/mediar-ai/screenpipe) Build agents powered by your screen history
-- [Ollamb](https://github.com/hengkysteen/ollamb) (Simple yet rich in features, cross-platform built with Flutter and designed for Ollama. Try the [web demo](https://hengkysteen.github.io/demo/ollamb/).)
-- [Writeopia](https://github.com/Writeopia/Writeopia) (Text editor with integration with Ollama)
 
 ### Cloud
 
@@ -437,9 +431,7 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [SwollamaCLI](https://github.com/marcusziade/Swollama) bundled with the Swollama Swift package. [Demo](https://github.com/marcusziade/Swollama?tab=readme-ov-file#cli-usage)
 - [aichat](https://github.com/sigoden/aichat) All-in-one LLM CLI tool featuring Shell Assistant, Chat-REPL, RAG, AI tools & agents, with access to OpenAI, Claude, Gemini, Ollama, Groq, and more.
 - [PowershAI](https://github.com/rrg92/powershai) PowerShell module that brings AI to terminal on Windows, including support for Ollama
-- [DeepShell](https://github.com/Abyss-c0re/deepshell) Your self-hosted AI assistant. Interactive Shell, Files and Folders analysis.
 - [orbiton](https://github.com/xyproto/orbiton) Configuration-free text editor and IDE with support for tab completion with Ollama.
-- [orca-cli](https://github.com/molbal/orca-cli) Ollama Registry CLI Application - Browse, pull and download models from Ollama Registry in your terminal.
 
 ### Apple Vision Pro
 
@@ -518,7 +510,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [Ollama for Zig](https://github.com/dravenk/ollama-zig)
 - [Abso](https://github.com/lunary-ai/abso) (OpenAI-compatible TypeScript SDK for any LLM provider)
 - [Nichey](https://github.com/goodreasonai/nichey) is a Python package for generating custom wikis for your research topic
-- [Ollama for D](https://github.com/kassane/ollama-d)
 
 ### Mobile
 

api/types.go

@@ -12,7 +12,6 @@ import (
 	"time"
 
 	"github.com/ollama/ollama/envconfig"
-	"github.com/ollama/ollama/types/model"
 )
 
 // StatusError is an error with an HTTP status code and message.
@@ -82,7 +81,7 @@ type GenerateRequest struct {
 
 	// Options lists model-specific options. For example, temperature can be
 	// set through this field, if the model supports it.
-	Options map[string]any `json:"options"`
+	Options map[string]interface{} `json:"options"`
 }
 
 // ChatRequest describes a request sent by [Client.Chat].
@@ -107,7 +106,7 @@ type ChatRequest struct {
 	Tools `json:"tools,omitempty"`
 
 	// Options lists model-specific options.
-	Options map[string]any `json:"options"`
+	Options map[string]interface{} `json:"options"`
 }
 
 type Tools []Tool
@@ -261,7 +260,7 @@ type EmbedRequest struct {
 	Truncate *bool `json:"truncate,omitempty"`
 
 	// Options lists model-specific options.
-	Options map[string]any `json:"options"`
+	Options map[string]interface{} `json:"options"`
 }
 
 // EmbedResponse is the response from [Client.Embed].
@@ -287,7 +286,7 @@ type EmbeddingRequest struct {
 	KeepAlive *Duration `json:"keep_alive,omitempty"`
 
 	// Options lists model-specific options.
-	Options map[string]any `json:"options"`
+	Options map[string]interface{} `json:"options"`
 }
 
 // EmbeddingResponse is the response from [Client.Embeddings].
@@ -333,7 +332,7 @@ type ShowRequest struct {
 	Template string `json:"template"`
 	Verbose  bool   `json:"verbose"`
 
-	Options map[string]any `json:"options"`
+	Options map[string]interface{} `json:"options"`
 
 	// Deprecated: set the model name with Model instead
 	Name string `json:"name"`
@@ -341,18 +340,16 @@ type ShowRequest struct {
 
 // ShowResponse is the response returned from [Client.Show].
 type ShowResponse struct {
-	License       string             `json:"license,omitempty"`
-	Modelfile     string             `json:"modelfile,omitempty"`
-	Parameters    string             `json:"parameters,omitempty"`
-	Template      string             `json:"template,omitempty"`
-	System        string             `json:"system,omitempty"`
-	Details       ModelDetails       `json:"details,omitempty"`
-	Messages      []Message          `json:"messages,omitempty"`
-	ModelInfo     map[string]any     `json:"model_info,omitempty"`
-	ProjectorInfo map[string]any     `json:"projector_info,omitempty"`
-	Tensors       []Tensor           `json:"tensors,omitempty"`
-	Capabilities  []model.Capability `json:"capabilities,omitempty"`
-	ModifiedAt    time.Time          `json:"modified_at,omitempty"`
+	License       string         `json:"license,omitempty"`
+	Modelfile     string         `json:"modelfile,omitempty"`
+	Parameters    string         `json:"parameters,omitempty"`
+	Template      string         `json:"template,omitempty"`
+	System        string         `json:"system,omitempty"`
+	Details       ModelDetails   `json:"details,omitempty"`
+	Messages      []Message      `json:"messages,omitempty"`
+	ModelInfo     map[string]any `json:"model_info,omitempty"`
+	ProjectorInfo map[string]any `json:"projector_info,omitempty"`
+	ModifiedAt    time.Time      `json:"modified_at,omitempty"`
 }
 
 // CopyRequest is the request passed to [Client.Copy].
@@ -470,13 +467,6 @@ type ModelDetails struct {
 	QuantizationLevel string   `json:"quantization_level"`
 }
 
-// Tensor describes the metadata for a given tensor.
-type Tensor struct {
-	Name  string   `json:"name"`
-	Type  string   `json:"type"`
-	Shape []uint64 `json:"shape"`
-}
-
 func (m *Metrics) Summary() {
 	if m.TotalDuration > 0 {
 		fmt.Fprintf(os.Stderr, "total duration:       %v\n", m.TotalDuration)
@@ -505,7 +495,7 @@ func (m *Metrics) Summary() {
 	}
 }
 
-func (opts *Options) FromMap(m map[string]any) error {
+func (opts *Options) FromMap(m map[string]interface{}) error {
 	valueOpts := reflect.ValueOf(opts).Elem() // names of the fields in the options struct
 	typeOpts := reflect.TypeOf(opts).Elem()   // types of the fields in the options struct
 
@@ -562,12 +552,12 @@ func (opts *Options) FromMap(m map[string]any) error {
 				}
 				field.SetString(val)
 			case reflect.Slice:
-				// JSON unmarshals to []any, not []string
-				val, ok := val.([]any)
+				// JSON unmarshals to []interface{}, not []string
+				val, ok := val.([]interface{})
 				if !ok {
 					return fmt.Errorf("option %q must be of type array", key)
 				}
-				// convert []any to []string
+				// convert []interface{} to []string
 				slice := make([]string, len(val))
 				for i, item := range val {
 					str, ok := item.(string)
@@ -674,7 +664,7 @@ func (d *Duration) UnmarshalJSON(b []byte) (err error) {
 }
 
 // FormatParams converts specified parameter options to their correct types
-func FormatParams(params map[string][]string) (map[string]any, error) {
+func FormatParams(params map[string][]string) (map[string]interface{}, error) {
 	opts := Options{}
 	valueOpts := reflect.ValueOf(&opts).Elem() // names of the fields in the options struct
 	typeOpts := reflect.TypeOf(opts)           // types of the fields in the options struct
@@ -688,7 +678,7 @@ func FormatParams(params map[string][]string) (map[string]any, error) {
 		}
 	}
 
-	out := make(map[string]any)
+	out := make(map[string]interface{})
 	// iterate params and set values based on json struct tags
 	for key, vals := range params {
 		if opt, ok := jsonOpts[key]; !ok {

api/types_test.go

@@ -134,7 +134,7 @@ func TestUseMmapParsingFromJSON(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			var oMap map[string]any
+			var oMap map[string]interface{}
 			err := json.Unmarshal([]byte(test.req), &oMap)
 			require.NoError(t, err)
 			opts := DefaultOptions()

benchmark/server_benchmark_test.go

@@ -1,178 +0,0 @@
-package benchmark
-
-import (
-	"context"
-	"flag"
-	"fmt"
-	"testing"
-	"time"
-
-	"github.com/ollama/ollama/api"
-)
-
-// Command line flags
-var modelFlag string
-
-func init() {
-	flag.StringVar(&modelFlag, "m", "", "Name of the model to benchmark")
-	flag.Lookup("m").DefValue = "model"
-}
-
-// modelName returns the model name from flags, failing the test if not set
-func modelName(b *testing.B) string {
-	if modelFlag == "" {
-		b.Fatal("Error: -m flag is required for benchmark tests")
-	}
-	return modelFlag
-}
-
-type TestCase struct {
-	name      string
-	prompt    string
-	maxTokens int
-}
-
-// runGenerateBenchmark contains the common generate and metrics logic
-func runGenerateBenchmark(b *testing.B, ctx context.Context, client *api.Client, req *api.GenerateRequest) {
-	start := time.Now()
-	var ttft time.Duration
-	var metrics api.Metrics
-
-	err := client.Generate(ctx, req, func(resp api.GenerateResponse) error {
-		if ttft == 0 && resp.Response != "" {
-			ttft = time.Since(start)
-		}
-		if resp.Done {
-			metrics = resp.Metrics
-		}
-		return nil
-	})
-
-	// Report custom metrics as part of the benchmark results
-	b.ReportMetric(float64(ttft.Milliseconds()), "ttft_ms")
-	b.ReportMetric(float64(metrics.LoadDuration.Milliseconds()), "load_ms")
-
-	// Token throughput metrics
-	promptThroughput := float64(metrics.PromptEvalCount) / metrics.PromptEvalDuration.Seconds()
-	genThroughput := float64(metrics.EvalCount) / metrics.EvalDuration.Seconds()
-	b.ReportMetric(promptThroughput, "prompt_tok/s")
-	b.ReportMetric(genThroughput, "gen_tok/s")
-
-	// Token counts
-	b.ReportMetric(float64(metrics.PromptEvalCount), "prompt_tokens")
-	b.ReportMetric(float64(metrics.EvalCount), "gen_tokens")
-	if err != nil {
-		b.Fatal(err)
-	}
-}
-
-// BenchmarkColdStart runs benchmarks with model loading from cold state
-func BenchmarkColdStart(b *testing.B) {
-	client := setup(b)
-	tests := []TestCase{
-		{"short_prompt", "Write a long story", 100},
-		{"medium_prompt", "Write a detailed economic analysis", 500},
-		{"long_prompt", "Write a comprehensive AI research paper", 1000},
-	}
-	m := modelName(b)
-
-	for _, tt := range tests {
-		b.Run(fmt.Sprintf("%s/cold/%s", m, tt.name), func(b *testing.B) {
-			ctx := context.Background()
-
-			// Set number of tokens as our throughput metric
-			b.SetBytes(int64(tt.maxTokens))
-
-			for b.Loop() {
-				b.StopTimer()
-				// Ensure model is unloaded before each iteration
-				unload(client, m, b)
-				b.StartTimer()
-
-				req := &api.GenerateRequest{
-					Model:   m,
-					Prompt:  tt.prompt,
-					Options: map[string]any{"num_predict": tt.maxTokens, "temperature": 0.1},
-				}
-
-				runGenerateBenchmark(b, ctx, client, req)
-			}
-		})
-	}
-}
-
-// BenchmarkWarmStart runs benchmarks with pre-loaded model
-func BenchmarkWarmStart(b *testing.B) {
-	client := setup(b)
-	tests := []TestCase{
-		{"short_prompt", "Write a long story", 100},
-		{"medium_prompt", "Write a detailed economic analysis", 500},
-		{"long_prompt", "Write a comprehensive AI research paper", 1000},
-	}
-	m := modelName(b)
-
-	for _, tt := range tests {
-		b.Run(fmt.Sprintf("%s/warm/%s", m, tt.name), func(b *testing.B) {
-			ctx := context.Background()
-
-			// Pre-warm the model
-			warmup(client, m, tt.prompt, b)
-
-			// Set number of tokens as our throughput metric
-			b.SetBytes(int64(tt.maxTokens))
-
-			for b.Loop() {
-				req := &api.GenerateRequest{
-					Model:   m,
-					Prompt:  tt.prompt,
-					Options: map[string]any{"num_predict": tt.maxTokens, "temperature": 0.1},
-				}
-
-				runGenerateBenchmark(b, ctx, client, req)
-			}
-		})
-	}
-}
-
-// setup verifies server and model availability
-func setup(b *testing.B) *api.Client {
-	client, err := api.ClientFromEnvironment()
-	if err != nil {
-		b.Fatal(err)
-	}
-	if _, err := client.Show(context.Background(), &api.ShowRequest{Model: modelName(b)}); err != nil {
-		b.Fatalf("Model unavailable: %v", err)
-	}
-
-	return client
-}
-
-// warmup ensures the model is loaded and warmed up
-func warmup(client *api.Client, model string, prompt string, b *testing.B) {
-	for range 3 {
-		err := client.Generate(
-			context.Background(),
-			&api.GenerateRequest{
-				Model:   model,
-				Prompt:  prompt,
-				Options: map[string]any{"num_predict": 50, "temperature": 0.1},
-			},
-			func(api.GenerateResponse) error { return nil },
-		)
-		if err != nil {
-			b.Logf("Error during model warm-up: %v", err)
-		}
-	}
-}
-
-// unload forces model unloading using KeepAlive: 0 parameter
-func unload(client *api.Client, model string, b *testing.B) {
-	req := &api.GenerateRequest{
-		Model:     model,
-		KeepAlive: &api.Duration{Duration: 0},
-	}
-	if err := client.Generate(context.Background(), req, func(api.GenerateResponse) error { return nil }); err != nil {
-		b.Logf("Unload error: %v", err)
-	}
-	time.Sleep(1 * time.Second)
-}

cmd/cmd.go

@@ -18,8 +18,6 @@ import (
 	"os/signal"
 	"path/filepath"
 	"runtime"
-	"slices"
-	"sort"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -268,7 +266,7 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 	opts := runOptions{
 		Model:    args[0],
 		WordWrap: os.Getenv("TERM") == "xterm-256color",
-		Options:  map[string]any{},
+		Options:  map[string]interface{}{},
 	}
 
 	format, err := cmd.Flags().GetString("format")
@@ -340,11 +338,6 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	opts.MultiModal = slices.Contains(info.Capabilities, model.CapabilityVision)
-
-	// TODO: remove the projector info and vision info checks below,
-	// these are left in for backwards compatibility with older servers
-	// that don't have the capabilities field in the model info
 	if len(info.ProjectorInfo) != 0 {
 		opts.MultiModal = true
 	}
@@ -575,9 +568,8 @@ func ShowHandler(cmd *cobra.Command, args []string) error {
 	parameters, errParams := cmd.Flags().GetBool("parameters")
 	system, errSystem := cmd.Flags().GetBool("system")
 	template, errTemplate := cmd.Flags().GetBool("template")
-	verbose, errVerbose := cmd.Flags().GetBool("verbose")
 
-	for _, boolErr := range []error{errLicense, errModelfile, errParams, errSystem, errTemplate, errVerbose} {
+	for _, boolErr := range []error{errLicense, errModelfile, errParams, errSystem, errTemplate} {
 		if boolErr != nil {
 			return errors.New("error retrieving flags")
 		}
@@ -615,7 +607,7 @@ func ShowHandler(cmd *cobra.Command, args []string) error {
 		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
 	}
 
-	req := api.ShowRequest{Name: args[0], Verbose: verbose}
+	req := api.ShowRequest{Name: args[0]}
 	resp, err := client.Show(cmd.Context(), &req)
 	if err != nil {
 		return err
@@ -638,10 +630,10 @@ func ShowHandler(cmd *cobra.Command, args []string) error {
 		return nil
 	}
 
-	return showInfo(resp, verbose, os.Stdout)
+	return showInfo(resp, os.Stdout)
 }
 
-func showInfo(resp *api.ShowResponse, verbose bool, w io.Writer) error {
+func showInfo(resp *api.ShowResponse, w io.Writer) error {
 	tableRender := func(header string, rows func() [][]string) {
 		fmt.Fprintln(w, " ", header)
 		table := tablewriter.NewWriter(w)
@@ -675,15 +667,6 @@ func showInfo(resp *api.ShowResponse, verbose bool, w io.Writer) error {
 		return
 	})
 
-	if len(resp.Capabilities) > 0 {
-		tableRender("Capabilities", func() (rows [][]string) {
-			for _, capability := range resp.Capabilities {
-				rows = append(rows, []string{"", capability.String()})
-			}
-			return
-		})
-	}
-
 	if resp.ProjectorInfo != nil {
 		tableRender("Projector", func() (rows [][]string) {
 			arch := resp.ProjectorInfo["general.architecture"].(string)
@@ -707,47 +690,6 @@ func showInfo(resp *api.ShowResponse, verbose bool, w io.Writer) error {
 		})
 	}
 
-	if resp.ModelInfo != nil && verbose {
-		tableRender("Metadata", func() (rows [][]string) {
-			keys := make([]string, 0, len(resp.ModelInfo))
-			for k := range resp.ModelInfo {
-				keys = append(keys, k)
-			}
-			sort.Strings(keys)
-
-			for _, k := range keys {
-				var v string
-				switch vData := resp.ModelInfo[k].(type) {
-				case bool:
-					v = fmt.Sprintf("%t", vData)
-				case string:
-					v = vData
-				case float64:
-					v = fmt.Sprintf("%g", vData)
-				case []any:
-					n := 3
-					if len(vData) < n {
-						n = len(vData)
-					}
-					v = fmt.Sprintf("%v", vData[:n])
-				default:
-					v = fmt.Sprintf("%T", vData)
-				}
-				rows = append(rows, []string{"", k, v})
-			}
-			return
-		})
-	}
-
-	if len(resp.Tensors) > 0 && verbose {
-		tableRender("Tensors", func() (rows [][]string) {
-			for _, t := range resp.Tensors {
-				rows = append(rows, []string{"", t.Name, t.Type, fmt.Sprint(t.Shape)})
-			}
-			return
-		})
-	}
-
 	head := func(s string, n int) (rows [][]string) {
 		scanner := bufio.NewScanner(strings.NewReader(s))
 		for scanner.Scan() && (len(rows) < n || n < 0) {
@@ -852,7 +794,7 @@ type runOptions struct {
 	Format      string
 	System      string
 	Images      []api.ImageData
-	Options     map[string]any
+	Options     map[string]interface{}
 	MultiModal  bool
 	KeepAlive   *api.Duration
 }
@@ -1254,7 +1196,6 @@ func NewCLI() *cobra.Command {
 	showCmd.Flags().Bool("parameters", false, "Show parameters of a model")
 	showCmd.Flags().Bool("template", false, "Show template of a model")
 	showCmd.Flags().Bool("system", false, "Show system message of a model")
-	showCmd.Flags().BoolP("verbose", "v", false, "Show detailed model information")
 
 	runCmd := &cobra.Command{
 		Use:     "run MODEL [PROMPT]",

cmd/cmd_test.go

@@ -16,7 +16,6 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/ollama/ollama/api"
-	"github.com/ollama/ollama/types/model"
 )
 
 func TestShowInfo(t *testing.T) {
@@ -28,7 +27,7 @@ func TestShowInfo(t *testing.T) {
 				ParameterSize:     "7B",
 				QuantizationLevel: "FP16",
 			},
-		}, false, &b); err != nil {
+		}, &b); err != nil {
 			t.Fatal(err)
 		}
 
@@ -58,7 +57,7 @@ func TestShowInfo(t *testing.T) {
 				ParameterSize:     "7B",
 				QuantizationLevel: "FP16",
 			},
-		}, false, &b); err != nil {
+		}, &b); err != nil {
 			t.Fatal(err)
 		}
 
@@ -69,60 +68,6 @@ func TestShowInfo(t *testing.T) {
     embedding length    0       
     quantization        FP16    
 
-`
-		if diff := cmp.Diff(expect, b.String()); diff != "" {
-			t.Errorf("unexpected output (-want +got):\n%s", diff)
-		}
-	})
-
-	t.Run("verbose model", func(t *testing.T) {
-		var b bytes.Buffer
-		if err := showInfo(&api.ShowResponse{
-			Details: api.ModelDetails{
-				Family:            "test",
-				ParameterSize:     "8B",
-				QuantizationLevel: "FP16",
-			},
-			Parameters: `
-			stop up`,
-			ModelInfo: map[string]any{
-				"general.architecture":    "test",
-				"general.parameter_count": float64(8_000_000_000),
-				"some.true_bool":          true,
-				"some.false_bool":         false,
-				"test.context_length":     float64(1000),
-				"test.embedding_length":   float64(11434),
-			},
-			Tensors: []api.Tensor{
-				{Name: "blk.0.attn_k.weight", Type: "BF16", Shape: []uint64{42, 3117}},
-				{Name: "blk.0.attn_q.weight", Type: "FP16", Shape: []uint64{3117, 42}},
-			},
-		}, true, &b); err != nil {
-			t.Fatal(err)
-		}
-
-		expect := `  Model
-    architecture        test     
-    parameters          8B       
-    context length      1000     
-    embedding length    11434    
-    quantization        FP16     
-
-  Parameters
-    stop    up    
-
-  Metadata
-    general.architecture       test     
-    general.parameter_count    8e+09    
-    some.false_bool            false    
-    some.true_bool             true     
-    test.context_length        1000     
-    test.embedding_length      11434    
-
-  Tensors
-    blk.0.attn_k.weight    BF16    [42 3117]    
-    blk.0.attn_q.weight    FP16    [3117 42]    
-
 `
 		if diff := cmp.Diff(expect, b.String()); diff != "" {
 			t.Errorf("unexpected output (-want +got):\n%s", diff)
@@ -144,7 +89,7 @@ func TestShowInfo(t *testing.T) {
 			stop you
 			stop up
 			temperature 99`,
-		}, false, &b); err != nil {
+		}, &b); err != nil {
 			t.Fatal(err)
 		}
 
@@ -181,7 +126,7 @@ func TestShowInfo(t *testing.T) {
 				"clip.vision.embedding_length": float64(0),
 				"clip.vision.projection_dim":   float64(0),
 			},
-		}, false, &b); err != nil {
+		}, &b); err != nil {
 			t.Fatal(err)
 		}
 
@@ -214,7 +159,7 @@ func TestShowInfo(t *testing.T) {
 Ahoy, matey!
 Weigh anchor!
 			`,
-		}, false, &b); err != nil {
+		}, &b); err != nil {
 			t.Fatal(err)
 		}
 
@@ -243,7 +188,7 @@ Weigh anchor!
 				QuantizationLevel: "FP16",
 			},
 			License: license,
-		}, false, &b); err != nil {
+		}, &b); err != nil {
 			t.Fatal(err)
 		}
 
@@ -261,34 +206,6 @@ Weigh anchor!
 			t.Errorf("unexpected output (-want +got):\n%s", diff)
 		}
 	})
-
-	t.Run("capabilities", func(t *testing.T) {
-		var b bytes.Buffer
-		if err := showInfo(&api.ShowResponse{
-			Details: api.ModelDetails{
-				Family:            "test",
-				ParameterSize:     "7B",
-				QuantizationLevel: "FP16",
-			},
-			Capabilities: []model.Capability{model.CapabilityVision, model.CapabilityTools},
-		}, false, &b); err != nil {
-			t.Fatal(err)
-		}
-
-		expect := "  Model\n" +
-			"    architecture    test    \n" +
-			"    parameters      7B      \n" +
-			"    quantization    FP16    \n" +
-			"\n" +
-			"  Capabilities\n" +
-			"    vision    \n" +
-			"    tools     \n" +
-			"\n"
-
-		if diff := cmp.Diff(expect, b.String()); diff != "" {
-			t.Errorf("unexpected output (-want +got):\n%s", diff)
-		}
-	})
 }
 
 func TestDeleteHandler(t *testing.T) {
@@ -790,132 +707,3 @@ func TestCreateHandler(t *testing.T) {
 		})
 	}
 }
-
-func TestNewCreateRequest(t *testing.T) {
-	tests := []struct {
-		name     string
-		from     string
-		opts     runOptions
-		expected *api.CreateRequest
-	}{
-		{
-			"basic test",
-			"newmodel",
-			runOptions{
-				Model:       "mymodel",
-				ParentModel: "",
-				Prompt:      "You are a fun AI agent",
-				Messages:    []api.Message{},
-				WordWrap:    true,
-			},
-			&api.CreateRequest{
-				From:  "mymodel",
-				Model: "newmodel",
-			},
-		},
-		{
-			"parent model test",
-			"newmodel",
-			runOptions{
-				Model:       "mymodel",
-				ParentModel: "parentmodel",
-				Messages:    []api.Message{},
-				WordWrap:    true,
-			},
-			&api.CreateRequest{
-				From:  "parentmodel",
-				Model: "newmodel",
-			},
-		},
-		{
-			"parent model as filepath test",
-			"newmodel",
-			runOptions{
-				Model:       "mymodel",
-				ParentModel: "/some/file/like/etc/passwd",
-				Messages:    []api.Message{},
-				WordWrap:    true,
-			},
-			&api.CreateRequest{
-				From:  "mymodel",
-				Model: "newmodel",
-			},
-		},
-		{
-			"parent model as windows filepath test",
-			"newmodel",
-			runOptions{
-				Model:       "mymodel",
-				ParentModel: "D:\\some\\file\\like\\etc\\passwd",
-				Messages:    []api.Message{},
-				WordWrap:    true,
-			},
-			&api.CreateRequest{
-				From:  "mymodel",
-				Model: "newmodel",
-			},
-		},
-		{
-			"options test",
-			"newmodel",
-			runOptions{
-				Model:       "mymodel",
-				ParentModel: "parentmodel",
-				Options: map[string]any{
-					"temperature": 1.0,
-				},
-			},
-			&api.CreateRequest{
-				From:  "parentmodel",
-				Model: "newmodel",
-				Parameters: map[string]any{
-					"temperature": 1.0,
-				},
-			},
-		},
-		{
-			"messages test",
-			"newmodel",
-			runOptions{
-				Model:       "mymodel",
-				ParentModel: "parentmodel",
-				System:      "You are a fun AI agent",
-				Messages: []api.Message{
-					{
-						Role:    "user",
-						Content: "hello there!",
-					},
-					{
-						Role:    "assistant",
-						Content: "hello to you!",
-					},
-				},
-				WordWrap: true,
-			},
-			&api.CreateRequest{
-				From:   "parentmodel",
-				Model:  "newmodel",
-				System: "You are a fun AI agent",
-				Messages: []api.Message{
-					{
-						Role:    "user",
-						Content: "hello there!",
-					},
-					{
-						Role:    "assistant",
-						Content: "hello to you!",
-					},
-				},
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			actual := NewCreateRequest(tt.from, tt.opts)
-			if !cmp.Equal(actual, tt.expected) {
-				t.Errorf("expected output %#v, got %#v", tt.expected, actual)
-			}
-		})
-	}
-}

cmd/interactive.go

@@ -18,7 +18,6 @@ import (
 	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/readline"
 	"github.com/ollama/ollama/types/errtypes"
-	"github.com/ollama/ollama/types/model"
 )
 
 type MultilineState int
@@ -348,7 +347,7 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 
 				switch args[1] {
 				case "info":
-					_ = showInfo(resp, false, os.Stderr)
+					_ = showInfo(resp, os.Stderr)
 				case "license":
 					if resp.License == "" {
 						fmt.Println("No license was specified for this model.")
@@ -460,16 +459,9 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 }
 
 func NewCreateRequest(name string, opts runOptions) *api.CreateRequest {
-	parentModel := opts.ParentModel
-
-	modelName := model.ParseName(parentModel)
-	if !modelName.IsValid() {
-		parentModel = ""
-	}
-
 	req := &api.CreateRequest{
-		Model: name,
-		From:  cmp.Or(parentModel, opts.Model),
+		Name: name,
+		From: cmp.Or(opts.ParentModel, opts.Model),
 	}
 
 	if opts.System != "" {

convert/convert.go

@@ -201,7 +201,7 @@ func ConvertModel(fsys fs.FS, ws io.WriteSeeker) error {
 	case "CohereForCausalLM":
 		conv = &commandrModel{}
 	default:
-		return fmt.Errorf("unsupported architecture %q", p.Architectures[0])
+		return errors.New("unsupported architecture")
 	}
 
 	if err := json.Unmarshal(bts, conv); err != nil {

convert/convert_gemma3.go

@@ -87,7 +87,7 @@ func (p *gemma3Model) KV(t *Tokenizer) ggml.KV {
 		kv["gemma3.embedding_length"] = p.HiddenSize
 		kv["gemma3.feed_forward_length"] = p.IntermediateSize
 	default:
-		kv["gemma3.context_length"] = cmp.Or(p.MaxPositionEmbeddings, 131072)
+		kv["gemma3.context_length"] = cmp.Or(p.MaxPositionEmbeddings, 8192)
 		kv["gemma3.embedding_length"] = p.TextModel.HiddenSize
 		kv["gemma3.feed_forward_length"] = p.TextModel.IntermediateSize
 		kv["gemma3.attention.sliding_window"] = p.TextModel.SlidingWindow

convert/sentencepiece/sentencepiece_model.pb.go

@@ -1360,7 +1360,7 @@ func file_sentencepiece_model_proto_rawDescGZIP() []byte {
 
 var file_sentencepiece_model_proto_enumTypes = make([]protoimpl.EnumInfo, 2)
 var file_sentencepiece_model_proto_msgTypes = make([]protoimpl.MessageInfo, 6)
-var file_sentencepiece_model_proto_goTypes = []any{
+var file_sentencepiece_model_proto_goTypes = []interface{}{
 	(TrainerSpec_ModelType)(0),         // 0: sentencepiece.TrainerSpec.ModelType
 	(ModelProto_SentencePiece_Type)(0), // 1: sentencepiece.ModelProto.SentencePiece.Type
 	(*TrainerSpec)(nil),                // 2: sentencepiece.TrainerSpec
@@ -1392,7 +1392,7 @@ func file_sentencepiece_model_proto_init() {
 		return
 	}
 	if !protoimpl.UnsafeEnabled {
-		file_sentencepiece_model_proto_msgTypes[0].Exporter = func(v any, i int) any {
+		file_sentencepiece_model_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
 			switch v := v.(*TrainerSpec); i {
 			case 0:
 				return &v.state
@@ -1406,7 +1406,7 @@ func file_sentencepiece_model_proto_init() {
 				return nil
 			}
 		}
-		file_sentencepiece_model_proto_msgTypes[1].Exporter = func(v any, i int) any {
+		file_sentencepiece_model_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
 			switch v := v.(*NormalizerSpec); i {
 			case 0:
 				return &v.state
@@ -1420,7 +1420,7 @@ func file_sentencepiece_model_proto_init() {
 				return nil
 			}
 		}
-		file_sentencepiece_model_proto_msgTypes[2].Exporter = func(v any, i int) any {
+		file_sentencepiece_model_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
 			switch v := v.(*SelfTestData); i {
 			case 0:
 				return &v.state
@@ -1434,7 +1434,7 @@ func file_sentencepiece_model_proto_init() {
 				return nil
 			}
 		}
-		file_sentencepiece_model_proto_msgTypes[3].Exporter = func(v any, i int) any {
+		file_sentencepiece_model_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
 			switch v := v.(*ModelProto); i {
 			case 0:
 				return &v.state
@@ -1448,7 +1448,7 @@ func file_sentencepiece_model_proto_init() {
 				return nil
 			}
 		}
-		file_sentencepiece_model_proto_msgTypes[4].Exporter = func(v any, i int) any {
+		file_sentencepiece_model_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
 			switch v := v.(*SelfTestData_Sample); i {
 			case 0:
 				return &v.state
@@ -1460,7 +1460,7 @@ func file_sentencepiece_model_proto_init() {
 				return nil
 			}
 		}
-		file_sentencepiece_model_proto_msgTypes[5].Exporter = func(v any, i int) any {
+		file_sentencepiece_model_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
 			switch v := v.(*ModelProto_SentencePiece); i {
 			case 0:
 				return &v.state

discover/cpu_common.go

@@ -12,7 +12,7 @@ func IsNUMA() bool {
 		// numa support in llama.cpp is linux only
 		return false
 	}
-	ids := map[string]any{}
+	ids := map[string]interface{}{}
 	packageIds, _ := filepath.Glob("/sys/devices/system/cpu/cpu*/topology/physical_package_id")
 	for _, packageId := range packageIds {
 		id, err := os.ReadFile(packageId)

discover/gpu_linux.go

@@ -111,7 +111,6 @@ func GetCPUDetails() ([]CPU, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer file.Close()
 	return linuxCPUDetails(file)
 }
 
@@ -169,11 +168,13 @@ func linuxCPUDetails(file io.Reader) ([]CPU, error) {
 	for id, s := range socketByID {
 		s.CoreCount = len(coreBySocket[id])
 		s.ThreadCount = 0
+		for _, tc := range threadsByCoreBySocket[id] {
+			s.ThreadCount += tc
+		}
 
 		// This only works if HT is enabled, consider a more reliable model, maybe cache size comparisons?
 		efficiencyCoreCount := 0
 		for _, threads := range threadsByCoreBySocket[id] {
-			s.ThreadCount += threads
 			if threads == 1 {
 				efficiencyCoreCount++
 			}

docs/api.md

@@ -558,10 +558,6 @@ Final response:
 {
   "model": "llama3.2",
   "created_at": "2023-08-04T19:22:45.499127Z",
-  "message": {
-    "role": "assistant",
-    "content": ""
-  },
   "done": true,
   "total_duration": 4883583458,
   "load_duration": 1334875,
@@ -1217,7 +1213,7 @@ Show information about a model including details, modelfile, template, parameter
 
 ```shell
 curl http://localhost:11434/api/show -d '{
-  "model": "llava"
+  "model": "llama3.2"
 }'
 ```
 
@@ -1260,11 +1256,7 @@ curl http://localhost:11434/api/show -d '{
     "tokenizer.ggml.pre": "llama-bpe",
     "tokenizer.ggml.token_type": [],        // populates if `verbose=true`
     "tokenizer.ggml.tokens": []             // populates if `verbose=true`
-  },
-  "capabilities": [
-    "completion",
-    "vision"
-  ],
+  }
 }
 ```
 

docs/benchmark.md

@@ -1,59 +0,0 @@
-# Benchmark
-
-Go benchmark tests that measure end-to-end performance of a running Ollama server. Run these tests to evaluate model inference performance on your hardware and measure the impact of code changes.
-
-## When to use
-
-Run these benchmarks when:
-- Making changes to the model inference engine
-- Modifying model loading/unloading logic
-- Changing prompt processing or token generation code
-- Implementing a new model architecture
-- Testing performance across different hardware setups
-
-## Prerequisites
-- Ollama server running locally with `ollama serve` on `127.0.0.1:11434`
-## Usage and Examples
-
->[!NOTE]
->All commands must be run from the root directory of the Ollama project.
-
-Basic syntax:
-```bash
-go test -bench=. ./benchmark/... -m $MODEL_NAME
-```
-
-Required flags:
-- `-bench=.`: Run all benchmarks
-- `-m`: Model name to benchmark
-
-Optional flags:
-- `-count N`: Number of times to run the benchmark (useful for statistical analysis)
-- `-timeout T`: Maximum time for the benchmark to run (e.g. "10m" for 10 minutes)
-
-Common usage patterns:
-
-Single benchmark run with a model specified:
-```bash
-go test -bench=. ./benchmark/... -m llama3.3
-```
-
-## Output metrics
-
-The benchmark reports several key metrics:
-
-- `gen_tok/s`: Generated tokens per second
-- `prompt_tok/s`: Prompt processing tokens per second
-- `ttft_ms`: Time to first token in milliseconds
-- `load_ms`: Model load time in milliseconds
-- `gen_tokens`: Total tokens generated
-- `prompt_tokens`: Total prompt tokens processed
-
-Each benchmark runs two scenarios:
-- Cold start: Model is loaded from disk for each test
-- Warm start: Model is pre-loaded in memory
-
-Three prompt lengths are tested for each scenario:
-- Short prompt (100 tokens)
-- Medium prompt (500 tokens)
-- Long prompt (1000 tokens)

docs/faq.md

@@ -20,13 +20,7 @@ Please refer to the [GPU docs](./gpu.md).
 
 ## How can I specify the context window size?
 
-By default, Ollama uses a context window size of 2048 tokens. 
-
-This can be overridden with the `OLLAMA_CONTEXT_LENGTH` environment variable. For example, to set the default context window to 8K, use: 
-
-```shell
-OLLAMA_CONTEXT_LENGTH=8192 ollama serve
-```
+By default, Ollama uses a context window size of 2048 tokens. This can be overridden with the `OLLAMA_CONTEXT_LENGTH` environment variable. For example, to set the default context length to 8K, use: `OLLAMA_CONTEXT_LENGTH=8192 ollama serve`.
 
 To change this when using `ollama run`, use `/set parameter`:
 
@@ -193,13 +187,6 @@ cloudflared tunnel --url http://localhost:11434 --http-host-header="localhost:11
 
 Ollama allows cross-origin requests from `127.0.0.1` and `0.0.0.0` by default. Additional origins can be configured with `OLLAMA_ORIGINS`.
 
-For browser extensions, you'll need to explicitly allow the extension's origin pattern. Set `OLLAMA_ORIGINS` to include `chrome-extension://*`, `moz-extension://*`, and `safari-web-extension://*` if you wish to allow all browser extensions access, or specific extensions as needed:
-
-```
-# Allow all Chrome, Firefox, and Safari extensions
-OLLAMA_ORIGINS=chrome-extension://*,moz-extension://*,safari-web-extension://* ollama serve
-```
-
 Refer to the section [above](#how-do-i-configure-ollama-server) for how to set environment variables on your platform.
 
 ## Where are models stored?

docs/troubleshooting.md

@@ -9,7 +9,7 @@ cat ~/.ollama/logs/server.log
 On **Linux** systems with systemd, the logs can be found with this command:
 
 ```shell
-journalctl -u ollama --no-pager --follow --pager-end 
+journalctl -u ollama --no-pager
 ```
 
 When you run Ollama in a **container**, the logs go to stdout/stderr in the container:

format/time_test.go

@@ -5,7 +5,7 @@ import (
 	"time"
 )
 
-func assertEqual(t *testing.T, a any, b any) {
+func assertEqual(t *testing.T, a interface{}, b interface{}) {
 	if a != b {
 		t.Errorf("Assert failed, expected %v, got %v", b, a)
 	}

fs/ggml/ggml.go

@@ -327,10 +327,6 @@ func (t Tensor) Size() uint64 {
 	return t.parameters() * t.typeSize() / t.blockSize()
 }
 
-func (t Tensor) Type() string {
-	return fileType(t.Kind).String()
-}
-
 type container interface {
 	Name() string
 	Decode(io.ReadSeeker) (model, error)
@@ -413,7 +409,7 @@ func Decode(rs io.ReadSeeker, maxArraySize int) (*GGML, int64, error) {
 	}, offset, nil
 }
 
-func (f GGML) GraphSize(context, batch uint64, numParallel int, kvCacheType string) (kv []uint64, partialOffload, fullOffload uint64) {
+func (f GGML) GraphSize(context, batch uint64, kvCacheType string) (kv, partialOffload, fullOffload uint64) {
 	embedding := f.KV().EmbeddingLength()
 	heads := f.KV().HeadCount()
 	headsKV := f.KV().HeadCountKV()
@@ -426,10 +422,7 @@ func (f GGML) GraphSize(context, batch uint64, numParallel int, kvCacheType stri
 	layers := f.Tensors().GroupLayers()
 
 	bytesPerElement := kvCacheBytesPerElement(kvCacheType)
-	kv = make([]uint64, f.KV().BlockCount())
-	for i := range kv {
-		kv[i] = uint64(float64(context*(embeddingHeadsK+embeddingHeadsV)*headsKV) * bytesPerElement)
-	}
+	kv = uint64(float64(context*f.KV().BlockCount()*(embeddingHeadsK+embeddingHeadsV)*headsKV) * bytesPerElement)
 
 	switch f.KV().Architecture() {
 	case "llama":
@@ -463,14 +456,16 @@ func (f GGML) GraphSize(context, batch uint64, numParallel int, kvCacheType stri
 	case "mllama":
 		var visionTokens, tiles uint64 = 1601, 4
 
-		crossAttentionLayers := f.KV().Uints("attention.cross_attention_layers")
-		for i := range kv {
-			if slices.Contains(crossAttentionLayers, uint32(i)) {
-				kv[i] = headsKV * (embeddingHeadsK + embeddingHeadsV) *
-					4 * // sizeof(float32)
-					visionTokens *
-					tiles
-			}
+		if crossAttentionLayers, ok := f.KV()["mllama.attention.cross_attention_layers"].(*array); ok {
+			kv = headsKV *
+				(embeddingHeadsK + embeddingHeadsV) * // one for K, one for V
+				(2* // sizeof(float16)
+					(f.KV().BlockCount()-uint64(crossAttentionLayers.size))* // num non-cross attention layers
+					context +
+					4* // sizeof(float32)
+						uint64(crossAttentionLayers.size)* // num cross attention layers
+						visionTokens*
+						tiles)
 		}
 
 		fullOffload = max(
@@ -506,20 +501,6 @@ func (f GGML) GraphSize(context, batch uint64, numParallel int, kvCacheType stri
 				4*embeddingHeadsK*context*8+
 				embedding*embeddingHeadsK*heads*9/16,
 		)
-
-		// Gemma2 also has sliding window attention but we only have an optimized implementation in the Ollama
-		// engine. Gemma3 always uses the Ollama engine.
-		if f.KV().Architecture() == "gemma3" {
-			const gemma3GlobalCacheCount = 6
-			slidingWindow := (uint64(numParallel) * uint64(f.KV().Uint("attention.sliding_window"))) + batch
-			for i := range kv {
-				// Every 6th layer is a global layer, which is the full context size that has already been set. The other
-				// layers are the smaller local (sliding) layers.
-				if (i+1)%gemma3GlobalCacheCount != 0 {
-					kv[i] = uint64(float64(slidingWindow*(embeddingHeadsK+embeddingHeadsV)*headsKV) * bytesPerElement)
-				}
-			}
-		}
 	case "command-r":
 		fullOffload = max(
 			4*batch*(embedding+vocab),
@@ -598,52 +579,39 @@ func (f GGML) GraphSize(context, batch uint64, numParallel int, kvCacheType stri
 }
 
 func (llm GGML) VisionGraphSize() (weights, graphSize uint64) {
-	if llm.KV().Uint("vision.block_count") == 0 {
-		return
-	}
-
-	for name, layer := range llm.Tensors().GroupLayers() {
-		if name == "v" || strings.HasPrefix(name, "v.") {
-			for _, tensor := range layer {
-				weights += tensor.Size()
-			}
-		}
-	}
-
-	imageSize := uint64(llm.KV().Uint("vision.image_size"))
-	patchSize := uint64(llm.KV().Uint("vision.patch_size"))
-	if patchSize == 0 {
-		slog.Warn("unknown patch size for vision model")
-		return
-	}
-
-	numChannels := uint64(llm.KV().Uint("vision.num_channels"))
-
-	numPatches := (imageSize / patchSize) * (imageSize / patchSize)
-	if _, ok := llm.Tensors().GroupLayers()["v"]["class_embd"]; ok {
-		numPatches++
-	}
-
-	headCount := uint64(llm.KV().Uint("vision.attention.head_count"))
-	embeddingLength := uint64(llm.KV().Uint("vision.embedding_length"))
-
 	switch llm.KV().Architecture() {
 	case "mllama":
+		for _, layer := range llm.Tensors().GroupLayers()["v"] {
+			weights += layer.Size()
+		}
+
+		kv := func(n string) uint64 {
+			if v, ok := llm.KV()["mllama.vision."+n].(uint32); ok {
+				return uint64(v)
+			}
+
+			return 0
+		}
+
+		imageSize := kv("image_size")
+
+		maxNumTiles := kv("max_num_tiles")
+		embeddingLength := kv("embedding_length")
+		headCount := kv("attention.head_count")
+
+		numPatches := (imageSize / kv("patch_size")) * (imageSize / kv("patch_size"))
+		if _, ok := llm.Tensors().GroupLayers()["v"]["class_embd"]; ok {
+			numPatches++
+		}
+
 		numPaddedPatches := numPatches + 8 - (numPatches%8)%8
 
-		maxNumTiles := uint64(llm.KV().Uint("vision.max_num_tiles"))
-
 		graphSize = 4 * (8 +
-			imageSize*imageSize*numChannels*maxNumTiles +
+			imageSize*imageSize*kv("num_channels")*maxNumTiles +
 			embeddingLength*numPatches*maxNumTiles +
 			9*embeddingLength*numPaddedPatches*maxNumTiles +
 			numPaddedPatches*maxNumTiles*numPaddedPatches*maxNumTiles*headCount)
-	case "gemma3":
-		graphSize = 4 * (imageSize*imageSize*numChannels +
-			embeddingLength*patchSize +
-			numPatches*numPatches*headCount)
 	}
-
 	return weights, graphSize
 }
 

integration/basic_test.go

@@ -22,7 +22,7 @@ func TestOrcaMiniBlueSky(t *testing.T) {
 		Model:  "orca-mini",
 		Prompt: "why is the sky blue?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 		},
@@ -39,7 +39,7 @@ func TestUnicode(t *testing.T) {
 		Model:  "deepseek-coder-v2:16b-lite-instruct-q2_K",
 		Prompt: "天空为什么是蓝色的?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 			// Workaround deepseek context shifting bug
@@ -61,7 +61,7 @@ func TestExtendedUnicodeOutput(t *testing.T) {
 		Model:  "gemma2:2b",
 		Prompt: "Output some smily face emoji",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 		},
@@ -96,7 +96,7 @@ func TestUnicodeModelDir(t *testing.T) {
 		Model:  "orca-mini",
 		Prompt: "why is the sky blue?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 		},

integration/concurrency_test.go

@@ -25,7 +25,7 @@ func TestMultiModelConcurrency(t *testing.T) {
 				Prompt:    "why is the ocean blue?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},
@@ -34,7 +34,7 @@ func TestMultiModelConcurrency(t *testing.T) {
 				Prompt:    "what is the origin of the us thanksgiving holiday?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},

integration/context_test.go

@@ -23,7 +23,7 @@ func TestLongInputContext(t *testing.T) {
 		Model:  "llama2",
 		Prompt: "Oh, don’t speak to me of Austria. Perhaps I don’t understand things, but Austria never has wished, and does not wish, for war. She is betraying us! Russia alone must save Europe. Our gracious sovereign recognizes his high vocation and will be true to it. That is the one thing I have faith in! Our good and wonderful sovereign has to perform the noblest role on earth, and he is so virtuous and noble that God will not forsake him. He will fulfill his vocation and crush the hydra of revolution, which has become more terrible than ever in the person of this murderer and villain! We alone must avenge the blood of the just one.... Whom, I ask you, can we rely on?... England with her commercial spirit will not and cannot understand the Emperor Alexander’s loftiness of soul. She has refused to evacuate Malta. She wanted to find, and still seeks, some secret motive in our actions. What answer did Novosíltsev get? None. The English have not understood and cannot understand the self-abnegation of our Emperor who wants nothing for himself, but only desires the good of mankind. And what have they promised? Nothing! And what little they have promised they will not perform! Prussia has always declared that Buonaparte is invincible, and that all Europe is powerless before him.... And I don’t believe a word that Hardenburg says, or Haugwitz either. This famous Prussian neutrality is just a trap. I have faith only in God and the lofty destiny of our adored monarch. He will save Europe! What country is this referring to?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 			"num_ctx":     128,
@@ -50,7 +50,7 @@ func TestContextExhaustion(t *testing.T) {
 		Model:  "llama2",
 		Prompt: "Write me a story with a ton of emojis?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 			"num_ctx":     128,

integration/llm_image_test.go

@@ -19,7 +19,7 @@ func TestIntegrationLlava(t *testing.T) {
 		Model:  "llava:7b",
 		Prompt: "what does the text in this image say?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"seed":        42,
 			"temperature": 0.0,
 		},
@@ -47,7 +47,7 @@ func TestIntegrationMllama(t *testing.T) {
 		Model:  "x/llama3.2-vision",
 		Prompt: "what does the text in this image say?",
 		Stream: &stream,
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"seed":        42,
 			"temperature": 0.0,
 		},
@@ -66,35 +66,6 @@ func TestIntegrationMllama(t *testing.T) {
 	DoGenerate(ctx, t, client, req, []string{resp}, 240*time.Second, 30*time.Second)
 }
 
-func TestIntegrationSplitBatch(t *testing.T) {
-	image, err := base64.StdEncoding.DecodeString(imageEncoding)
-	require.NoError(t, err)
-	req := api.GenerateRequest{
-		Model: "gemma3:4b",
-		// Fill up a chunk of the batch so the image will partially spill over into the next one
-		System: "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed aliquet, justo in malesuada lobortis, odio ligula volutpat quam, quis faucibus ipsum magna quis sapien. Aliquam in venenatis diam, eu viverra magna. Phasellus imperdiet hendrerit volutpat. Vivamus sem ex, facilisis placerat felis non, dictum elementum est. Phasellus aliquam imperdiet lacus, eget placerat ligula sodales vel. Pellentesque nec auctor mi. Curabitur arcu nisi, faucibus eget nunc id, viverra interdum mi. Curabitur ornare ipsum ex, ac euismod ex aliquam in. Vestibulum id magna at purus accumsan fermentum. Proin scelerisque posuere nunc quis interdum. Maecenas sed mollis nisl. Etiam vitae ipsum interdum, placerat est quis, tincidunt velit. Nullam tempor nibh non lorem volutpat efficitur. Cras laoreet diam imperdiet ipsum auctor bibendum. Suspendisse ultrices urna sed metus sagittis suscipit. Quisque ullamcorper aliquam nibh ut mollis. Aenean dapibus mauris pharetra, venenatis elit ac, hendrerit odio. Cras vestibulum erat tempor, lobortis justo eu, lobortis ipsum. Nam laoreet dapibus sem. Proin vel diam ultrices, elementum ante et, ornare lectus. Proin eu accumsan nisl. Praesent ac ex vitae ipsum vulputate tristique facilisis sit amet lacus. Nullam faucibus magna a pellentesque pretium. Nunc lacinia ullamcorper sollicitudin. Donec vitae accumsan turpis, sed porttitor est. Donec porttitor mi vitae augue faucibus, vel mollis diam tincidunt.",
-		Prompt: "what does the text in this image say?",
-		Stream: &stream,
-		Options: map[string]any{
-			"seed":        42,
-			"temperature": 0.0,
-		},
-		Images: []api.ImageData{
-			image,
-		},
-	}
-
-	// Note: sometimes it returns "the ollamas" sometimes "the ollams"
-	resp := "the ollam"
-	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
-	defer cancel()
-	client, _, cleanup := InitServerConnection(ctx, t)
-	defer cleanup()
-	require.NoError(t, PullIfMissing(ctx, client, req.Model))
-	// llava models on CPU can be quite slow to start,
-	DoGenerate(ctx, t, client, req, []string{resp}, 120*time.Second, 30*time.Second)
-}
-
 const imageEncoding = `iVBORw0KGgoAAAANSUhEUgAAANIAAAB4CAYAAACHHqzKAAAAAXNSR0IArs4c6QAAAIRlWElmTU0AKgAAAAgABQESAAMAAAABAAEAAAEaAAUAAAABAAAASgEb
 AAUAAAABAAAAUgEoAAMAAAABAAIAAIdpAAQAAAABAAAAWgAAAAAAAABIAAAAAQAAAEgAAAABAAOgAQADAAAAAQABAACgAgAEAAAAAQAAANKgAwAEAAAAAQAA
 AHgAAAAAXdsepgAAAAlwSFlzAAALEwAACxMBAJqcGAAAAVlpVFh0WE1MOmNvbS5hZG9iZS54bXAAAAAAADx4OnhtcG1ldGEgeG1sbnM6eD0iYWRvYmU6bnM6

integration/llm_test.go

@@ -20,7 +20,7 @@ var (
 			Model:  "orca-mini",
 			Prompt: "why is the ocean blue?",
 			Stream: &stream,
-			Options: map[string]any{
+			Options: map[string]interface{}{
 				"seed":        42,
 				"temperature": 0.0,
 			},
@@ -28,7 +28,7 @@ var (
 			Model:  "orca-mini",
 			Prompt: "what is the origin of the us thanksgiving holiday?",
 			Stream: &stream,
-			Options: map[string]any{
+			Options: map[string]interface{}{
 				"seed":        42,
 				"temperature": 0.0,
 			},

integration/max_queue_test.go

@@ -32,7 +32,7 @@ func TestMaxQueue(t *testing.T) {
 	req := api.GenerateRequest{
 		Model:  "orca-mini",
 		Prompt: "write a long historical fiction story about christopher columbus.  use at least 10 facts from his actual journey",
-		Options: map[string]any{
+		Options: map[string]interface{}{
 			"seed":        42,
 			"temperature": 0.0,
 		},

integration/utils_test.go

@@ -291,7 +291,7 @@ func GenerateRequests() ([]api.GenerateRequest, [][]string) {
 				Prompt:    "why is the ocean blue?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},
@@ -300,7 +300,7 @@ func GenerateRequests() ([]api.GenerateRequest, [][]string) {
 				Prompt:    "why is the color of dirt brown?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},
@@ -309,7 +309,7 @@ func GenerateRequests() ([]api.GenerateRequest, [][]string) {
 				Prompt:    "what is the origin of the us thanksgiving holiday?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},
@@ -318,7 +318,7 @@ func GenerateRequests() ([]api.GenerateRequest, [][]string) {
 				Prompt:    "what is the origin of independence day?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},
@@ -327,7 +327,7 @@ func GenerateRequests() ([]api.GenerateRequest, [][]string) {
 				Prompt:    "what is the composition of air?",
 				Stream:    &stream,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
-				Options: map[string]any{
+				Options: map[string]interface{}{
 					"seed":        42,
 					"temperature": 0.0,
 				},

kvcache/cache.go

@@ -43,13 +43,8 @@ type Cache interface {
 
 	// ** cache management **
 
-	// Init sets up runtime parameters.
-	// backend: Used to allocate cache data storage and execute management operations (such as defrag)
-	// dtype: The data type for storing cache entries
-	// maxSequences: The maximum number of sequences stored in the cache - across all batches
-	// capacity: The number of cache entries to store, per sequence
-	// maxBatch: The maximum number of tokens that can occur in a single batch
-	Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int)
+	// Init sets up runtime parameters
+	Init(backend ml.Backend, dtype ml.DType, capacity int32)
 
 	// Close closes the cache and frees resources associated with it
 	Close()
@@ -57,16 +52,11 @@ type Cache interface {
 	// StartForward is called before the start of the model's forward pass.
 	// For each token in the coming batch, there must be a corresponding
 	// entry in positions and seqs.
-	StartForward(ctx ml.Context, batch input.Batch) error
+	StartForward(ctx ml.Context, opts input.Options) error
 
 	// CopyPrefix copies tokens in the range [0, len) from srcSeq to dstSeq
 	CopyPrefix(srcSeq, dstSeq int, len int32)
 
-	// CanResume returns true if the cache can continue with the next token at
-	// the given position and sequence. Assumes that the caller has already
-	// verified the contents of the cache.
-	CanResume(seq int, pos int32) bool
-
 	// Remove deletes tokens in the range [beginIndex, endIndex) from seq. Set
 	// endIndex to math.MaxInt32 to remove everything starting at beginIndex.
 	//

kvcache/causal.go

@@ -20,6 +20,7 @@ type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, e
 // The mask is of shape history size, batch size
 type Causal struct {
 	DType      ml.DType
+	Capacity   int32
 	windowSize int32
 
 	opts CausalOptions
@@ -97,7 +98,7 @@ func NewSWACache(windowSize int32, shift shiftFn) *Causal {
 	}
 }
 
-func (c *Causal) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
+func (c *Causal) Init(backend ml.Backend, dtype ml.DType, capacity int32) {
 	if c.config == nil {
 		var config ml.CacheConfig
 		if cc, ok := backend.(ml.BackendCacheConfig); ok {
@@ -118,16 +119,9 @@ func (c *Causal) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity
 		c.config.MaskDType = ml.DTypeF32
 	}
 
-	var cacheSize int
-	if c.windowSize == math.MaxInt32 || capacity < int(c.windowSize) {
-		cacheSize = maxSequences * capacity
-	} else {
-		cacheSize = (maxSequences * int(c.windowSize)) + maxBatch
-	}
-	cacheSize = roundUp(cacheSize, c.config.CachePadding)
-	c.cells = make([]cacheCell, cacheSize)
-
 	c.DType = dtype
+	c.Capacity = int32(roundUp(int(capacity), c.config.CachePadding))
+	c.cells = make([]cacheCell, c.Capacity)
 	c.cellRanges = make(map[int]cellRange)
 	c.backend = backend
 }
@@ -146,14 +140,12 @@ func (c *Causal) Close() {
 	}
 }
 
-func (c *Causal) StartForward(ctx ml.Context, batch input.Batch) error {
-	c.curBatchSize = len(batch.Positions)
-	c.curSequences = batch.Sequences
-	c.curPositions = batch.Positions
+func (c *Causal) StartForward(ctx ml.Context, opts input.Options) error {
+	c.curBatchSize = len(opts.Positions)
+	c.curSequences = opts.Sequences
+	c.curPositions = opts.Positions
 	c.opts.Except = nil
 
-	c.updateSlidingWindow()
-
 	var err error
 	c.curLoc, err = c.findStartLoc()
 	if errors.Is(err, ErrKvCacheFull) {
@@ -165,8 +157,8 @@ func (c *Causal) StartForward(ctx ml.Context, batch input.Batch) error {
 	}
 
 	c.curCellRange = newRange()
-	for i, pos := range batch.Positions {
-		seq := batch.Sequences[i]
+	for i, pos := range opts.Positions {
+		seq := opts.Sequences[i]
 
 		c.cells[c.curLoc+i] = cacheCell{pos: pos, sequences: []int{seq}}
 
@@ -218,51 +210,7 @@ func (c *Causal) findStartLoc() (int, error) {
 		}
 	}
 
-	return 0, fmt.Errorf("%w (length: %v)", ErrKvCacheFull, len(c.cells))
-}
-
-func (c *Causal) updateSlidingWindow() {
-	if c.windowSize == math.MaxInt32 {
-		return
-	}
-
-	// create a map of unique sequences to the lowest position in that sequence
-	lowestPos := make(map[int]int32)
-	for i := range c.curPositions {
-		seq := c.curSequences[i]
-
-		pos, ok := lowestPos[seq]
-		if !ok {
-			pos = c.curPositions[i]
-		} else if c.curPositions[i] < pos {
-			pos = c.curPositions[i]
-		}
-
-		lowestPos[seq] = pos
-	}
-
-	// delete any entries that are beyond the window of the oldest position in the sequence
-	for seq, pos := range lowestPos {
-		oldRange, ok := c.cellRanges[seq]
-		if !ok {
-			continue
-		}
-
-		newRange := newRange()
-
-		for i := oldRange.min; i <= oldRange.max; i++ {
-			if slices.Contains(c.cells[i].sequences, seq) {
-				if c.cells[i].pos < pos-c.windowSize {
-					c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == seq })
-				} else {
-					newRange.min = min(newRange.min, i)
-					newRange.max = max(newRange.max, i)
-				}
-			}
-		}
-
-		c.cellRanges[seq] = newRange
-	}
+	return 0, fmt.Errorf("%w (length: %v)", ErrKvCacheFull, c.Capacity)
 }
 
 func roundDown(length, pad int) int {
@@ -317,7 +265,7 @@ func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) {
 	return maskTensor, nil
 }
 
-func (c *Causal) moveCells(ctx ml.Context, src, dst, length int) {
+func (c *Causal) moveCells(ctx ml.Context, src, dst, len int) {
 	for i, key := range c.keys {
 		if key == nil {
 			continue
@@ -327,8 +275,8 @@ func (c *Causal) moveCells(ctx ml.Context, src, dst, length int) {
 		numKVHeads := key.Dim(1)
 		rowSize := key.Stride(2)
 
-		kSrcView := key.View(ctx, rowSize*src, kHeadDim*numKVHeads*length)
-		kDstView := key.View(ctx, rowSize*dst, kHeadDim*numKVHeads*length)
+		kSrcView := key.View(ctx, rowSize*src, kHeadDim*numKVHeads*len)
+		kDstView := key.View(ctx, rowSize*dst, kHeadDim*numKVHeads*len)
 
 		value := c.values[i]
 		var vSrcView, vDstView ml.Tensor
@@ -336,14 +284,14 @@ func (c *Causal) moveCells(ctx ml.Context, src, dst, length int) {
 			vHeadDim := value.Dim(1)
 			elemSize := value.Stride(0)
 
-			vSrcView = value.View(ctx, elemSize*src, length, len(c.cells)*elemSize, vHeadDim*numKVHeads)
-			vDstView = value.View(ctx, elemSize*dst, length, len(c.cells)*elemSize, vHeadDim*numKVHeads)
+			vSrcView = value.View(ctx, elemSize*src, len, int(c.Capacity)*elemSize, vHeadDim*numKVHeads)
+			vDstView = value.View(ctx, elemSize*dst, len, int(c.Capacity)*elemSize, vHeadDim*numKVHeads)
 		} else {
 			vHeadDim := value.Dim(0)
 			rowSize := value.Stride(2)
 
-			vSrcView = value.View(ctx, rowSize*src, vHeadDim*numKVHeads*length)
-			vDstView = value.View(ctx, rowSize*dst, vHeadDim*numKVHeads*length)
+			vSrcView = value.View(ctx, rowSize*src, vHeadDim*numKVHeads*len)
+			vDstView = value.View(ctx, rowSize*dst, vHeadDim*numKVHeads*len)
 		}
 
 		ctx.Forward(
@@ -373,8 +321,7 @@ func (c *Causal) defrag() {
 	ctx := c.backend.NewContext()
 
 	// For every move, 6 tensors are required per layer (2 views and a
-	// copy for each of k and v). We also need to refer to the original
-	// k and v cache tensors - once per layer, not per move.
+	// copy for each of k and v).
 	layers := 0
 	for _, key := range c.keys {
 		if key == nil {
@@ -383,7 +330,7 @@ func (c *Causal) defrag() {
 		layers++
 	}
 
-	maxMoves := (ctx.MaxGraphNodes() - 2*layers) / (6 * layers)
+	maxMoves := ctx.MaxGraphNodes() / (6 * layers)
 	moves := 0
 
 	var pendingSrc, pendingDst, pendingLen int
@@ -532,14 +479,14 @@ func (c *Causal) Put(ctx ml.Context, key, value ml.Tensor) {
 	}
 
 	if _, ok := c.keys[c.curLayer]; !ok {
-		c.keys[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, kHeadDim, numKVHeads, len(c.cells))
+		c.keys[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, kHeadDim, numKVHeads, int(c.Capacity))
 	}
 
 	if _, ok := c.values[c.curLayer]; !ok {
 		if c.config.PermutedV {
-			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, len(c.cells), vHeadDim, numKVHeads)
+			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, int(c.Capacity), vHeadDim, numKVHeads)
 		} else {
-			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, vHeadDim, numKVHeads, len(c.cells))
+			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, vHeadDim, numKVHeads, int(c.Capacity))
 		}
 	}
 
@@ -550,7 +497,7 @@ func (c *Causal) Put(ctx ml.Context, key, value ml.Tensor) {
 		elemSize := c.values[c.curLayer].Stride(0)
 
 		value = value.Permute(ctx, 1, 2, 0, 3)
-		ctx.Forward(value.Copy(ctx, c.values[c.curLayer].View(ctx, elemSize*c.curLoc, batchSize, len(c.cells)*elemSize, vHeadDim*numKVHeads)))
+		ctx.Forward(value.Copy(ctx, c.values[c.curLayer].View(ctx, elemSize*c.curLoc, batchSize, int(c.Capacity)*elemSize, vHeadDim*numKVHeads)))
 	} else {
 		rowSize := c.values[c.curLayer].Stride(2)
 
@@ -581,35 +528,6 @@ func (c *Causal) CopyPrefix(srcSeq, dstSeq int, len int32) {
 	c.cellRanges[dstSeq] = seqRange
 }
 
-func (c *Causal) CanResume(seq int, pos int32) bool {
-	if c.windowSize == math.MaxInt32 {
-		return true
-	}
-
-	seqRange, ok := c.cellRanges[seq]
-	if !ok {
-		return false
-	}
-
-	// for sliding window, check that the window of the new sequence is contained in
-	// the window of what we are storing
-	var last int32 = -1
-	for i := seqRange.min; i <= seqRange.max; i++ {
-		if slices.Contains(c.cells[i].sequences, seq) {
-			last = max(last, c.cells[i].pos)
-		}
-	}
-
-	if last == -1 {
-		return false
-	}
-
-	lastWindowStart := max(0, last-c.windowSize)
-	posWindowStart := max(0, pos-c.windowSize)
-
-	return posWindowStart >= lastWindowStart
-}
-
 func (c *Causal) shift(seq int, beginIndex, offset int32) error {
 	if c.shiftFn == nil {
 		return ErrNotSupported
@@ -664,12 +582,6 @@ func (c *Causal) shift(seq int, beginIndex, offset int32) error {
 }
 
 func (c *Causal) Remove(seq int, beginIndex, endIndex int32) error {
-	// TODO(jessegross): We should check to see if removing the middle of the sequence will
-	// cause the sliding window to encompass tokens that we no longer have. If so, then we
-	// should return an error, which will trigger the runner to evaluate the full history and
-	// rebuild the window. However, if we have multimodal inputs in our history, this reuse
-	// results in use after free, so we don't do it for now.
-
 	var offset int32
 	if endIndex != math.MaxInt32 {
 		offset = beginIndex - endIndex
@@ -684,7 +596,8 @@ func (c *Causal) Remove(seq int, beginIndex, endIndex int32) error {
 			} else {
 				if c.cells[i].pos >= endIndex {
 					if slices.ContainsFunc(c.cells[i].sequences, func(s int) bool { return s != seq }) {
-						return errors.New("shifting cells shared by multiple sequences not supported")
+						// TODO(jessegross): Need to be careful about data shared between sequences
+						return errors.New("shifting on cells shared by multiple sequences not yet implemented")
 					}
 
 					c.cells[i].pos += offset

kvcache/causal_test.go

@@ -25,7 +25,7 @@ func TestStore(t *testing.T) {
 	cache := NewCausalCache(nil)
 	defer cache.Close()
 
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
+	cache.Init(backend, ml.DTypeF16, 16)
 
 	tests := []testCase{
 		{
@@ -58,11 +58,11 @@ func TestSWA(t *testing.T) {
 	cache := NewSWACache(1, nil)
 	defer cache.Close()
 
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
+	cache.Init(backend, ml.DTypeF32, 16)
 
 	tests := []testCase{
 		{
-			name:          "FirstBatch",
+			name:          "SlidingWindow",
 			in:            []float32{1, 2, 3, 4},
 			inShape:       []int{1, 1, 4},
 			seqs:          []int{0, 0, 0, 0},
@@ -71,16 +71,6 @@ func TestSWA(t *testing.T) {
 			expectedShape: []int{1, 1, 4},
 			expectedMask:  []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0},
 		},
-		{
-			name:          "SecondBatch",
-			in:            []float32{5, 6},
-			inShape:       []int{1, 1, 2},
-			seqs:          []int{0, 0},
-			pos:           []int32{4, 5},
-			expected:      []float32{5, 6, 3, 4},
-			expectedShape: []int{1, 1, 4},
-			expectedMask:  []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1))},
-		},
 	}
 
 	testCache(t, backend, cache, tests)
@@ -91,7 +81,7 @@ func TestSequences(t *testing.T) {
 	cache := NewCausalCache(nil)
 	defer cache.Close()
 
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
+	cache.Init(backend, ml.DTypeF16, 16)
 
 	tests := []testCase{
 		{
@@ -126,7 +116,7 @@ func TestRemove(t *testing.T) {
 	})
 	defer cache.Close()
 
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
+	cache.Init(backend, ml.DTypeF16, 16)
 
 	tests := []testCase{
 		{
@@ -191,7 +181,7 @@ func TestDefrag(t *testing.T) {
 	})
 	defer cache.Close()
 
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
+	cache.Init(backend, ml.DTypeF16, 16)
 
 	tests := []testCase{
 		{
@@ -239,7 +229,7 @@ func TestCopy(t *testing.T) {
 	cache := NewCausalCache(func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { return key, nil })
 	defer cache.Close()
 
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
+	cache.Init(backend, ml.DTypeF16, 16)
 
 	tests := []testCase{
 		{
@@ -280,7 +270,7 @@ func testCache(t *testing.T, backend ml.Backend, cache Cache, tests []testCase)
 			context := backend.NewContext()
 			defer context.Close()
 
-			err := cache.StartForward(context, input.Batch{Positions: test.pos, Sequences: test.seqs})
+			err := cache.StartForward(context, input.Options{Positions: test.pos, Sequences: test.seqs})
 			if err != nil {
 				panic(err)
 			}
@@ -300,77 +290,6 @@ func testCache(t *testing.T, backend ml.Backend, cache Cache, tests []testCase)
 	}
 }
 
-func TestCanResume(t *testing.T) {
-	backend := &testBackend{}
-	windowSize := int32(4)
-	cache := NewSWACache(windowSize, nil)
-	defer cache.Close()
-
-	cache.Init(backend, ml.DTypeF16, 1, 16, 16)
-
-	context := backend.NewContext()
-	defer context.Close()
-
-	err := cache.StartForward(context, input.Batch{
-		Positions: []int32{0, 1, 2, 3},
-		Sequences: []int{0, 0, 0, 0},
-	})
-	if err != nil {
-		t.Fatalf("StartForward failed: %v", err)
-	}
-
-	cache.SetLayer(0)
-	tensor, _ := context.FromFloatSlice([]float32{1, 2, 3, 4}, 1, 1, 4)
-	cache.Put(context, tensor, tensor)
-
-	// with window size 4, nothing has slid out of the window yet
-	if !cache.CanResume(0, 0) {
-		t.Errorf("CanResume(0, 0) = false, want true (within window)")
-	}
-	if !cache.CanResume(0, 1) {
-		t.Errorf("CanResume(0, 1) = false, want true (within window)")
-	}
-	if !cache.CanResume(0, 2) {
-		t.Errorf("CanResume(0, 2) = false, want true (within window)")
-	}
-	if !cache.CanResume(0, 3) {
-		t.Errorf("CanResume(0, 3) = false, want true (latest position)")
-	}
-
-	// shift window by adding position 4
-	err = cache.StartForward(context, input.Batch{
-		Positions: []int32{4, 5},
-		Sequences: []int{0, 0},
-	})
-	if err != nil {
-		t.Fatalf("StartForward failed: %v", err)
-	}
-
-	cache.SetLayer(0)
-	tensor, _ = context.FromFloatSlice([]float32{5, 6}, 1, 1, 2)
-	cache.Put(context, tensor, tensor)
-
-	// only the latest position has overlapping windows
-	if cache.CanResume(0, 0) {
-		t.Errorf("after shift: CanResume(0, 0) = true, want false (outside window)")
-	}
-	if cache.CanResume(0, 1) {
-		t.Errorf("after shift: CanResume(0, 1) = true, want false (outside window)")
-	}
-	if cache.CanResume(0, 2) {
-		t.Errorf("after shift: CanResume(0, 2) = true, want false (outside window)")
-	}
-	if cache.CanResume(0, 3) {
-		t.Errorf("after shift: CanResume(0, 3) = true, want false (outside window)")
-	}
-	if cache.CanResume(0, 4) {
-		t.Errorf("after shift: CanResume(0, 4) = true, want false (outside window)")
-	}
-	if !cache.CanResume(0, 5) {
-		t.Errorf("after shift: CanResume(0, 5) = false, want true (latest position)")
-	}
-}
-
 type testBackend struct{}
 
 func (b *testBackend) Config() ml.Config {
@@ -433,6 +352,7 @@ func (c *testContext) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
 }
 
 func (c *testContext) Input() ml.Context    { return c }
+func (c *testContext) Output() ml.Context   { return c }
 func (c *testContext) Layer(int) ml.Context { return c }
 
 func (c *testContext) Forward(...ml.Tensor) ml.Context { return c }

kvcache/encoder.go

@@ -49,7 +49,7 @@ func NewEncoderCache() *EncoderCache {
 	}
 }
 
-func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
+func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, capacity int32) {
 	if c.config == nil {
 		var config ml.CacheConfig
 		if cc, ok := backend.(ml.BackendCacheConfig); ok {
@@ -58,10 +58,6 @@ func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, ca
 		c.config = &config
 	}
 
-	if maxSequences > 1 {
-		panic(fmt.Errorf("encoder cache does not support multiple sequences; requested: %v", maxSequences))
-	}
-
 	if c.config.CachePadding != 0 && c.config.CachePadding != 1 {
 		panic(fmt.Errorf("encoder cache is unable to enforce requested CachePadding (%v)", c.config.CachePadding))
 	}
@@ -83,10 +79,10 @@ func (c *EncoderCache) Close() {
 	}
 }
 
-func (c *EncoderCache) StartForward(ctx ml.Context, batch input.Batch) error {
+func (c *EncoderCache) StartForward(ctx ml.Context, opts input.Options) error {
 	// We work with the most recent image
-	if len(batch.Multimodal) > 0 {
-		c.curPos = batch.Positions[batch.Multimodal[len(batch.Multimodal)-1].Index]
+	if len(opts.Multimodal) > 0 {
+		c.curPos = opts.Positions[opts.Multimodal[len(opts.Multimodal)-1].Index]
 	}
 
 	return nil
@@ -134,10 +130,6 @@ func (c *EncoderCache) CopyPrefix(srcSeq, dstSeq int, len int32) {
 	panic("encoder cache does not support multiple sequences")
 }
 
-func (c *EncoderCache) CanResume(seq int, pos int32) bool {
-	return true
-}
-
 func (c *EncoderCache) Remove(seq int, beginIndex, endIndex int32) error {
 	if c.encoderPos >= beginIndex && c.encoderPos < endIndex {
 		c.encoderCached = false

kvcache/wrapper.go

@@ -23,9 +23,9 @@ func NewWrapperCache(caches ...Cache) *WrapperCache {
 	}
 }
 
-func (c *WrapperCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {
+func (c *WrapperCache) Init(backend ml.Backend, dtype ml.DType, capacity int32) {
 	for _, cache := range c.caches {
-		cache.Init(backend, dtype, maxSequences, capacity, maxBatch)
+		cache.Init(backend, dtype, capacity)
 	}
 }
 
@@ -41,14 +41,14 @@ func (c *WrapperCache) Close() {
 	}
 }
 
-func (c *WrapperCache) StartForward(ctx ml.Context, batch input.Batch) error {
+func (c *WrapperCache) StartForward(ctx ml.Context, opts input.Options) error {
 	for i, cache := range c.caches {
-		err := cache.StartForward(ctx, batch)
+		err := cache.StartForward(ctx, opts)
 		if err != nil {
 			// unwind on error - Remove with endIndex set to math.MaxInt32 does not fail
 			for j := i - 1; j >= 0; j-- {
-				for k := range batch.Positions {
-					_ = c.caches[j].Remove(batch.Sequences[k], batch.Positions[k], math.MaxInt32)
+				for k := range opts.Positions {
+					_ = c.caches[j].Remove(opts.Sequences[k], opts.Positions[k], math.MaxInt32)
 				}
 			}
 			return err
@@ -87,16 +87,6 @@ func (c *WrapperCache) CopyPrefix(srcSeq, dstSeq int, len int32) {
 	}
 }
 
-func (c *WrapperCache) CanResume(seq int, pos int32) bool {
-	for _, cache := range c.caches {
-		if !cache.CanResume(seq, pos) {
-			return false
-		}
-	}
-
-	return true
-}
-
 func (c *WrapperCache) Remove(seq int, beginIndex, endIndex int32) error {
 	// If the one of these fails, the caller is supposed to retry with endIndex set to math.MaxInt32, which should not fail
 	for _, cache := range c.caches {

llama/llama.cpp/src/llama-arch.cpp

@@ -37,7 +37,6 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_MINICPM3,         "minicpm3"         },
     { LLM_ARCH_GEMMA,            "gemma"            },
     { LLM_ARCH_GEMMA2,           "gemma2"           },
-    { LLM_ARCH_GEMMA3,           "gemma3"           },
     { LLM_ARCH_STARCODER2,       "starcoder2"       },
     { LLM_ARCH_MAMBA,            "mamba"            },
     { LLM_ARCH_XVERSE,           "xverse"           },
@@ -805,24 +804,6 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
             { LLM_TENSOR_FFN_POST_NORM,   "blk.%d.post_ffw_norm" },
         },
     },
-    {
-        LLM_ARCH_GEMMA3,
-        {
-            { LLM_TENSOR_TOKEN_EMBD,      "token_embd" },
-            { LLM_TENSOR_OUTPUT_NORM,     "output_norm" },
-            { LLM_TENSOR_ATTN_NORM,       "blk.%d.attn_norm" },
-            { LLM_TENSOR_ATTN_Q,          "blk.%d.attn_q" },
-            { LLM_TENSOR_ATTN_K,          "blk.%d.attn_k" },
-            { LLM_TENSOR_ATTN_V,          "blk.%d.attn_v" },
-            { LLM_TENSOR_ATTN_OUT,        "blk.%d.attn_output" },
-            { LLM_TENSOR_ATTN_POST_NORM,  "blk.%d.post_attention_norm" },
-            { LLM_TENSOR_FFN_NORM,        "blk.%d.ffn_norm" },
-            { LLM_TENSOR_FFN_GATE,        "blk.%d.ffn_gate" },
-            { LLM_TENSOR_FFN_DOWN,        "blk.%d.ffn_down" },
-            { LLM_TENSOR_FFN_UP,          "blk.%d.ffn_up" },
-            { LLM_TENSOR_FFN_POST_NORM,   "blk.%d.post_ffw_norm" },
-        },
-    },
     {
         LLM_ARCH_STARCODER2,
         {

llama/llama.cpp/src/llama-arch.h

@@ -41,7 +41,6 @@ enum llm_arch {
     LLM_ARCH_MINICPM3,
     LLM_ARCH_GEMMA,
     LLM_ARCH_GEMMA2,
-    LLM_ARCH_GEMMA3,
     LLM_ARCH_STARCODER2,
     LLM_ARCH_MAMBA,
     LLM_ARCH_XVERSE,

llama/llama.cpp/src/llama-model.cpp

@@ -878,9 +878,6 @@ void llama_model::load_hparams(llama_model_loader & ml) {
                     default: type = LLM_TYPE_UNKNOWN;
                }
             } break;
-        case LLM_ARCH_GEMMA3:
-            {
-            } break;
         case LLM_ARCH_STARCODER2:
             {
                 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
@@ -2540,9 +2537,6 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
                         layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
                     }
                 } break;
-            case LLM_ARCH_GEMMA3:
-                {
-                } break;
             case LLM_ARCH_STARCODER2:
                 {
                     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
@@ -4035,7 +4029,6 @@ enum llama_rope_type llama_model_rope_type(const struct llama_model * model) {
         case LLM_ARCH_PHIMOE:
         case LLM_ARCH_GEMMA:
         case LLM_ARCH_GEMMA2:
-        case LLM_ARCH_GEMMA3:
         case LLM_ARCH_STARCODER2:
         case LLM_ARCH_OPENELM:
         case LLM_ARCH_GPTNEOX:

llama/llama.cpp/src/llama-quant.cpp

@@ -737,15 +737,6 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
         // This used to be a regex, but <regex> has an extreme cost to compile times.
         bool quantize = name.rfind("weight") == name.size() - 6; // ends with 'weight'?
 
-        // don't quantize vision stuff
-        quantize &= name.find("v.blk.") == std::string::npos;
-
-        quantize &= name.find("mm.mm_input_projection.weight") == std::string::npos;
-        quantize &= name.find("mm.mm_soft_emb_norm.weight") == std::string::npos;
-        quantize &= name.find("v.patch_embedding.weight") == std::string::npos;
-        quantize &= name.find("v.position_embedding.weight") == std::string::npos;
-        quantize &= name.find("v.post_layernorm.weight") == std::string::npos;
-
         // quantize only 2D and 3D tensors (experts)
         quantize &= (ggml_n_dims(tensor) >= 2);
 

llama/llama.go

@@ -166,10 +166,6 @@ func (c *Context) KvCacheDefrag() {
 	C.llama_kv_cache_defrag(c.c)
 }
 
-func (c *Context) KvCacheCanShift() bool {
-	return bool(C.llama_kv_cache_can_shift(c.c))
-}
-
 // Get the embeddings for a sequence id
 func (c *Context) GetEmbeddingsSeq(seqId int) []float32 {
 	e := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId)))

llama/patches/0021-gemma3-quantization.patch

@@ -1,113 +0,0 @@
-From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
-From: Patrick Devine <patrick@infrahq.com>
-Date: Fri, 14 Mar 2025 16:33:23 -0700
-Subject: [PATCH] gemma3 quantization
-
----
- src/llama-arch.cpp  | 19 +++++++++++++++++++
- src/llama-arch.h    |  1 +
- src/llama-model.cpp |  7 +++++++
- src/llama-quant.cpp |  9 +++++++++
- 4 files changed, 36 insertions(+)
-
-diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
-index b6f20286..b443fcd3 100644
---- a/src/llama-arch.cpp
-+++ b/src/llama-arch.cpp
-@@ -37,6 +37,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
-     { LLM_ARCH_MINICPM3,         "minicpm3"         },
-     { LLM_ARCH_GEMMA,            "gemma"            },
-     { LLM_ARCH_GEMMA2,           "gemma2"           },
-+    { LLM_ARCH_GEMMA3,           "gemma3"           },
-     { LLM_ARCH_STARCODER2,       "starcoder2"       },
-     { LLM_ARCH_MAMBA,            "mamba"            },
-     { LLM_ARCH_XVERSE,           "xverse"           },
-@@ -804,6 +805,24 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
-             { LLM_TENSOR_FFN_POST_NORM,   "blk.%d.post_ffw_norm" },
-         },
-     },
-+    {
-+        LLM_ARCH_GEMMA3,
-+        {
-+            { LLM_TENSOR_TOKEN_EMBD,      "token_embd" },
-+            { LLM_TENSOR_OUTPUT_NORM,     "output_norm" },
-+            { LLM_TENSOR_ATTN_NORM,       "blk.%d.attn_norm" },
-+            { LLM_TENSOR_ATTN_Q,          "blk.%d.attn_q" },
-+            { LLM_TENSOR_ATTN_K,          "blk.%d.attn_k" },
-+            { LLM_TENSOR_ATTN_V,          "blk.%d.attn_v" },
-+            { LLM_TENSOR_ATTN_OUT,        "blk.%d.attn_output" },
-+            { LLM_TENSOR_ATTN_POST_NORM,  "blk.%d.post_attention_norm" },
-+            { LLM_TENSOR_FFN_NORM,        "blk.%d.ffn_norm" },
-+            { LLM_TENSOR_FFN_GATE,        "blk.%d.ffn_gate" },
-+            { LLM_TENSOR_FFN_DOWN,        "blk.%d.ffn_down" },
-+            { LLM_TENSOR_FFN_UP,          "blk.%d.ffn_up" },
-+            { LLM_TENSOR_FFN_POST_NORM,   "blk.%d.post_ffw_norm" },
-+        },
-+    },
-     {
-         LLM_ARCH_STARCODER2,
-         {
-diff --git a/src/llama-arch.h b/src/llama-arch.h
-index ec742224..aad92a5d 100644
---- a/src/llama-arch.h
-+++ b/src/llama-arch.h
-@@ -41,6 +41,7 @@ enum llm_arch {
-     LLM_ARCH_MINICPM3,
-     LLM_ARCH_GEMMA,
-     LLM_ARCH_GEMMA2,
-+    LLM_ARCH_GEMMA3,
-     LLM_ARCH_STARCODER2,
-     LLM_ARCH_MAMBA,
-     LLM_ARCH_XVERSE,
-diff --git a/src/llama-model.cpp b/src/llama-model.cpp
-index ab1a07d1..70183041 100644
---- a/src/llama-model.cpp
-+++ b/src/llama-model.cpp
-@@ -878,6 +878,9 @@ void llama_model::load_hparams(llama_model_loader & ml) {
-                     default: type = LLM_TYPE_UNKNOWN;
-                }
-             } break;
-+        case LLM_ARCH_GEMMA3:
-+            {
-+            } break;
-         case LLM_ARCH_STARCODER2:
-             {
-                 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
-@@ -2537,6 +2540,9 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
-                         layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
-                     }
-                 } break;
-+            case LLM_ARCH_GEMMA3:
-+                {
-+                } break;
-             case LLM_ARCH_STARCODER2:
-                 {
-                     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
-@@ -4029,6 +4035,7 @@ enum llama_rope_type llama_model_rope_type(const struct llama_model * model) {
-         case LLM_ARCH_PHIMOE:
-         case LLM_ARCH_GEMMA:
-         case LLM_ARCH_GEMMA2:
-+        case LLM_ARCH_GEMMA3:
-         case LLM_ARCH_STARCODER2:
-         case LLM_ARCH_OPENELM:
-         case LLM_ARCH_GPTNEOX:
-diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp
-index 6eb1da08..d2f3a510 100644
---- a/src/llama-quant.cpp
-+++ b/src/llama-quant.cpp
-@@ -737,6 +737,15 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
-         // This used to be a regex, but <regex> has an extreme cost to compile times.
-         bool quantize = name.rfind("weight") == name.size() - 6; // ends with 'weight'?
- 
-+        // don't quantize vision stuff
-+        quantize &= name.find("v.blk.") == std::string::npos;
-+
-+        quantize &= name.find("mm.mm_input_projection.weight") == std::string::npos;
-+        quantize &= name.find("mm.mm_soft_emb_norm.weight") == std::string::npos;
-+        quantize &= name.find("v.patch_embedding.weight") == std::string::npos;
-+        quantize &= name.find("v.position_embedding.weight") == std::string::npos;
-+        quantize &= name.find("v.post_layernorm.weight") == std::string::npos;
-+
-         // quantize only 2D and 3D tensors (experts)
-         quantize &= (ggml_n_dims(tensor) >= 2);
- 

llama/patches/0022-add-rdna4-support.patch

@@ -1,103 +0,0 @@
-From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
-From: Saman <saman.khatir@amd.com>
-Date: Wed, 19 Mar 2025 14:02:26 -0700
-Subject: [PATCH] add rdna4 support
-
----
- ggml/src/ggml-cuda/common.cuh    | 6 ++++--
- ggml/src/ggml-cuda/mmq.cu        | 2 +-
- ggml/src/ggml-cuda/mmq.cuh       | 4 ++--
- ggml/src/ggml-cuda/mmvq.cu       | 4 ++--
- ggml/src/ggml-cuda/vendors/hip.h | 4 ++++
- 5 files changed, 13 insertions(+), 7 deletions(-)
-
-diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
-index adf0d3ec..b24593fc 100644
---- a/ggml/src/ggml-cuda/common.cuh
-+++ b/ggml/src/ggml-cuda/common.cuh
-@@ -61,11 +61,13 @@
- #define GGML_CUDA_CC_RDNA1      (GGML_CUDA_CC_OFFSET_AMD + 0x1010) // RX 5000
- #define GGML_CUDA_CC_RDNA2      (GGML_CUDA_CC_OFFSET_AMD + 0x1030) // RX 6000, minimum for dp4a
- #define GGML_CUDA_CC_RDNA3      (GGML_CUDA_CC_OFFSET_AMD + 0x1100) // RX 7000, minimum for WMMA
-+#define GGML_CUDA_CC_RDNA4      (GGML_CUDA_CC_OFFSET_AMD + 0x1200) // RX 9000
- 
- #define GGML_CUDA_CC_IS_RDNA(cc)  (cc >= GGML_CUDA_CC_RDNA1)
- #define GGML_CUDA_CC_IS_RDNA1(cc) (cc >= GGML_CUDA_CC_RDNA1 && cc < GGML_CUDA_CC_RDNA2)
- #define GGML_CUDA_CC_IS_RDNA2(cc) (cc >= GGML_CUDA_CC_RDNA2 && cc < GGML_CUDA_CC_RDNA3)
--#define GGML_CUDA_CC_IS_RDNA3(cc) (cc >= GGML_CUDA_CC_RDNA3)
-+#define GGML_CUDA_CC_IS_RDNA3(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA4)
-+#define GGML_CUDA_CC_IS_RDNA4(cc) (cc >= GGML_CUDA_CC_RDNA4)
- #define GGML_CUDA_CC_IS_GCN(cc)   (cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA)
- #define GGML_CUDA_CC_IS_CDNA(cc)  (cc >= GGML_CUDA_CC_CDNA && cc < GGML_CUDA_CC_RDNA1)
- 
-@@ -386,7 +388,7 @@ static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, i
- #if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__)
- #if defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(RDNA2)
-     c = __builtin_amdgcn_sdot4(a, b, c, false);
--#elif defined(RDNA3)
-+#elif defined(RDNA3) || defined(RDNA4)
-     c = __builtin_amdgcn_sudot4( true, a, true, b, c, false);
- #elif defined(__gfx1010__) || defined(__gfx900__)
-     int tmp1;
-diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu
-index 10f2ebb1..933d945c 100644
---- a/ggml/src/ggml-cuda/mmq.cu
-+++ b/ggml/src/ggml-cuda/mmq.cu
-@@ -149,5 +149,5 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11) {
-         return !fp16_mma_hardware_available(cc) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
-     }
- 
--    return (!GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
-+    return (!GGML_CUDA_CC_IS_RDNA4(cc) && !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
- }
-diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh
-index 0451c65f..66ce2bc9 100644
---- a/ggml/src/ggml-cuda/mmq.cuh
-+++ b/ggml/src/ggml-cuda/mmq.cuh
-@@ -2577,9 +2577,9 @@ static __device__ void mul_mat_q_process_tile(
- 
- template <ggml_type type, int mmq_x, int nwarps, bool need_check>
- #if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__)
--#if defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
-+#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
-     __launch_bounds__(WARP_SIZE*nwarps, 2)
--#endif // defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
-+#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
- #else
- #if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
-     __launch_bounds__(WARP_SIZE*nwarps, 1)
-diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu
-index 4fb466ca..23ae7abc 100644
---- a/ggml/src/ggml-cuda/mmvq.cu
-+++ b/ggml/src/ggml-cuda/mmvq.cu
-@@ -62,13 +62,13 @@ static __global__ void mul_mat_vec_q(
- 
-     constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
- 
--#if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && (defined(RDNA2) || defined(RDNA3))
-+#if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4))
-     constexpr int nwarps              = 1;
-     constexpr int rows_per_cuda_block = 1;
- #else
-     constexpr int nwarps              = ncols_y <= 4 ? 4 : 2;
-     constexpr int rows_per_cuda_block = ncols_y == 1 ? 1 : 2;
--#endif // defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && !defined(RDNA2) && !defined(RDNA3)
-+#endif // defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && !defined(RDNA2) && !defined(RDNA3) && !defined(RDNA4)
- 
-     const     int tid = WARP_SIZE*threadIdx.y + threadIdx.x;
-     const     int row0 = rows_per_cuda_block*blockIdx.x;
-diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h
-index 81964611..a62544b5 100644
---- a/ggml/src/ggml-cuda/vendors/hip.h
-+++ b/ggml/src/ggml-cuda/vendors/hip.h
-@@ -150,6 +150,10 @@
- #define CDNA
- #endif
- 
-+#if defined(__gfx1200__) || defined(__gfx1201__)
-+#define RDNA4
-+#endif
-+
- #if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \
-     defined(__gfx1150__) || defined(__gfx1151__)
- #define RDNA3

llm/memory.go

@@ -15,12 +15,12 @@ import (
 )
 
 // This algorithm looks for a complete fit to determine if we need to unload other models
-func PredictServerFit(allGpus discover.GpuInfoList, f *ggml.GGML, adapters, projectors []string, opts api.Options, numParallel int) (bool, uint64) {
+func PredictServerFit(allGpus discover.GpuInfoList, f *ggml.GGML, adapters, projectors []string, opts api.Options) (bool, uint64) {
 	// Split up the GPUs by type and try them
 	var estimatedVRAM uint64
 	for _, gpus := range allGpus.ByLibrary() {
 		var layerCount int
-		estimate := EstimateGPULayers(gpus, f, projectors, opts, numParallel)
+		estimate := EstimateGPULayers(gpus, f, projectors, opts)
 		layerCount, estimatedVRAM = estimate.Layers, estimate.VRAMSize
 		if opts.NumGPU < 0 {
 			if layerCount > 0 && layerCount >= int(f.KV().BlockCount()+1) {
@@ -71,7 +71,7 @@ type MemoryEstimate struct {
 
 // Given a model and one or more GPU targets, predict how many layers and bytes we can load, and the total size
 // The GPUs provided must all be the same Library
-func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []string, opts api.Options, numParallel int) MemoryEstimate {
+func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []string, opts api.Options) MemoryEstimate {
 	// Graph size for a partial offload, applies to all GPUs
 	var graphPartialOffload uint64
 
@@ -137,19 +137,13 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 		}
 	}
 
-	kv, graphPartialOffload, graphFullOffload := f.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)), numParallel, kvct)
+	kv, graphPartialOffload, graphFullOffload := f.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)), kvct)
 
-	if len(kv) > 0 {
-		layerSize += kv[0]
-	}
-
-	var kvTotal uint64
-	for _, kvLayer := range kv {
-		kvTotal += kvLayer
-	}
+	// KV is proportional to the number of layers
+	layerSize += kv / f.KV().BlockCount()
 
 	if graphPartialOffload == 0 {
-		graphPartialOffload = f.KV().GQA() * kvTotal / 6
+		graphPartialOffload = f.KV().GQA() * kv / 6
 	}
 	if graphFullOffload == 0 {
 		graphFullOffload = graphPartialOffload
@@ -223,9 +217,9 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 		// Some models have inconsistent layer sizes
 		if blk, ok := layers[fmt.Sprintf("blk.%d", i)]; ok {
 			layerSize = blk.Size()
-			layerSize += kv[i]
-			memoryWeights += blk.Size()
+			layerSize += kv / f.KV().BlockCount()
 		}
+		memoryWeights += layerSize
 
 		if opts.NumGPU >= 0 && layerCount >= opts.NumGPU {
 			// Stop allocating on GPU(s) once we hit the users target NumGPU
@@ -321,7 +315,7 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 		layersRequested:     opts.NumGPU,
 		layersModel:         int(f.KV().BlockCount()) + 1,
 		availableList:       availableList,
-		kv:                  kvTotal,
+		kv:                  kv,
 		allocationsList:     allocationsList,
 		memoryWeights:       memoryWeights,
 		memoryLayerOutput:   memoryLayerOutput,
@@ -380,9 +374,9 @@ func (m MemoryEstimate) LogValue() slog.Value {
 			slog.Group(
 				"weights",
 				// memory of the weights
-				"total", format.HumanBytes2(m.memoryWeights+m.memoryLayerOutput),
+				"total", format.HumanBytes2(m.memoryWeights),
 				// memory of repeating layers
-				"repeating", format.HumanBytes2(m.memoryWeights),
+				"repeating", format.HumanBytes2(m.memoryWeights-m.memoryLayerOutput),
 				// memory of non-repeating layers
 				"nonrepeating", format.HumanBytes2(m.memoryLayerOutput),
 			),

llm/memory_test.go

@@ -61,7 +61,7 @@ func TestEstimateGPULayers(t *testing.T) {
 	projectors := []string{}
 	opts := api.DefaultOptions()
 	t.Run("cpu", func(t *testing.T) {
-		estimate := EstimateGPULayers(gpus, ggml, projectors, opts, 1)
+		estimate := EstimateGPULayers(gpus, ggml, projectors, opts)
 		assert.Equal(t, 0, estimate.Layers)
 		assert.Equal(t, uint64(0), estimate.Graph)
 	})
@@ -112,7 +112,7 @@ func TestEstimateGPULayers(t *testing.T) {
 			gpus[1].FreeMemory += gpuMinimumMemory + layerSize + s.layer1*layerSize + 1
 			gpus[0].FreeMemory += max(graphFullOffload, graphPartialOffload)
 			gpus[1].FreeMemory += max(graphFullOffload, graphPartialOffload)
-			estimate := EstimateGPULayers(gpus, ggml, projectors, opts, 1)
+			estimate := EstimateGPULayers(gpus, ggml, projectors, opts)
 			assert.Equal(t, int(s.expect0+s.expect1), estimate.Layers, "scenario %d: %v", i, s)
 			assert.Equal(t, fmt.Sprintf("%d,%d", s.expect0, s.expect1), estimate.TensorSplit, "scenario %d: %v", i, s)
 			var layerSums uint64

llm/server.go

@@ -109,7 +109,7 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 		gpus = discover.GetCPUInfo()
 	}
 
-	estimate := EstimateGPULayers(gpus, f, projectors, opts, numParallel)
+	estimate := EstimateGPULayers(gpus, f, projectors, opts)
 	if len(gpus) > 1 || gpus[0].Library != "cpu" {
 		switch {
 		case gpus[0].Library == "metal" && estimate.VRAMSize > systemTotalMemory:
@@ -402,7 +402,7 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 			s.cmd.Env = append(s.cmd.Env, visibleDevicesEnv+"="+visibleDevicesEnvVal)
 		}
 
-		slog.Info("starting llama server", "cmd", s.cmd)
+		slog.Info("starting llama server", "cmd", s.cmd.String())
 		if envconfig.Debug() {
 			filteredEnv := []string{}
 			for _, ev := range s.cmd.Env {
@@ -470,7 +470,7 @@ const ( // iota is reset to 0
 	ServerStatusError
 )
 
-func (s ServerStatus) String() string {
+func (s ServerStatus) ToString() string {
 	switch s {
 	case ServerStatusReady:
 		return "llm server ready"
@@ -485,9 +485,12 @@ func (s ServerStatus) String() string {
 	}
 }
 
-type ServerStatusResponse struct {
-	Status   ServerStatus `json:"status"`
-	Progress float32      `json:"progress"`
+type ServerStatusResp struct {
+	Status          string  `json:"status"`
+	SlotsIdle       int     `json:"slots_idle"`
+	SlotsProcessing int     `json:"slots_processing"`
+	Error           string  `json:"error"`
+	Progress        float32 `json:"progress"`
 }
 
 func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
@@ -499,7 +502,7 @@ func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
 		}
 		if s.cmd.ProcessState.ExitCode() == -1 {
 			// Most likely a signal killed it, log some more details to try to help troubleshoot
-			slog.Warn("llama runner process no longer running", "sys", s.cmd.ProcessState.Sys(), "string", s.cmd.ProcessState)
+			slog.Warn("llama runner process no longer running", "sys", s.cmd.ProcessState.Sys(), "string", s.cmd.ProcessState.String())
 		}
 		return ServerStatusError, fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
 	}
@@ -524,19 +527,21 @@ func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
 		return ServerStatusError, fmt.Errorf("read health request: %w", err)
 	}
 
-	var ssr ServerStatusResponse
-	if err := json.Unmarshal(body, &ssr); err != nil {
+	var status ServerStatusResp
+	if err := json.Unmarshal(body, &status); err != nil {
 		return ServerStatusError, fmt.Errorf("health unmarshal encode response: %w", err)
 	}
 
-	switch ssr.Status {
-	case ServerStatusLoadingModel:
-		s.loadProgress = ssr.Progress
-		return ssr.Status, nil
-	case ServerStatusReady, ServerStatusNoSlotsAvailable:
-		return ssr.Status, nil
+	switch status.Status {
+	case "ok":
+		return ServerStatusReady, nil
+	case "no slot available":
+		return ServerStatusNoSlotsAvailable, nil
+	case "loading model":
+		s.loadProgress = status.Progress
+		return ServerStatusLoadingModel, nil
 	default:
-		return ssr.Status, fmt.Errorf("server error: %+v", ssr)
+		return ServerStatusError, fmt.Errorf("server error: %+v", status)
 	}
 }
 
@@ -611,7 +616,7 @@ func (s *llmServer) WaitUntilRunning(ctx context.Context) error {
 		status, _ := s.getServerStatus(ctx)
 		if lastStatus != status && status != ServerStatusReady {
 			// Only log on status changes
-			slog.Info("waiting for server to become available", "status", status)
+			slog.Info("waiting for server to become available", "status", status.ToString())
 		}
 		switch status {
 		case ServerStatusReady:
@@ -625,7 +630,7 @@ func (s *llmServer) WaitUntilRunning(ctx context.Context) error {
 				slog.Debug(fmt.Sprintf("model load progress %0.2f", s.loadProgress))
 				stallTimer = time.Now().Add(stallDuration)
 			} else if !fullyLoaded && int(s.loadProgress*100.0) >= 100 {
-				slog.Debug("model load completed, waiting for server to become available", "status", status)
+				slog.Debug("model load completed, waiting for server to become available", "status", status.ToString())
 				stallTimer = time.Now().Add(stallDuration)
 				fullyLoaded = true
 			}
@@ -666,26 +671,63 @@ type ImageData struct {
 	AspectRatioID int    `json:"aspect_ratio_id"`
 }
 
+type completion struct {
+	Content      string `json:"content"`
+	Model        string `json:"model"`
+	Prompt       string `json:"prompt"`
+	Stop         bool   `json:"stop"`
+	StoppedLimit bool   `json:"stopped_limit"`
+
+	Timings struct {
+		PredictedN  int     `json:"predicted_n"`
+		PredictedMS float64 `json:"predicted_ms"`
+		PromptN     int     `json:"prompt_n"`
+		PromptMS    float64 `json:"prompt_ms"`
+	}
+}
+
 type CompletionRequest struct {
 	Prompt  string
 	Format  json.RawMessage
 	Images  []ImageData
 	Options *api.Options
-
-	Grammar string // set before sending the request to the subprocess
 }
 
 type CompletionResponse struct {
-	Content            string        `json:"content"`
-	DoneReason         string        `json:"done_reason"`
-	Done               bool          `json:"done"`
-	PromptEvalCount    int           `json:"prompt_eval_count"`
-	PromptEvalDuration time.Duration `json:"prompt_eval_duration"`
-	EvalCount          int           `json:"eval_count"`
-	EvalDuration       time.Duration `json:"eval_duration"`
+	Content            string
+	DoneReason         string
+	Done               bool
+	PromptEvalCount    int
+	PromptEvalDuration time.Duration
+	EvalCount          int
+	EvalDuration       time.Duration
 }
 
 func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
+	request := map[string]any{
+		"prompt":            req.Prompt,
+		"stream":            true,
+		"n_predict":         req.Options.NumPredict,
+		"n_keep":            req.Options.NumKeep,
+		"main_gpu":          req.Options.MainGPU,
+		"temperature":       req.Options.Temperature,
+		"top_k":             req.Options.TopK,
+		"top_p":             req.Options.TopP,
+		"min_p":             req.Options.MinP,
+		"typical_p":         req.Options.TypicalP,
+		"repeat_last_n":     req.Options.RepeatLastN,
+		"repeat_penalty":    req.Options.RepeatPenalty,
+		"presence_penalty":  req.Options.PresencePenalty,
+		"frequency_penalty": req.Options.FrequencyPenalty,
+		"mirostat":          req.Options.Mirostat,
+		"mirostat_tau":      req.Options.MirostatTau,
+		"mirostat_eta":      req.Options.MirostatEta,
+		"seed":              req.Options.Seed,
+		"stop":              req.Options.Stop,
+		"image_data":        req.Images,
+		"cache_prompt":      true,
+	}
+
 	if len(req.Format) > 0 {
 		switch string(req.Format) {
 		case `null`, `""`:
@@ -693,7 +735,7 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 			// these as "not set".
 			break
 		case `"json"`:
-			req.Grammar = grammarJSON
+			request["grammar"] = grammarJSON
 		default:
 			if req.Format[0] != '{' {
 				return fmt.Errorf("invalid format: %q; expected \"json\" or a valid JSON Schema object", req.Format)
@@ -704,15 +746,10 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 			if g == nil {
 				return fmt.Errorf("invalid JSON schema in format")
 			}
-			req.Grammar = string(g)
+			request["grammar"] = string(g)
 		}
 	}
 
-	if req.Options == nil {
-		opts := api.DefaultOptions()
-		req.Options = &opts
-	}
-
 	if err := s.sem.Acquire(ctx, 1); err != nil {
 		if errors.Is(err, context.Canceled) {
 			slog.Info("aborting completion request due to client closing the connection")
@@ -733,7 +770,7 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 	if err != nil {
 		return err
 	} else if status != ServerStatusReady {
-		return fmt.Errorf("unexpected server status: %s", status)
+		return fmt.Errorf("unexpected server status: %s", status.ToString())
 	}
 
 	// Handling JSON marshaling with special characters unescaped.
@@ -741,7 +778,7 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 	enc := json.NewEncoder(buffer)
 	enc.SetEscapeHTML(false)
 
-	if err := enc.Encode(req); err != nil {
+	if err := enc.Encode(request); err != nil {
 		return fmt.Errorf("failed to marshal data: %v", err)
 	}
 
@@ -792,7 +829,7 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 				evt = line
 			}
 
-			var c CompletionResponse
+			var c completion
 			if err := json.Unmarshal(evt, &c); err != nil {
 				return fmt.Errorf("error unmarshalling llm prediction response: %v", err)
 			}
@@ -816,8 +853,20 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 				})
 			}
 
-			if c.Done {
-				fn(c)
+			if c.Stop {
+				doneReason := "stop"
+				if c.StoppedLimit {
+					doneReason = "length"
+				}
+
+				fn(CompletionResponse{
+					Done:               true,
+					DoneReason:         doneReason,
+					PromptEvalCount:    c.Timings.PromptN,
+					PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
+					EvalCount:          c.Timings.PredictedN,
+					EvalDuration:       parseDurationMs(c.Timings.PredictedMS),
+				})
 				return nil
 			}
 		}
@@ -865,7 +914,7 @@ func (s *llmServer) Embedding(ctx context.Context, input string) ([]float32, err
 	if err != nil {
 		return nil, err
 	} else if status != ServerStatusReady {
-		return nil, fmt.Errorf("unexpected server status: %s", status)
+		return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
 	}
 
 	data, err := json.Marshal(EmbeddingRequest{Content: input})
@@ -1010,3 +1059,12 @@ func (s *llmServer) EstimatedVRAMByGPU(gpuID string) uint64 {
 	}
 	return 0
 }
+
+func parseDurationMs(ms float64) time.Duration {
+	dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
+	if err != nil {
+		panic(err)
+	}
+
+	return dur
+}

ml/backend.go

@@ -2,7 +2,6 @@ package ml
 
 import (
 	"bytes"
-	"context"
 	"encoding/binary"
 	"fmt"
 	"os"
@@ -61,10 +60,6 @@ type CacheConfig struct {
 
 // BackendParams controls how the backend loads and executes models
 type BackendParams struct {
-	// Progress is a callback function that allows reporting percentage completion
-	// of model loading
-	Progress func(float32)
-
 	// NumThreads sets the number of threads to use if running on the CPU
 	NumThreads int
 
@@ -81,9 +76,9 @@ type BackendParams struct {
 	FlashAttention bool
 }
 
-var backends = make(map[string]func(context.Context, *os.File, BackendParams) (Backend, error))
+var backends = make(map[string]func(*os.File, BackendParams) (Backend, error))
 
-func RegisterBackend(name string, f func(context.Context, *os.File, BackendParams) (Backend, error)) {
+func RegisterBackend(name string, f func(*os.File, BackendParams) (Backend, error)) {
 	if _, ok := backends[name]; ok {
 		panic("backend: backend already registered")
 	}
@@ -91,9 +86,9 @@ func RegisterBackend(name string, f func(context.Context, *os.File, BackendParam
 	backends[name] = f
 }
 
-func NewBackend(ctx context.Context, f *os.File, params BackendParams) (Backend, error) {
+func NewBackend(f *os.File, params BackendParams) (Backend, error) {
 	if backend, ok := backends["ggml"]; ok {
-		return backend(ctx, f, params)
+		return backend(f, params)
 	}
 
 	return nil, fmt.Errorf("unsupported backend")
@@ -110,10 +105,12 @@ type Context interface {
 	MaxGraphNodes() int
 	Close()
 
-	// Input returns a context appropriate for creating tensors that are
-	// inputs to the model (which includes things like output locations)
+	// Input returns a context appropriate for creating input tensors
 	Input() Context
 
+	// Output returns a context appropriate for creating output tensors
+	Output() Context
+
 	// Layer returns a context appropriate for creating intermediate tensors
 	Layer(int) Context
 }

ml/backend/ggml/ggml.go

@@ -9,17 +9,15 @@ package ggml
 import "C"
 
 import (
-	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"maps"
 	"os"
-	"runtime"
 	"slices"
 	"strconv"
 	"strings"
-	"sync/atomic"
 	"unicode"
 	"unsafe"
 
@@ -48,6 +46,9 @@ type Backend struct {
 	// input is the backend used for inputs
 	input *C.struct_ggml_backend_buffer_type
 
+	// output is the backend used for outputs
+	output *C.struct_ggml_backend_buffer_type
+
 	// layers is the backend used for repeating layers
 	layers map[int]*C.struct_ggml_backend_buffer_type
 
@@ -57,7 +58,7 @@ type Backend struct {
 	maxGraphNodes int
 }
 
-func New(ctx context.Context, r *os.File, params ml.BackendParams) (ml.Backend, error) {
+func New(r *os.File, params ml.BackendParams) (ml.Backend, error) {
 	meta, n, err := fs.Decode(r, -1)
 	if err != nil {
 		return nil, err
@@ -296,16 +297,12 @@ func New(ctx context.Context, r *os.File, params ml.BackendParams) (ml.Backend,
 		}
 	}
 
-	var doneBytes atomic.Uint64
-	totalBytes := uint64(n) - meta.Tensors().Offset
-
-	g, ctx := errgroup.WithContext(ctx)
-	g.SetLimit(runtime.GOMAXPROCS(0))
+	// concurrently read in tensor data. uses a section reader which is safe for concurrent reads
+	sr := io.NewSectionReader(r, int64(meta.Tensors().Offset), n-int64(meta.Tensors().Offset))
+	var g errgroup.Group
 	for _, t := range meta.Tensors().Items() {
-		g.Go(func() error {
-			tts := make([]*C.struct_ggml_tensor, max(1, len(targets[t.Name])))
-			for i := range tts {
-				target := targets[t.Name][i]
+		for _, target := range targets[t.Name] {
+			g.Go(func() error {
 				if target == "" {
 					target = t.Name
 				}
@@ -315,44 +312,23 @@ func New(ctx context.Context, r *os.File, params ml.BackendParams) (ml.Backend,
 					return fmt.Errorf("unassigned tensor: %s", t.Name)
 				}
 
-				tts[i] = tt
-			}
-
-			sr := io.NewSectionReader(r, int64(meta.Tensors().Offset+t.Offset), int64(t.Size()))
-			bts := make([]byte, 128*format.KibiByte)
-
-			var s uint64
-			for s < t.Size() {
-				n, err := io.ReadFull(sr, bts[:min(len(bts), int(t.Size()-s))])
+				bts := make([]byte, t.Size())
+				n, err := io.ReadFull(io.NewSectionReader(sr, int64(t.Offset), int64(t.Size())), bts)
 				if err != nil {
 					return err
 				}
 
-				for _, tt := range tts {
-					C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), C.size_t(s), C.size_t(n))
+				if n != len(bts) {
+					return errors.New("short read")
 				}
 
-				s += uint64(n)
-
-				if params.Progress != nil {
-					done := doneBytes.Add(uint64(n))
-					params.Progress(float32(done) / float32(totalBytes))
-				}
-			}
-
-			return nil
-		})
+				C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), 0, C.size_t(t.Size()))
+				return nil
+			})
+		}
 	}
 
-	// start a goroutine to cancel the errgroup if the parent context is done
-	go func() {
-		<-ctx.Done()
-		g.Go(func() error {
-			return ctx.Err()
-		})
-	}()
-
-	if err := g.Wait(); err != nil {
+	if g.Wait() != nil {
 		return nil, err
 	}
 
@@ -395,9 +371,10 @@ func New(ctx context.Context, r *os.File, params ml.BackendParams) (ml.Backend,
 			(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&schedBufts[0])),
 			C.int(len(schedBackends)),
 			C.size_t(maxGraphNodes),
-			C._Bool(len(gpus) > 1 && slices.Contains(gpus, output.d)),
+			true,
 		),
-		input: deviceBufferTypes[input.d],
+		input:  deviceBufferTypes[input.d],
+		output: deviceBufferTypes[output.d],
 		layers: func() map[int]*C.struct_ggml_backend_buffer_type {
 			m := make(map[int]*C.struct_ggml_backend_buffer_type)
 			for i, layer := range layers {
@@ -478,6 +455,19 @@ func (c Context) Input() ml.Context {
 	return &c
 }
 
+func (c Context) Output() ml.Context {
+	if c.b.output != nil {
+		return &Context{
+			b:             c.b,
+			ctx:           c.ctx,
+			buft:          c.b.output,
+			maxGraphNodes: c.maxGraphNodes,
+		}
+	}
+
+	return &c
+}
+
 func (c Context) Layer(i int) ml.Context {
 	if buft, ok := c.b.layers[i]; ok {
 		return &Context{

ml/backend/ggml/ggml/src/ggml-cuda/common.cuh

@@ -61,13 +61,11 @@
 #define GGML_CUDA_CC_RDNA1      (GGML_CUDA_CC_OFFSET_AMD + 0x1010) // RX 5000
 #define GGML_CUDA_CC_RDNA2      (GGML_CUDA_CC_OFFSET_AMD + 0x1030) // RX 6000, minimum for dp4a
 #define GGML_CUDA_CC_RDNA3      (GGML_CUDA_CC_OFFSET_AMD + 0x1100) // RX 7000, minimum for WMMA
-#define GGML_CUDA_CC_RDNA4      (GGML_CUDA_CC_OFFSET_AMD + 0x1200) // RX 9000
 
 #define GGML_CUDA_CC_IS_RDNA(cc)  (cc >= GGML_CUDA_CC_RDNA1)
 #define GGML_CUDA_CC_IS_RDNA1(cc) (cc >= GGML_CUDA_CC_RDNA1 && cc < GGML_CUDA_CC_RDNA2)
 #define GGML_CUDA_CC_IS_RDNA2(cc) (cc >= GGML_CUDA_CC_RDNA2 && cc < GGML_CUDA_CC_RDNA3)
-#define GGML_CUDA_CC_IS_RDNA3(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA4)
-#define GGML_CUDA_CC_IS_RDNA4(cc) (cc >= GGML_CUDA_CC_RDNA4)
+#define GGML_CUDA_CC_IS_RDNA3(cc) (cc >= GGML_CUDA_CC_RDNA3)
 #define GGML_CUDA_CC_IS_GCN(cc)   (cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA)
 #define GGML_CUDA_CC_IS_CDNA(cc)  (cc >= GGML_CUDA_CC_CDNA && cc < GGML_CUDA_CC_RDNA1)
 
@@ -388,7 +386,7 @@ static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, i
 #if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__)
 #if defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(RDNA2)
     c = __builtin_amdgcn_sdot4(a, b, c, false);
-#elif defined(RDNA3) || defined(RDNA4)
+#elif defined(RDNA3)
     c = __builtin_amdgcn_sudot4( true, a, true, b, c, false);
 #elif defined(__gfx1010__) || defined(__gfx900__)
     int tmp1;

ml/backend/ggml/ggml/src/ggml-cuda/mmq.cu

@@ -149,5 +149,5 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11) {
         return !fp16_mma_hardware_available(cc) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
     }
 
-    return (!GGML_CUDA_CC_IS_RDNA4(cc) && !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
+    return (!GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
 }

ml/backend/ggml/ggml/src/ggml-cuda/mmq.cuh

@@ -2577,9 +2577,9 @@ static __device__ void mul_mat_q_process_tile(
 
 template <ggml_type type, int mmq_x, int nwarps, bool need_check>
 #if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__)
-#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
+#if defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
     __launch_bounds__(WARP_SIZE*nwarps, 2)
-#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
+#endif // defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN)
 #else
 #if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
     __launch_bounds__(WARP_SIZE*nwarps, 1)

ml/backend/ggml/ggml/src/ggml-cuda/mmvq.cu

@@ -62,13 +62,13 @@ static __global__ void mul_mat_vec_q(
 
     constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
 
-#if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4))
+#if defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && (defined(RDNA2) || defined(RDNA3))
     constexpr int nwarps              = 1;
     constexpr int rows_per_cuda_block = 1;
 #else
     constexpr int nwarps              = ncols_y <= 4 ? 4 : 2;
     constexpr int rows_per_cuda_block = ncols_y == 1 ? 1 : 2;
-#endif // defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && !defined(RDNA2) && !defined(RDNA3) && !defined(RDNA4)
+#endif // defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__) && !defined(RDNA2) && !defined(RDNA3)
 
     const     int tid = WARP_SIZE*threadIdx.y + threadIdx.x;
     const     int row0 = rows_per_cuda_block*blockIdx.x;

ml/backend/ggml/ggml/src/ggml-cuda/vendors/hip.h

@@ -150,10 +150,6 @@
 #define CDNA
 #endif
 
-#if defined(__gfx1200__) || defined(__gfx1201__)
-#define RDNA4
-#endif
-
 #if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || \
     defined(__gfx1150__) || defined(__gfx1151__)
 #define RDNA3

ml/backend/ggml/ggml/src/ollama-debug.c

@@ -1,5 +1,4 @@
 #include <string.h>
-#include <inttypes.h>
 
 #include "ollama-debug.h"
 
@@ -25,7 +24,7 @@ static void print_tensor(const void *tensor, void (*cb)(const void *, int),
     fprintf(stderr, "[");
     for (int i = 0; i < dims[0]; i++) {
         if (i >= nitems && i < dims[0] - nitems) {
-            fprintf(stderr, "... (%" PRIi64 " more), ", dims[0] - 2 * nitems);
+            fprintf(stderr, "... (%lld more), ", dims[0] - 2 * nitems);
             int skip = dims[0] - 2 * nitems;
             if (ndims > 1) {
                 stride += mul(dims + 1, ndims - 1) * skip;
@@ -68,7 +67,7 @@ static void print_tensor_i32(const void *tensor, int i) {
 }
 
 static void ollama_debug_tensor(const struct ggml_tensor *tensor, bool verbose, const char *prefix, int indent) {
-    fprintf(stderr, "%s%s %s (%s): [%" PRIi64 " %" PRIi64 " %" PRIi64 " %" PRIi64 "]\n", prefix, tensor->name,
+    fprintf(stderr, "%s%s %s (%s): [%lld %lld %lld %lld]\n", prefix, tensor->name,
             ggml_op_name(tensor->op), ggml_type_name(tensor->type), tensor->ne[0],
             tensor->ne[1], tensor->ne[2], tensor->ne[3]);
 

model/input/input.go

@@ -1,7 +1,5 @@
 package input
 
-import "github.com/ollama/ollama/ml"
-
 // Input represents one token in the input stream
 type Input struct {
 	// Token is a single element of text.
@@ -17,12 +15,6 @@ type Input struct {
 	// stored in Multimodal, used for caching and comparing
 	// equality.
 	MultimodalHash uint64
-
-	// SameBatch forces the following number of tokens to be processed
-	// in a single batch, breaking and extending batches as needed.
-	// Useful for things like images that must be processed in one
-	// shot.
-	SameBatch int
 }
 
 // MultimodalIndex is a multimodal element (such as an image)
@@ -35,24 +27,11 @@ type MultimodalIndex struct {
 	Multimodal any
 }
 
-// Batch contains the inputs for a model forward pass
-type Batch struct {
-	// Inputs is the input tokens, including placeholders for multimodal inputs.
-	Inputs ml.Tensor
-
-	// Multimodal is a set of multimodal embeddings previously created by
-	// EncodeMultimodal, along with an index into Inputs. Unused for text-only
-	// models or for batches without multimodal elements.
+// Options contains the inputs for a model forward pass
+type Options struct {
+	Inputs     []int32
 	Multimodal []MultimodalIndex
-
-	// Positions is the position for each Input, relative to its sequence. Equal
-	// in length to Inputs.
-	Positions []int32
-
-	// Sequences is the sequence for each Input. Equal in length to Inputs.
-	Sequences []int
-
-	// Outputs are the set of indicies into Inputs for which output data should
-	// be returned.
-	Outputs []int32
+	Positions  []int32
+	Sequences  []int
+	Outputs    []int32
 }

model/model.go

@@ -1,7 +1,6 @@
 package model
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	_ "image/jpeg"
@@ -23,11 +22,9 @@ import (
 	"github.com/ollama/ollama/model/input"
 )
 
-var ErrNoVisionModel = errors.New("this model is missing data required for image input")
-
 // Model implements a specific model architecture, defining the forward pass and any model-specific configuration
 type Model interface {
-	Forward(ml.Context, input.Batch) (ml.Tensor, error)
+	Forward(ml.Context, input.Options) (ml.Tensor, error)
 
 	Backend() ml.Backend
 	Config() config
@@ -61,7 +58,7 @@ type MultimodalProcessor interface {
 	// This function is also responsible for updating MultimodalHash for any Multimodal
 	// that is modified to ensure that there is a unique hash value that accurately
 	// represents the contents.
-	PostTokenize([]input.Input) ([]input.Input, error)
+	PostTokenize(ml.Context, []input.Input) ([]input.Input, error)
 }
 
 // Base implements the common fields and methods for all models
@@ -95,14 +92,14 @@ func Register(name string, f func(ml.Config) (Model, error)) {
 }
 
 // New initializes a new model instance with the provided configuration based on the metadata in the model file
-func New(ctx context.Context, modelPath string, params ml.BackendParams) (Model, error) {
+func New(modelPath string, params ml.BackendParams) (Model, error) {
 	r, err := os.Open(modelPath)
 	if err != nil {
 		return nil, err
 	}
 	defer r.Close()
 
-	b, err := ml.NewBackend(ctx, r, params)
+	b, err := ml.NewBackend(r, params)
 	if err != nil {
 		return nil, err
 	}
@@ -281,30 +278,24 @@ func canNil(t reflect.Type) bool {
 		t.Kind() == reflect.Slice
 }
 
-func Forward(ctx ml.Context, m Model, inputs []int32, batch input.Batch) (ml.Tensor, error) {
-	if len(batch.Positions) != len(batch.Sequences) {
-		return nil, fmt.Errorf("length of positions (%v) must match length of seqs (%v)", len(batch.Positions), len(batch.Sequences))
+func Forward(ctx ml.Context, m Model, opts input.Options) (ml.Tensor, error) {
+	if len(opts.Positions) != len(opts.Sequences) {
+		return nil, fmt.Errorf("length of positions (%v) must match length of seqs (%v)", len(opts.Positions), len(opts.Sequences))
 	}
 
-	if len(batch.Positions) < 1 {
+	if len(opts.Positions) < 1 {
 		return nil, errors.New("batch size cannot be less than 1")
 	}
 
-	var err error
-	batch.Inputs, err = ctx.Input().FromIntSlice(inputs, len(inputs))
-	if err != nil {
-		return nil, err
-	}
-
 	cache := m.Config().Cache
 	if cache != nil {
-		err := cache.StartForward(ctx, batch)
+		err := cache.StartForward(ctx, opts)
 		if err != nil {
 			return nil, err
 		}
 	}
 
-	t, err := m.Forward(ctx, batch)
+	t, err := m.Forward(ctx, opts)
 	if err != nil {
 		return nil, err
 	}

model/model_test.go

@@ -163,7 +163,7 @@ func TestGetTextProcessor(t *testing.T) {
 
 type notTextProcessorModel struct{}
 
-func (notTextProcessorModel) Forward(ml.Context, input.Batch) (ml.Tensor, error) {
+func (notTextProcessorModel) Forward(ml.Context, input.Options) (ml.Tensor, error) {
 	panic("unimplemented")
 }
 

model/models/gemma2/model.go

@@ -38,6 +38,7 @@ const (
 func New(c ml.Config) (model.Model, error) {
 	m := Model{
 		SentencePieceModel: model.NewSentencePieceModel(
+			c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Scores: c.Floats("tokenizer.ggml.scores"),
@@ -167,18 +168,23 @@ func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs, outputs ml.Ten
 	return hiddenState.Add(ctx, residual)
 }
 
-func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
-	positions, err := ctx.Input().FromIntSlice(batch.Positions, len(batch.Positions))
+func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
+	inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs))
 	if err != nil {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(batch.Outputs, len(batch.Outputs))
+	positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions))
 	if err != nil {
 		return nil, err
 	}
 
-	hiddenState := m.TokenEmbedding.Forward(ctx, batch.Inputs)
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	if err != nil {
+		return nil, err
+	}
+
+	hiddenState := m.TokenEmbedding.Forward(ctx, inputs)
 	hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(m.Options.hiddenSize)))
 
 	if len(m.Layers) == gemma27BLayerCount {
@@ -205,7 +211,8 @@ func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
 	// final logit softcap
 	hiddenState = hiddenState.Scale(ctx, 1.0/float64(m.Options.finalLogitSoftcap))
 	hiddenState = hiddenState.Tanh(ctx)
-	return hiddenState.Scale(ctx, float64(m.Options.finalLogitSoftcap)), nil
+	hiddenState = hiddenState.Scale(ctx, float64(m.Options.finalLogitSoftcap))
+	return hiddenState.Rows(ctx, outputs), nil
 }
 
 func init() {

model/models/gemma3/model.go

@@ -2,9 +2,10 @@ package gemma3
 
 import (
 	"bytes"
+	"encoding/binary"
+	"hash/fnv"
 	"image"
 	"math"
-	"slices"
 
 	"github.com/ollama/ollama/kvcache"
 	"github.com/ollama/ollama/ml"
@@ -55,6 +56,7 @@ func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, i
 func New(c ml.Config) (model.Model, error) {
 	m := Model{
 		SentencePieceModel: model.NewSentencePieceModel(
+			c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Scores: c.Floats("tokenizer.ggml.scores"),
@@ -82,10 +84,6 @@ func New(c ml.Config) (model.Model, error) {
 }
 
 func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, error) {
-	if len(m.VisionModel.Layers) == 0 {
-		return nil, model.ErrNoVisionModel
-	}
-
 	image, _, err := image.Decode(bytes.NewReader(multimodalData))
 	if err != nil {
 		return nil, err
@@ -110,23 +108,36 @@ func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, er
 	return visionOutputs, nil
 }
 
-func (m *Model) PostTokenize(inputs []input.Input) ([]input.Input, error) {
+type imageToken struct {
+	embedding ml.Tensor
+	index     int
+}
+
+func (m *Model) PostTokenize(ctx ml.Context, inputs []input.Input) ([]input.Input, error) {
 	var result []input.Input
+	fnvHash := fnv.New64a()
 
 	for _, inp := range inputs {
 		if inp.Multimodal == nil {
 			result = append(result, inp)
 		} else {
+			imageInputs := []input.Input{
+				{Token: 108},    // "\n\n"
+				{Token: 255999}, // "<start_of_image>""
+			}
+			result = append(result, imageInputs...)
+
+			// add image embeddings
 			inputMultimodal := inp.Multimodal.(ml.Tensor)
 
-			result = append(result,
-				input.Input{Token: 108, SameBatch: inputMultimodal.Dim(1) + 3},               // "\n\n"
-				input.Input{Token: 255999},                                                   // "<start_of_image>""
-				input.Input{Multimodal: inputMultimodal, MultimodalHash: inp.MultimodalHash}, // image data is on the first placeholder
-			)
+			for i := range inputMultimodal.Dim(1) {
+				fnvHash.Reset()
+				binary.Write(fnvHash, binary.NativeEndian, inp.MultimodalHash)
+				fnvHash.Write([]byte{byte(i)})
 
-			// add image token placeholders
-			result = append(result, slices.Repeat([]input.Input{{Token: 0}}, inputMultimodal.Dim(1)-1)...)
+				imageToken := imageToken{embedding: inputMultimodal, index: i}
+				result = append(result, input.Input{Multimodal: imageToken, MultimodalHash: fnvHash.Sum64()})
+			}
 
 			result = append(result,
 				input.Input{Token: 256000}, // <end_of_image>
@@ -138,18 +149,23 @@ func (m *Model) PostTokenize(inputs []input.Input) ([]input.Input, error) {
 	return result, nil
 }
 
-func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
-	positions, err := ctx.Input().FromIntSlice(batch.Positions, len(batch.Positions))
+func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
+	inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs))
 	if err != nil {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(batch.Outputs, len(batch.Outputs))
+	positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions))
 	if err != nil {
 		return nil, err
 	}
 
-	return m.TextModel.Forward(ctx, batch.Inputs, positions, outputs, batch, m.Cache), nil
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	if err != nil {
+		return nil, err
+	}
+
+	return m.TextModel.Forward(ctx, inputs, positions, outputs, opts, m.Cache), nil
 }
 
 func init() {

model/models/gemma3/model_text.go

@@ -45,6 +45,7 @@ func newTextModel(c ml.Config) *TextModel {
 
 	m := TextModel{
 		SentencePieceModel: model.NewSentencePieceModel(
+			c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Scores: c.Floats("tokenizer.ggml.scores"),
@@ -170,20 +171,53 @@ func (l *TextLayer) Forward(ctx ml.Context, layer int, hiddenState, positionIDs,
 	return hiddenState.Add(ctx, residual)
 }
 
-func (m *TextModel) Forward(ctx ml.Context, inputs, positions, outputs ml.Tensor, batch input.Batch, cache kvcache.Cache) ml.Tensor {
+func setImageEmbeddings(ctx ml.Context, hiddenState ml.Tensor, multimodal []input.MultimodalIndex) []int {
+	var embedding ml.Tensor
+	var src, dst, length int
+	var except []int
+
+	for _, image := range multimodal {
+		imageToken := image.Multimodal.(imageToken)
+		imageSrc := imageToken.index
+		imageDst := image.Index
+
+		if embedding == nil {
+			embedding = imageToken.embedding
+			src = imageSrc
+			dst = imageDst
+			length = 1
+		} else if embedding == imageToken.embedding && imageSrc+1 == src && imageDst+1 == dst {
+			src = imageSrc
+			dst = imageDst
+			length++
+		} else if embedding == imageToken.embedding && src+length == imageSrc && dst+length == imageDst {
+			length++
+		} else {
+			visionOutputs := embedding.View(ctx, src*embedding.Stride(1), length*embedding.Dim(0))
+			ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, dst*hiddenState.Stride(1), length*hiddenState.Dim(0))))
+
+			embedding = imageToken.embedding
+			src = imageSrc
+			dst = imageDst
+			length = 1
+		}
+
+		except = append(except, imageDst)
+	}
+
+	if embedding != nil {
+		visionOutputs := embedding.View(ctx, src*embedding.Stride(1), length*embedding.Dim(0))
+		ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, dst*hiddenState.Stride(1), length*hiddenState.Dim(0))))
+	}
+
+	return except
+}
+
+func (m *TextModel) Forward(ctx ml.Context, inputs, positions, outputs ml.Tensor, opts input.Options, cache kvcache.Cache) ml.Tensor {
 	hiddenState := m.TokenEmbedding.Forward(ctx, inputs)
 	hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(m.TextOptions.hiddenSize)))
 
-	// set image embeddings
-	var except []int
-	for _, image := range batch.Multimodal {
-		visionOutputs := image.Multimodal.(ml.Tensor)
-		ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, image.Index*hiddenState.Stride(1), visionOutputs.Dim(0)*visionOutputs.Dim(1))))
-
-		for i := range visionOutputs.Dim(1) {
-			except = append(except, image.Index+i)
-		}
-	}
+	except := setImageEmbeddings(ctx, hiddenState, opts.Multimodal)
 
 	for i, layer := range m.Layers {
 		// gemma alternates between the sliding window (local) and causal (global)

model/models/llama/model.go

@@ -139,18 +139,23 @@ func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs, outputs ml.Ten
 	return hiddenState.Add(ctx, residual)
 }
 
-func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
-	positions, err := ctx.Input().FromIntSlice(batch.Positions, len(batch.Positions))
+func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
+	inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs))
 	if err != nil {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(batch.Outputs, len(batch.Outputs))
+	positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions))
 	if err != nil {
 		return nil, err
 	}
 
-	hiddenState := m.TokenEmbedding.Forward(ctx, batch.Inputs)
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	if err != nil {
+		return nil, err
+	}
+
+	hiddenState := m.TokenEmbedding.Forward(ctx, inputs)
 
 	for i, layer := range m.Layers {
 		m.Cache.SetLayer(i)

model/models/mllama/model.go

@@ -63,10 +63,6 @@ func New(c ml.Config) (model.Model, error) {
 }
 
 func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, error) {
-	if len(m.VisionModel.Transformer.Layers) == 0 || len(m.GlobalTransformer.Layers) == 0 {
-		return nil, model.ErrNoVisionModel
-	}
-
 	image, _, err := image.Decode(bytes.NewReader(multimodalData))
 	if err != nil {
 		return nil, err
@@ -106,17 +102,17 @@ func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, er
 	return m.Projector.Forward(ctx, crossAttentionStates), nil
 }
 
-func (m *Model) PostTokenize(inputs []input.Input) ([]input.Input, error) {
+func (m *Model) PostTokenize(ctx ml.Context, inputs []input.Input) ([]input.Input, error) {
 	var images []input.Input
 	fnvHash := fnv.New64a()
 
 	for i := range inputs {
 		if inputs[i].Multimodal == nil {
 			if len(images) > 0 {
-				inputs[i].Multimodal = []ml.Tensor{images[0].Multimodal.(ml.Tensor)}
+				inputs[i].Multimodal = images[0].Multimodal
 				inputs[i].MultimodalHash = images[0].MultimodalHash
 				for j := 1; j < len(images); j++ {
-					inputs[i].Multimodal = append(inputs[i].Multimodal.([]ml.Tensor), images[0].Multimodal.(ml.Tensor))
+					inputs[i].Multimodal = inputs[i].Multimodal.(ml.Tensor).Concat(ctx, images[j].Multimodal.(ml.Tensor), 3)
 					fnvHash.Reset()
 					binary.Write(fnvHash, binary.NativeEndian, inputs[i].MultimodalHash)
 					binary.Write(fnvHash, binary.NativeEndian, inputs[j].MultimodalHash)
@@ -135,27 +131,29 @@ func (m *Model) PostTokenize(inputs []input.Input) ([]input.Input, error) {
 	return inputs, nil
 }
 
-func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
+func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 	var crossAttentionStates ml.Tensor
-	if len(batch.Multimodal) > 0 {
-		images := batch.Multimodal[len(batch.Multimodal)-1].Multimodal.([]ml.Tensor)
-		if len(images) > 0 {
-			crossAttentionStates = images[len(images)-1]
-		}
+	if len(opts.Multimodal) > 0 {
+		crossAttentionStates = opts.Multimodal[len(opts.Multimodal)-1].Multimodal.(ml.Tensor)
 	}
 
-	positions, err := ctx.Input().FromIntSlice(batch.Positions, len(batch.Positions))
+	inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs))
 	if err != nil {
 		return nil, err
 	}
 
-	outputs, err := ctx.Input().FromIntSlice(batch.Outputs, len(batch.Outputs))
+	positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions))
+	if err != nil {
+		return nil, err
+	}
+
+	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}
 
 	// TODO: attention mask, cross attention mask
-	return m.TextModel.Forward(ctx, batch.Inputs, positions, outputs, nil, crossAttentionStates, nil, m.Cache.(*kvcache.WrapperCache)), nil
+	return m.TextModel.Forward(ctx, inputs, positions, outputs, nil, crossAttentionStates, nil, m.Cache.(*kvcache.WrapperCache)), nil
 }
 
 func init() {

model/process_text_spm.go

@@ -1,23 +1,29 @@
 package model
 
 import (
-	"container/heap"
-	"fmt"
+	"iter"
 	"log/slog"
-	"strconv"
 	"strings"
+
+	"github.com/dlclark/regexp2"
+	queue "github.com/emirpasic/gods/v2/queues/priorityqueue"
 )
 
 const spmWhitespaceSep = "▁"
 
+func replaceWhitespaceBySeperator(s string) string {
+	return strings.ReplaceAll(s, " ", spmWhitespaceSep)
+}
+
 type SentencePieceModel struct {
 	maxTokenLen int
+	pre         *regexp2.Regexp
 	vocab       *Vocabulary
 }
 
 var _ TextProcessor = (*SentencePieceModel)(nil)
 
-func NewSentencePieceModel(vocab *Vocabulary) SentencePieceModel {
+func NewSentencePieceModel(pre string, vocab *Vocabulary) SentencePieceModel {
 	slog.Debug("Tokens", "num tokens", len(vocab.Values), "vals", vocab.Values[:5], "scores", vocab.Scores[:5], "types", vocab.Types[:5])
 
 	counter := map[int]int{}
@@ -38,6 +44,7 @@ func NewSentencePieceModel(vocab *Vocabulary) SentencePieceModel {
 
 	return SentencePieceModel{
 		maxTokenLen: maxTokenLen,
+		pre:         regexp2.MustCompile(pre, regexp2.Unicode|regexp2.RE2),
 		vocab:       vocab,
 	}
 }
@@ -46,9 +53,20 @@ func (spm SentencePieceModel) Is(id int32, special Special) bool {
 	return spm.vocab.Is(id, special)
 }
 
+func (spm *SentencePieceModel) split(s string) iter.Seq[string] {
+	return func(yield func(string) bool) {
+		for m, _ := spm.pre.FindStringMatch(s); m != nil; m, _ = spm.pre.FindNextMatch(m) {
+			if !yield(m.String()) {
+				break
+			}
+		}
+	}
+}
+
 func (spm SentencePieceModel) Encode(s string, addSpecial bool) ([]int32, error) {
 	fragments := []fragment{{value: s}}
 	for _, special := range spm.vocab.SpecialVocabulary() {
+		// TODO: process special tokens concurrently
 		id := spm.vocab.Encode(special)
 		for i := 0; i < len(fragments); i++ {
 			frag := fragments[i]
@@ -73,6 +91,7 @@ func (spm SentencePieceModel) Encode(s string, addSpecial bool) ([]int32, error)
 			fragments = append(fragments[:i], append(middle, fragments[i+1:]...)...)
 		}
 	}
+	slog.Debug("fragments", "frags", fragments)
 
 	var ids []int32
 	for _, frag := range fragments {
@@ -81,96 +100,105 @@ func (spm SentencePieceModel) Encode(s string, addSpecial bool) ([]int32, error)
 			continue
 		}
 
-		text := strings.ReplaceAll(frag.value, " ", spmWhitespaceSep)
+		for split := range spm.split(frag.value) {
+			split = replaceWhitespaceBySeperator(split)
 
-		if id := spm.vocab.Encode(text); id >= 0 {
-			ids = append(ids, id)
-			continue
-		}
-
-		q := &queue{}
-		heap.Init(q)
-
-		runes := []rune(text)
-		merges := make([]merge, len(runes))
-		for r := range runes {
-			merges[r] = merge{
-				p:     r - 1,
-				n:     r + 1,
-				runes: []rune{runes[r]},
-			}
-		}
-
-		pairwise := func(a, b int) *candidate {
-			if a < 0 || b >= len(runes) {
-				return nil
-			}
-
-			left, right := string(merges[a].runes), string(merges[b].runes)
-			if id := spm.vocab.Encode(left + right); id >= 0 {
-				return &candidate{
-					a:     a,
-					b:     b,
-					score: spm.vocab.Scores[id],
-					size:  len(left) + len(right),
-				}
-			}
-
-			return nil
-		}
-
-		for i := range len(runes) - 1 {
-			if pair := pairwise(i, i+1); pair != nil {
-				heap.Push(q, pair)
-			}
-		}
-
-		for q.Len() > 0 {
-			pair := heap.Pop(q).(*candidate)
-			left, right := merges[pair.a], merges[pair.b]
-
-			if string(left.runes) == "" || string(right.runes) == "" || len(string(left.runes))+len(string(right.runes)) != pair.size {
+			var sb strings.Builder
+			sb.Write([]byte(split))
+			if id := spm.vocab.Encode(sb.String()); id >= 0 {
+				ids = append(ids, id)
 				continue
 			}
 
-			merges[pair.a].runes = append(left.runes, right.runes...)
-			merges[pair.b].runes = nil
-			merges[pair.a].n = right.n
-			if right.n < len(merges) {
-				merges[right.n].p = pair.a
+			runes := []rune(sb.String())
+			pq := queue.NewWith(func(a, b any) int {
+				priA := a.(*candidate)
+				priB := b.(*candidate)
+				if priA.score > priB.score || (priA.score == priB.score && priA.a < priB.a) {
+					return -1
+				}
+				return 1
+			})
+
+			merges := make([]merge, len(runes))
+			for r := range runes {
+				merges[r] = merge{
+					p:     r - 1,
+					n:     r + 1,
+					runes: []rune{runes[r]},
+				}
 			}
 
-			if pair := pairwise(merges[pair.a].p, pair.a); pair != nil {
-				heap.Push(q, pair)
+			slog.Debug("tokenizer", "merges", merges)
+
+			pairwise := func(a, b int) *candidate {
+				if a < 0 || b >= len(runes) {
+					return nil
+				}
+
+				left, right := string(merges[a].runes), string(merges[b].runes)
+				if id := spm.vocab.Encode(left + right); id >= 0 {
+					return &candidate{
+						a:     a,
+						b:     b,
+						score: spm.vocab.Scores[id],
+					}
+				}
+				return nil
 			}
 
-			if pair := pairwise(pair.a, merges[pair.a].n); pair != nil {
-				heap.Push(q, pair)
+			for i := range len(runes) - 1 {
+				if pair := pairwise(i, i+1); pair != nil {
+					pq.Enqueue(pair)
+				}
 			}
-		}
 
-		for _, merge := range merges {
-			if token := string(merge.runes); token != "" {
-				id := spm.vocab.Encode(token)
+			pqv := pq.Values()
+			for _, v := range pqv {
+				e := v.(*candidate)
+				slog.Debug("candidate", "candidate", e)
+			}
 
-				if id >= 0 {
-					ids = append(ids, id)
+			for !pq.Empty() {
+				v, _ := pq.Dequeue()
+				pair := v.(*candidate)
+				left, right := merges[pair.a], merges[pair.b]
+
+				slog.Debug("pair", "left", left, "right", right)
+				if len(left.runes) == 0 || len(right.runes) == 0 {
 					continue
 				}
 
-				// Fallback to byte tokenization
-				var result []int32
-				for _, b := range []byte(token) {
-					byteToken := fmt.Sprintf("<0x%02X>", b)
-					unknownID := spm.vocab.Encode(byteToken)
-					if unknownID >= 0 {
-						result = append(result, unknownID)
-					} else {
-						slog.Debug("unknown byte token", "byte", b, "token", byteToken)
-					}
+				if id := spm.vocab.Encode(string(left.runes) + string(right.runes)); id < 0 {
+					continue
 				}
 
-				ids = append(ids, result...)
+				merges[pair.a].runes = append(left.runes, right.runes...)
+				merges[pair.b].runes = nil
+				merges[pair.a].n = right.n
+				if right.n < len(merges) {
+					merges[right.n].p = pair.a
+				}
+
+				if pair := pairwise(merges[pair.a].p, pair.a); pair != nil {
+					pq.Enqueue(pair)
+				}
+
+				if pair := pairwise(pair.a, merges[pair.a].n); pair != nil {
+					pq.Enqueue(pair)
+				}
+			}
+
+			slog.Debug("merges", "merges", merges)
+
+			for _, merge := range merges {
+				if len(merge.runes) > 0 {
+					if id := spm.vocab.Encode(string(merge.runes)); id >= 0 {
+						ids = append(ids, id)
+					} else {
+						slog.Debug("missing token", "token", string(merge.runes))
+					}
+				}
 			}
 		}
 	}
@@ -201,30 +229,6 @@ func (spm SentencePieceModel) Encode(s string, addSpecial bool) ([]int32, error)
 type candidate struct {
 	a, b  int
 	score float32
-	size  int
-}
-
-type queue []*candidate
-
-func (q queue) Len() int { return len(q) }
-
-func (q queue) Less(i, j int) bool {
-	return (q[i].score > q[j].score) || (q[i].score == q[j].score && q[i].a < q[j].a)
-}
-
-func (q queue) Swap(i, j int) { q[i], q[j] = q[j], q[i] }
-
-func (q *queue) Push(x interface{}) {
-	item := x.(*candidate)
-	*q = append(*q, item)
-}
-
-func (q *queue) Pop() interface{} {
-	old := *q
-	n := len(old)
-	item := old[n-1]
-	*q = old[0 : n-1]
-	return item
 }
 
 func (spm SentencePieceModel) Decode(ids []int32) (string, error) {
@@ -232,26 +236,11 @@ func (spm SentencePieceModel) Decode(ids []int32) (string, error) {
 	for _, id := range ids {
 		data := spm.vocab.Decode(id)
 		data = strings.ReplaceAll(data, spmWhitespaceSep, " ")
-
-		// For tokenizers that use byte tokens like "<0xEA>"
-		// convert them to the partial unicode character
-		// so they are buffered correctly by the runner instead
-		// of being sent back to the api as "<0xEA>"
-		if len(data) == 6 && strings.HasPrefix(data, "<0x") && strings.HasSuffix(data, ">") {
-			byteVal, err := strconv.ParseUint(data[1:5], 0, 8)
-			if err != nil {
-				return "", fmt.Errorf("failed to parse hex byte: %v", err)
-			}
-
-			if err := sb.WriteByte(byte(byteVal)); err != nil {
-				return "", err
-			}
-		} else {
-			if _, err := sb.WriteString(data); err != nil {
-				return "", err
-			}
+		if _, err := sb.WriteString(data); err != nil {
+			return "", err
 		}
 	}
 
+	slog.Debug("decoded", "ids", ids, "text", sb.String())
 	return sb.String(), nil
 }

model/process_text_spm_test.go

@@ -25,6 +25,8 @@ func loadSentencePieceVocab(t *testing.T) SentencePieceModel {
 		t.Fatal(err)
 	}
 
+	preTokenizer := `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`
+
 	var v Vocabulary
 
 	for _, piece := range spm.GetPieces() {
@@ -45,7 +47,7 @@ func loadSentencePieceVocab(t *testing.T) SentencePieceModel {
 		}
 	}
 
-	return NewSentencePieceModel(&v)
+	return NewSentencePieceModel(preTokenizer, &v)
 }
 
 func TestSentencePieceEncode(t *testing.T) {
@@ -114,59 +116,3 @@ func TestSentencePieceEncode(t *testing.T) {
 		}
 	})
 }
-
-func TestSentencePieceModelDecodeByteTokens(t *testing.T) {
-	vocab := &Vocabulary{
-		Values: []string{
-			"normal",
-			"<0xEA>",
-			"<0x41>",
-			"<0xC3>",
-			"<0xA3>",
-		},
-		Types: []uint32{
-			TOKEN_TYPE_NORMAL,
-			TOKEN_TYPE_BYTE,
-			TOKEN_TYPE_BYTE,
-			TOKEN_TYPE_BYTE,
-			TOKEN_TYPE_BYTE,
-		},
-		Scores: []float32{0, 0, 0, 0, 0},
-	}
-
-	spm := NewSentencePieceModel(vocab)
-
-	tests := []struct {
-		name     string
-		ids      []int32
-		expected string
-	}{
-		{
-			name:     "single byte token",
-			ids:      []int32{1},
-			expected: "\xea",
-		},
-		{
-			name:     "ASCII byte token",
-			ids:      []int32{2},
-			expected: "A",
-		},
-		{
-			name:     "multiple byte tokens forming UTF-8 character",
-			ids:      []int32{3, 4},
-			expected: "ã",
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result, err := spm.Decode(tt.ids)
-			if err != nil {
-				t.Errorf("failed to decode token IDs %v: %v", tt.ids, err)
-			}
-			if result != tt.expected {
-				t.Errorf("got %q, want %q", result, tt.expected)
-			}
-		})
-	}
-}

openai/openai.go

@@ -23,10 +23,10 @@ import (
 var finishReasonToolCalls = "tool_calls"
 
 type Error struct {
-	Message string  `json:"message"`
-	Type    string  `json:"type"`
-	Param   any     `json:"param"`
-	Code    *string `json:"code"`
+	Message string      `json:"message"`
+	Type    string      `json:"type"`
+	Param   interface{} `json:"param"`
+	Code    *string     `json:"code"`
 }
 
 type ErrorResponse struct {
@@ -465,7 +465,7 @@ func fromChatRequest(r ChatCompletionRequest) (*api.ChatRequest, error) {
 		}
 	}
 
-	options := make(map[string]any)
+	options := make(map[string]interface{})
 
 	switch stop := r.Stop.(type) {
 	case string:

openai/openai_test.go

@@ -219,7 +219,7 @@ func TestChatMiddleware(t *testing.T) {
 							{
 								Function: api.ToolCallFunction{
 									Name: "get_current_weather",
-									Arguments: map[string]any{
+									Arguments: map[string]interface{}{
 										"location": "Paris, France",
 										"format":   "celsius",
 									},

runner/llamarunner/cache.go

@@ -213,16 +213,8 @@ func (c *InputCache) ShiftDiscard(inputLen int, numKeep int) int {
 	return discard
 }
 
-type ErrReprocessInputs struct {
-	Inputs []input
-}
-
-func (e *ErrReprocessInputs) Error() string {
-	return fmt.Sprintf("kv cache shift not supported, inputs need reprocessing (input count: %v)", len(e.Inputs))
-}
-
-// ShiftCacheSlot frees up space in the KV cache by deleting the oldest half of history
-// and shifting the newest half into that space (saving numKeep inputs at the beginning).
+// Frees up space in the KV cache by deleting the oldest half of history and shifting
+// the newest half into that space (saving numKeep inputs at the beginning).
 //
 // Assumes that at least 1 entry can be freed up by shifting (i.e. numKeep < numCtx)
 func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int) error {
@@ -230,8 +222,7 @@ func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int) error {
 		return fmt.Errorf("unable to shift context - keep exceeds context (keep: %v context: %v)", numKeep, c.numCtx)
 	}
 
-	inputLen := len(slot.Inputs)
-	discard := c.ShiftDiscard(inputLen, numKeep)
+	discard := c.ShiftDiscard(len(slot.Inputs), numKeep)
 
 	if discard <= 0 {
 		return nil
@@ -240,42 +231,16 @@ func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int) error {
 	slog.Debug("context limit hit - shifting", "id", slot.Id, "limit", c.numCtx, "input", len(slot.Inputs),
 		"keep", numKeep, "discard", discard)
 
-	var shiftFailed bool
-
-	if c.lc.KvCacheCanShift() {
-		// For models that support shifting, attempt to shift the KV cache
-		if !c.lc.KvCacheSeqRm(slot.Id, numKeep, numKeep+discard) {
-			shiftFailed = true
-			slog.Debug("kv cache removal not supported, clearing cache and returning inputs for reprocessing", "id", slot.Id)
-		} else {
-			c.lc.KvCacheSeqAdd(slot.Id, numKeep+discard, inputLen, -discard)
-		}
-	} else {
-		// For models that don't support shifting
-		shiftFailed = true
-		slog.Debug("kv cache cannot shift, clearing cache and returning inputs for reprocessing", "id", slot.Id)
+	// TODO (jessegross): KV cache removal can fail for certain types of models
+	if !c.lc.KvCacheSeqRm(slot.Id, numKeep, numKeep+discard) {
+		return fmt.Errorf("unable to remove old kv cache entries (id: %v, keep: %v discard: %v)", slot.Id, numKeep, discard)
 	}
+	c.lc.KvCacheSeqAdd(slot.Id, numKeep+discard, len(slot.Inputs), -discard)
 
-	if shiftFailed {
-		// Create new input slice with preserved tokens (numKeep + remaining tokens after discard)
-		newInputs := make([]input, numKeep+inputLen-(numKeep+discard))
-		copy(newInputs[:numKeep], slot.Inputs[:numKeep])
-		copy(newInputs[numKeep:], slot.Inputs[numKeep+discard:])
-
-		// Clear the entire KV cache
-		_ = c.lc.KvCacheSeqRm(slot.Id, 0, -1)
-		// Reset the slot inputs since we've cleared the cache
-		slot.Inputs = []input{}
-
-		// Return error with inputs that need to be reprocessed
-		return &ErrReprocessInputs{Inputs: newInputs}
-	}
-
-	// Standard shift succeeded - update input array
-	for i := numKeep + discard; i < inputLen; i++ {
+	for i := numKeep + discard; i < len(slot.Inputs); i++ {
 		slot.Inputs[i-discard] = slot.Inputs[i]
 	}
-	slot.Inputs = slot.Inputs[:inputLen-discard]
+	slot.Inputs = slot.Inputs[:len(slot.Inputs)-discard]
 
 	return nil
 }

runner/llamarunner/runner.go

@@ -24,7 +24,6 @@ import (
 
 	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/llama"
-	"github.com/ollama/ollama/llm"
 	"github.com/ollama/ollama/runner/common"
 )
 
@@ -100,7 +99,7 @@ type NewSequenceParams struct {
 	embedding      bool
 }
 
-func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
+func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequenceParams) (*Sequence, error) {
 	s.ready.Wait()
 
 	startTime := time.Now()
@@ -164,7 +163,7 @@ func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSe
 // inputs processes the prompt and images into a list of inputs
 // by splitting the prompt on [img-<n>] tags, tokenizing text and
 // generating image embeddings for each image
-func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input, error) {
+func (s *Server) inputs(prompt string, images []ImageData) ([]input, error) {
 	var inputs []input
 	var parts []string
 	var matches [][]string
@@ -230,7 +229,7 @@ type Server struct {
 	image *ImageContext
 
 	// status for external health reporting - loading, ready to serve, etc.
-	status llm.ServerStatus
+	status ServerStatus
 
 	// current progress on loading the model
 	progress float32
@@ -389,15 +388,7 @@ func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch)
 				if len(seq.pendingInputs) == 0 {
 					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
 					if err != nil {
-						var reprocess *ErrReprocessInputs
-						if errors.As(err, &reprocess) {
-							// Prepend these inputs to the sequence's inputs queue for reprocessing
-							seq.inputs = append(reprocess.Inputs, seq.inputs...)
-							// Continue processing as normal
-							continue
-						} else {
-							return err
-						}
+						return err
 					}
 				} else {
 					break
@@ -550,18 +541,75 @@ func (s *Server) processBatch(tokenBatch *llama.Batch, embedBatch *llama.Batch)
 	return nil
 }
 
+// TODO (jmorganca): use structs from the api package to avoid duplication
+// this way the api acts as a proxy instead of using a different api for the
+// runner
+type Options struct {
+	api.Runner
+
+	NumKeep          int      `json:"n_keep"`
+	Seed             int      `json:"seed"`
+	NumPredict       int      `json:"n_predict"`
+	TopK             int      `json:"top_k"`
+	TopP             float32  `json:"top_p"`
+	MinP             float32  `json:"min_p"`
+	TypicalP         float32  `json:"typical_p"`
+	RepeatLastN      int      `json:"repeat_last_n"`
+	Temperature      float32  `json:"temperature"`
+	RepeatPenalty    float32  `json:"repeat_penalty"`
+	PresencePenalty  float32  `json:"presence_penalty"`
+	FrequencyPenalty float32  `json:"frequency_penalty"`
+	Mirostat         int      `json:"mirostat"`
+	MirostatTau      float32  `json:"mirostat_tau"`
+	MirostatEta      float32  `json:"mirostat_eta"`
+	Stop             []string `json:"stop"`
+}
+
+type ImageData struct {
+	Data          []byte `json:"data"`
+	ID            int    `json:"id"`
+	AspectRatioID int    `json:"aspect_ratio_id"`
+}
+
+type CompletionRequest struct {
+	Prompt      string      `json:"prompt"`
+	Images      []ImageData `json:"image_data"`
+	Grammar     string      `json:"grammar"`
+	CachePrompt bool        `json:"cache_prompt"`
+
+	Options
+}
+
+type Timings struct {
+	PredictedN  int     `json:"predicted_n"`
+	PredictedMS float64 `json:"predicted_ms"`
+	PromptN     int     `json:"prompt_n"`
+	PromptMS    float64 `json:"prompt_ms"`
+}
+
+type CompletionResponse struct {
+	Content string `json:"content"`
+	Stop    bool   `json:"stop"`
+
+	Model        string  `json:"model,omitempty"`
+	Prompt       string  `json:"prompt,omitempty"`
+	StoppedLimit bool    `json:"stopped_limit,omitempty"`
+	PredictedN   int     `json:"predicted_n,omitempty"`
+	PredictedMS  float64 `json:"predicted_ms,omitempty"`
+	PromptN      int     `json:"prompt_n,omitempty"`
+	PromptMS     float64 `json:"prompt_ms,omitempty"`
+
+	Timings Timings `json:"timings"`
+}
+
 func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
-	var req llm.CompletionRequest
+	var req CompletionRequest
+	req.Options = Options(api.DefaultOptions())
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 		http.Error(w, "Bad request", http.StatusBadRequest)
 		return
 	}
 
-	if req.Options == nil {
-		opts := api.DefaultOptions()
-		req.Options = &opts
-	}
-
 	// Set the headers to indicate streaming
 	w.Header().Set("Content-Type", "application/json")
 	w.Header().Set("Transfer-Encoding", "chunked")
@@ -572,28 +620,26 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Extract options from the CompletionRequest
-	samplingParams := llama.SamplingParams{
-		TopK:           req.Options.TopK,
-		TopP:           req.Options.TopP,
-		MinP:           req.Options.MinP,
-		TypicalP:       req.Options.TypicalP,
-		Temp:           req.Options.Temperature,
-		RepeatLastN:    req.Options.RepeatLastN,
-		PenaltyRepeat:  req.Options.RepeatPenalty,
-		PenaltyFreq:    req.Options.FrequencyPenalty,
-		PenaltyPresent: req.Options.PresencePenalty,
-		Mirostat:       req.Options.Mirostat,
-		MirostatTau:    req.Options.MirostatTau,
-		MirostatEta:    req.Options.MirostatEta,
-		Seed:           uint32(req.Options.Seed),
-		Grammar:        req.Grammar,
-	}
+	var samplingParams llama.SamplingParams
+	samplingParams.TopK = req.TopK
+	samplingParams.TopP = req.TopP
+	samplingParams.MinP = req.MinP
+	samplingParams.TypicalP = req.TypicalP
+	samplingParams.Temp = req.Temperature
+	samplingParams.RepeatLastN = req.RepeatLastN
+	samplingParams.PenaltyRepeat = req.RepeatPenalty
+	samplingParams.PenaltyFreq = req.FrequencyPenalty
+	samplingParams.PenaltyPresent = req.PresencePenalty
+	samplingParams.Mirostat = req.Mirostat
+	samplingParams.MirostatTau = req.MirostatTau
+	samplingParams.MirostatEta = req.MirostatEta
+	samplingParams.Seed = uint32(req.Seed)
+	samplingParams.Grammar = req.Grammar
 
 	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
-		numPredict:     req.Options.NumPredict,
-		stop:           req.Options.Stop,
-		numKeep:        req.Options.NumKeep,
+		numPredict:     req.NumPredict,
+		stop:           req.Stop,
+		numKeep:        req.NumKeep,
 		samplingParams: &samplingParams,
 		embedding:      false,
 	})
@@ -607,7 +653,7 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 		if errors.Is(err, context.Canceled) {
 			slog.Info("aborting completion request due to client closing the connection")
 		} else {
-			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
+			slog.Error("Failed to acquire semaphore", "error", err)
 		}
 		return
 	}
@@ -616,10 +662,9 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	found := false
 	for i, sq := range s.seqs {
 		if sq == nil {
-			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, true)
+			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
 			if err != nil {
 				s.mu.Unlock()
-				s.seqsSem.Release(1)
 				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
 				return
 			}
@@ -635,7 +680,6 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	s.mu.Unlock()
 
 	if !found {
-		s.seqsSem.Release(1)
 		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
 		return
 	}
@@ -647,7 +691,7 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 			return
 		case content, ok := <-seq.responses:
 			if ok {
-				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
+				if err := json.NewEncoder(w).Encode(&CompletionResponse{
 					Content: content,
 				}); err != nil {
 					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
@@ -658,17 +702,15 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 				flusher.Flush()
 			} else {
 				// Send the final response
-				doneReason := "stop"
-				if seq.doneReason == "limit" {
-					doneReason = "length"
-				}
-				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
-					Done:               true,
-					DoneReason:         doneReason,
-					PromptEvalCount:    seq.numPromptInputs,
-					PromptEvalDuration: seq.startGenerationTime.Sub(seq.startProcessingTime),
-					EvalCount:          seq.numDecoded,
-					EvalDuration:       time.Since(seq.startGenerationTime),
+				if err := json.NewEncoder(w).Encode(&CompletionResponse{
+					Stop:         true,
+					StoppedLimit: seq.doneReason == "limit",
+					Timings: Timings{
+						PromptN:     seq.numPromptInputs,
+						PromptMS:    float64(seq.startGenerationTime.Sub(seq.startProcessingTime).Milliseconds()),
+						PredictedN:  seq.numDecoded,
+						PredictedMS: float64(time.Since(seq.startGenerationTime).Milliseconds()),
+					},
 				}); err != nil {
 					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
 				}
@@ -679,8 +721,17 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+type EmbeddingRequest struct {
+	Content     string `json:"content"`
+	CachePrompt bool   `json:"cache_prompt"`
+}
+
+type EmbeddingResponse struct {
+	Embedding []float32 `json:"embedding"`
+}
+
 func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
-	var req llm.EmbeddingRequest
+	var req EmbeddingRequest
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 		http.Error(w, fmt.Sprintf("bad request: %s", err), http.StatusBadRequest)
 		return
@@ -701,7 +752,7 @@ func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
 		if errors.Is(err, context.Canceled) {
 			slog.Info("aborting embeddings request due to client closing the connection")
 		} else {
-			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
+			slog.Error("Failed to acquire semaphore", "error", err)
 		}
 		return
 	}
@@ -710,10 +761,9 @@ func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
 	found := false
 	for i, sq := range s.seqs {
 		if sq == nil {
-			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, false)
+			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
 			if err != nil {
 				s.mu.Unlock()
-				s.seqsSem.Release(1)
 				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
 				return
 			}
@@ -726,24 +776,47 @@ func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
 	s.mu.Unlock()
 
 	if !found {
-		s.seqsSem.Release(1)
 		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
 		return
 	}
 
 	embedding := <-seq.embedding
 
-	if err := json.NewEncoder(w).Encode(&llm.EmbeddingResponse{
+	if err := json.NewEncoder(w).Encode(&EmbeddingResponse{
 		Embedding: embedding,
 	}); err != nil {
 		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
 	}
 }
 
+type HealthResponse struct {
+	Status   string  `json:"status"`
+	Progress float32 `json:"progress"`
+}
+
+type ServerStatus int
+
+const (
+	ServerStatusReady ServerStatus = iota
+	ServerStatusLoadingModel
+	ServerStatusError
+)
+
+func (s ServerStatus) ToString() string {
+	switch s {
+	case ServerStatusReady:
+		return "ok"
+	case ServerStatusLoadingModel:
+		return "loading model"
+	default:
+		return "server error"
+	}
+}
+
 func (s *Server) health(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
-	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
-		Status:   s.status,
+	if err := json.NewEncoder(w).Encode(&HealthResponse{
+		Status:   s.status.ToString(),
 		Progress: s.progress,
 	}); err != nil {
 		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
@@ -806,7 +879,7 @@ func (s *Server) loadModel(
 		panic(err)
 	}
 
-	s.status = llm.ServerStatusReady
+	s.status = ServerStatusReady
 	s.ready.Done()
 }
 
@@ -864,7 +937,7 @@ func Execute(args []string) error {
 		parallel:  *parallel,
 		seqs:      make([]*Sequence, *parallel),
 		seqsSem:   semaphore.NewWeighted(int64(*parallel)),
-		status:    llm.ServerStatusLoadingModel,
+		status:    ServerStatusLoadingModel,
 	}
 
 	var tensorSplitFloats []float32

runner/ollamarunner/cache.go

@@ -31,10 +31,8 @@ type InputCache struct {
 	cache kvcache.Cache
 }
 
-func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots int, batchSize int, multiUserCache bool) (*InputCache, error) {
-	numCtx := kvSize / int32(numSlots)
-
-	if numCtx < 1 {
+func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots int, multiUserCache bool) (*InputCache, error) {
+	if kvSize/int32(numSlots) < 1 {
 		return nil, fmt.Errorf("must have at least one kv cache entry per parallel sequence (kv: %v parallel: %v)", kvSize, numSlots)
 	}
 
@@ -46,11 +44,11 @@ func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots
 
 	cache := model.Config().Cache
 	if cache != nil {
-		cache.Init(model.Backend(), kvCacheTypeFromStr(kvCacheType), numSlots, int(numCtx), batchSize)
+		cache.Init(model.Backend(), kvCacheTypeFromStr(kvCacheType), kvSize)
 	}
 
 	return &InputCache{
-		numCtx:         numCtx,
+		numCtx:         kvSize / int32(numSlots),
 		enabled:        cache != nil,
 		slots:          slots,
 		multiUserCache: multiUserCache,
@@ -91,7 +89,7 @@ type InputCacheSlot struct {
 	lastUsed time.Time
 }
 
-func (c *InputCache) LoadCacheSlot(prompt []input.Input) (*InputCacheSlot, []input.Input, error) {
+func (c *InputCache) LoadCacheSlot(prompt []input.Input, cachePrompt bool) (*InputCacheSlot, []input.Input, error) {
 	var slot *InputCacheSlot
 	var numPast int32
 	var err error
@@ -109,6 +107,10 @@ func (c *InputCache) LoadCacheSlot(prompt []input.Input) (*InputCacheSlot, []inp
 		return nil, nil, err
 	}
 
+	if !cachePrompt {
+		numPast = 0
+	}
+
 	slot.InUse = true
 	slot.lastUsed = time.Now()
 
@@ -118,10 +120,6 @@ func (c *InputCache) LoadCacheSlot(prompt []input.Input) (*InputCacheSlot, []inp
 	}
 
 	if c.cache != nil {
-		if numPast > 0 && !c.cache.CanResume(slot.Id, numPast) {
-			numPast = 0
-		}
-
 		err = c.cache.Remove(slot.Id, numPast, math.MaxInt32)
 		if err != nil {
 			// Some models don't support partial erasure
@@ -229,8 +227,6 @@ func countCommonPrefix(a []input.Input, b []input.Input) int32 {
 	return count
 }
 
-// TODO(jessegross): If we need to reprocess the inputs we should ensure that
-// we don't split up a SameBatch
 func (c *InputCache) ShiftDiscard(inputLen int32, numKeep int32) int32 {
 	targetFree := (c.numCtx - numKeep) / 2
 	targetFree = max(targetFree, 1)
@@ -245,14 +241,6 @@ func (c *InputCache) ShiftDiscard(inputLen int32, numKeep int32) int32 {
 	return discard
 }
 
-type ErrReprocessInputs struct {
-	Inputs []input.Input
-}
-
-func (e *ErrReprocessInputs) Error() string {
-	return fmt.Sprintf("kv cache shift not supported, inputs need reprocessing (input count: %v)", len(e.Inputs))
-}
-
 // Frees up space in the KV cache by deleting the oldest half of history and shifting
 // the newest half into that space (saving numKeep inputs at the beginning).
 //
@@ -272,23 +260,11 @@ func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int32) error {
 	slog.Debug("context limit hit - shifting", "id", slot.Id, "limit", c.numCtx, "input", len(slot.Inputs),
 		"keep", numKeep, "discard", discard)
 
+	// TODO (jessegross): KV cache removal can fail for certain types of models
 	if c.cache != nil {
 		err := c.cache.Remove(slot.Id, numKeep, numKeep+discard)
 		if err != nil {
-			slog.Debug("kv cache removal unsupported, clearing cache and returning inputs for reprocessing",
-				"id", slot.Id, "error", err)
-
-			// Create new input slice with preserved tokens (numKeep + remaining tokens after discard)
-			newInputs := make([]input.Input, numKeep+inputLen-(numKeep+discard))
-			copy(newInputs[:numKeep], slot.Inputs[:numKeep])
-			copy(newInputs[numKeep:], slot.Inputs[numKeep+discard:])
-
-			// Reset the cache
-			_ = c.cache.Remove(slot.Id, 0, -1)
-			slot.Inputs = []input.Input{}
-
-			// Return error with inputs that need to be reprocessed
-			return &ErrReprocessInputs{Inputs: newInputs}
+			return fmt.Errorf("unable to remove old kv cache entries (id: %v, keep: %v discard: %v): %w", slot.Id, numKeep, discard, err)
 		}
 	}
 

runner/ollamarunner/cache_test.go

@@ -1,13 +1,10 @@
 package ollamarunner
 
 import (
-	"errors"
-	"fmt"
 	"image"
 	"testing"
 	"time"
 
-	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/model/input"
 )
 
@@ -300,220 +297,3 @@ func TestShiftDiscard(t *testing.T) {
 		})
 	}
 }
-
-func TestLoadCacheSlot(t *testing.T) {
-	tests := []struct {
-		name           string
-		cache          InputCache
-		prompt         []input.Input
-		wantErr        bool
-		expectedSlotId int
-		expectedPrompt int // expected length of remaining prompt
-	}{
-		{
-			name: "Basic cache hit - single user",
-			cache: InputCache{
-				multiUserCache: false,
-				slots: []InputCacheSlot{
-					{
-						Id:       0,
-						Inputs:   []input.Input{{Token: 1}, {Token: 2}},
-						InUse:    false,
-						lastUsed: time.Now().Add(-time.Second),
-					},
-					{
-						Id:       1,
-						Inputs:   []input.Input{},
-						InUse:    false,
-						lastUsed: time.Now().Add(-2 * time.Second),
-					},
-				},
-			},
-			prompt:         []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
-			wantErr:        false,
-			expectedSlotId: 0,
-			expectedPrompt: 1, // Only token 3 remains
-		},
-		{
-			name: "Basic cache hit - multi user",
-			cache: InputCache{
-				multiUserCache: true,
-				slots: []InputCacheSlot{
-					{
-						Id:       0,
-						Inputs:   []input.Input{{Token: 1}, {Token: 2}},
-						InUse:    false,
-						lastUsed: time.Now().Add(-time.Second),
-					},
-					{
-						Id:       1,
-						Inputs:   []input.Input{},
-						InUse:    false,
-						lastUsed: time.Now().Add(-2 * time.Second),
-					},
-				},
-			},
-			prompt:         []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
-			wantErr:        false,
-			expectedSlotId: 0,
-			expectedPrompt: 1, // Only token 3 remains
-		},
-		{
-			name: "Exact match - leave one input",
-			cache: InputCache{
-				multiUserCache: false,
-				slots: []InputCacheSlot{
-					{
-						Id:       0,
-						Inputs:   []input.Input{{Token: 1}, {Token: 2}},
-						InUse:    false,
-						lastUsed: time.Now().Add(-time.Second),
-					},
-				},
-			},
-			prompt:         []input.Input{{Token: 1}, {Token: 2}},
-			wantErr:        false,
-			expectedSlotId: 0,
-			expectedPrompt: 1, // Should leave 1 token for sampling
-		},
-		{
-			name: "No available slots",
-			cache: InputCache{
-				multiUserCache: false,
-				slots: []InputCacheSlot{
-					{
-						Id:       0,
-						Inputs:   []input.Input{{Token: 1}, {Token: 2}},
-						InUse:    true,
-						lastUsed: time.Now().Add(-time.Second),
-					},
-				},
-			},
-			prompt:         []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
-			wantErr:        true,
-			expectedSlotId: -1,
-			expectedPrompt: -1,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			slot, remainingPrompt, err := tt.cache.LoadCacheSlot(tt.prompt)
-
-			// Check error state
-			if (err != nil) != tt.wantErr {
-				t.Errorf("LoadCacheSlot() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-
-			if tt.wantErr {
-				return // Skip further checks if we expected an error
-			}
-
-			// Verify slot ID
-			if slot.Id != tt.expectedSlotId {
-				t.Errorf("LoadCacheSlot() slot ID = %v, expected %v", slot.Id, tt.expectedSlotId)
-			}
-
-			// Verify slot is now marked in use
-			if !slot.InUse {
-				t.Errorf("LoadCacheSlot() slot not marked InUse")
-			}
-
-			// Verify remaining prompt length
-			if len(remainingPrompt) != tt.expectedPrompt {
-				t.Errorf("LoadCacheSlot() remaining prompt length = %v, expected %v",
-					len(remainingPrompt), tt.expectedPrompt)
-			}
-		})
-	}
-}
-
-// Mock implementation of the Cache interface
-type mockCache struct {
-	shouldFail bool
-}
-
-// Implement only the methods needed for the test
-func (m *mockCache) Remove(seq int, beginIndex, endIndex int32) error {
-	if m.shouldFail {
-		return fmt.Errorf("mock cache removal error")
-	}
-	return nil
-}
-
-// Stub implementations for other interface methods
-func (m *mockCache) SetLayer(layer int)                                                            {}
-func (m *mockCache) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor)                          { return nil, nil, nil }
-func (m *mockCache) Put(ctx ml.Context, key, value ml.Tensor)                                      {}
-func (m *mockCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) {}
-func (m *mockCache) Close()                                                                        {}
-func (m *mockCache) StartForward(ctx ml.Context, batch input.Batch) error                          { return nil }
-func (m *mockCache) CopyPrefix(srcSeq, dstSeq int, len int32)                                      {}
-func (m *mockCache) SetConfig(ml.CacheConfig)                                                      {}
-func (m *mockCache) CanResume(seq int, pos int32) bool                                             { return true }
-
-func TestShiftCacheSlot(t *testing.T) {
-	tests := []struct {
-		name          string
-		numCtx        int32
-		inputs        []input.Input
-		numKeep       int32
-		cacheErr      bool
-		wantErr       any
-		wantInputsLen int
-	}{
-		{
-			name:          "Normal shift",
-			numCtx:        10,
-			inputs:        []input.Input{{Token: 1}, {Token: 2}, {Token: 3}, {Token: 4}, {Token: 5}, {Token: 6}, {Token: 7}, {Token: 8}, {Token: 9}, {Token: 10}},
-			numKeep:       2,
-			cacheErr:      false, // No error
-			wantErr:       nil,
-			wantInputsLen: 6, // After discarding 4 tokens
-		},
-		{
-			name:          "Cache removal fails",
-			numCtx:        10,
-			inputs:        []input.Input{{Token: 1}, {Token: 2}, {Token: 3}, {Token: 4}, {Token: 5}, {Token: 6}, {Token: 7}, {Token: 8}, {Token: 9}, {Token: 10}},
-			numKeep:       2,
-			cacheErr:      true,
-			wantErr:       &ErrReprocessInputs{},
-			wantInputsLen: 0, // Original inputs should be cleared
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			mock := &mockCache{shouldFail: tt.cacheErr}
-			c := InputCache{
-				numCtx: tt.numCtx,
-				cache:  mock,
-			}
-			slot := &InputCacheSlot{
-				Id:     123,
-				Inputs: make([]input.Input, len(tt.inputs)),
-			}
-			copy(slot.Inputs, tt.inputs)
-
-			err := c.ShiftCacheSlot(slot, tt.numKeep)
-
-			if tt.wantErr != nil {
-				if err == nil {
-					t.Errorf("Expected error but got nil")
-					return
-				}
-
-				if !errors.As(err, &tt.wantErr) {
-					t.Errorf("Expected error of type %T but got %T: %v", tt.wantErr, err, err)
-				}
-			} else if err != nil {
-				t.Errorf("Unexpected error: %v", err)
-			}
-
-			if len(slot.Inputs) != tt.wantInputsLen {
-				t.Errorf("Slot inputs length after operation: got %v, want %v", len(slot.Inputs), tt.wantInputsLen)
-			}
-		})
-	}
-}

runner/ollamarunner/runner.go

@@ -24,7 +24,6 @@ import (
 	"golang.org/x/sync/semaphore"
 
 	"github.com/ollama/ollama/api"
-	"github.com/ollama/ollama/llm"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/model"
 	"github.com/ollama/ollama/model/input"
@@ -34,14 +33,10 @@ import (
 	_ "github.com/ollama/ollama/model/models"
 )
 
-type contextList struct {
-	list []ml.Context
-}
-
 type Sequence struct {
-	// ctxs are used for allocating tensors that last the lifetime of the sequence, such as
+	// ctx for allocating tensors that last the lifetime of the sequence, such as
 	// multimodal embeddings
-	ctxs *contextList
+	ctx ml.Context
 
 	// batch index
 	iBatch int
@@ -99,12 +94,13 @@ type NewSequenceParams struct {
 	embedding  bool
 }
 
-func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSequenceParams) (*Sequence, error) {
+func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequenceParams) (*Sequence, error) {
 	s.ready.Wait()
 
 	startTime := time.Now()
+	ctx := s.model.Backend().NewContext()
 
-	inputs, ctxs, err := s.inputs(prompt, images)
+	inputs, err := s.inputs(ctx, prompt, images)
 	if err != nil {
 		return nil, fmt.Errorf("failed to process inputs: %w", err)
 	} else if len(inputs) == 0 {
@@ -120,36 +116,8 @@ func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSe
 
 	if int32(len(inputs)) > s.cache.numCtx {
 		discard := int32(len(inputs)) - s.cache.numCtx
-		promptStart := params.numKeep + discard
-
-		// If we need to truncate in the middle of a unbreakable batch, remove the entire batch
-		sameBatch := 0
-		for i, inp := range inputs {
-			if sameBatch > 0 {
-				sameBatch--
-
-				if promptStart == int32(i) {
-					promptStart++
-				}
-			} else if promptStart == int32(i) {
-				break
-			}
-
-			if inp.SameBatch != 0 {
-				if int32(i) < params.numKeep {
-					return nil, fmt.Errorf("SameBatch may not be specified within numKeep (index: %v numKeep: %v SameBatch: %v)", i, params.numKeep, inp.SameBatch)
-				}
-
-				sameBatch = inp.SameBatch
-			}
-		}
-
-		if promptStart >= int32(len(inputs)) {
-			return nil, errors.New("entire prompt removed by truncation")
-		}
-
 		newInputs := inputs[:params.numKeep]
-		newInputs = append(newInputs, inputs[promptStart:]...)
+		newInputs = append(newInputs, inputs[params.numKeep+discard:]...)
 
 		slog.Warn("truncating input prompt", "limit", s.cache.numCtx, "prompt", len(inputs), "keep", params.numKeep, "new", len(newInputs))
 		inputs = newInputs
@@ -158,7 +126,7 @@ func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSe
 	// TODO(jessegross): Ingest cached history for grammar
 
 	return &Sequence{
-		ctxs:                ctxs,
+		ctx:                 ctx,
 		inputs:              inputs,
 		numPromptInputs:     len(inputs),
 		startProcessingTime: startTime,
@@ -177,7 +145,7 @@ func (s *Server) NewSequence(prompt string, images []llm.ImageData, params NewSe
 // inputs processes the prompt and images into a list of inputs
 // by splitting the prompt on [img-<n>] tags, tokenizing text and
 // decoding images
-func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input.Input, *contextList, error) {
+func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]input.Input, error) {
 	var inputs []input.Input
 	var parts []string
 	var matches [][]string
@@ -192,19 +160,12 @@ func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input.Input, *
 		parts = []string{prompt}
 	}
 
-	var contexts contextList
-	runtime.AddCleanup(&contexts, func(ctxs []ml.Context) {
-		for _, ctx := range ctxs {
-			ctx.Close()
-		}
-	}, contexts.list)
-
 	postTokenize := false
 	for i, part := range parts {
 		// text - tokenize
 		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 
 		for _, t := range tokens {
@@ -224,14 +185,12 @@ func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input.Input, *
 			}
 
 			if imageIndex < 0 {
-				return nil, nil, fmt.Errorf("invalid image index: %d", n)
+				return nil, fmt.Errorf("invalid image index: %d", n)
 			}
 
-			ctx := s.model.Backend().NewContext()
-			contexts.list = append(contexts.list, ctx)
 			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 
 			s.multimodalHash.Reset()
@@ -245,13 +204,13 @@ func (s *Server) inputs(prompt string, images []llm.ImageData) ([]input.Input, *
 
 	if visionModel && postTokenize {
 		var err error
-		inputs, err = multimodalProcessor.PostTokenize(inputs)
+		inputs, err = multimodalProcessor.PostTokenize(ctx, inputs)
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 	}
 
-	return inputs, &contexts, nil
+	return inputs, nil
 }
 
 type Server struct {
@@ -263,7 +222,7 @@ type Server struct {
 	model model.Model
 
 	// status for external health reporting - loading, ready to serve, etc.
-	status llm.ServerStatus
+	status ServerStatus
 
 	// current progress on loading the model
 	progress float32
@@ -292,9 +251,6 @@ type Server struct {
 	// KV cache
 	cache *InputCache
 
-	// next sequence for prompt processing to avoid starvation
-	nextSeq int
-
 	// multimodalHash generates hashes for comparing equality
 	// of non-text data
 	multimodalHash maphash.Hash
@@ -349,6 +305,7 @@ func (s *Server) removeSequence(seqIndex int, reason string) {
 	close(seq.responses)
 	close(seq.embedding)
 	seq.cache.InUse = false
+	seq.ctx.Close()
 	s.seqs[seqIndex] = nil
 	s.seqsSem.Release(1)
 }
@@ -376,22 +333,16 @@ func (s *Server) processBatch() error {
 	}
 	defer s.mu.Unlock()
 
-	var batchInputs []int32
-	var batch input.Batch
-
-	resumeSeq := -1
-	seqIdx := s.nextSeq - 1
-	for range s.seqs {
-		seqIdx = (seqIdx + 1) % len(s.seqs)
-		seq := s.seqs[seqIdx]
+	var options input.Options
 
+	for i, seq := range s.seqs {
 		if seq == nil {
 			continue
 		}
 
 		// if past the num predict limit
 		if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict {
-			s.removeSequence(seqIdx, "limit")
+			s.removeSequence(i, "limit")
 			continue
 		}
 
@@ -400,61 +351,33 @@ func (s *Server) processBatch() error {
 			seq.cache.Inputs = []input.Input{}
 		}
 
-		batchSize := s.batchSize
-
-		for i, inp := range seq.inputs {
-			// If we are required to put following inputs into a single batch then extend the
-			// batch size. Since we are only extending the size the minimum amount possible, this
-			// will cause a break if we have existing inputs.
-			minBatch := 1 + inp.SameBatch
-			if minBatch > batchSize {
-				batchSize = minBatch
+		for j, inp := range seq.inputs {
+			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx {
+				if len(seq.pendingInputs) == 0 {
+					err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
+					if err != nil {
+						return err
+					}
+				} else {
+					break
+				}
 			}
 
-			// Stop if the required batch would put us over the total batch size (including tokens
-			// added by other sequences). If we haven't been able to add anything yet then pick up
-			// here again for the next batch to avoid starvation, though we can opportunistically
-			// check if other sequences can still squeeze something in.
-			if len(batchInputs)+minBatch > batchSize {
-				if len(seq.pendingInputs) == 0 && resumeSeq == -1 {
-					resumeSeq = seqIdx
-				}
+			if j >= s.batchSize {
 				break
 			}
 
-			// If the sum of our working set (already processed tokens, tokens we added to this
-			// batch, required following tokens) exceeds the context size, then trigger a shift
-			// now so we don't have to do one later when we can't break the batch.
-			if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+minBatch) > s.cache.numCtx {
-				if len(seq.pendingInputs) != 0 {
-					break
-				}
-
-				err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep)
-				if err != nil {
-					var reprocess *ErrReprocessInputs
-					if errors.As(err, &reprocess) {
-						// Prepend these inputs to the sequence's inputs queue for reprocessing
-						seq.inputs = append(reprocess.Inputs, seq.inputs...)
-						// Skip this sequence but continue processing the rest
-						continue
-					} else {
-						return err
-					}
-				}
-			}
-
-			batchInputs = append(batchInputs, inp.Token)
+			options.Inputs = append(options.Inputs, inp.Token)
 			if inp.Multimodal != nil {
-				batch.Multimodal = append(batch.Multimodal, input.MultimodalIndex{Index: len(batchInputs) - 1, Multimodal: inp.Multimodal})
+				options.Multimodal = append(options.Multimodal, input.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: inp.Multimodal})
 			}
 
-			batch.Positions = append(batch.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
-			batch.Sequences = append(batch.Sequences, seq.cache.Id)
+			options.Positions = append(options.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
+			options.Sequences = append(options.Sequences, seq.cache.Id)
 
-			seq.iBatch = len(batch.Outputs)
-			if i+1 == len(seq.inputs) {
-				batch.Outputs = append(batch.Outputs, int32(len(batchInputs)-1))
+			seq.iBatch = len(options.Outputs)
+			if j+1 == len(seq.inputs) {
+				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
 			}
 			seq.pendingInputs = append(seq.pendingInputs, inp)
 		}
@@ -462,20 +385,14 @@ func (s *Server) processBatch() error {
 		seq.inputs = seq.inputs[len(seq.pendingInputs):]
 	}
 
-	if resumeSeq != -1 {
-		s.nextSeq = resumeSeq
-	} else {
-		s.nextSeq = seqIdx + 1
-	}
-
-	if len(batchInputs) == 0 {
+	if len(options.Inputs) == 0 {
 		return nil
 	}
 
 	ctx := s.model.Backend().NewContext()
 	defer ctx.Close()
 
-	modelOutput, err := model.Forward(ctx, s.model, batchInputs, batch)
+	modelOutput, err := model.Forward(ctx, s.model, options)
 	if err != nil {
 		return fmt.Errorf("failed to decode batch: %w", err)
 	}
@@ -515,7 +432,7 @@ func (s *Server) processBatch() error {
 		}
 
 		// sample a token
-		vocabSize := len(logits) / len(batch.Outputs)
+		vocabSize := len(logits) / len(options.Outputs)
 
 		token, err := seq.sampler.Sample(logits[seq.iBatch*vocabSize : (seq.iBatch+1)*vocabSize])
 		if err != nil {
@@ -584,18 +501,75 @@ func (s *Server) processBatch() error {
 	return nil
 }
 
+// TODO (jmorganca): use structs from the api package to avoid duplication
+// this way the api acts as a proxy instead of using a different api for the
+// runner
+type Options struct {
+	api.Runner
+
+	NumKeep          int      `json:"n_keep"`
+	Seed             int      `json:"seed"`
+	NumPredict       int      `json:"n_predict"`
+	TopK             int      `json:"top_k"`
+	TopP             float32  `json:"top_p"`
+	MinP             float32  `json:"min_p"`
+	TypicalP         float32  `json:"typical_p"`
+	RepeatLastN      int      `json:"repeat_last_n"`
+	Temperature      float32  `json:"temperature"`
+	RepeatPenalty    float32  `json:"repeat_penalty"`
+	PresencePenalty  float32  `json:"presence_penalty"`
+	FrequencyPenalty float32  `json:"frequency_penalty"`
+	Mirostat         int      `json:"mirostat"`
+	MirostatTau      float32  `json:"mirostat_tau"`
+	MirostatEta      float32  `json:"mirostat_eta"`
+	Stop             []string `json:"stop"`
+}
+
+type ImageData struct {
+	Data          []byte `json:"data"`
+	ID            int    `json:"id"`
+	AspectRatioID int    `json:"aspect_ratio_id"`
+}
+
+type CompletionRequest struct {
+	Prompt      string      `json:"prompt"`
+	Images      []ImageData `json:"image_data"`
+	Grammar     string      `json:"grammar"`
+	CachePrompt bool        `json:"cache_prompt"`
+
+	Options
+}
+
+type Timings struct {
+	PredictedN  int     `json:"predicted_n"`
+	PredictedMS float64 `json:"predicted_ms"`
+	PromptN     int     `json:"prompt_n"`
+	PromptMS    float64 `json:"prompt_ms"`
+}
+
+type CompletionResponse struct {
+	Content string `json:"content"`
+	Stop    bool   `json:"stop"`
+
+	Model        string  `json:"model,omitempty"`
+	Prompt       string  `json:"prompt,omitempty"`
+	StoppedLimit bool    `json:"stopped_limit,omitempty"`
+	PredictedN   int     `json:"predicted_n,omitempty"`
+	PredictedMS  float64 `json:"predicted_ms,omitempty"`
+	PromptN      int     `json:"prompt_n,omitempty"`
+	PromptMS     float64 `json:"prompt_ms,omitempty"`
+
+	Timings Timings `json:"timings"`
+}
+
 func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
-	var req llm.CompletionRequest
+	var req CompletionRequest
+	req.Options = Options(api.DefaultOptions())
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 		http.Error(w, "Bad request", http.StatusBadRequest)
 		return
 	}
 
-	if req.Options == nil {
-		opts := api.DefaultOptions()
-		req.Options = &opts
-	}
-
 	// Set the headers to indicate streaming
 	w.Header().Set("Content-Type", "application/json")
 	w.Header().Set("Transfer-Encoding", "chunked")
@@ -617,18 +591,18 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	}
 
 	sampler := sample.NewSampler(
-		req.Options.Temperature,
-		req.Options.TopK,
-		req.Options.TopP,
-		req.Options.MinP,
-		req.Options.Seed,
+		req.Temperature,
+		req.TopK,
+		req.TopP,
+		req.MinP,
+		req.Seed,
 		grammar,
 	)
 
 	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
-		numPredict: req.Options.NumPredict,
-		stop:       req.Options.Stop,
-		numKeep:    int32(req.Options.NumKeep),
+		numPredict: req.NumPredict,
+		stop:       req.Stop,
+		numKeep:    int32(req.NumKeep),
 		sampler:    sampler,
 		embedding:  false,
 	})
@@ -642,7 +616,7 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 		if errors.Is(err, context.Canceled) {
 			slog.Info("aborting completion request due to client closing the connection")
 		} else {
-			http.Error(w, fmt.Sprintf("Failed to acquire semaphore: %v", err), http.StatusInternalServerError)
+			slog.Error("Failed to acquire semaphore", "error", err)
 		}
 		return
 	}
@@ -651,10 +625,9 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	found := false
 	for i, sq := range s.seqs {
 		if sq == nil {
-			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs)
+			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
 			if err != nil {
 				s.mu.Unlock()
-				s.seqsSem.Release(1)
 				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
 				return
 			}
@@ -668,7 +641,6 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	s.mu.Unlock()
 
 	if !found {
-		s.seqsSem.Release(1)
 		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
 		return
 	}
@@ -680,7 +652,7 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 			return
 		case content, ok := <-seq.responses:
 			if ok {
-				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
+				if err := json.NewEncoder(w).Encode(&CompletionResponse{
 					Content: content,
 				}); err != nil {
 					http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
@@ -691,17 +663,15 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 				flusher.Flush()
 			} else {
 				// Send the final response
-				doneReason := "stop"
-				if seq.doneReason == "limit" {
-					doneReason = "length"
-				}
-				if err := json.NewEncoder(w).Encode(&llm.CompletionResponse{
-					Done:               true,
-					DoneReason:         doneReason,
-					PromptEvalCount:    seq.numPromptInputs,
-					PromptEvalDuration: seq.startGenerationTime.Sub(seq.startProcessingTime),
-					EvalCount:          seq.numPredicted,
-					EvalDuration:       time.Since(seq.startGenerationTime),
+				if err := json.NewEncoder(w).Encode(&CompletionResponse{
+					Stop:         true,
+					StoppedLimit: seq.doneReason == "limit",
+					Timings: Timings{
+						PromptN:     seq.numPromptInputs,
+						PromptMS:    float64(seq.startGenerationTime.Sub(seq.startProcessingTime).Milliseconds()),
+						PredictedN:  seq.numPredicted,
+						PredictedMS: float64(time.Since(seq.startGenerationTime).Milliseconds()),
+					},
 				}); err != nil {
 					http.Error(w, fmt.Sprintf("failed to encode final response: %v", err), http.StatusInternalServerError)
 				}
@@ -712,10 +682,102 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+type EmbeddingRequest struct {
+	Content     string `json:"content"`
+	CachePrompt bool   `json:"cache_prompt"`
+}
+
+type EmbeddingResponse struct {
+	Embedding []float32 `json:"embedding"`
+}
+
+func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
+	var req EmbeddingRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		http.Error(w, fmt.Sprintf("bad request: %s", err), http.StatusBadRequest)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+
+	slog.Debug("embedding request", "content", req.Content)
+
+	seq, err := s.NewSequence(req.Content, nil, NewSequenceParams{embedding: true})
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Failed to create new sequence: %v", err), http.StatusInternalServerError)
+		return
+	}
+
+	// Ensure there is a place to put the sequence, released when removed from s.seqs
+	if err := s.seqsSem.Acquire(r.Context(), 1); err != nil {
+		if errors.Is(err, context.Canceled) {
+			slog.Info("aborting embeddings request due to client closing the connection")
+		} else {
+			slog.Error("Failed to acquire semaphore", "error", err)
+		}
+		return
+	}
+
+	s.mu.Lock()
+	found := false
+	for i, sq := range s.seqs {
+		if sq == nil {
+			seq.cache, seq.inputs, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
+			if err != nil {
+				s.mu.Unlock()
+				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
+				return
+			}
+			s.seqs[i] = seq
+			s.cond.Signal()
+			found = true
+			break
+		}
+	}
+	s.mu.Unlock()
+
+	if !found {
+		http.Error(w, "could not find an available sequence", http.StatusInternalServerError)
+		return
+	}
+
+	embedding := <-seq.embedding
+
+	if err := json.NewEncoder(w).Encode(&EmbeddingResponse{
+		Embedding: embedding,
+	}); err != nil {
+		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
+	}
+}
+
+type HealthResponse struct {
+	Status   string  `json:"status"`
+	Progress float32 `json:"progress"`
+}
+
+type ServerStatus int
+
+const (
+	ServerStatusReady ServerStatus = iota
+	ServerStatusLoadingModel
+	ServerStatusError
+)
+
+func (s ServerStatus) ToString() string {
+	switch s {
+	case ServerStatusReady:
+		return "ok"
+	case ServerStatusLoadingModel:
+		return "loading model"
+	default:
+		return "server error"
+	}
+}
+
 func (s *Server) health(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
-	if err := json.NewEncoder(w).Encode(&llm.ServerStatusResponse{
-		Status:   s.status,
+	if err := json.NewEncoder(w).Encode(&HealthResponse{
+		Status:   s.status.ToString(),
 		Progress: s.progress,
 	}); err != nil {
 		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
@@ -734,7 +796,6 @@ func (m *multiLPath) String() string {
 }
 
 func (s *Server) loadModel(
-	ctx context.Context,
 	mpath string,
 	params ml.BackendParams,
 	lpath multiLPath,
@@ -744,7 +805,7 @@ func (s *Server) loadModel(
 	multiUserCache bool,
 ) {
 	var err error
-	s.model, err = model.New(ctx, mpath, params)
+	s.model, err = model.New(mpath, params)
 	if err != nil {
 		panic(err)
 	}
@@ -756,7 +817,7 @@ func (s *Server) loadModel(
 		panic("loras are not yet implemented")
 	}
 
-	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, s.batchSize, multiUserCache)
+	s.cache, err = NewInputCache(s.model, kvCacheType, int32(kvSize), parallel, multiUserCache)
 	if err != nil {
 		panic(err)
 	}
@@ -770,7 +831,7 @@ func (s *Server) loadModel(
 	s.seqs = make([]*Sequence, s.parallel)
 	s.seqsSem = semaphore.NewWeighted(int64(s.parallel))
 
-	s.status = llm.ServerStatusReady
+	s.status = ServerStatusReady
 	s.ready.Done()
 }
 
@@ -822,7 +883,7 @@ func Execute(args []string) error {
 
 	server := &Server{
 		batchSize: *batchSize,
-		status:    llm.ServerStatusLoadingModel,
+		status:    ServerStatusLoadingModel,
 	}
 
 	// TODO(jessegross): Parameters that need to be implemented:
@@ -840,9 +901,6 @@ func Execute(args []string) error {
 	}
 
 	params := ml.BackendParams{
-		Progress: func(progress float32) {
-			server.progress = progress
-		},
 		NumThreads:     *threads,
 		NumGPULayers:   *numGPULayers,
 		MainGPU:        *mainGPU,
@@ -851,13 +909,13 @@ func Execute(args []string) error {
 	}
 
 	server.ready.Add(1)
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	go server.loadModel(ctx, *mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
+	go server.loadModel(*mpath, params, lpaths, *parallel, *kvCacheType, *kvSize, *multiUserCache)
 
 	server.cond = sync.NewCond(&server.mu)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	go server.run(ctx)
 
 	addr := "127.0.0.1:" + strconv.Itoa(*port)
@@ -869,13 +927,9 @@ func Execute(args []string) error {
 	defer listener.Close()
 
 	mux := http.NewServeMux()
-	// TODO: support embeddings
-	mux.HandleFunc("POST /embedding", func(w http.ResponseWriter, r *http.Request) {
-		http.Error(w, "this model does not support embeddings", http.StatusNotImplemented)
-	})
-
-	mux.HandleFunc("POST /completion", server.completion)
-	mux.HandleFunc("GET /health", server.health)
+	mux.HandleFunc("/embedding", server.embeddings)
+	mux.HandleFunc("/completion", server.completion)
+	mux.HandleFunc("/health", server.health)
 
 	httpServer := http.Server{
 		Handler: mux,

sample/samplers.go

@@ -1,11 +1,10 @@
 package sample
 
 import (
-	"errors"
 	"math"
-	"math/rand/v2"
-	"slices"
+	"math/rand"
 	"sync"
+	"time"
 
 	"github.com/ollama/ollama/llama"
 )
@@ -26,10 +25,6 @@ type Sampler struct {
 }
 
 func (s *Sampler) Sample(logits []float32) (int32, error) {
-	if len(logits) == 0 {
-		return -1, errors.New("sample: no logits provided to sample")
-	}
-
 	tokens := make([]token, len(logits))
 	for i := range logits {
 		tokens[i].id = int32(i)
@@ -91,50 +86,53 @@ func (s *Sampler) sample(tokens []token) (token, error) {
 	// topK also sorts the tokens in descending order of logits
 	tokens = topK(tokens, s.topK)
 
-	// scale and normalize the tokens in place
-	temperature(tokens, s.temperature)
-	softmax(tokens)
-
 	tokens = topP(tokens, s.topP)
 	tokens = minP(tokens, s.minP)
 
-	var r float32
-	if s.rng != nil {
-		r = s.rng.Float32()
-	} else {
-		r = rand.Float32()
-	}
+	// token logit values are updated to probabilities
+	temperature(tokens, s.temperature)
+	softmax(tokens)
+	return tokens[dist(tokens, s.rng.Int63())], nil
 
-	// Calculate cumulative sum of probabilities
-	var sum float32
-	for i := range tokens {
-		sum += tokens[i].value
-		tokens[i].value = sum
-	}
-	r *= tokens[len(tokens)-1].value
+	// // TODO: this should fall back to greedy sampling
+	// // or topP, topK values etc should be such that
+	// // there are always tokens to sample from
+	// if len(tokens) == 0 {
+	// 	return token{}, errors.New("no tokens to sample from")
+	// }
 
-	idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
-		if token.value < target {
-			return -1
-		}
-		return 1
-	})
+	// var r float32
+	// if s.rng != nil {
+	// 	r = s.rng.Float32()
+	// } else {
+	// 	r = rand.Float32()
+	// }
 
-	if math.IsNaN(float64(sum)) {
-		return token{}, errors.New("sample: logits sum to NaN, check model output")
-	}
-	return tokens[idx], nil
+	// // Calculate cumulative sum of probabilities
+	// var sum float32
+	// for i := range tokens {
+	// 	sum += tokens[i].value
+	// 	tokens[i].value = sum
+	// }
+	// r *= tokens[len(tokens)-1].value
+
+	// idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int {
+	// 	if token.value < target {
+	// 		return -1
+	// 	}
+	// 	return 1
+	// })
+
+	// return tokens[idx], nil
 }
 
 // TODO(parthsareen): update sampler interface to use json unmarshal https://github.com/ollama/ollama/issues/9278
 func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int, grammar *Grammar) Sampler {
 	var rng *rand.Rand
 	if seed != -1 {
-		// PCG requires two parameters: sequence and stream
-		// Use original seed for sequence
-		sequence := uint64(seed)
-		// Use golden ratio hash to generate statistically independent seeds
-		rng = rand.New(rand.NewPCG(sequence, sequence^0x9E3779B9))
+		rng = rand.New(rand.NewSource(int64(seed)))
+	} else {
+		rng = rand.New(rand.NewSource(time.Now().UnixNano()))
 	}
 	if temperature < 0.0 {
 		temperature = 0.0

sample/samplers_test.go

@@ -1,7 +1,6 @@
 package sample
 
 import (
-	"math"
 	"math/rand/v2"
 	"testing"
 )
@@ -30,29 +29,6 @@ func TestWeighted(t *testing.T) {
 	if want != got {
 		t.Errorf("index mismatch: want %d, got %d", want, got)
 	}
-
-	// Test very high p
-	logits = []float32{1.0, 0.9999999999999999, 0.5, 0.1}
-	// Use extremely small topP to filter out all tokens
-	sampler = NewSampler(1.0, 0, 1e-10, 0, 0, nil)
-	got, err = sampler.Sample(logits)
-	if err != nil {
-		t.Error(err)
-		return
-	}
-	// Should get the token with the highest logit
-	want = int32(0)
-	if want != got {
-		t.Errorf("index mismatch: want %d, got %d", want, got)
-	}
-
-	logits = []float32{float32(math.NaN()), float32(math.NaN()), float32(math.NaN())}
-	sampler = NewSampler(1, 0, 0.95, 0.05, 0, nil)
-	got, err = sampler.Sample(logits)
-	if err == nil {
-		t.Errorf("expected error, got %d", got)
-		return
-	}
 }
 
 func BenchmarkSample(b *testing.B) {

sample/transforms.go

@@ -3,6 +3,7 @@ package sample
 import (
 	"container/heap"
 	"math"
+	"math/rand"
 	"slices"
 )
 
@@ -25,38 +26,6 @@ func (h *tokenHeap) Pop() any {
 	return x
 }
 
-// temperature applies scaling to the logits
-func temperature(ts []token, temp float32) {
-	// Ensure temperature clipping near 0 to avoid numerical instability
-	temp = max(temp, 1e-7)
-	for i := range ts {
-		ts[i].value = ts[i].value / temp
-	}
-}
-
-// softmax applies normalization to the logits
-func softmax(ts []token) {
-	// Find max logit for numerical stability
-	maxLogit := float32(math.Inf(-1))
-	for _, t := range ts {
-		if t.value > maxLogit {
-			maxLogit = t.value
-		}
-	}
-
-	// Compute exp(x - max)
-	var sum float32
-	for i, v := range ts {
-		ts[i].value = float32(math.Exp(float64(v.value - maxLogit)))
-		sum += ts[i].value
-	}
-
-	// exp(x - max) / sum(exp(x - max))
-	for i := range ts {
-		ts[i].value /= sum
-	}
-}
-
 // topK limits the number of tokens considered to the k highest logits
 func topK(ts []token, k int) []token {
 	if k >= len(ts) || k <= 0 {
@@ -96,7 +65,6 @@ func topK(ts []token, k int) []token {
 }
 
 // topP limits tokens to those with cumulative probability p
-// requires ts to be sorted in descending order of probabilities
 func topP(ts []token, p float32) []token {
 	if p == 1.0 {
 		return ts
@@ -107,24 +75,93 @@ func topP(ts []token, p float32) []token {
 	for i, t := range ts {
 		sum += t.value
 		if sum > float32(p) {
-			return ts[:i+1]
+			ts = ts[:i+1]
+			return ts
 		}
 	}
 
 	return ts
 }
 
-// minP filters tokens with probabilities >= p * max_prob
-// requires ts to be sorted in descending order of probabilities
+// minP limits tokens to those with cumulative probability p
 func minP(ts []token, p float32) []token {
-	maxProb := ts[0].value
+	if p == 1.0 {
+		return ts
+	}
 
-	threshold := maxProb * p
-
-	for i, t := range ts {
-		if t.value < threshold {
-			return ts[:i]
+	maxProb := float32(math.Inf(-1))
+	for _, token := range ts {
+		if token.value > maxProb {
+			maxProb = token.value
 		}
 	}
+
+	threshold := maxProb * float32(p)
+
+	// Filter tokens in-place
+	validTokens := ts[:0]
+	for i, token := range ts {
+		if token.value >= threshold {
+			validTokens = append(validTokens, ts[i])
+		}
+	}
+
+	ts = validTokens
 	return ts
 }
+
+func temperature(ts []token, temp float32) {
+	for i := range ts {
+		ts[i].value /= temp
+	}
+}
+
+func softmax(ts []token) {
+	if len(ts) == 0 {
+		return
+	}
+
+	// Find max logit for numerical stability
+	maxLogit := ts[0].value
+	for _, t := range ts {
+		if t.value > maxLogit {
+			maxLogit = t.value
+		}
+	}
+
+	// Compute exp(logit - maxLogit) and sum them
+	var sumExp float32
+	for i, t := range ts {
+		expVal := float32(math.Exp(float64(t.value - maxLogit)))
+		ts[i].value = expVal
+		sumExp += expVal
+	}
+
+	// Normalize probabilities
+	for i := range ts {
+		ts[i].value /= sumExp
+	}
+}
+
+// applyDist selects a token based on probabilities and seed
+func dist(ts []token, seed int64) int {
+	rng := rand.New(rand.NewSource(seed))
+
+	cdf := make([]float32, len(ts))
+	var cumSum float32
+	for i, t := range ts {
+		cumSum += t.value
+		cdf[i] = cumSum
+	}
+
+	r := rng.Float32() * cumSum
+
+	// Select token based on CDF
+	for i, probSum := range cdf {
+		if r < probSum {
+			return i
+		}
+	}
+
+	return len(ts) - 1
+}

sample/transforms_test.go

@@ -1,8 +1,13 @@
 package sample
 
 import (
+	"encoding/binary"
+	"errors"
 	"math"
 	"math/rand/v2"
+	"os"
+	"path/filepath"
+	"runtime"
 	"testing"
 )
 
@@ -32,132 +37,64 @@ func compareLogits(t *testing.T, name string, want []float32, got []token) {
 	}
 }
 
-func TestTemperature(t *testing.T) {
-	input := []float32{1.0, 4.0, -2.0, 0.0}
-	tokens := toTokens(input)
-	temperature(tokens, 0.5)
-	want := []float32{2.0, 8.0, -4.0, 0.0}
-	compareLogits(t, "temperature(0.5)", want, tokens)
+func TestTemperatureAndSoftmax(t *testing.T) {
+	input := []float32{1, 4, -2, 0}
+	got := temperature(toTokens(input), 0.5)
 
-	input = []float32{1.0, 4.0, -2.0, 0.0}
-	tokens = toTokens(input)
-	temperature(tokens, 1.0)
-	want = []float32{1.0, 4.0, -2.0, 0.0}
-	compareLogits(t, "temperature(1)", want, tokens)
-
-	input = []float32{1.0, 4.0, -2.0, 0.0}
-	tokens = toTokens(input)
-	temperature(tokens, 0.0)
-	want = []float32{1e7, 4e7, -2e7, 0.0}
-	compareLogits(t, "temperature(0)", want, tokens)
-}
-
-func TestSoftmax(t *testing.T) {
-	tests := []struct {
-		name     string
-		input    []float32
-		expected []float32
-	}{
-		{
-			name:     "correctness softmax",
-			input:    []float32{1, -2, 3, 0},
-			expected: []float32{0.113550, 0.005653, 0.839024, 0.041773},
-		},
-		{
-			name:  "normal distribution",
-			input: []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367},
-		},
-		{
-			name:  "single value",
-			input: []float32{1.0},
-		},
-		{
-			name:  "identical values",
-			input: []float32{0.9, 0.9, 0.9},
-		},
-		{
-			name:  "large values",
-			input: []float32{1000.0, 2000.0, 3000.0},
-		},
-		{
-			name:  "small values",
-			input: []float32{1e-6, 2e-6, 3e-6},
-		},
-		{
-			name:  "negative values",
-			input: []float32{-1.0, -2.0, -3.0},
-		},
-		{
-			name:  "mixed values",
-			input: []float32{-100.0, 0.0, 100.0},
-		},
+	// Check probabilities sum to 1
+	var sum float32
+	for _, token := range got {
+		sum += token.value
+	}
+	if math.Abs(float64(sum-1.0)) > 1e-6 {
+		t.Errorf("probabilities don't sum to 1: got %f", sum)
 	}
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			tokens := toTokens(tt.input)
-			softmax(tokens)
-
-			if tt.expected != nil {
-				compareLogits(t, tt.name, tt.expected, tokens)
-				return
-			}
-
-			// Check probabilities sum to 1
-			var sum float32
-			for _, token := range tokens {
-				sum += token.value
-				if token.value < 0 || token.value > 1 {
-					t.Errorf("probability out of range [0,1]: got %f", token.value)
-				}
-			}
-			if math.Abs(float64(sum-1.0)) > 1e-6 {
-				t.Errorf("probabilities don't sum to 1: got %f", sum)
-			}
-		})
+	got = temperature(toTokens(input), 1)
+	// Check probabilities sum to 1
+	sum = 0.0
+	for _, token := range got {
+		sum += token.value
+	}
+	if math.Abs(float64(sum-1.0)) > 1e-6 {
+		t.Errorf("probabilities don't sum to 1: got %f", sum)
 	}
 }
 
 func TestTopK(t *testing.T) {
 	input := []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
-	tokens := toTokens(input)
-	tokens = topK(tokens, 5)
-	if len(tokens) != 5 {
-		t.Errorf("topK(5): wrong length: want 5, got %d", len(tokens))
+
+	// Test k=5
+	got := topK(toTokens(input), 5)
+	if len(got) != 5 {
+		t.Errorf("topK(5): wrong length: want 5, got %d", len(got))
 	}
+	// Should keep highest 3 values in descending order
 	want := []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154}
-	compareLogits(t, "topK(3)", want, tokens)
+	compareLogits(t, "topK(3)", want, got)
 
-	tokens = toTokens(input)
-	tokens = topK(tokens, 20)
-	if len(tokens) != len(input) {
-		t.Errorf("topK(20): wrong length: want %d, got %d", len(input), len(tokens))
+	got = topK(toTokens(input), 20)
+	if len(got) != len(input) {
+		t.Errorf("topK(20): wrong length: want %d, got %d", len(input), len(got))
 	}
 
+	// Test k=-1
 	input = []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
 	want = []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154, 0.043722924, 0.036774673, 0.026986899, 0.01681367, 0.0046718004, 0.00412893, 0.0030491839}
-	tokens = toTokens(input)
-	tokens = topK(tokens, -1)
-	if len(tokens) != len(input) {
-		t.Errorf("topK(-1): wrong length: want %d, got %d", len(input), len(tokens))
+	got = topK(toTokens(input), -1)
+	if len(got) != len(input) {
+		t.Errorf("topK(-1): wrong length: want %d, got %d", len(input), len(got))
 	}
-	compareLogits(t, "topK(-1)", want, tokens)
+	compareLogits(t, "topK(-1)", want, got)
 
+	// Test k=0
 	input = []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
 	want = []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154, 0.043722924, 0.036774673, 0.026986899, 0.01681367, 0.0046718004, 0.00412893, 0.0030491839}
-	tokens = toTokens(input)
-	tokens = topK(tokens, 0)
-	if len(tokens) != len(input) {
-		t.Errorf("topK(-1): wrong length: want %d, got %d", len(input), len(tokens))
-	}
-	compareLogits(t, "topK(-1)", want, tokens)
-
-	input = []float32{-1e7, -2e7, -3e7, -4e7}
-	tokens = toTokens(input)
-	tokens = topK(tokens, 1)
-	if len(tokens) < 1 {
-		t.Error("topK should keep at least one token")
+	got = topK(toTokens(input), 0)
+	if len(got) != len(input) {
+		t.Errorf("topK(-1): wrong length: want %d, got %d", len(input), len(got))
 	}
+	compareLogits(t, "topK(-1)", want, got)
 }
 
 func TestTopP(t *testing.T) {
@@ -165,134 +102,142 @@ func TestTopP(t *testing.T) {
 	tokens := toTokens(input)
 
 	// First apply temperature and softmax to get probabilities
-	softmax(tokens)
+	tokens = temperature(tokens, 1)
 	tokens = topK(tokens, 20)
 
-	// Test with very high p value
-	got := topP(tokens, 1.0)
-
-	// Should keep all tokens since p is 1
-	if len(got) != len(input) {
-		t.Errorf("topP(1.0): should keep all tokens, got %d, want %d", len(got), len(input))
-	}
-
-	// Test with normal p value
-	got = topP(tokens, 0.95)
+	// Then apply topP
+	got := topP(tokens, 0.95)
 
+	// Should keep tokens until cumsum > 0.95
 	if len(got) > 3 {
-		t.Errorf("topP(0.95): kept too many tokens: got %d", len(tokens))
-		t.Logf("got: %v", got)
-	}
-
-	// Test edge case - ensure at least one token remains
-	input = []float32{-1e6, -1e6, -1e7}
-	tokens = toTokens(input)
-	tokens = topK(tokens, 20)
-	softmax(tokens)
-	got = topP(tokens, 0.0)
-	if len(got) < 1 {
-		t.Error("topP should keep at least one token")
-	}
-
-	// Test with zero p value
-	got = topP(tokens, 0.0)
-
-	// Should keep only the highest probability token
-	if len(got) != 1 {
-		t.Errorf("topP(0.0): should keep only one token, got %d", len(got))
-		t.Logf("got: %v", got)
-	}
-
-	tokens = toTokens(input)
-	tokens = topK(tokens, 20)
-	softmax(tokens)
-	got = topP(tokens, 1e-10)
-	if len(got) == 0 {
-		t.Errorf("topP(1e-10): should keep at least one token, got %d", len(got))
+		t.Errorf("topP(0.95): kept too many tokens: got %d", len(got))
 		t.Logf("got: %v", got)
 	}
 }
 
 func TestMinP(t *testing.T) {
-	input := []float32{-2, 0, -1, -3, 2, 1, 4, 3}
+	input := []float32{-3, -2, -1, 0, 1, 2, 4, 3}
 	tokens := toTokens(input)
 
 	// First apply temperature and softmax
-	tokens = topK(tokens, 20)
-	softmax(tokens)
+	tokens = temperature(tokens, 1)
 
-	tokens = minP(tokens, 1.0)
-
-	if len(tokens) != 1 {
-		t.Errorf("minP(1.0): should keep all tokens, got %d, want %d", len(tokens), len(tokens))
-	}
-
-	// Test with normal p value
-	tokens = toTokens(input) // Reset tokens
-	tokens = topK(tokens, 20)
-	softmax(tokens)
-	tokens = minP(tokens, 0.2)
+	// Then apply minP
+	got := minP(tokens, 0.2)
 
 	// Should keep tokens with prob >= 0.2 * max_prob
-	if len(tokens) > 3 {
-		t.Errorf("minP(0.2): kept too many tokens: got %d", len(tokens))
-		t.Logf("got: %v", tokens)
+	if len(got) > 3 {
+		t.Errorf("minP(0.2): kept too many tokens: got %d", len(got))
 	}
+}
+
+func TestSortLogits(t *testing.T) {
+	input := []float32{0.026986899, 0.043722924, 0.036774673, 0.27755088, 0.0046718004, 0.08582123, 0.20409796, 0.00412893, 0.15720603, 0.045046154, 0.0030491839, 0.01681367}
+	tokens := toTokens(input)
 
-	// Test with zero p value
-	tokens = toTokens(input) // Reset tokens
 	tokens = topK(tokens, 20)
-	softmax(tokens)
-	tokens = minP(tokens, 0.0)
 
-	// Should keep only the highest probability token
-	if len(tokens) != len(input) {
-		t.Errorf("minP(0.0): should keep only one token, got %d", len(tokens))
-		t.Logf("got: %v", tokens)
-	}
-
-	// Test with single token
-	tokens = toTokens(input[:1])
-	tokens = topK(tokens, 20)
-	softmax(tokens)
-	tokens = minP(tokens, 0.1)
-
-	// Should keep only the highest probability token
-	if len(tokens) != 1 {
-		t.Errorf("minP(0.1): should return single token, got %d", len(tokens))
-		t.Logf("got: %v", tokens)
-	}
-
-	input = []float32{1e-10, 1e-10, 1e-10}
-	tokens = toTokens(input)
-	softmax(tokens)
-	tokens = minP(tokens, 1.0)
-	if len(tokens) < 1 {
-		t.Error("minP should keep at least one token even with extreme probabilities")
-		got := minP(tokens, 1.0)
-
-		if len(got) != 1 {
-			t.Errorf("minP(1.0): should keep all tokens, got %d, want %d", len(got), len(tokens))
-		}
-
-		// Test with normal p value
-		got = minP(tokens, 0.2)
-
-		// Should keep tokens with prob >= 0.2 * max_prob
-		if len(got) > 3 {
-			t.Errorf("minP(0.2): kept too many tokens: got %d", len(got))
-			t.Logf("got: %v", got)
-		}
-
-		// Test with zero p value
-		got = minP(tokens, 0.0)
-
-		// Should keep only the highest probability token
-		if len(got) != len(tokens) {
-			t.Errorf("minP(0.0): should keep only one token, got %d", len(got))
-			t.Logf("got: %v", got)
+	for i := 1; i < len(tokens); i++ {
+		if tokens[i].value > tokens[i-1].value {
+			t.Errorf("sortLogits: tokens not sorted in descending order at index %d: %f > %f",
+				i, tokens[i].value, tokens[i-1].value)
 		}
 	}
+
+	want := []float32{0.27755088, 0.20409796, 0.15720603, 0.08582123, 0.045046154, 0.043722924, 0.036774673, 0.026986899, 0.01681367, 0.0046718004, 0.00412893, 0.0030491839}
+	compareLogits(t, "sortLogits", want, tokens)
+}
+
+// TestSortLogitsWithRealData tests sorting behavior using real model logit distributions
+func TestSortLogitsWithRealData(t *testing.T) {
+	// This will be populated from testdata/logits.bin
+	// Format: 32-bit float array in binary format
+	logits, err := loadTestLogits(t)
+	if err != nil {
+		t.Skipf("Skipping real logit test: %v", err)
+		return
+	}
+
+	tokens := toTokens(logits)
+	sortLogits(tokens)
+
+	// Calculate n for verification
+	n := int(math.Sqrt(float64(len(tokens)))) + 1
+	if n > 1000 {
+		n = 1000
+	} else if n < 100 {
+		n = 100
+	}
+
+	t.Logf("Testing with %d tokens, partial sorting top %d", len(tokens), n)
+
+	// Only verify the top n elements are sorted (which is what we guarantee)
+	// This is much faster than checking the entire array
+	topN := tokens[:n]
+	for i := 1; i < len(topN); i++ {
+		if topN[i].value > topN[i-1].value {
+			t.Fatalf("top %d tokens not properly sorted at index %d: %.15f > %.15f",
+				n, i, topN[i].value, topN[i-1].value)
+		}
+	}
+
+	// Verify we didn't lose any high value tokens by checking that
+	// all tokens after position n are <= the nth token
+	// Do this in chunks to avoid timeouts on large arrays
+	nthValue := tokens[n-1].value
+	const chunkSize = 1000
+
+	for start := n; start < len(tokens); start += chunkSize {
+		end := min(start+chunkSize, len(tokens))
+		for i := start; i < end; i++ {
+			if tokens[i].value > nthValue {
+				t.Fatalf("found higher value token after position %d: tokens[%d].value = %.15f > %.15f",
+					n, i, tokens[i].value, nthValue)
+			}
+		}
+	}
+}
+
+// loadTestLogits loads logit test data from testdata/logits.bin
+func loadTestLogits(t *testing.T) ([]float32, error) {
+	t.Helper()
+
+	_, currFile, _, ok := runtime.Caller(0)
+	if !ok {
+		return nil, errors.New("could not determine test file path")
+	}
+	testDataPath := filepath.Join(filepath.Dir(currFile), "testdata", "logits.bin")
+
+	file, err := os.Open(testDataPath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	stat, err := file.Stat()
+	if err != nil {
+		return nil, err
+	}
+
+	numFloats := stat.Size() / 4 // each float32 is 4 bytes
+	if numFloats*4 != stat.Size() {
+		return nil, errors.New("logits.bin has invalid size: not a multiple of 4 bytes")
+	}
+
+	logits := make([]float32, numFloats)
+	for i := range logits {
+		var val uint32
+		if err := binary.Read(file, binary.LittleEndian, &val); err != nil {
+			return nil, err
+		}
+		logits[i] = math.Float32frombits(val)
+	}
+
+	if len(logits) == 0 {
+		return nil, errors.New("logits.bin is empty")
+	}
+
+	return logits, nil
 }
 
 func BenchmarkTransforms(b *testing.B) {
@@ -315,19 +260,11 @@ func BenchmarkTransforms(b *testing.B) {
 		}
 	})
 
-	b.Run("Softmax", func(b *testing.B) {
-		b.ResetTimer()
-		for b.Loop() {
-			copy(tokensCopy, tokens)
-			softmax(tokensCopy)
-		}
-	})
-
 	b.Run("TopK", func(b *testing.B) {
 		b.ResetTimer()
 		for b.Loop() {
 			copy(tokensCopy, tokens)
-			tokens = topK(tokensCopy, 10)
+			topK(tokensCopy, 10)
 		}
 	})
 
@@ -335,7 +272,7 @@ func BenchmarkTransforms(b *testing.B) {
 		b.ResetTimer()
 		for b.Loop() {
 			copy(tokensCopy, tokens)
-			tokens = topP(tokensCopy, 0.9)
+			topP(tokensCopy, 0.9)
 		}
 	})
 
@@ -343,7 +280,7 @@ func BenchmarkTransforms(b *testing.B) {
 		b.ResetTimer()
 		for b.Loop() {
 			copy(tokensCopy, tokens)
-			tokens = minP(tokensCopy, 0.2)
+			minP(tokensCopy, 0.2)
 		}
 	})
 
@@ -351,7 +288,7 @@ func BenchmarkTransforms(b *testing.B) {
 		b.ResetTimer()
 		for b.Loop() {
 			copy(tokensCopy, tokens)
-			tokens = topK(tokensCopy, 200000)
+			topK(tokensCopy, 200000)
 		}
 	})
 }

scripts/build_darwin.sh

@@ -8,7 +8,7 @@ usage() {
     exit 1
 }
 
-export VERSION=${VERSION:-$(git describe --tags --first-parent --abbrev=7 --long --dirty --always | sed -e "s/^v//g")}
+export VERSION=${VERSION:-$(git describe --tags --dirty)}
 export GOFLAGS="'-ldflags=-w -s \"-X=github.com/ollama/ollama/version.Version=${VERSION#v}\" \"-X=github.com/ollama/ollama/server.mode=release\"'"
 export CGO_CPPFLAGS='-mmacosx-version-min=11.3'
 

server/download.go

@@ -29,9 +29,8 @@ import (
 const maxRetries = 6
 
 var (
-	errMaxRetriesExceeded   = errors.New("max retries exceeded")
-	errPartStalled          = errors.New("part stalled")
-	errMaxRedirectsExceeded = errors.New("maximum redirects exceeded (10) for directURL")
+	errMaxRetriesExceeded = errors.New("max retries exceeded")
+	errPartStalled        = errors.New("part stalled")
 )
 
 var blobDownloadManager sync.Map
@@ -237,7 +236,7 @@ func (b *blobDownload) run(ctx context.Context, requestURL *url.URL, opts *regis
 
 			newOpts.CheckRedirect = func(req *http.Request, via []*http.Request) error {
 				if len(via) > 10 {
-					return errMaxRedirectsExceeded
+					return errors.New("maximum redirects exceeded (10) for directURL")
 				}
 
 				// if the hostname is the same, allow the redirect

server/images.go

@@ -35,9 +35,14 @@ var (
 	errCapabilityCompletion = errors.New("completion")
 	errCapabilityTools      = errors.New("tools")
 	errCapabilityInsert     = errors.New("insert")
-	errCapabilityVision     = errors.New("vision")
-	errCapabilityEmbedding  = errors.New("embedding")
-	errInsecureProtocol     = errors.New("insecure protocol http")
+)
+
+type Capability string
+
+const (
+	CapabilityCompletion = Capability("completion")
+	CapabilityTools      = Capability("tools")
+	CapabilityInsert     = Capability("insert")
 )
 
 type registryOptions struct {
@@ -60,83 +65,52 @@ type Model struct {
 	System         string
 	License        []string
 	Digest         string
-	Options        map[string]any
+	Options        map[string]interface{}
 	Messages       []api.Message
 
 	Template *template.Template
 }
 
-// Capabilities returns the capabilities that the model supports
-func (m *Model) Capabilities() []model.Capability {
-	capabilities := []model.Capability{}
-
-	// Check for completion capability
-	r, err := os.Open(m.ModelPath)
-	if err == nil {
-		defer r.Close()
-
-		f, _, err := ggml.Decode(r, 0)
-		if err == nil {
-			if _, ok := f.KV()[fmt.Sprintf("%s.pooling_type", f.KV().Architecture())]; ok {
-				capabilities = append(capabilities, model.CapabilityEmbedding)
-			} else {
-				capabilities = append(capabilities, model.CapabilityCompletion)
-			}
-			if _, ok := f.KV()[fmt.Sprintf("%s.vision.block_count", f.KV().Architecture())]; ok {
-				capabilities = append(capabilities, model.CapabilityVision)
-			}
-		} else {
-			slog.Error("couldn't decode ggml", "error", err)
-		}
-	} else {
-		slog.Error("couldn't open model file", "error", err)
-	}
-
-	if m.Template == nil {
-		return capabilities
-	}
-
-	// Check for tools capability
-	if slices.Contains(m.Template.Vars(), "tools") {
-		capabilities = append(capabilities, model.CapabilityTools)
-	}
-
-	// Check for insert capability
-	if slices.Contains(m.Template.Vars(), "suffix") {
-		capabilities = append(capabilities, model.CapabilityInsert)
-	}
-
-	return capabilities
-}
-
 // CheckCapabilities checks if the model has the specified capabilities returning an error describing
 // any missing or unknown capabilities
-func (m *Model) CheckCapabilities(want ...model.Capability) error {
-	available := m.Capabilities()
+func (m *Model) CheckCapabilities(caps ...Capability) error {
 	var errs []error
+	for _, cap := range caps {
+		switch cap {
+		case CapabilityCompletion:
+			r, err := os.Open(m.ModelPath)
+			if err != nil {
+				slog.Error("couldn't open model file", "error", err)
+				continue
+			}
+			defer r.Close()
 
-	// Map capabilities to their corresponding error
-	capToErr := map[model.Capability]error{
-		model.CapabilityCompletion: errCapabilityCompletion,
-		model.CapabilityTools:      errCapabilityTools,
-		model.CapabilityInsert:     errCapabilityInsert,
-		model.CapabilityVision:     errCapabilityVision,
-		model.CapabilityEmbedding:  errCapabilityEmbedding,
-	}
+			// TODO(mxyng): decode the GGML into model to avoid doing this multiple times
+			f, _, err := ggml.Decode(r, 0)
+			if err != nil {
+				slog.Error("couldn't decode ggml", "error", err)
+				continue
+			}
 
-	for _, cap := range want {
-		err, ok := capToErr[cap]
-		if !ok {
+			if _, ok := f.KV()[fmt.Sprintf("%s.pooling_type", f.KV().Architecture())]; ok {
+				errs = append(errs, errCapabilityCompletion)
+			}
+		case CapabilityTools:
+			if !slices.Contains(m.Template.Vars(), "tools") {
+				errs = append(errs, errCapabilityTools)
+			}
+		case CapabilityInsert:
+			vars := m.Template.Vars()
+			if !slices.Contains(vars, "suffix") {
+				errs = append(errs, errCapabilityInsert)
+			}
+		default:
 			slog.Error("unknown capability", "capability", cap)
 			return fmt.Errorf("unknown capability: %s", cap)
 		}
-
-		if !slices.Contains(available, cap) {
-			errs = append(errs, err)
-		}
 	}
 
-	if len(errs) > 0 {
+	if err := errors.Join(errs...); err != nil {
 		return fmt.Errorf("%w %w", errCapabilities, errors.Join(errs...))
 	}
 
@@ -505,7 +479,7 @@ func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
 	fn(api.ProgressResponse{Status: "retrieving manifest"})
 
 	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
-		return errInsecureProtocol
+		return errors.New("insecure protocol http")
 	}
 
 	manifest, _, err := GetManifest(mp)
@@ -569,7 +543,7 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
 	}
 
 	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
-		return errInsecureProtocol
+		return errors.New("insecure protocol http")
 	}
 
 	fn(api.ProgressResponse{Status: "pulling manifest"})

server/images_test.go

@@ -1,360 +0,0 @@
-package server
-
-import (
-	"bytes"
-	"encoding/binary"
-	"os"
-	"path/filepath"
-	"strings"
-	"testing"
-
-	"github.com/ollama/ollama/template"
-	"github.com/ollama/ollama/types/model"
-)
-
-// Constants for GGUF magic bytes and version
-var (
-	ggufMagic = []byte{0x47, 0x47, 0x55, 0x46} // "GGUF"
-	ggufVer   = uint32(3)                      // Version 3
-)
-
-// Helper function to create mock GGUF data
-func createMockGGUFData(architecture string, vision bool) []byte {
-	var buf bytes.Buffer
-
-	// Write GGUF header
-	buf.Write(ggufMagic)
-	binary.Write(&buf, binary.LittleEndian, ggufVer)
-
-	// Write tensor count (0 for our test)
-	var numTensors uint64 = 0
-	binary.Write(&buf, binary.LittleEndian, numTensors)
-
-	// Calculate number of metadata entries
-	numMetaEntries := uint64(1) // architecture entry
-	if vision {
-		numMetaEntries++
-	}
-	// Add embedding entry if architecture is "bert"
-	if architecture == "bert" {
-		numMetaEntries++
-	}
-	binary.Write(&buf, binary.LittleEndian, numMetaEntries)
-
-	// Write architecture metadata
-	archKey := "general.architecture"
-	keyLen := uint64(len(archKey))
-	binary.Write(&buf, binary.LittleEndian, keyLen)
-	buf.WriteString(archKey)
-
-	// String type (8)
-	var strType uint32 = 8
-	binary.Write(&buf, binary.LittleEndian, strType)
-
-	// String length
-	strLen := uint64(len(architecture))
-	binary.Write(&buf, binary.LittleEndian, strLen)
-	buf.WriteString(architecture)
-
-	if vision {
-		visionKey := architecture + ".vision.block_count"
-		keyLen = uint64(len(visionKey))
-		binary.Write(&buf, binary.LittleEndian, keyLen)
-		buf.WriteString(visionKey)
-
-		// uint32 type (4)
-		var uint32Type uint32 = 4
-		binary.Write(&buf, binary.LittleEndian, uint32Type)
-
-		// uint32 value (1)
-		var countVal uint32 = 1
-		binary.Write(&buf, binary.LittleEndian, countVal)
-	}
-	// Write embedding metadata if architecture is "bert"
-	if architecture == "bert" {
-		poolKey := architecture + ".pooling_type"
-		keyLen = uint64(len(poolKey))
-		binary.Write(&buf, binary.LittleEndian, keyLen)
-		buf.WriteString(poolKey)
-
-		// uint32 type (4)
-		var uint32Type uint32 = 4
-		binary.Write(&buf, binary.LittleEndian, uint32Type)
-
-		// uint32 value (1)
-		var poolingVal uint32 = 1
-		binary.Write(&buf, binary.LittleEndian, poolingVal)
-	}
-
-	return buf.Bytes()
-}
-
-func TestModelCapabilities(t *testing.T) {
-	// Create a temporary directory for test files
-	tempDir, err := os.MkdirTemp("", "model_capabilities_test")
-	if err != nil {
-		t.Fatalf("Failed to create temp directory: %v", err)
-	}
-	defer os.RemoveAll(tempDir)
-
-	// Create different types of mock model files
-	completionModelPath := filepath.Join(tempDir, "model.bin")
-	visionModelPath := filepath.Join(tempDir, "vision_model.bin")
-	embeddingModelPath := filepath.Join(tempDir, "embedding_model.bin")
-	// Create a simple model file for tests that don't depend on GGUF content
-	simpleModelPath := filepath.Join(tempDir, "simple_model.bin")
-
-	err = os.WriteFile(completionModelPath, createMockGGUFData("llama", false), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create completion model file: %v", err)
-	}
-	err = os.WriteFile(visionModelPath, createMockGGUFData("llama", true), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create completion model file: %v", err)
-	}
-	err = os.WriteFile(embeddingModelPath, createMockGGUFData("bert", false), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create embedding model file: %v", err)
-	}
-	err = os.WriteFile(simpleModelPath, []byte("dummy model data"), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create simple model file: %v", err)
-	}
-
-	toolsInsertTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}")
-	if err != nil {
-		t.Fatalf("Failed to parse template: %v", err)
-	}
-	chatTemplate, err := template.Parse("{{ .prompt }}")
-	if err != nil {
-		t.Fatalf("Failed to parse template: %v", err)
-	}
-	toolsTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}")
-	if err != nil {
-		t.Fatalf("Failed to parse template: %v", err)
-	}
-
-	testModels := []struct {
-		name         string
-		model        Model
-		expectedCaps []model.Capability
-	}{
-		{
-			name: "model with completion capability",
-			model: Model{
-				ModelPath: completionModelPath,
-				Template:  chatTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityCompletion},
-		},
-
-		{
-			name: "model with completion, tools, and insert capability",
-			model: Model{
-				ModelPath: completionModelPath,
-				Template:  toolsInsertTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityInsert},
-		},
-		{
-			name: "model with tools and insert capability",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  toolsInsertTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityTools, model.CapabilityInsert},
-		},
-		{
-			name: "model with tools capability",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  toolsTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityTools},
-		},
-		{
-			name: "model with vision capability",
-			model: Model{
-				ModelPath: visionModelPath,
-				Template:  chatTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision},
-		},
-		{
-			name: "model with vision, tools, and insert capability",
-			model: Model{
-				ModelPath: visionModelPath,
-				Template:  toolsInsertTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityInsert},
-		},
-		{
-			name: "model with embedding capability",
-			model: Model{
-				ModelPath: embeddingModelPath,
-				Template:  chatTemplate,
-			},
-			expectedCaps: []model.Capability{model.CapabilityEmbedding},
-		},
-	}
-
-	// compare two slices of model.Capability regardless of order
-	compareCapabilities := func(a, b []model.Capability) bool {
-		if len(a) != len(b) {
-			return false
-		}
-
-		aCount := make(map[model.Capability]int)
-		for _, cap := range a {
-			aCount[cap]++
-		}
-
-		bCount := make(map[model.Capability]int)
-		for _, cap := range b {
-			bCount[cap]++
-		}
-
-		for cap, count := range aCount {
-			if bCount[cap] != count {
-				return false
-			}
-		}
-
-		return true
-	}
-
-	for _, tt := range testModels {
-		t.Run(tt.name, func(t *testing.T) {
-			// Test Capabilities method
-			caps := tt.model.Capabilities()
-			if !compareCapabilities(caps, tt.expectedCaps) {
-				t.Errorf("Expected capabilities %v, got %v", tt.expectedCaps, caps)
-			}
-		})
-	}
-}
-
-func TestModelCheckCapabilities(t *testing.T) {
-	// Create a temporary directory for test files
-	tempDir, err := os.MkdirTemp("", "model_check_capabilities_test")
-	if err != nil {
-		t.Fatalf("Failed to create temp directory: %v", err)
-	}
-	defer os.RemoveAll(tempDir)
-
-	visionModelPath := filepath.Join(tempDir, "vision_model.bin")
-	simpleModelPath := filepath.Join(tempDir, "model.bin")
-	embeddingModelPath := filepath.Join(tempDir, "embedding_model.bin")
-
-	err = os.WriteFile(simpleModelPath, []byte("dummy model data"), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create simple model file: %v", err)
-	}
-	err = os.WriteFile(visionModelPath, createMockGGUFData("llama", true), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create vision model file: %v", err)
-	}
-	err = os.WriteFile(embeddingModelPath, createMockGGUFData("bert", false), 0o644)
-	if err != nil {
-		t.Fatalf("Failed to create embedding model file: %v", err)
-	}
-
-	toolsInsertTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}")
-	if err != nil {
-		t.Fatalf("Failed to parse template: %v", err)
-	}
-	chatTemplate, err := template.Parse("{{ .prompt }}")
-	if err != nil {
-		t.Fatalf("Failed to parse template: %v", err)
-	}
-	toolsTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}")
-	if err != nil {
-		t.Fatalf("Failed to parse template: %v", err)
-	}
-
-	tests := []struct {
-		name           string
-		model          Model
-		checkCaps      []model.Capability
-		expectedErrMsg string
-	}{
-		{
-			name: "completion model without tools capability",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  chatTemplate,
-			},
-			checkCaps:      []model.Capability{model.CapabilityTools},
-			expectedErrMsg: "does not support tools",
-		},
-		{
-			name: "model with all needed capabilities",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  toolsInsertTemplate,
-			},
-			checkCaps: []model.Capability{model.CapabilityTools, model.CapabilityInsert},
-		},
-		{
-			name: "model missing insert capability",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  toolsTemplate,
-			},
-			checkCaps:      []model.Capability{model.CapabilityInsert},
-			expectedErrMsg: "does not support insert",
-		},
-		{
-			name: "model missing vision capability",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  toolsTemplate,
-			},
-			checkCaps:      []model.Capability{model.CapabilityVision},
-			expectedErrMsg: "does not support vision",
-		},
-		{
-			name: "model with vision capability",
-			model: Model{
-				ModelPath: visionModelPath,
-				Template:  chatTemplate,
-			},
-			checkCaps: []model.Capability{model.CapabilityVision},
-		},
-		{
-			name: "model with embedding capability",
-			model: Model{
-				ModelPath: embeddingModelPath,
-				Template:  chatTemplate,
-			},
-			checkCaps: []model.Capability{model.CapabilityEmbedding},
-		},
-		{
-			name: "unknown capability",
-			model: Model{
-				ModelPath: simpleModelPath,
-				Template:  chatTemplate,
-			},
-			checkCaps:      []model.Capability{"unknown"},
-			expectedErrMsg: "unknown capability",
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			// Test CheckCapabilities method
-			err := tt.model.CheckCapabilities(tt.checkCaps...)
-			if tt.expectedErrMsg == "" {
-				if err != nil {
-					t.Errorf("Expected no error, got: %v", err)
-				}
-			} else {
-				if err == nil {
-					t.Errorf("Expected error containing %q, got nil", tt.expectedErrMsg)
-				} else if !strings.Contains(err.Error(), tt.expectedErrMsg) {
-					t.Errorf("Expected error containing %q, got: %v", tt.expectedErrMsg, err)
-				}
-			}
-		})
-	}
-}

server/internal/cache/blob/cache.go

@@ -146,7 +146,7 @@ func debugger(err *error) func(step string) {
 // be in either of the following forms:
 //
 //	@<digest>
-//	<name>@<digest>
+//	<name>
 //	<name>
 //
 // If a digest is provided, it is returned as is and nothing else happens.
@@ -160,6 +160,8 @@ func debugger(err *error) func(step string) {
 // hashed is passed to a PutBytes call to ensure that the manifest is in the
 // blob store. This is done to ensure that future calls to [Get] succeed in
 // these cases.
+//
+// TODO(bmizerany): Move Links/Resolve/etc. out of this package.
 func (c *DiskCache) Resolve(name string) (Digest, error) {
 	name, digest := splitNameDigest(name)
 	if digest != "" {
@@ -277,6 +279,18 @@ func (c *DiskCache) Get(d Digest) (Entry, error) {
 // It returns an error if either the name or digest is invalid, or if link
 // creation encounters any issues.
 func (c *DiskCache) Link(name string, d Digest) error {
+	// TODO(bmizerany): Move link handling from cache to registry.
+	//
+	// We originally placed links in the cache due to its storage
+	// knowledge. However, the registry likely offers better context for
+	// naming concerns, and our API design shouldn't be tightly coupled to
+	// our on-disk format.
+	//
+	// Links work effectively when independent from physical location -
+	// they can reference content with matching SHA regardless of storage
+	// location. In an upcoming change, we plan to shift this
+	// responsibility to the registry where it better aligns with the
+	// system's conceptual model.
 	manifest, err := c.manifestPath(name)
 	if err != nil {
 		return err
@@ -327,9 +341,7 @@ func (c *DiskCache) GetFile(d Digest) string {
 	return absJoin(c.dir, "blobs", filename)
 }
 
-// Links returns a sequence of link names. The sequence is in lexical order.
-// Names are converted from their relative path form to their name form but are
-// not guaranteed to be valid. Callers should validate the names before using.
+// Links returns a sequence of links in the cache in lexical order.
 func (c *DiskCache) Links() iter.Seq2[string, error] {
 	return func(yield func(string, error) bool) {
 		for path, err := range c.links() {
@@ -402,14 +414,12 @@ func (c *DiskCache) links() iter.Seq2[string, error] {
 }
 
 type checkWriter struct {
-	size int64
 	d    Digest
-	f    *os.File
+	size int64
+	n    int64
 	h    hash.Hash
-
-	w   io.Writer // underlying writer; set by creator
-	n   int64
-	err error
+	f    *os.File
+	err  error
 
 	testHookBeforeFinalWrite func(*os.File)
 }
@@ -425,10 +435,6 @@ func (w *checkWriter) seterr(err error) error {
 // underlying writer is guaranteed to be the last byte of p as verified by the
 // hash.
 func (w *checkWriter) Write(p []byte) (int, error) {
-	if w.err != nil {
-		return 0, w.err
-	}
-
 	_, err := w.h.Write(p)
 	if err != nil {
 		return 0, w.seterr(err)
@@ -447,7 +453,7 @@ func (w *checkWriter) Write(p []byte) (int, error) {
 	if nextSize > w.size {
 		return 0, w.seterr(fmt.Errorf("content exceeds expected size: %d > %d", nextSize, w.size))
 	}
-	n, err := w.w.Write(p)
+	n, err := w.f.Write(p)
 	w.n += int64(n)
 	return n, w.seterr(err)
 }
@@ -487,12 +493,10 @@ func (c *DiskCache) copyNamedFile(name string, file io.Reader, out Digest, size
 
 	// Copy file to f, but also into h to double-check hash.
 	cw := &checkWriter{
-		d:    out,
-		size: size,
-		h:    sha256.New(),
-		f:    f,
-		w:    f,
-
+		d:                        out,
+		size:                     size,
+		h:                        sha256.New(),
+		f:                        f,
 		testHookBeforeFinalWrite: c.testHookBeforeFinalWrite,
 	}
 	n, err := io.Copy(cw, file)
@@ -528,6 +532,11 @@ func splitNameDigest(s string) (name, digest string) {
 var errInvalidName = errors.New("invalid name")
 
 func nameToPath(name string) (_ string, err error) {
+	if strings.Contains(name, "@") {
+		// TODO(bmizerany): HACK: Fix names.Parse to validate.
+		// TODO(bmizerany): merge with default parts (maybe names.Merge(a, b))
+		return "", errInvalidName
+	}
 	n := names.Parse(name)
 	if !n.IsFullyQualified() {
 		return "", errInvalidName
@@ -538,7 +547,8 @@ func nameToPath(name string) (_ string, err error) {
 func absJoin(pp ...string) string {
 	abs, err := filepath.Abs(filepath.Join(pp...))
 	if err != nil {
-		panic(err) // this should never happen
+		// Likely a bug bug or a bad OS problem. Just panic.
+		panic(err)
 	}
 	return abs
 }

server/internal/cache/blob/chunked.go

@@ -1,73 +0,0 @@
-package blob
-
-import (
-	"crypto/sha256"
-	"errors"
-	"io"
-	"os"
-)
-
-// Chunk represents a range of bytes in a blob.
-type Chunk struct {
-	Start int64
-	End   int64
-}
-
-// Size returns end minus start plus one.
-func (c Chunk) Size() int64 {
-	return c.End - c.Start + 1
-}
-
-// Chunker writes to a blob in chunks.
-// Its zero value is invalid. Use [DiskCache.Chunked] to create a new Chunker.
-type Chunker struct {
-	digest Digest
-	size   int64
-	f      *os.File // nil means pre-validated
-}
-
-// Chunked returns a new Chunker, ready for use storing a blob of the given
-// size in chunks.
-//
-// Use [Chunker.Put] to write data to the blob at specific offsets.
-func (c *DiskCache) Chunked(d Digest, size int64) (*Chunker, error) {
-	name := c.GetFile(d)
-	info, err := os.Stat(name)
-	if err == nil && info.Size() == size {
-		return &Chunker{}, nil
-	}
-	f, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY, 0o666)
-	if err != nil {
-		return nil, err
-	}
-	return &Chunker{digest: d, size: size, f: f}, nil
-}
-
-// Put copies chunk.Size() bytes from r to the blob at the given offset,
-// merging the data with the existing blob. It returns an error if any. As a
-// special case, if r has less than chunk.Size() bytes, Put returns
-// io.ErrUnexpectedEOF.
-func (c *Chunker) Put(chunk Chunk, d Digest, r io.Reader) error {
-	if c.f == nil {
-		return nil
-	}
-
-	cw := &checkWriter{
-		d:    d,
-		size: chunk.Size(),
-		h:    sha256.New(),
-		f:    c.f,
-		w:    io.NewOffsetWriter(c.f, chunk.Start),
-	}
-
-	_, err := io.CopyN(cw, r, chunk.Size())
-	if err != nil && errors.Is(err, io.EOF) {
-		return io.ErrUnexpectedEOF
-	}
-	return err
-}
-
-// Close closes the underlying file.
-func (c *Chunker) Close() error {
-	return c.f.Close()
-}

server/internal/cache/blob/digest.go

@@ -63,10 +63,6 @@ func (d Digest) Short() string {
 	return fmt.Sprintf("%x", d.sum[:4])
 }
 
-func (d Digest) Sum() [32]byte {
-	return d.sum
-}
-
 func (d Digest) Compare(other Digest) int {
 	return slices.Compare(d.sum[:], other.sum[:])
 }

server/internal/chunks/chunks.go

@@ -0,0 +1,78 @@
+package chunks
+
+import (
+	"fmt"
+	"iter"
+	"strconv"
+	"strings"
+)
+
+type Chunk struct {
+	Start, End int64
+}
+
+func New(start, end int64) Chunk {
+	return Chunk{start, end}
+}
+
+// ParseRange parses a string in the form "unit=range" where unit is a string
+// and range is a string in the form "start-end". It returns the unit and the
+// range as a Chunk.
+func ParseRange(s string) (unit string, _ Chunk, _ error) {
+	unit, r, _ := strings.Cut(s, "=")
+	if r == "" {
+		return unit, Chunk{}, nil
+	}
+	c, err := Parse(r)
+	if err != nil {
+		return "", Chunk{}, err
+	}
+	return unit, c, err
+}
+
+// Parse parses a string in the form "start-end" and returns the Chunk.
+func Parse(s string) (Chunk, error) {
+	startStr, endStr, _ := strings.Cut(s, "-")
+	start, err := strconv.ParseInt(startStr, 10, 64)
+	if err != nil {
+		return Chunk{}, fmt.Errorf("invalid start: %v", err)
+	}
+	end, err := strconv.ParseInt(endStr, 10, 64)
+	if err != nil {
+		return Chunk{}, fmt.Errorf("invalid end: %v", err)
+	}
+	if start > end {
+		return Chunk{}, fmt.Errorf("invalid range %d-%d: start > end", start, end)
+	}
+	return Chunk{start, end}, nil
+}
+
+// Of returns a sequence of contiguous Chunks of size chunkSize that cover
+// the range [0, size), in order.
+func Of(size, chunkSize int64) iter.Seq[Chunk] {
+	return func(yield func(Chunk) bool) {
+		for start := int64(0); start < size; start += chunkSize {
+			end := min(start+chunkSize-1, size-1)
+			if !yield(Chunk{start, end}) {
+				break
+			}
+		}
+	}
+}
+
+// Count returns the number of Chunks of size chunkSize needed to cover the
+// range [0, size).
+func Count(size, chunkSize int64) int64 {
+	return (size + chunkSize - 1) / chunkSize
+}
+
+// Size returns end minus start plus one.
+func (c Chunk) Size() int64 {
+	return c.End - c.Start + 1
+}
+
+// String returns the string representation of the Chunk in the form
+// "{start}-{end}".
+func (c Chunk) String() string {
+	return fmt.Sprintf("%d-%d", c.Start, c.End)
+}

server/internal/chunks/chunks_test.go

@@ -0,0 +1,65 @@
+package chunks
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestOf(t *testing.T) {
+	cases := []struct {
+		total     int64
+		chunkSize int64
+		want      []Chunk
+	}{
+		{0, 1, nil},
+		{1, 1, []Chunk{{0, 0}}},
+		{1, 2, []Chunk{{0, 0}}},
+		{2, 1, []Chunk{{0, 0}, {1, 1}}},
+		{10, 9, []Chunk{{0, 8}, {9, 9}}},
+	}
+
+	for _, tt := range cases {
+		got := slices.Collect(Of(tt.total, tt.chunkSize))
+		if !slices.Equal(got, tt.want) {
+			t.Errorf("[%d/%d]: got %v; want %v", tt.total, tt.chunkSize, got, tt.want)
+		}
+	}
+}
+
+func TestSize(t *testing.T) {
+	cases := []struct {
+		c    Chunk
+		want int64
+	}{
+		{Chunk{0, 0}, 1},
+		{Chunk{0, 1}, 2},
+		{Chunk{3, 4}, 2},
+	}
+
+	for _, tt := range cases {
+		got := tt.c.Size()
+		if got != tt.want {
+			t.Errorf("%v: got %d; want %d", tt.c, got, tt.want)
+		}
+	}
+}
+
+func TestCount(t *testing.T) {
+	cases := []struct {
+		total     int64
+		chunkSize int64
+		want      int64
+	}{
+		{0, 1, 0},
+		{1, 1, 1},
+		{1, 2, 1},
+		{2, 1, 2},
+		{10, 9, 2},
+	}
+	for _, tt := range cases {
+		got := Count(tt.total, tt.chunkSize)
+		if got != tt.want {
+			t.Errorf("[%d/%d]: got %d; want %d", tt.total, tt.chunkSize, got, tt.want)
+		}
+	}
+}

server/internal/client/ollama/registry.go

@@ -19,13 +19,11 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
-	"iter"
 	"log/slog"
 	"net/http"
 	"os"
 	"path/filepath"
 	"runtime"
-	"runtime/debug"
 	"slices"
 	"strconv"
 	"strings"
@@ -37,7 +35,10 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/ollama/ollama/server/internal/cache/blob"
+	"github.com/ollama/ollama/server/internal/chunks"
+	"github.com/ollama/ollama/server/internal/internal/backoff"
 	"github.com/ollama/ollama/server/internal/internal/names"
+	"github.com/ollama/ollama/server/internal/internal/syncs"
 
 	_ "embed"
 )
@@ -59,18 +60,18 @@ var (
 	// ErrCached is passed to [Trace.PushUpdate] when a layer already
 	// exists. It is a non-fatal error and is never returned by [Registry.Push].
 	ErrCached = errors.New("cached")
-
-	// ErrIncomplete is returned by [Registry.Pull] when a model pull was
-	// incomplete due to one or more layer download failures. Users that
-	// want specific errors should use [WithTrace].
-	ErrIncomplete = errors.New("incomplete")
 )
 
 // Defaults
 const (
 	// DefaultChunkingThreshold is the threshold at which a layer should be
 	// split up into chunks when downloading.
-	DefaultChunkingThreshold = 64 << 20
+	DefaultChunkingThreshold = 128 << 20
+
+	// DefaultMaxChunkSize is the default maximum size of a chunk to
+	// download. It is configured based on benchmarks and aims to strike a
+	// balance between download speed and memory usage.
+	DefaultMaxChunkSize = 8 << 20
 )
 
 var defaultCache = sync.OnceValues(func() (*blob.DiskCache, error) {
@@ -210,13 +211,20 @@ type Registry struct {
 	// pushing or pulling models. If zero, the number of streams is
 	// determined by [runtime.GOMAXPROCS].
 	//
-	// A negative value means no limit.
+	// Clients that want "unlimited" streams should set this to a large
+	// number.
 	MaxStreams int
 
 	// ChunkingThreshold is the maximum size of a layer to download in a single
 	// request. If zero, [DefaultChunkingThreshold] is used.
 	ChunkingThreshold int64
 
+	// MaxChunkSize is the maximum size of a chunk to download. If zero,
+	// the default is [DefaultMaxChunkSize].
+	//
+	// It is only used when a layer is larger than [MaxChunkingThreshold].
+	MaxChunkSize int64
+
 	// Mask, if set, is the name used to convert non-fully qualified names
 	// to fully qualified names. If empty, [DefaultMask] is used.
 	Mask string
@@ -258,7 +266,6 @@ func DefaultRegistry() (*Registry, error) {
 	}
 
 	var rc Registry
-	rc.UserAgent = UserAgent()
 	rc.Key, err = ssh.ParseRawPrivateKey(keyPEM)
 	if err != nil {
 		return nil, err
@@ -274,35 +281,25 @@ func DefaultRegistry() (*Registry, error) {
 	return &rc, nil
 }
 
-func UserAgent() string {
-	buildinfo, _ := debug.ReadBuildInfo()
-
-	version := buildinfo.Main.Version
-	if version == "(devel)" {
-		// When using `go run .` the version is "(devel)". This is seen
-		// as an invalid version by ollama.com and so it defaults to
-		// "needs upgrade" for some requests, such as pulls. These
-		// checks can be skipped by using the special version "v0.0.0",
-		// so we set it to that here.
-		version = "v0.0.0"
-	}
-
-	return fmt.Sprintf("ollama/%s (%s %s) Go/%s",
-		version,
-		runtime.GOARCH,
-		runtime.GOOS,
-		runtime.Version(),
-	)
-}
-
 func (r *Registry) maxStreams() int {
-	return cmp.Or(r.MaxStreams, runtime.GOMAXPROCS(0))
+	n := cmp.Or(r.MaxStreams, runtime.GOMAXPROCS(0))
+
+	// Large downloads require a writter stream, so ensure we have at least
+	// two streams to avoid a deadlock.
+	return max(n, 2)
 }
 
 func (r *Registry) maxChunkingThreshold() int64 {
 	return cmp.Or(r.ChunkingThreshold, DefaultChunkingThreshold)
 }
 
+// chunkSizeFor returns the chunk size for a layer of the given size. If the
+// size is less than or equal to the max chunking threshold, the size is
+// returned; otherwise, the max chunk size is returned.
+func (r *Registry) maxChunkSize() int64 {
+	return cmp.Or(r.MaxChunkSize, DefaultMaxChunkSize)
+}
+
 type PushParams struct {
 	// From is an optional destination name for the model. If empty, the
 	// destination name is the same as the source name.
@@ -421,20 +418,12 @@ func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error {
 	return err
 }
 
-// trackingReader is an io.Reader that tracks the number of bytes read and
-// calls the update function with the layer, the number of bytes read.
-//
-// It always calls update with a nil error.
-type trackingReader struct {
-	l      *Layer
-	r      io.Reader
-	update func(l *Layer, n int64, err error)
-}
-
-func (r *trackingReader) Read(p []byte) (n int, err error) {
-	n, err = r.r.Read(p)
-	r.update(r.l, int64(n), nil)
-	return
+func canRetry(err error) bool {
+	var re *Error
+	if !errors.As(err, &re) {
+		return false
+	}
+	return re.Status >= 500
 }
 
 // Pull pulls the model with the given name from the remote registry into the
@@ -445,15 +434,15 @@ func (r *trackingReader) Read(p []byte) (n int, err error) {
 // typically slower than splitting the model up across layers, and is mostly
 // utilized for layers of type equal to "application/vnd.ollama.image".
 func (r *Registry) Pull(ctx context.Context, name string) error {
-	m, err := r.Resolve(ctx, name)
+	scheme, n, _, err := r.parseNameExtended(name)
 	if err != nil {
 		return err
 	}
 
-	// TODO(bmizerany): decide if this should be considered valid. Maybe
-	// server-side we special case '{}' to have some special meaning? Maybe
-	// "archiving" a tag (which is how we reason about it in the registry
-	// already, just with a different twist).
+	m, err := r.Resolve(ctx, name)
+	if err != nil {
+		return err
+	}
 	if len(m.Layers) == 0 {
 		return fmt.Errorf("%w: no layers", ErrManifestInvalid)
 	}
@@ -463,132 +452,142 @@ func (r *Registry) Pull(ctx context.Context, name string) error {
 		return err
 	}
 
-	// TODO(bmizerany): work to remove the need to do this
+	exists := func(l *Layer) bool {
+		info, err := c.Get(l.Digest)
+		return err == nil && info.Size == l.Size
+	}
+
+	t := traceFromContext(ctx)
+
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(r.maxStreams())
+
 	layers := m.Layers
 	if m.Config != nil && m.Config.Digest.IsValid() {
 		layers = append(layers, m.Config)
 	}
 
-	// Send initial layer trace events to allow clients to have an
-	// understanding of work to be done before work starts.
-	var expected int64
-	t := traceFromContext(ctx)
 	for _, l := range layers {
-		t.update(l, 0, nil)
-		expected += l.Size
-	}
-
-	var received atomic.Int64
-	var g errgroup.Group
-	g.SetLimit(r.maxStreams())
-	for _, l := range layers {
-		info, err := c.Get(l.Digest)
-		if err == nil && info.Size == l.Size {
-			received.Add(l.Size)
+		if exists(l) {
 			t.update(l, l.Size, ErrCached)
 			continue
 		}
 
-		var wg sync.WaitGroup
-		chunked, err := c.Chunked(l.Digest, l.Size)
+		blobURL := fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", scheme, n.Host(), n.Namespace(), n.Model(), l.Digest)
+		req, err := r.newRequest(ctx, "GET", blobURL, nil)
 		if err != nil {
 			t.update(l, 0, err)
 			continue
 		}
 
-		for cs, err := range r.chunksums(ctx, name, l) {
-			if err != nil {
-				// Chunksum stream interrupted. Note in trace
-				// log and let in-flight downloads complete.
-				// This will naturally trigger ErrIncomplete
-				// since received < expected bytes.
-				t.update(l, 0, err)
-				break
-			}
+		t.update(l, 0, nil)
 
-			cacheKey := fmt.Sprintf(
-				"v1 pull chunksum %s %s %d-%d",
-				l.Digest,
-				cs.Digest,
-				cs.Chunk.Start,
-				cs.Chunk.End,
-			)
-			cacheKeyDigest := blob.DigestFromBytes(cacheKey)
-			_, err := c.Get(cacheKeyDigest)
-			if err == nil {
-				received.Add(cs.Chunk.Size())
-				t.update(l, cs.Chunk.Size(), ErrCached)
-				continue
-			}
-
-			wg.Add(1)
-			g.Go(func() (err error) {
-				defer func() {
-					if err == nil {
-						// Ignore cache key write errors for now. We've already
-						// reported to trace that the chunk is complete.
-						//
-						// Ideally, we should only report completion to trace
-						// after successful cache commit. This current approach
-						// works but could trigger unnecessary redownloads if
-						// the checkpoint key is missing on next pull.
-						//
-						// Not incorrect, just suboptimal - fix this in a
-						// future update.
-						_ = blob.PutBytes(c, cacheKeyDigest, cacheKey)
-
-						received.Add(cs.Chunk.Size())
-					} else {
-						t.update(l, 0, err)
-					}
-					wg.Done()
-				}()
-
-				req, err := http.NewRequestWithContext(ctx, "GET", cs.URL, nil)
-				if err != nil {
-					return err
-				}
-				req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", cs.Chunk.Start, cs.Chunk.End))
+		if l.Size <= r.maxChunkingThreshold() {
+			g.Go(func() error {
+				// TODO(bmizerany): retry/backoff like below in
+				// the chunking case
 				res, err := sendRequest(r.client(), req)
 				if err != nil {
 					return err
 				}
 				defer res.Body.Close()
-
-				body := &trackingReader{l: l, r: res.Body, update: t.update}
-				return chunked.Put(cs.Chunk, cs.Digest, body)
+				err = c.Put(l.Digest, res.Body, l.Size)
+				if err == nil {
+					t.update(l, l.Size, nil)
+				}
+				return err
 			})
-		}
+		} else {
+			q := syncs.NewRelayReader()
 
-		// Close writer immediately after downloads finish, not at Pull
-		// exit. Using defer would keep file descriptors open until all
-		// layers complete, potentially exhausting system limits with
-		// many layers.
-		//
-		// The WaitGroup tracks when all chunks finish downloading,
-		// allowing precise writer closure in a background goroutine.
-		// Each layer briefly uses one extra goroutine while at most
-		// maxStreams()-1 chunks download in parallel.
-		//
-		// This caps file descriptors at maxStreams() instead of
-		// growing with layer count.
-		g.Go(func() error {
-			wg.Wait()
-			chunked.Close()
-			return nil
-		})
+			g.Go(func() (err error) {
+				defer func() { q.CloseWithError(err) }()
+				return c.Put(l.Digest, q, l.Size)
+			})
+
+			var progress atomic.Int64
+
+			// We want to avoid extra round trips per chunk due to
+			// redirects from the registry to the blob store, so
+			// fire an initial request to get the final URL and
+			// then use that URL for the chunk requests.
+			req.Header.Set("Range", "bytes=0-0")
+			res, err := sendRequest(r.client(), req)
+			if err != nil {
+				return err
+			}
+			res.Body.Close()
+			req = res.Request.WithContext(req.Context())
+
+			wp := writerPool{size: r.maxChunkSize()}
+
+			for chunk := range chunks.Of(l.Size, r.maxChunkSize()) {
+				if ctx.Err() != nil {
+					break
+				}
+
+				ticket := q.Take()
+				g.Go(func() (err error) {
+					defer func() {
+						if err != nil {
+							q.CloseWithError(err)
+						}
+						ticket.Close()
+						t.update(l, progress.Load(), err)
+					}()
+
+					for _, err := range backoff.Loop(ctx, 3*time.Second) {
+						if err != nil {
+							return err
+						}
+						err := func() error {
+							req := req.Clone(req.Context())
+							req.Header.Set("Range", fmt.Sprintf("bytes=%s", chunk))
+							res, err := sendRequest(r.client(), req)
+							if err != nil {
+								return err
+							}
+							defer res.Body.Close()
+
+							tw := wp.get()
+							tw.Reset(ticket)
+							defer wp.put(tw)
+
+							_, err = io.CopyN(tw, res.Body, chunk.Size())
+							if err != nil {
+								return maybeUnexpectedEOF(err)
+							}
+							if err := tw.Flush(); err != nil {
+								return err
+							}
+
+							total := progress.Add(chunk.Size())
+							if total >= l.Size {
+								q.Close()
+							}
+							return nil
+						}()
+						if !canRetry(err) {
+							return err
+						}
+					}
+					return nil
+				})
+			}
+		}
 	}
+
 	if err := g.Wait(); err != nil {
 		return err
 	}
-	if received.Load() != expected {
-		return fmt.Errorf("%w: received %d/%d bytes", ErrIncomplete, received.Load(), expected)
-	}
 
+	// store the manifest blob
 	md := blob.DigestFromBytes(m.Data)
 	if err := blob.PutBytes(c, md, m.Data); err != nil {
 		return err
 	}
+
+	// commit the manifest with a link
 	return c.Link(m.Name, md)
 }
 
@@ -616,6 +615,8 @@ type Manifest struct {
 	Config *Layer `json:"config"`
 }
 
+var emptyDigest, _ = blob.ParseDigest("sha256:0000000000000000000000000000000000000000000000000000000000000000")
+
 // Layer returns the layer with the given
 // digest, or nil if not found.
 func (m *Manifest) Layer(d blob.Digest) *Layer {
@@ -627,30 +628,6 @@ func (m *Manifest) Layer(d blob.Digest) *Layer {
 	return nil
 }
 
-func (m *Manifest) All() iter.Seq[*Layer] {
-	return func(yield func(*Layer) bool) {
-		if !yield(m.Config) {
-			return
-		}
-		for _, l := range m.Layers {
-			if !yield(l) {
-				return
-			}
-		}
-	}
-}
-
-func (m *Manifest) Size() int64 {
-	var size int64
-	if m.Config != nil {
-		size += m.Config.Size
-	}
-	for _, l := range m.Layers {
-		size += l.Size
-	}
-	return size
-}
-
 // MarshalJSON implements json.Marshaler.
 //
 // NOTE: It adds an empty config object to the manifest, which is required by
@@ -666,9 +643,10 @@ func (m Manifest) MarshalJSON() ([]byte, error) {
 		// last phase of the commit which expects it, but does nothing
 		// with it. This will be fixed in a future release of
 		// ollama.com.
-		Config Layer `json:"config"`
+		Config *Layer `json:"config"`
 	}{
-		M: M(m),
+		M:      M(m),
+		Config: &Layer{Digest: emptyDigest},
 	}
 	return json.Marshal(v)
 }
@@ -758,135 +736,6 @@ func (r *Registry) Resolve(ctx context.Context, name string) (*Manifest, error)
 	return m, nil
 }
 
-type chunksum struct {
-	URL    string
-	Chunk  blob.Chunk
-	Digest blob.Digest
-}
-
-// chunksums returns a sequence of chunksums for the given layer. If the layer is under the
-// chunking threshold, a single chunksum is returned that covers the entire layer. If the layer
-// is over the chunking threshold, the chunksums are read from the chunksums endpoint.
-func (r *Registry) chunksums(ctx context.Context, name string, l *Layer) iter.Seq2[chunksum, error] {
-	return func(yield func(chunksum, error) bool) {
-		scheme, n, _, err := r.parseNameExtended(name)
-		if err != nil {
-			yield(chunksum{}, err)
-			return
-		}
-
-		if l.Size < r.maxChunkingThreshold() {
-			// any layer under the threshold should be downloaded
-			// in one go.
-			cs := chunksum{
-				URL: fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s",
-					scheme,
-					n.Host(),
-					n.Namespace(),
-					n.Model(),
-					l.Digest,
-				),
-				Chunk:  blob.Chunk{Start: 0, End: l.Size - 1},
-				Digest: l.Digest,
-			}
-			yield(cs, nil)
-			return
-		}
-
-		// The response is a sequence of chunksums.
-		//
-		// Chunksums are chunks of a larger blob that can be
-		// downloaded and verified independently.
-		//
-		// The chunksums endpoint is a GET request that returns a
-		// sequence of chunksums in the following format:
-		//
-		//     > GET /v2/<namespace>/<model>/chunksums/<digest>
-		//
-		//     < HTTP/1.1 200 OK
-		//     < Content-Location: <blobURL>
-		//     <
-		//     < <digest> <start>-<end>
-		//     < ...
-		//
-		// The <blobURL> is the URL to download the chunks from and
-		// each <digest> is the digest of the chunk, and <start>-<end>
-		// is the range the chunk in the blob.
-		//
-		// Ranges may be used directly in Range headers like
-		// "bytes=<start>-<end>".
-		//
-		// The chunksums returned are guaranteed to be contiguous and
-		// include all bytes of the layer. If the stream is cut short,
-		// clients should retry.
-
-		chunksumsURL := fmt.Sprintf("%s://%s/v2/%s/%s/chunksums/%s",
-			scheme,
-			n.Host(),
-			n.Namespace(),
-			n.Model(),
-			l.Digest,
-		)
-
-		req, err := r.newRequest(ctx, "GET", chunksumsURL, nil)
-		if err != nil {
-			yield(chunksum{}, err)
-			return
-		}
-		res, err := sendRequest(r.client(), req)
-		if err != nil {
-			yield(chunksum{}, err)
-			return
-		}
-		defer res.Body.Close()
-		if res.StatusCode != 200 {
-			err := fmt.Errorf("chunksums: unexpected status code %d", res.StatusCode)
-			yield(chunksum{}, err)
-			return
-		}
-		blobURL := res.Header.Get("Content-Location")
-
-		s := bufio.NewScanner(res.Body)
-		s.Split(bufio.ScanWords)
-		for {
-			if !s.Scan() {
-				if s.Err() != nil {
-					yield(chunksum{}, s.Err())
-				}
-				return
-			}
-			d, err := blob.ParseDigest(s.Bytes())
-			if err != nil {
-				yield(chunksum{}, fmt.Errorf("invalid digest: %q", s.Bytes()))
-				return
-			}
-
-			if !s.Scan() {
-				err := s.Err()
-				if err == nil {
-					err = fmt.Errorf("missing chunk range for digest %s", d)
-				}
-				yield(chunksum{}, err)
-				return
-			}
-			chunk, err := parseChunk(s.Bytes())
-			if err != nil {
-				yield(chunksum{}, fmt.Errorf("invalid chunk range for digest %s: %q", d, s.Bytes()))
-				return
-			}
-
-			cs := chunksum{
-				URL:    blobURL,
-				Chunk:  chunk,
-				Digest: d,
-			}
-			if !yield(cs, nil) {
-				return
-			}
-		}
-	}
-}
-
 func (r *Registry) client() *http.Client {
 	if r.HTTPClient != nil {
 		return r.HTTPClient
@@ -1049,6 +898,13 @@ func checkData(url string) string {
 	return fmt.Sprintf("GET,%s,%s", url, zeroSum)
 }
 
+func maybeUnexpectedEOF(err error) error {
+	if errors.Is(err, io.EOF) {
+		return io.ErrUnexpectedEOF
+	}
+	return err
+}
+
 type publicError struct {
 	wrapped error
 	message string
@@ -1135,22 +991,27 @@ func splitExtended(s string) (scheme, name, digest string) {
 	return scheme, s, digest
 }
 
-// parseChunk parses a string in the form "start-end" and returns the Chunk.
-func parseChunk[S ~string | ~[]byte](s S) (blob.Chunk, error) {
-	startPart, endPart, found := strings.Cut(string(s), "-")
-	if !found {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: missing '-'", s)
-	}
-	start, err := strconv.ParseInt(startPart, 10, 64)
-	if err != nil {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid start to %q: %v", s, err)
-	}
-	end, err := strconv.ParseInt(endPart, 10, 64)
-	if err != nil {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid end to %q: %v", s, err)
-	}
-	if start > end {
-		return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: start > end", s)
-	}
-	return blob.Chunk{Start: start, End: end}, nil
+type writerPool struct {
+	size int64 // set by the caller
+
+	mu sync.Mutex
+	ws []*bufio.Writer
+}
+
+func (p *writerPool) get() *bufio.Writer {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if len(p.ws) == 0 {
+		return bufio.NewWriterSize(nil, int(p.size))
+	}
+	w := p.ws[len(p.ws)-1]
+	p.ws = p.ws[:len(p.ws)-1]
+	return w
+}
+
+func (p *writerPool) put(w *bufio.Writer) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	w.Reset(nil)
+	p.ws = append(p.ws, w)
 }

server/internal/client/ollama/registry_test.go

@@ -9,41 +9,22 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
-	"net"
+	"math/rand/v2"
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"path"
 	"reflect"
+	"slices"
 	"strings"
-	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/ollama/ollama/server/internal/cache/blob"
+	"github.com/ollama/ollama/server/internal/chunks"
 	"github.com/ollama/ollama/server/internal/testutil"
 )
 
-func ExampleRegistry_cancelOnFirstError() {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	ctx = WithTrace(ctx, &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			if err != nil {
-				// Discontinue pulling layers if there is an
-				// error instead of continuing to pull more
-				// data.
-				cancel()
-			}
-		},
-	})
-
-	var r Registry
-	if err := r.Pull(ctx, "model"); err != nil {
-		// panic for demo purposes
-		panic(err)
-	}
-}
-
 func TestManifestMarshalJSON(t *testing.T) {
 	// All manifests should contain an "empty" config object.
 	var m Manifest
@@ -76,21 +57,21 @@ func (rr recordRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
 
 // newClient constructs a cache with predefined manifests for testing. The manifests are:
 //
-//	empty:         no data
-//	zero:          no layers
-//	single:        one layer with the contents "exists"
-//	multiple:      two layers with the contents "exists" and "here"
-//	notfound:      a layer that does not exist in the cache
-//	null:          one null layer (e.g. [null])
-//	sizemismatch:  one valid layer, and one with a size mismatch (file size is less than the reported size)
-//	invalid:       a layer with invalid JSON data
+//	empty: no data
+//	zero: no layers
+//	single: one layer with the contents "exists"
+//	multiple: two layers with the contents "exists" and "here"
+//	notfound: a layer that does not exist in the cache
+//	null: one null layer (e.g. [null])
+//	sizemismatch: one valid layer, and one with a size mismatch (file size is less than the reported size)
+//	invalid: a layer with invalid JSON data
 //
 // Tests that want to ensure the client does not communicate with the upstream
 // registry should pass a nil handler, which will cause a panic if
 // communication is attempted.
 //
 // To simulate a network error, pass a handler that returns a 499 status code.
-func newClient(t *testing.T, upstreamRegistry http.HandlerFunc) (*Registry, *blob.DiskCache) {
+func newClient(t *testing.T, h http.HandlerFunc) (*Registry, *blob.DiskCache) {
 	t.Helper()
 
 	c, err := blob.Open(t.TempDir())
@@ -108,7 +89,7 @@ func newClient(t *testing.T, upstreamRegistry http.HandlerFunc) (*Registry, *blo
 	r := &Registry{
 		Cache: c,
 		HTTPClient: &http.Client{
-			Transport: recordRoundTripper(upstreamRegistry),
+			Transport: recordRoundTripper(h),
 		},
 	}
 
@@ -335,8 +316,15 @@ func TestPushCommitRoundtripError(t *testing.T) {
 	}
 }
 
+func checkNotExist(t *testing.T, err error) {
+	t.Helper()
+	if !errors.Is(err, fs.ErrNotExist) {
+		t.Fatalf("err = %v; want fs.ErrNotExist", err)
+	}
+}
+
 func TestRegistryPullInvalidName(t *testing.T) {
-	rc, _ := newRegistryClient(t, nil)
+	rc, _ := newClient(t, nil)
 	err := rc.Pull(t.Context(), "://")
 	if !errors.Is(err, ErrNameInvalid) {
 		t.Errorf("err = %v; want %v", err, ErrNameInvalid)
@@ -352,16 +340,245 @@ func TestRegistryPullInvalidManifest(t *testing.T) {
 	}
 
 	for _, resp := range cases {
-		rc, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
+		rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
 			io.WriteString(w, resp)
 		})
-		err := rc.Pull(t.Context(), "http://example.com/a/b")
+		err := rc.Pull(t.Context(), "x")
 		if !errors.Is(err, ErrManifestInvalid) {
 			t.Errorf("err = %v; want invalid manifest", err)
 		}
 	}
 }
 
+func TestRegistryPullNotCached(t *testing.T) {
+	check := testutil.Checker(t)
+
+	var c *blob.DiskCache
+	var rc *Registry
+
+	d := blob.DigestFromBytes("some data")
+	rc, c = newClient(t, func(w http.ResponseWriter, r *http.Request) {
+		if strings.Contains(r.URL.Path, "/blobs/") {
+			io.WriteString(w, "some data")
+			return
+		}
+		fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":9}]}`, d)
+	})
+
+	// Confirm that the layer does not exist locally
+	_, err := rc.ResolveLocal("model")
+	checkNotExist(t, err)
+
+	_, err = c.Get(d)
+	checkNotExist(t, err)
+
+	err = rc.Pull(t.Context(), "model")
+	check(err)
+
+	mw, err := rc.Resolve(t.Context(), "model")
+	check(err)
+	mg, err := rc.ResolveLocal("model")
+	check(err)
+	if !reflect.DeepEqual(mw, mg) {
+		t.Errorf("mw = %v; mg = %v", mw, mg)
+	}
+
+	// Confirm successful download
+	info, err := c.Get(d)
+	check(err)
+	if info.Digest != d {
+		t.Errorf("info.Digest = %v; want %v", info.Digest, d)
+	}
+	if info.Size != 9 {
+		t.Errorf("info.Size = %v; want %v", info.Size, 9)
+	}
+
+	data, err := os.ReadFile(c.GetFile(d))
+	check(err)
+	if string(data) != "some data" {
+		t.Errorf("data = %q; want %q", data, "exists")
+	}
+}
+
+func TestRegistryPullCached(t *testing.T) {
+	cached := blob.DigestFromBytes("exists")
+	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
+		if strings.Contains(r.URL.Path, "/blobs/") {
+			w.WriteHeader(499) // should not be called
+			return
+		}
+		if strings.Contains(r.URL.Path, "/manifests/") {
+			fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":6}]}`, cached)
+		}
+	})
+
+	var errs []error
+	var reads []int64
+	ctx := WithTrace(t.Context(), &Trace{
+		Update: func(d *Layer, n int64, err error) {
+			t.Logf("update %v %d %v", d, n, err)
+			reads = append(reads, n)
+			errs = append(errs, err)
+		},
+	})
+
+	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+	defer cancel()
+
+	err := rc.Pull(ctx, "single")
+	testutil.Check(t, err)
+
+	want := []int64{6}
+	if !errors.Is(errors.Join(errs...), ErrCached) {
+		t.Errorf("errs = %v; want %v", errs, ErrCached)
+	}
+	if !slices.Equal(reads, want) {
+		t.Errorf("pairs = %v; want %v", reads, want)
+	}
+}
+
+func TestRegistryPullManifestNotFound(t *testing.T) {
+	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	})
+	err := rc.Pull(t.Context(), "notfound")
+	checkErrCode(t, err, 404, "")
+}
+
+func TestRegistryPullResolveRemoteError(t *testing.T) {
+	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		io.WriteString(w, `{"errors":[{"code":"an_error"}]}`)
+	})
+	err := rc.Pull(t.Context(), "single")
+	checkErrCode(t, err, 500, "an_error")
+}
+
+func TestRegistryPullResolveRoundtripError(t *testing.T) {
+	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
+		if strings.Contains(r.URL.Path, "/manifests/") {
+			w.WriteHeader(499) // force RoundTrip error
+			return
+		}
+	})
+	err := rc.Pull(t.Context(), "single")
+	if !errors.Is(err, errRoundTrip) {
+		t.Errorf("err = %v; want %v", err, errRoundTrip)
+	}
+}
+
+// TestRegistryPullMixedCachedNotCached tests that cached layers do not
+// interfere with pulling layers that are not cached
+func TestRegistryPullMixedCachedNotCached(t *testing.T) {
+	x := blob.DigestFromBytes("xxxxxx")
+	e := blob.DigestFromBytes("exists")
+	y := blob.DigestFromBytes("yyyyyy")
+
+	for i := range 10 {
+		t.Logf("iteration %d", i)
+
+		digests := []blob.Digest{x, e, y}
+
+		rand.Shuffle(len(digests), func(i, j int) {
+			digests[i], digests[j] = digests[j], digests[i]
+		})
+
+		manifest := fmt.Sprintf(`{
+			"layers": [
+				{"digest":"%s","size":6},
+				{"digest":"%s","size":6},
+				{"digest":"%s","size":6}
+			]
+		}`, digests[0], digests[1], digests[2])
+
+		rc, c := newClient(t, func(w http.ResponseWriter, r *http.Request) {
+			switch path.Base(r.URL.Path) {
+			case "latest":
+				io.WriteString(w, manifest)
+			case x.String():
+				io.WriteString(w, "xxxxxx")
+			case e.String():
+				io.WriteString(w, "exists")
+			case y.String():
+				io.WriteString(w, "yyyyyy")
+			default:
+				panic(fmt.Sprintf("unexpected request: %v", r))
+			}
+		})
+
+		ctx := WithTrace(t.Context(), &Trace{
+			Update: func(l *Layer, n int64, err error) {
+				t.Logf("update %v %d %v", l, n, err)
+			},
+		})
+
+		// Check that we pull all layers that we can.
+
+		err := rc.Pull(ctx, "mixed")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		for _, d := range digests {
+			info, err := c.Get(d)
+			if err != nil {
+				t.Fatalf("Get(%v): %v", d, err)
+			}
+			if info.Size != 6 {
+				t.Errorf("info.Size = %v; want %v", info.Size, 6)
+			}
+		}
+	}
+}
+
+func TestRegistryPullChunking(t *testing.T) {
+	rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) {
+		t.Log("request:", r.URL.Host, r.Method, r.URL.Path, r.Header.Get("Range"))
+		if r.URL.Host != "blob.store" {
+			// The production registry redirects to the blob store.
+			http.Redirect(w, r, "http://blob.store"+r.URL.Path, http.StatusFound)
+			return
+		}
+		if strings.Contains(r.URL.Path, "/blobs/") {
+			rng := r.Header.Get("Range")
+			if rng == "" {
+				http.Error(w, "missing range", http.StatusBadRequest)
+				return
+			}
+			_, c, err := chunks.ParseRange(r.Header.Get("Range"))
+			if err != nil {
+				panic(err)
+			}
+			io.WriteString(w, "remote"[c.Start:c.End+1])
+			return
+		}
+		fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":6}]}`, blob.DigestFromBytes("remote"))
+	})
+
+	// Force chunking by setting the threshold to less than the size of the
+	// layer.
+	rc.ChunkingThreshold = 3
+	rc.MaxChunkSize = 3
+
+	var reads []int64
+	ctx := WithTrace(t.Context(), &Trace{
+		Update: func(d *Layer, n int64, err error) {
+			if err != nil {
+				t.Errorf("update %v %d %v", d, n, err)
+			}
+			reads = append(reads, n)
+		},
+	})
+
+	err := rc.Pull(ctx, "remote")
+	testutil.Check(t, err)
+
+	want := []int64{0, 3, 6}
+	if !slices.Equal(reads, want) {
+		t.Errorf("reads = %v; want %v", reads, want)
+	}
+}
+
 func TestRegistryResolveByDigest(t *testing.T) {
 	check := testutil.Checker(t)
 
@@ -399,6 +616,26 @@ func TestInsecureSkipVerify(t *testing.T) {
 	testutil.Check(t, err)
 }
 
+func TestCanRetry(t *testing.T) {
+	cases := []struct {
+		err  error
+		want bool
+	}{
+		{nil, false},
+		{errors.New("x"), false},
+		{ErrCached, false},
+		{ErrManifestInvalid, false},
+		{ErrNameInvalid, false},
+		{&Error{Status: 100}, false},
+		{&Error{Status: 500}, true},
+	}
+	for _, tt := range cases {
+		if got := canRetry(tt.err); got != tt.want {
+			t.Errorf("CanRetry(%v) = %v; want %v", tt.err, got, tt.want)
+		}
+	}
+}
+
 func TestErrorUnmarshal(t *testing.T) {
 	cases := []struct {
 		name    string
@@ -550,23 +787,17 @@ func TestParseNameExtended(t *testing.T) {
 
 func TestUnlink(t *testing.T) {
 	t.Run("found by name", func(t *testing.T) {
-		check := testutil.Checker(t)
-
-		rc, _ := newRegistryClient(t, nil)
-		// make a blob and link it
-		d := blob.DigestFromBytes("{}")
-		err := blob.PutBytes(rc.Cache, d, "{}")
-		check(err)
-		err = rc.Cache.Link("registry.ollama.ai/library/single:latest", d)
-		check(err)
+		rc, _ := newClient(t, nil)
 
 		// confirm linked
-		_, err = rc.ResolveLocal("single")
-		check(err)
+		_, err := rc.ResolveLocal("single")
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
 
 		// unlink
 		_, err = rc.Unlink("single")
-		check(err)
+		testutil.Check(t, err)
 
 		// confirm unlinked
 		_, err = rc.ResolveLocal("single")
@@ -575,7 +806,7 @@ func TestUnlink(t *testing.T) {
 		}
 	})
 	t.Run("not found by name", func(t *testing.T) {
-		rc, _ := newRegistryClient(t, nil)
+		rc, _ := newClient(t, nil)
 		ok, err := rc.Unlink("manifestNotFound")
 		if err != nil {
 			t.Fatal(err)
@@ -585,369 +816,3 @@ func TestUnlink(t *testing.T) {
 		}
 	})
 }
-
-// Many tests from here out, in this file are based on a single blob, "abc",
-// with the checksum of its sha256 hash. The checksum is:
-//
-//	"abc" -> sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad
-//
-// Using the literal value instead of a constant with fmt.Xprintf calls proved
-// to be the most readable and maintainable approach. The sum is consistently
-// used in the tests and unique so searches do not yield false positives.
-
-func checkRequest(t *testing.T, req *http.Request, method, path string) {
-	t.Helper()
-	if got := req.URL.Path; got != path {
-		t.Errorf("URL = %q, want %q", got, path)
-	}
-	if req.Method != method {
-		t.Errorf("Method = %q, want %q", req.Method, method)
-	}
-}
-
-func newRegistryClient(t *testing.T, h http.HandlerFunc) (*Registry, context.Context) {
-	s := httptest.NewServer(h)
-	t.Cleanup(s.Close)
-	cache, err := blob.Open(t.TempDir())
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-		},
-	})
-
-	rc := &Registry{
-		Cache: cache,
-		HTTPClient: &http.Client{Transport: &http.Transport{
-			Dial: func(network, addr string) (net.Conn, error) {
-				return net.Dial(network, s.Listener.Addr().String())
-			},
-		}},
-	}
-	return rc, ctx
-}
-
-func TestPullChunked(t *testing.T) {
-	var steps atomic.Int64
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch steps.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab"))
-			fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c"))
-		case 3, 4:
-			checkRequest(t, r, "GET", "/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			switch rng := r.Header.Get("Range"); rng {
-			case "bytes=0-1":
-				io.WriteString(w, "ab")
-			case "bytes=2-2":
-				t.Logf("writing c")
-				io.WriteString(w, "c")
-			default:
-				t.Errorf("unexpected range %q", rng)
-			}
-		default:
-			t.Errorf("unexpected steps %d: %v", steps.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.ChunkingThreshold = 1 // force chunking
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	testutil.Check(t, err)
-
-	_, err = c.Cache.Resolve("o.com/library/abc:latest")
-	testutil.Check(t, err)
-
-	if g := steps.Load(); g != 4 {
-		t.Fatalf("got %d steps, want 4", g)
-	}
-}
-
-func TestPullCached(t *testing.T) {
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-		io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-	})
-
-	check := testutil.Checker(t)
-
-	// Premeptively cache the blob
-	d, err := blob.ParseDigest("sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-	check(err)
-	err = blob.PutBytes(c.Cache, d, []byte("abc"))
-	check(err)
-
-	// Pull only the manifest, which should be enough to resolve the cached blob
-	err = c.Pull(ctx, "http://o.com/library/abc")
-	check(err)
-}
-
-func TestPullManifestError(t *testing.T) {
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-		w.WriteHeader(http.StatusNotFound)
-		io.WriteString(w, `{"errors":[{"code":"MANIFEST_UNKNOWN"}]}`)
-	})
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	if err == nil {
-		t.Fatalf("expected error")
-	}
-	var got *Error
-	if !errors.Is(err, ErrModelNotFound) {
-		t.Fatalf("err = %v, want %v", got, ErrModelNotFound)
-	}
-}
-
-func TestPullLayerError(t *testing.T) {
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-		io.WriteString(w, `!`)
-	})
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	if err == nil {
-		t.Fatalf("expected error")
-	}
-	var want *json.SyntaxError
-	if !errors.As(err, &want) {
-		t.Fatalf("err = %T, want %T", err, want)
-	}
-}
-
-func TestPullLayerChecksumError(t *testing.T) {
-	var step atomic.Int64
-	c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab"))
-			fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c"))
-		case 3:
-			w.WriteHeader(http.StatusNotFound)
-			io.WriteString(w, `{"errors":[{"code":"BLOB_UNKNOWN"}]}`)
-		case 4:
-			io.WriteString(w, "c")
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.MaxStreams = 1
-	c.ChunkingThreshold = 1 // force chunking
-
-	var written atomic.Int64
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-			written.Add(n)
-		},
-	})
-
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	var got *Error
-	if !errors.As(err, &got) || got.Code != "BLOB_UNKNOWN" {
-		t.Fatalf("err = %v, want %v", err, got)
-	}
-
-	if g := written.Load(); g != 1 {
-		t.Fatalf("wrote %d bytes, want 1", g)
-	}
-}
-
-func TestPullChunksumStreamError(t *testing.T) {
-	var step atomic.Int64
-	c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-
-			// Write one valid chunksum and one invalid chunksum
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) // valid
-			fmt.Fprint(w, "sha256:!")                              // invalid
-		case 3:
-			io.WriteString(w, "ab")
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.ChunkingThreshold = 1 // force chunking
-
-	got := c.Pull(ctx, "http://o.com/library/abc")
-	if !errors.Is(got, ErrIncomplete) {
-		t.Fatalf("err = %v, want %v", got, ErrIncomplete)
-	}
-}
-
-type flushAfterWriter struct {
-	w io.Writer
-}
-
-func (f *flushAfterWriter) Write(p []byte) (n int, err error) {
-	n, err = f.w.Write(p)
-	f.w.(http.Flusher).Flush() // panic if not a flusher
-	return
-}
-
-func TestPullChunksumStreaming(t *testing.T) {
-	csr, csw := io.Pipe()
-	defer csw.Close()
-
-	var step atomic.Int64
-	c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fw := &flushAfterWriter{w} // ensure client gets data as it arrives by aggressively flushing
-			_, err := io.Copy(fw, csr)
-			if err != nil {
-				t.Errorf("copy: %v", err)
-			}
-		case 3:
-			io.WriteString(w, "ab")
-		case 4:
-			io.WriteString(w, "c")
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.ChunkingThreshold = 1 // force chunking
-
-	update := make(chan int64, 1)
-	ctx := WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-			if n > 0 {
-				update <- n
-			}
-		},
-	})
-
-	errc := make(chan error, 1)
-	go func() {
-		errc <- c.Pull(ctx, "http://o.com/library/abc")
-	}()
-
-	// Send first chunksum and ensure it kicks off work immediately
-	fmt.Fprintf(csw, "%s 0-1\n", blob.DigestFromBytes("ab"))
-	if g := <-update; g != 2 {
-		t.Fatalf("got %d, want 2", g)
-	}
-
-	// now send the second chunksum and ensure it kicks off work immediately
-	fmt.Fprintf(csw, "%s 2-2\n", blob.DigestFromBytes("c"))
-	if g := <-update; g != 1 {
-		t.Fatalf("got %d, want 1", g)
-	}
-	csw.Close()
-	testutil.Check(t, <-errc)
-}
-
-func TestPullChunksumsCached(t *testing.T) {
-	var step atomic.Int64
-	c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) {
-		switch step.Add(1) {
-		case 1:
-			checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest")
-			io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`)
-		case 2:
-			w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
-			fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab"))
-			fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c"))
-		case 3, 4:
-			switch rng := r.Header.Get("Range"); rng {
-			case "bytes=0-1":
-				io.WriteString(w, "ab")
-			case "bytes=2-2":
-				io.WriteString(w, "c")
-			default:
-				t.Errorf("unexpected range %q", rng)
-			}
-		default:
-			t.Errorf("unexpected steps %d: %v", step.Load(), r)
-			http.Error(w, "unexpected steps", http.StatusInternalServerError)
-		}
-	})
-
-	c.MaxStreams = 1        // force serial processing of chunksums
-	c.ChunkingThreshold = 1 // force chunking
-
-	ctx, cancel := context.WithCancel(t.Context())
-	defer cancel()
-
-	// Cancel the pull after the first chunksum is processed, but before
-	// the second chunksum is processed (which is waiting because
-	// MaxStreams=1). This should cause the second chunksum to error out
-	// leaving the blob incomplete.
-	ctx = WithTrace(ctx, &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			if n > 0 {
-				cancel()
-			}
-		},
-	})
-	err := c.Pull(ctx, "http://o.com/library/abc")
-	if !errors.Is(err, context.Canceled) {
-		t.Fatalf("err = %v, want %v", err, context.Canceled)
-	}
-
-	_, err = c.Cache.Resolve("o.com/library/abc:latest")
-	if !errors.Is(err, fs.ErrNotExist) {
-		t.Fatalf("err = %v, want nil", err)
-	}
-
-	// Reset state and pull again to ensure the blob chunks that should
-	// have been cached are, and the remaining chunk was downloaded, making
-	// the blob complete.
-	step.Store(0)
-	var written atomic.Int64
-	var cached atomic.Int64
-	ctx = WithTrace(t.Context(), &Trace{
-		Update: func(l *Layer, n int64, err error) {
-			t.Log("trace:", l.Digest.Short(), n, err)
-			if errors.Is(err, ErrCached) {
-				cached.Add(n)
-			}
-			written.Add(n)
-		},
-	})
-
-	check := testutil.Checker(t)
-
-	err = c.Pull(ctx, "http://o.com/library/abc")
-	check(err)
-
-	_, err = c.Cache.Resolve("o.com/library/abc:latest")
-	check(err)
-
-	if g := written.Load(); g != 3 {
-		t.Fatalf("wrote %d bytes, want 3", g)
-	}
-	if g := cached.Load(); g != 2 { // "ab" should have been cached
-		t.Fatalf("cached %d bytes, want 3", g)
-	}
-}

server/internal/cmd/oppbench/oppbench.go

@@ -0,0 +1,11 @@
+package main
+
+import (
+	"fmt"
+	"os"
+)
+
+func main() {
+	fmt.Println("Run as 'go test -bench=.' to run the benchmarks")
+	os.Exit(1)
+}

server/internal/cmd/oppbench/oppbench_test.go

@@ -0,0 +1,107 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/ollama/ollama/server/internal/chunks"
+	"golang.org/x/sync/errgroup"
+)
+
+func BenchmarkDownload(b *testing.B) {
+	run := func(fileSize, chunkSize int64) {
+		name := fmt.Sprintf("size=%d/chunksize=%d", fileSize, chunkSize)
+		b.Run(name, func(b *testing.B) { benchmarkDownload(b, fileSize, chunkSize) })
+	}
+
+	run(100<<20, 8<<20)
+	run(100<<20, 16<<20)
+	run(100<<20, 32<<20)
+	run(100<<20, 64<<20)
+	run(100<<20, 128<<20) // 1 chunk
+}
+
+func run(ctx context.Context, c *http.Client, chunk chunks.Chunk) error {
+	const blobURL = "https://ollama.com/v2/x/x/blobs/sha256-4824460d29f2058aaf6e1118a63a7a197a09bed509f0e7d4e2efb1ee273b447d"
+	req, err := http.NewRequestWithContext(ctx, "GET", blobURL, nil)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Range", fmt.Sprintf("bytes=%s", chunk))
+	res, err := c.Do(req)
+	if err != nil {
+		return err
+	}
+	defer res.Body.Close()
+
+	_, err = io.CopyN(io.Discard, res.Body, chunk.Size()) // will io.EOF on short read
+	return err
+}
+
+var sleepTime atomic.Int64
+
+func benchmarkDownload(b *testing.B, fileSize, chunkSize int64) {
+	client := &http.Client{
+		Transport: func() http.RoundTripper {
+			tr := http.DefaultTransport.(*http.Transport).Clone()
+			tr.DisableKeepAlives = true
+			return tr
+		}(),
+	}
+	defer client.CloseIdleConnections()
+
+	// warm up the client
+	run(context.Background(), client, chunks.New(0, 1<<20))
+
+	b.SetBytes(fileSize)
+	b.ReportAllocs()
+
+	// Give our CDN a min to breathe between benchmarks.
+	time.Sleep(time.Duration(sleepTime.Swap(3)))
+
+	for b.Loop() {
+		g, ctx := errgroup.WithContext(b.Context())
+		g.SetLimit(runtime.GOMAXPROCS(0))
+		for chunk := range chunks.Of(fileSize, chunkSize) {
+			g.Go(func() error { return run(ctx, client, chunk) })
+		}
+		if err := g.Wait(); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkWrite(b *testing.B) {
+	b.Run("chunksize=1MiB", func(b *testing.B) { benchmarkWrite(b, 1<<20) })
+}
+
+func benchmarkWrite(b *testing.B, chunkSize int) {
+	b.ReportAllocs()
+
+	dir := b.TempDir()
+	f, err := os.Create(filepath.Join(dir, "write-single"))
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer f.Close()
+
+	data := make([]byte, chunkSize)
+	b.SetBytes(int64(chunkSize))
+	r := bytes.NewReader(data)
+	for b.Loop() {
+		r.Reset(data)
+		_, err := io.Copy(f, r)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}

server/internal/registry/server.go

@@ -1,5 +1,6 @@
-// Package registry implements an http.Handler for handling local Ollama API
-// model management requests. See [Local] for details.
+// Package registry provides an http.Handler for handling local Ollama API
+// requests for performing tasks related to the ollama.com model registry and
+// the local disk cache.
 package registry
 
 import (
@@ -9,7 +10,6 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
-	"maps"
 	"net/http"
 	"sync"
 	"time"
@@ -18,11 +18,16 @@ import (
 	"github.com/ollama/ollama/server/internal/client/ollama"
 )
 
-// Local implements an http.Handler for handling local Ollama API model
-// management requests, such as pushing, pulling, and deleting models.
+// Local is an http.Handler for handling local Ollama API requests for
+// performing tasks related to the ollama.com model registry combined with the
+// local disk cache.
 //
-// It can be arranged for all unknown requests to be passed through to a
-// fallback handler, if one is provided.
+// It is not concern of Local, or this package, to handle model creation, which
+// proceeds any registry operations for models it produces.
+//
+// NOTE: The package built for dealing with model creation should use
+// [DefaultCache] to access the blob store and not attempt to read or write
+// directly to the blob disk cache.
 type Local struct {
 	Client *ollama.Registry // required
 	Logger *slog.Logger     // required
@@ -58,7 +63,6 @@ func (e serverError) Error() string {
 var (
 	errMethodNotAllowed = &serverError{405, "method_not_allowed", "method not allowed"}
 	errNotFound         = &serverError{404, "not_found", "not found"}
-	errModelNotFound    = &serverError{404, "not_found", "model not found"}
 	errInternalError    = &serverError{500, "internal_error", "internal server error"}
 )
 
@@ -171,16 +175,8 @@ func (s *Local) serveHTTP(rec *statusCodeRecorder, r *http.Request) {
 }
 
 type params struct {
-	// DeprecatedName is the name of the model to push, pull, or delete,
-	// but is deprecated. New clients should use [Model] instead.
-	//
-	// Use [model()] to get the model name for both old and new API requests.
-	DeprecatedName string `json:"name"`
-
-	// Model is the name of the model to push, pull, or delete.
-	//
-	// Use [model()] to get the model name for both old and new API requests.
-	Model string `json:"model"`
+	DeprecatedName string `json:"name"`  // Use [params.model]
+	Model          string `json:"model"` // Use [params.model]
 
 	// AllowNonTLS is a flag that indicates a client using HTTP
 	// is doing so, deliberately.
@@ -193,18 +189,9 @@ type params struct {
 	// confusing flags such as this.
 	AllowNonTLS bool `json:"insecure"`
 
-	// Stream, if true, will make the server send progress updates in a
-	// streaming of JSON objects. If false, the server will send a single
-	// JSON object with the final status as "success", or an error object
-	// if an error occurred.
-	//
-	// Unfortunately, this API was designed to be a bit awkward. Stream is
-	// defined to default to true if not present, so we need a way to check
-	// if the client decisively set it to false. So, we use a pointer to a
-	// bool. Gross.
-	//
-	// Use [stream()] to get the correct value for this field.
-	Stream *bool `json:"stream"`
+	// ProgressStream is a flag that indicates the client is expecting a stream of
+	// progress updates.
+	ProgressStream bool `json:"stream"`
 }
 
 // model returns the model name for both old and new API requests.
@@ -212,13 +199,6 @@ func (p params) model() string {
 	return cmp.Or(p.Model, p.DeprecatedName)
 }
 
-func (p params) stream() bool {
-	if p.Stream == nil {
-		return true
-	}
-	return *p.Stream
-}
-
 func (s *Local) handleDelete(_ http.ResponseWriter, r *http.Request) error {
 	if r.Method != "DELETE" {
 		return errMethodNotAllowed
@@ -232,16 +212,16 @@ func (s *Local) handleDelete(_ http.ResponseWriter, r *http.Request) error {
 		return err
 	}
 	if !ok {
-		return errModelNotFound
+		return &serverError{404, "not_found", "model not found"}
 	}
-	if s.Prune != nil {
-		return s.Prune()
+	if s.Prune == nil {
+		return nil
 	}
-	return nil
+	return s.Prune()
 }
 
 type progressUpdateJSON struct {
-	Status    string      `json:"status,omitempty,omitzero"`
+	Status    string      `json:"status"`
 	Digest    blob.Digest `json:"digest,omitempty,omitzero"`
 	Total     int64       `json:"total,omitempty,omitzero"`
 	Completed int64       `json:"completed,omitempty,omitzero"`
@@ -257,17 +237,6 @@ func (s *Local) handlePull(w http.ResponseWriter, r *http.Request) error {
 		return err
 	}
 
-	enc := json.NewEncoder(w)
-	if !p.stream() {
-		if err := s.Client.Pull(r.Context(), p.model()); err != nil {
-			if errors.Is(err, ollama.ErrModelNotFound) {
-				return errModelNotFound
-			}
-			return err
-		}
-		return enc.Encode(progressUpdateJSON{Status: "success"})
-	}
-
 	maybeFlush := func() {
 		fl, _ := w.(http.Flusher)
 		if fl != nil {
@@ -277,74 +246,69 @@ func (s *Local) handlePull(w http.ResponseWriter, r *http.Request) error {
 	defer maybeFlush()
 
 	var mu sync.Mutex
-	progress := make(map[*ollama.Layer]int64)
+	enc := json.NewEncoder(w)
+	enc.Encode(progressUpdateJSON{Status: "pulling manifest"})
 
-	progressCopy := make(map[*ollama.Layer]int64, len(progress))
-	flushProgress := func() {
-		defer maybeFlush()
+	ctx := ollama.WithTrace(r.Context(), &ollama.Trace{
+		Update: func(l *ollama.Layer, n int64, err error) {
+			mu.Lock()
+			defer mu.Unlock()
 
-		// TODO(bmizerany): Flushing every layer in one update doesn't
-		// scale well. We could flush only the modified layers or track
-		// the full download. Needs further consideration, though it's
-		// fine for now.
-		mu.Lock()
-		maps.Copy(progressCopy, progress)
-		mu.Unlock()
-		for l, n := range progressCopy {
+			// TODO(bmizerany): coalesce these updates; writing per
+			// update is expensive
 			enc.Encode(progressUpdateJSON{
 				Digest:    l.Digest,
+				Status:    "pulling",
 				Total:     l.Size,
 				Completed: n,
 			})
-		}
-	}
-	defer flushProgress()
-
-	t := time.NewTicker(1000 * time.Hour) // "unstarted" timer
-	start := sync.OnceFunc(func() {
-		flushProgress() // flush initial state
-		t.Reset(100 * time.Millisecond)
-	})
-	ctx := ollama.WithTrace(r.Context(), &ollama.Trace{
-		Update: func(l *ollama.Layer, n int64, err error) {
-			if n > 0 {
-				// Block flushing progress updates until every
-				// layer is accounted for. Clients depend on a
-				// complete model size to calculate progress
-				// correctly; if they use an incomplete total,
-				// progress indicators would erratically jump
-				// as new layers are registered.
-				start()
-			}
-			mu.Lock()
-			progress[l] += n
-			mu.Unlock()
 		},
 	})
 
 	done := make(chan error, 1)
 	go func() {
+		// TODO(bmizerany): continue to support non-streaming responses
 		done <- s.Client.Pull(ctx, p.model())
 	}()
 
-	for {
-		select {
-		case <-t.C:
-			flushProgress()
-		case err := <-done:
-			flushProgress()
-			if err != nil {
-				var status string
-				if errors.Is(err, ollama.ErrModelNotFound) {
-					status = fmt.Sprintf("error: model %q not found", p.model())
-				} else {
-					status = fmt.Sprintf("error: %v", err)
+	func() {
+		t := time.NewTicker(100 * time.Millisecond)
+		defer t.Stop()
+		for {
+			select {
+			case <-t.C:
+				mu.Lock()
+				maybeFlush()
+				mu.Unlock()
+			case err := <-done:
+				if err != nil {
+					var status string
+					if errors.Is(err, ollama.ErrModelNotFound) {
+						status = fmt.Sprintf("error: model %q not found", p.model())
+						enc.Encode(progressUpdateJSON{Status: status})
+					} else {
+						status = fmt.Sprintf("error: %v", err)
+						enc.Encode(progressUpdateJSON{Status: status})
+					}
+					return
 				}
-				enc.Encode(progressUpdateJSON{Status: status})
+
+				// These final updates are not strictly necessary, because they have
+				// already happened at this point. Our pull handler code used to do
+				// these steps after, not during, the pull, and they were slow, so we
+				// wanted to provide feedback to users what was happening. For now, we
+				// keep them to not jar users who are used to seeing them. We can phase
+				// them out with a new and nicer UX later. One without progress bars
+				// and digests that no one cares about.
+				enc.Encode(progressUpdateJSON{Status: "verifying layers"})
+				enc.Encode(progressUpdateJSON{Status: "writing manifest"})
+				enc.Encode(progressUpdateJSON{Status: "success"})
+				return
 			}
-			return nil
 		}
-	}
+	}()
+
+	return nil
 }
 
 func decodeUserJSON[T any](r io.Reader) (T, error) {

server/internal/registry/server_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"io/fs"
 	"net"
@@ -159,6 +160,7 @@ var registryFS = sync.OnceValue(func() fs.FS {
 	// to \n when parsing the txtar on Windows.
 	data := bytes.ReplaceAll(registryTXT, []byte("\r\n"), []byte("\n"))
 	a := txtar.Parse(data)
+	fmt.Printf("%q\n", a.Comment)
 	fsys, err := txtar.FS(a)
 	if err != nil {
 		panic(err)
@@ -177,7 +179,7 @@ func TestServerPull(t *testing.T) {
 			w.WriteHeader(404)
 			io.WriteString(w, `{"errors": [{"code": "MANIFEST_UNKNOWN", "message": "manifest unknown"}]}`)
 		default:
-			t.Logf("serving blob: %s", r.URL.Path)
+			t.Logf("serving file: %s", r.URL.Path)
 			modelsHandler.ServeHTTP(w, r)
 		}
 	})
@@ -186,7 +188,7 @@ func TestServerPull(t *testing.T) {
 		t.Helper()
 
 		if got.Code != 200 {
-			t.Errorf("Code = %d; want 200", got.Code)
+			t.Fatalf("Code = %d; want 200", got.Code)
 		}
 		gotlines := got.Body.String()
 		t.Logf("got:\n%s", gotlines)
@@ -195,29 +197,35 @@ func TestServerPull(t *testing.T) {
 			want, unwanted := strings.CutPrefix(want, "!")
 			want = strings.TrimSpace(want)
 			if !unwanted && !strings.Contains(gotlines, want) {
-				t.Errorf("! missing %q in body", want)
+				t.Fatalf("! missing %q in body", want)
 			}
 			if unwanted && strings.Contains(gotlines, want) {
-				t.Errorf("! unexpected %q in body", want)
+				t.Fatalf("! unexpected %q in body", want)
 			}
 		}
 	}
 
 	got := s.send(t, "POST", "/api/pull", `{"model": "BOOM"}`)
 	checkResponse(got, `
+		{"status":"pulling manifest"}
 		{"status":"error: request error https://example.com/v2/library/BOOM/manifests/latest: registry responded with status 999: boom"}
 	`)
 
 	got = s.send(t, "POST", "/api/pull", `{"model": "smol"}`)
 	checkResponse(got, `
-		{"digest":"sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312","total":5}
-		{"digest":"sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356","total":3}
-		{"digest":"sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312","total":5,"completed":5}
-		{"digest":"sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356","total":3,"completed":3}
+		{"status":"pulling manifest"}
+		{"status":"pulling","digest":"sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312","total":5}
+		{"status":"pulling","digest":"sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356","total":3}
+		{"status":"pulling","digest":"sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312","total":5,"completed":5}
+		{"status":"pulling","digest":"sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356","total":3,"completed":3}
+		{"status":"verifying layers"}
+		{"status":"writing manifest"}
+		{"status":"success"}
 	`)
 
 	got = s.send(t, "POST", "/api/pull", `{"model": "unknown"}`)
 	checkResponse(got, `
+		{"status":"pulling manifest"}
 		{"status":"error: model \"unknown\" not found"}
 	`)
 
@@ -232,39 +240,19 @@ func TestServerPull(t *testing.T) {
 
 	got = s.send(t, "POST", "/api/pull", `{"model": "://"}`)
 	checkResponse(got, `
+		{"status":"pulling manifest"}
 		{"status":"error: invalid or missing name: \"\""}
-	`)
 
-	// Non-streaming pulls
-	got = s.send(t, "POST", "/api/pull", `{"model": "://", "stream": false}`)
-	checkErrorResponse(t, got, 400, "bad_request", "invalid or missing name")
-	got = s.send(t, "POST", "/api/pull", `{"model": "smol", "stream": false}`)
-	checkResponse(got, `
-		{"status":"success"}
-		!digest
-		!total
-		!completed
+		!verifying
+		!writing
+		!success
 	`)
-	got = s.send(t, "POST", "/api/pull", `{"model": "unknown", "stream": false}`)
-	checkErrorResponse(t, got, 404, "not_found", "model not found")
 }
 
 func TestServerUnknownPath(t *testing.T) {
 	s := newTestServer(t, nil)
 	got := s.send(t, "DELETE", "/api/unknown", `{}`)
 	checkErrorResponse(t, got, 404, "not_found", "not found")
-
-	var fellback bool
-	s.Fallback = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		fellback = true
-	})
-	got = s.send(t, "DELETE", "/api/unknown", `{}`)
-	if !fellback {
-		t.Fatal("expected Fallback to be called")
-	}
-	if got.Code != 200 {
-		t.Fatalf("Code = %d; want 200", got.Code)
-	}
 }
 
 func checkErrorResponse(t *testing.T, got *httptest.ResponseRecorder, status int, code, msg string) {

server/model.go

@@ -82,7 +82,7 @@ func detectChatTemplate(layers []*layerGGML) ([]*layerGGML, error) {
 	for _, layer := range layers {
 		if s := layer.GGML.KV().ChatTemplate(); s != "" {
 			if t, err := template.Named(s); err != nil {
-				slog.Debug("template detection", "error", err, "template", s)
+				slog.Debug("template detection", "error", err)
 			} else {
 				layer, err := NewLayer(t.Reader(), "application/vnd.ollama.image.template")
 				if err != nil {

server/modelpath.go

@@ -31,10 +31,9 @@ const (
 
 var (
 	ErrInvalidImageFormat  = errors.New("invalid image format")
-	ErrInvalidDigestFormat = errors.New("invalid digest format")
 	ErrInvalidProtocol     = errors.New("invalid protocol scheme")
 	ErrInsecureProtocol    = errors.New("insecure protocol http")
-	ErrModelPathInvalid    = errors.New("invalid model path")
+	ErrInvalidDigestFormat = errors.New("invalid digest format")
 )
 
 func ParseModelPath(name string) ModelPath {
@@ -74,6 +73,8 @@ func ParseModelPath(name string) ModelPath {
 	return mp
 }
 
+var errModelPathInvalid = errors.New("invalid model path")
+
 func (mp ModelPath) GetNamespaceRepository() string {
 	return fmt.Sprintf("%s/%s", mp.Namespace, mp.Repository)
 }

server/prompt.go

@@ -26,6 +26,7 @@ func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api.
 	var system []api.Message
 
 	isMllama := checkMllamaModelFamily(m)
+	isGemma3 := checkGemma3ModelFamily(m)
 
 	var imageNumTokens int
 	// TODO: Ideally we would compute this from the projector metadata but some pieces are implementation dependent
@@ -40,7 +41,7 @@ func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api.
 	n := len(msgs) - 1
 	// in reverse, find all messages that fit into context window
 	for i := n; i >= 0; i-- {
-		if isMllama && len(msgs[i].Images) > 1 {
+		if (isMllama || isGemma3) && len(msgs[i].Images) > 1 {
 			return "", nil, errTooManyImages
 		}
 
@@ -157,3 +158,12 @@ func checkMllamaModelFamily(m *Model) bool {
 	}
 	return false
 }
+
+func checkGemma3ModelFamily(m *Model) bool {
+	for _, arch := range m.Config.ModelFamilies {
+		if arch == "gemma3" {
+			return true
+		}
+	}
+	return false
+}

server/routes.go

@@ -72,7 +72,7 @@ var (
 	errBadTemplate = errors.New("template error")
 )
 
-func modelOptions(model *Model, requestOpts map[string]any) (api.Options, error) {
+func modelOptions(model *Model, requestOpts map[string]interface{}) (api.Options, error) {
 	opts := api.DefaultOptions()
 	if err := opts.FromMap(model.Options); err != nil {
 		return api.Options{}, err
@@ -87,7 +87,7 @@ func modelOptions(model *Model, requestOpts map[string]any) (api.Options, error)
 
 // scheduleRunner schedules a runner after validating inputs such as capabilities and model options.
 // It returns the allocated runner, model instance, and consolidated options if successful and error otherwise.
-func (s *Server) scheduleRunner(ctx context.Context, name string, caps []model.Capability, requestOpts map[string]any, keepAlive *api.Duration) (llm.LlamaServer, *Model, *api.Options, error) {
+func (s *Server) scheduleRunner(ctx context.Context, name string, caps []Capability, requestOpts map[string]any, keepAlive *api.Duration) (llm.LlamaServer, *Model, *api.Options, error) {
 	if name == "" {
 		return nil, nil, nil, fmt.Errorf("model %w", errRequired)
 	}
@@ -144,7 +144,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		return
 	}
 
-	m, err := GetModel(name.String())
+	model, err := GetModel(name.String())
 	if err != nil {
 		switch {
 		case errors.Is(err, fs.ErrNotExist):
@@ -159,7 +159,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 
 	// expire the runner
 	if req.Prompt == "" && req.KeepAlive != nil && int(req.KeepAlive.Seconds()) == 0 {
-		s.sched.expireRunner(m)
+		s.sched.expireRunner(model)
 
 		c.JSON(http.StatusOK, api.GenerateResponse{
 			Model:      req.Model,
@@ -176,9 +176,9 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		return
 	}
 
-	caps := []model.Capability{model.CapabilityCompletion}
+	caps := []Capability{CapabilityCompletion}
 	if req.Suffix != "" {
-		caps = append(caps, model.CapabilityInsert)
+		caps = append(caps, CapabilityInsert)
 	}
 
 	r, m, opts, err := s.scheduleRunner(c.Request.Context(), name.String(), caps, req.Options, req.KeepAlive)
@@ -203,7 +203,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		return
 	}
 
-	isMllama := checkMllamaModelFamily(m)
+	isMllama := checkMllamaModelFamily(model)
 	if isMllama && len(req.Images) > 1 {
 		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "this model only supports one image: more than one image sent"})
 		return
@@ -211,7 +211,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 
 	images := make([]llm.ImageData, len(req.Images))
 	for i := range req.Images {
-		if isMllama && len(m.ProjectorPaths) > 0 {
+		if isMllama && len(model.ProjectorPaths) > 0 {
 			data, opts, err := mllama.Preprocess(bytes.NewReader(req.Images[i]))
 			if err != nil {
 				c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "error processing image"})
@@ -422,7 +422,7 @@ func (s *Server) EmbedHandler(c *gin.Context) {
 		return
 	}
 
-	r, m, opts, err := s.scheduleRunner(c.Request.Context(), name.String(), []model.Capability{}, req.Options, req.KeepAlive)
+	r, m, opts, err := s.scheduleRunner(c.Request.Context(), name.String(), []Capability{}, req.Options, req.KeepAlive)
 	if err != nil {
 		handleScheduleError(c, req.Model, err)
 		return
@@ -435,7 +435,7 @@ func (s *Server) EmbedHandler(c *gin.Context) {
 		return
 	}
 
-	kvData, _, err := getModelData(m.ModelPath, false)
+	kvData, err := getKVData(m.ModelPath, false)
 	if err != nil {
 		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
 		return
@@ -483,7 +483,8 @@ func (s *Server) EmbedHandler(c *gin.Context) {
 	}
 
 	if err := g.Wait(); err != nil {
-		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": strings.TrimSpace(err.Error())})
+		slog.Error("embedding generation failed", "error", err)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Errorf("failed to generate embeddings: %v", err)})
 		return
 	}
 
@@ -530,7 +531,7 @@ func (s *Server) EmbeddingsHandler(c *gin.Context) {
 		return
 	}
 
-	r, _, _, err := s.scheduleRunner(c.Request.Context(), name.String(), []model.Capability{}, req.Options, req.KeepAlive)
+	r, _, _, err := s.scheduleRunner(c.Request.Context(), name.String(), []Capability{}, req.Options, req.KeepAlive)
 	if err != nil {
 		handleScheduleError(c, req.Model, err)
 		return
@@ -544,7 +545,8 @@ func (s *Server) EmbeddingsHandler(c *gin.Context) {
 
 	embedding, err := r.Embedding(c.Request.Context(), req.Prompt)
 	if err != nil {
-		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": strings.TrimSpace(err.Error())})
+		slog.Info(fmt.Sprintf("embedding generation failed: %v", err))
+		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Errorf("failed to generate embedding: %v", err)})
 		return
 	}
 
@@ -777,7 +779,7 @@ func (s *Server) ShowHandler(c *gin.Context) {
 func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 	name := model.ParseName(req.Model)
 	if !name.IsValid() {
-		return nil, ErrModelPathInvalid
+		return nil, errModelPathInvalid
 	}
 	name, err := getExistingName(name)
 	if err != nil {
@@ -813,20 +815,19 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 	}
 
 	resp := &api.ShowResponse{
-		License:      strings.Join(m.License, "\n"),
-		System:       m.System,
-		Template:     m.Template.String(),
-		Details:      modelDetails,
-		Messages:     msgs,
-		Capabilities: m.Capabilities(),
-		ModifiedAt:   manifest.fi.ModTime(),
+		License:    strings.Join(m.License, "\n"),
+		System:     m.System,
+		Template:   m.Template.String(),
+		Details:    modelDetails,
+		Messages:   msgs,
+		ModifiedAt: manifest.fi.ModTime(),
 	}
 
 	var params []string
 	cs := 30
 	for k, v := range m.Options {
 		switch val := v.(type) {
-		case []any:
+		case []interface{}:
 			for _, nv := range val {
 				params = append(params, fmt.Sprintf("%-*s %#v", cs, k, nv))
 			}
@@ -849,23 +850,16 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 	fmt.Fprint(&sb, m.String())
 	resp.Modelfile = sb.String()
 
-	kvData, tensors, err := getModelData(m.ModelPath, req.Verbose)
+	kvData, err := getKVData(m.ModelPath, req.Verbose)
 	if err != nil {
 		return nil, err
 	}
-
 	delete(kvData, "general.name")
 	delete(kvData, "tokenizer.chat_template")
 	resp.ModelInfo = kvData
 
-	tensorData := make([]api.Tensor, len(tensors.Items()))
-	for cnt, t := range tensors.Items() {
-		tensorData[cnt] = api.Tensor{Name: t.Name, Type: t.Type(), Shape: t.Shape}
-	}
-	resp.Tensors = tensorData
-
 	if len(m.ProjectorPaths) > 0 {
-		projectorData, _, err := getModelData(m.ProjectorPaths[0], req.Verbose)
+		projectorData, err := getKVData(m.ProjectorPaths[0], req.Verbose)
 		if err != nil {
 			return nil, err
 		}
@@ -875,17 +869,17 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
 	return resp, nil
 }
 
-func getModelData(digest string, verbose bool) (ggml.KV, ggml.Tensors, error) {
+func getKVData(digest string, verbose bool) (ggml.KV, error) {
 	maxArraySize := 0
 	if verbose {
 		maxArraySize = -1
 	}
-	data, err := llm.LoadModel(digest, maxArraySize)
+	kvData, err := llm.LoadModel(digest, maxArraySize)
 	if err != nil {
-		return nil, ggml.Tensors{}, err
+		return nil, err
 	}
 
-	kv := data.KV()
+	kv := kvData.KV()
 
 	if !verbose {
 		for k := range kv {
@@ -895,7 +889,7 @@ func getModelData(digest string, verbose bool) (ggml.KV, ggml.Tensors, error) {
 		}
 	}
 
-	return kv, data.Tensors(), nil
+	return kv, nil
 }
 
 func (s *Server) ListHandler(c *gin.Context) {
@@ -1336,7 +1330,7 @@ func Serve(ln net.Listener) error {
 	return nil
 }
 
-func waitForStream(c *gin.Context, ch chan any) {
+func waitForStream(c *gin.Context, ch chan interface{}) {
 	c.Header("Content-Type", "application/json")
 	for resp := range ch {
 		switch r := resp.(type) {
@@ -1469,9 +1463,9 @@ func (s *Server) ChatHandler(c *gin.Context) {
 		return
 	}
 
-	caps := []model.Capability{model.CapabilityCompletion}
+	caps := []Capability{CapabilityCompletion}
 	if len(req.Tools) > 0 {
-		caps = append(caps, model.CapabilityTools)
+		caps = append(caps, CapabilityTools)
 	}
 
 	name := model.ParseName(req.Model)

server/sched.go

@@ -20,7 +20,6 @@ import (
 	"github.com/ollama/ollama/format"
 	"github.com/ollama/ollama/fs/ggml"
 	"github.com/ollama/ollama/llm"
-	"github.com/ollama/ollama/types/model"
 )
 
 type LlmRequest struct {
@@ -38,7 +37,7 @@ type Scheduler struct {
 	pendingReqCh  chan *LlmRequest
 	finishedReqCh chan *LlmRequest
 	expiredCh     chan *runnerRef
-	unloadedCh    chan any
+	unloadedCh    chan interface{}
 
 	loaded   map[string]*runnerRef
 	loadedMu sync.Mutex
@@ -68,7 +67,7 @@ func InitScheduler(ctx context.Context) *Scheduler {
 		pendingReqCh:  make(chan *LlmRequest, maxQueue),
 		finishedReqCh: make(chan *LlmRequest, maxQueue),
 		expiredCh:     make(chan *runnerRef, maxQueue),
-		unloadedCh:    make(chan any, maxQueue),
+		unloadedCh:    make(chan interface{}, maxQueue),
 		loaded:        make(map[string]*runnerRef),
 		newServerFn:   llm.NewLlamaServer,
 		getGpuFn:      discover.GetGPUInfo,
@@ -196,7 +195,7 @@ func (s *Scheduler) processPending(ctx context.Context) {
 					}
 
 					// Embedding models should always be loaded with parallel=1
-					if pending.model.CheckCapabilities(model.CapabilityCompletion) != nil {
+					if pending.model.CheckCapabilities(CapabilityCompletion) != nil {
 						numParallel = 1
 					}
 
@@ -618,8 +617,8 @@ func (runner *runnerRef) needsReload(ctx context.Context, req *LlmRequest) bool
 // a before and after GPU memory allocation.  The returned channel
 // will be notified when we're done waiting, or have timed out and should
 // proceed anyway
-func (runner *runnerRef) waitForVRAMRecovery() chan any {
-	finished := make(chan any, 1)
+func (runner *runnerRef) waitForVRAMRecovery() chan interface{} {
+	finished := make(chan interface{}, 1)
 
 	// CPU or Metal don't need checking, so no waiting required
 	// windows can page VRAM, only cuda currently can report accurate used vram usage
@@ -712,7 +711,7 @@ func pickBestFullFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.GpuIn
 			req.opts.NumCtx = req.origNumCtx * p
 			if !envconfig.SchedSpread() {
 				for _, g := range sgl {
-					if ok, estimatedVRAM = llm.PredictServerFit([]discover.GpuInfo{g}, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, p); ok {
+					if ok, estimatedVRAM = llm.PredictServerFit([]discover.GpuInfo{g}, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts); ok {
 						slog.Info("new model will fit in available VRAM in single GPU, loading", "model", req.model.ModelPath, "gpu", g.ID, "parallel", p, "available", g.FreeMemory, "required", format.HumanBytes2(estimatedVRAM))
 						*numParallel = p
 						return []discover.GpuInfo{g}
@@ -728,7 +727,7 @@ func pickBestFullFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.GpuIn
 		// Now try all the GPUs
 		for _, p := range numParallelToTry {
 			req.opts.NumCtx = req.origNumCtx * p
-			if ok, estimatedVRAM = llm.PredictServerFit(sgl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, p); ok {
+			if ok, estimatedVRAM = llm.PredictServerFit(sgl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts); ok {
 				slog.Info("new model will fit in available VRAM, loading", "model", req.model.ModelPath, "library", sgl[0].Library, "parallel", p, "required", format.HumanBytes2(estimatedVRAM))
 				*numParallel = p
 				return sgl
@@ -751,7 +750,7 @@ func pickBestPartialFitByLibrary(req *LlmRequest, f *ggml.GGML, gpus discover.Gp
 	var bestEstimate uint64
 	var bestFit int
 	for i, gl := range byLibrary {
-		_, estimatedVRAM := llm.PredictServerFit(gl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts, *numParallel)
+		_, estimatedVRAM := llm.PredictServerFit(gl, f, req.model.AdapterPaths, req.model.ProjectorPaths, req.opts)
 		if estimatedVRAM > bestEstimate {
 			bestEstimate = estimatedVRAM
 			bestFit = i
@@ -826,7 +825,7 @@ func (s *Scheduler) expireRunner(model *Model) {
 // If not, pick a runner to unload, else return nil and the request can be loaded
 func (s *Scheduler) maybeFindCPURunnerToUnload(req *LlmRequest, f *ggml.GGML, gpus discover.GpuInfoList) *runnerRef {
 	slog.Debug("evaluating if CPU model load will fit in available system memory")
-	estimate := llm.EstimateGPULayers(gpus, f, req.model.ProjectorPaths, req.opts, req.opts.NumCtx/req.origNumCtx)
+	estimate := llm.EstimateGPULayers(gpus, f, req.model.ProjectorPaths, req.opts)
 	if estimate.TotalSize <= gpus[0].FreeMemory {
 		slog.Debug("cpu inference mode, model fits in available system memory", "model", format.HumanBytes2(estimate.TotalSize), "available", format.HumanBytes2(gpus[0].FreeMemory))
 		return nil

template/gemma3-instruct.gotmpl

@@ -1,13 +0,0 @@
-{{- range $i, $_ := .Messages }}
-{{- $last := eq (len (slice $.Messages $i)) 1 }}
-{{- if eq .Role "user" }}<start_of_turn>user
-{{- if and (eq $i 1) $.System }}
-{{ $.System }}
-{{ end }}
-{{ .Content }}<end_of_turn>
-{{ else if eq .Role "assistant" }}<start_of_turn>model
-{{ .Content }}<end_of_turn>
-{{ end }}
-{{- if $last }}<start_of_turn>model
-{{ end }}
-{{- end }}

template/gemma3-instruct.json

@@ -1,6 +0,0 @@
-{
-  "stop": [
-    "<end_of_turn>"
-  ],
-  "temperature": 0.1
-}

template/index.json

@@ -87,10 +87,6 @@
     "template": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}",
     "name": "gemma-instruct"
   },
-  {
-    "template": "{{ bos_token }}\n{%- if messages[0]['role'] == 'system' -%}\n    {%- if messages[0]['content'] is string -%}\n        {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}\n    {%- else -%}\n        {%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}\n    {%- endif -%}\n    {%- set loop_messages = messages[1:] -%}\n{%- else -%}\n    {%- set first_user_prefix = \"\" -%}\n    {%- set loop_messages = messages -%}\n{%- endif -%}\n{%- for message in loop_messages -%}\n    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}\n        {{ raise_exception(\"Conversation roles must alternate user/assistant/user/assistant/...\") }}\n    {%- endif -%}\n    {%- if (message['role'] == 'assistant') -%}\n        {%- set role = \"model\" -%}\n    {%- else -%}\n        {%- set role = message['role'] -%}\n    {%- endif -%}\n    {{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else \"\") }}\n    {%- if message['content'] is string -%}\n        {{ message['content'] | trim }}\n    {%- elif message['content'] is iterable -%}\n        {%- for item in message['content'] -%}\n            {%- if item['type'] == 'image' -%}\n                {{ '<start_of_image>' }}\n            {%- elif item['type'] == 'text' -%}\n                {{ item['text'] | trim }}\n            {%- endif -%}\n        {%- endfor -%}\n    {%- else -%}\n        {{ raise_exception(\"Invalid content type\") }}\n    {%- endif -%}\n    {{ '<end_of_turn>\n' }}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n    {{'<start_of_turn>model\n'}}\n{%- endif -%}\n",
-    "name": "gemma3-instruct"
-  },
   {
     "template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
     "name": "llama3-instruct"

template/testdata/gemma3-instruct.gotmpl/system-user-assistant-user

@@ -1,10 +0,0 @@
-<start_of_turn>user
-You are a helpful assistant.
-
-Hello, how are you?<end_of_turn>
-<start_of_turn>model
-I'm doing great. How can I help you today?<end_of_turn>
-<start_of_turn>user
-I'd like to show off how chat templating works!<end_of_turn>
-<start_of_turn>model
-

template/testdata/gemma3-instruct.gotmpl/user

@@ -1,4 +0,0 @@
-<start_of_turn>user
-Hello, how are you?<end_of_turn>
-<start_of_turn>model
-

template/testdata/gemma3-instruct.gotmpl/user-assistant-user

@@ -1,8 +0,0 @@
-<start_of_turn>user
-Hello, how are you?<end_of_turn>
-<start_of_turn>model
-I'm doing great. How can I help you today?<end_of_turn>
-<start_of_turn>user
-I'd like to show off how chat templating works!<end_of_turn>
-<start_of_turn>model
-

types/model/capability.go

@@ -1,15 +0,0 @@
-package model
-
-type Capability string
-
-const (
-	CapabilityCompletion = Capability("completion")
-	CapabilityTools      = Capability("tools")
-	CapabilityInsert     = Capability("insert")
-	CapabilityVision     = Capability("vision")
-	CapabilityEmbedding  = Capability("embedding")
-)
-
-func (c Capability) String() string {
-	return string(c)
-}