server/internal/client/ollama: hold DiskCache on Registry

Previously, clients of a Registry had to carry around a DiskCache to use
it. This change makes the DiskCache an optional field on the Registry
struct.

This also changes DefaultCache to initialize one on first use. This
prevents overhead of building the cache if it is never used, or per
Registry request that involves use of DefaultCache.

Also, slip in some minor docs on Trace.

Dockerfile

@@ -12,7 +12,7 @@ FROM --platform=linux/amd64 rocm/dev-almalinux-8:${ROCMVERSION}-complete AS base
 RUN yum install -y yum-utils \
     && yum-config-manager --add-repo https://dl.rockylinux.org/vault/rocky/8.5/AppStream/\$basearch/os/ \
     && rpm --import https://dl.rockylinux.org/pub/rocky/RPM-GPG-KEY-Rocky-8 \
-    && dnf install -y yum-utils ccache gcc-toolset-10-gcc-10.2.1-8.2.el8 gcc-toolset-10-gcc-c++-10.2.1-8.2.el8 gcc-toolset-10-binutils-2.35-11.el8 \
+    && dnf install -y yum-utils ccache gcc-toolset-10-gcc-10.2.1-8.2.el8 gcc-toolset-10-gcc-c++-10.2.1-8.2.el8 \
     && yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
 ENV PATH=/opt/rh/gcc-toolset-10/root/usr/bin:$PATH
 
@@ -86,11 +86,10 @@ RUN --mount=type=cache,target=/root/.ccache \
         && cmake --install build --component CUDA --strip --parallel 8
 
 FROM base AS build
-WORKDIR /go/src/github.com/ollama/ollama
-COPY go.mod go.sum .
-RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux-$(case $(uname -m) in x86_64) echo amd64 ;; aarch64) echo arm64 ;; esac).tar.gz | tar xz -C /usr/local
+ARG GOVERSION=1.23.4
+RUN curl -fsSL https://golang.org/dl/go${GOVERSION}.linux-$(case $(uname -m) in x86_64) echo amd64 ;; aarch64) echo arm64 ;; esac).tar.gz | tar xz -C /usr/local
 ENV PATH=/usr/local/go/bin:$PATH
-RUN go mod download
+WORKDIR /go/src/github.com/ollama/ollama
 COPY . .
 ARG GOFLAGS="'-ldflags=-w -s'"
 ENV CGO_ENABLED=1

README.md

@@ -1,5 +1,5 @@
 <div align="center">
-  <a href="https://ollama.com">
+  <a href="https://ollama.com" />
     <img alt="ollama" height="200px" src="https://github.com/ollama/ollama/assets/3325447/0d0b44e2-8f4a-4e99-9b52-a5c1c741c8f7">
   </a>
 </div>
@@ -54,7 +54,6 @@ Here are some example models that can be downloaded:
 
 | Model              | Parameters | Size  | Download                         |
 | ------------------ | ---------- | ----- | -------------------------------- |
-| QwQ                | 32B        | 20GB  | `ollama run qwq`                 |
 | DeepSeek-R1        | 7B         | 4.7GB | `ollama run deepseek-r1`         |
 | DeepSeek-R1        | 671B       | 404GB | `ollama run deepseek-r1:671b`    |
 | Llama 3.3          | 70B        | 43GB  | `ollama run llama3.3`            |
@@ -65,7 +64,7 @@ Here are some example models that can be downloaded:
 | Llama 3.1          | 8B         | 4.7GB | `ollama run llama3.1`            |
 | Llama 3.1          | 405B       | 231GB | `ollama run llama3.1:405b`       |
 | Phi 4              | 14B        | 9.1GB | `ollama run phi4`                |
-| Phi 4 Mini         | 3.8B       | 2.5GB | `ollama run phi4-mini`           |
+| Phi 3 Mini         | 3.8B       | 2.3GB | `ollama run phi3`                |
 | Gemma 2            | 2B         | 1.6GB | `ollama run gemma2:2b`           |
 | Gemma 2            | 9B         | 5.5GB | `ollama run gemma2`              |
 | Gemma 2            | 27B        | 16GB  | `ollama run gemma2:27b`          |
@@ -76,7 +75,7 @@ Here are some example models that can be downloaded:
 | Code Llama         | 7B         | 3.8GB | `ollama run codellama`           |
 | Llama 2 Uncensored | 7B         | 3.8GB | `ollama run llama2-uncensored`   |
 | LLaVA              | 7B         | 4.5GB | `ollama run llava`               |
-| Granite-3.2         | 8B         | 4.9GB | `ollama run granite3.2`          |
+| Solar              | 10.7B      | 6.1GB | `ollama run solar`               |
 
 > [!NOTE]
 > You should have at least 8 GB of RAM available to run the 7B models, 16 GB to run the 13B models, and 32 GB to run the 33B models.
@@ -276,7 +275,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 ### Web & Desktop
 
 - [Open WebUI](https://github.com/open-webui/open-webui)
-- [SwiftChat (macOS with ReactNative)](https://github.com/aws-samples/swift-chat)
 - [Enchanted (macOS native)](https://github.com/AugustDev/enchanted)
 - [Hollama](https://github.com/fmaclen/hollama)
 - [Lollms-Webui](https://github.com/ParisNeo/lollms-webui)
@@ -390,7 +388,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [LangBot](https://github.com/RockChinQ/LangBot) (LLM-based instant messaging bots platform, with Agents, RAG features, supports multiple platforms)
 - [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.)
 
 ### Cloud
 
@@ -434,7 +431,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 
 ### Apple Vision Pro
 
-- [SwiftChat](https://github.com/aws-samples/swift-chat) (Cross-platform AI chat app supporting Apple Vision Pro via "Designed for iPad")
 - [Enchanted](https://github.com/AugustDev/enchanted)
 
 ### Database
@@ -512,13 +508,10 @@ See the [API documentation](./docs/api.md) for all endpoints.
 
 ### Mobile
 
-- [SwiftChat](https://github.com/aws-samples/swift-chat) (Lightning-fast Cross-platform AI chat app with native UI for Android, iOS and iPad)
 - [Enchanted](https://github.com/AugustDev/enchanted)
 - [Maid](https://github.com/Mobile-Artificial-Intelligence/maid)
 - [Ollama App](https://github.com/JHubi1/ollama-app) (Modern and easy-to-use multi-platform client for Ollama)
 - [ConfiChat](https://github.com/1runeberg/confichat) (Lightweight, standalone, multi-platform, and privacy focused LLM chat interface with optional encryption)
-- [Ollama Android Chat](https://github.com/sunshine0523/OllamaServer) (No need for Termux, start the Ollama service with one click on an Android device)
-- [Reins](https://github.com/ibrahimcetin/reins) (Easily tweak parameters, customize system prompts per chat, and enhance your AI experiments with reasoning model support.)
 
 ### Extensions & Plugins
 
@@ -564,7 +557,6 @@ See the [API documentation](./docs/api.md) for all endpoints.
 - [TextLLaMA](https://github.com/adarshM84/TextLLaMA) A Chrome Extension that helps you write emails, correct grammar, and translate into any language
 - [Simple-Discord-AI](https://github.com/zyphixor/simple-discord-ai)
 - [LLM Telegram Bot](https://github.com/innightwolfsleep/llm_telegram_bot) (telegram bot, primary for RP. Oobabooga-like buttons, [A1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui) API integration e.t.c)
-- [mcp-llm](https://github.com/sammcj/mcp-llm) (MCP Server to allow LLMs to call other LLMs)
 
 ### Supported backends
 

api/types.go

@@ -361,9 +361,9 @@ type CopyRequest struct {
 // PullRequest is the request passed to [Client.Pull].
 type PullRequest struct {
 	Model    string `json:"model"`
-	Insecure bool   `json:"insecure,omitempty"` // Deprecated: ignored
-	Username string `json:"username"`           // Deprecated: ignored
-	Password string `json:"password"`           // Deprecated: ignored
+	Insecure bool   `json:"insecure,omitempty"`
+	Username string `json:"username"`
+	Password string `json:"password"`
 	Stream   *bool  `json:"stream,omitempty"`
 
 	// Deprecated: set the model name with Model instead

cmd/cmd.go

@@ -34,6 +34,7 @@ import (
 	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/format"
+	"github.com/ollama/ollama/llama"
 	"github.com/ollama/ollama/parser"
 	"github.com/ollama/ollama/progress"
 	"github.com/ollama/ollama/runner"
@@ -255,7 +256,6 @@ func StopHandler(cmd *cobra.Command, args []string) error {
 		if strings.Contains(err.Error(), "not found") {
 			return fmt.Errorf("couldn't find model \"%s\" to stop", args[0])
 		}
-		return err
 	}
 	return nil
 }
@@ -338,16 +338,10 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if len(info.ProjectorInfo) != 0 {
-		opts.MultiModal = true
-	}
-	for k := range info.ModelInfo {
-		if strings.Contains(k, ".vision.") {
-			opts.MultiModal = true
-			break
-		}
-	}
-
+	// TODO(jessegross): We should either find another way to know if this is
+	// a vision model or remove the logic. Also consider that other modalities will
+	// need different behavior anyways.
+	opts.MultiModal = len(info.ProjectorInfo) != 0 || envconfig.NewEngine()
 	opts.ParentModel = info.Details.ParentModel
 
 	if interactive {
@@ -1280,6 +1274,7 @@ func NewCLI() *cobra.Command {
 
 	runnerCmd := &cobra.Command{
 		Use:    "runner",
+		Short:  llama.PrintSystemInfo(),
 		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			return runner.Execute(os.Args[1:])

docs/development.md

@@ -118,35 +118,6 @@ To run tests, use `go test`:
 go test ./...
 ```
 
-> NOTE: In rare cirumstances, you may nedd to change a package using the new
-> "synctest" package in go1.24.
->
-> If you do not have the "synctest" package enabled, you will not see build or
-> test failures resulting from your change(s), if any, locally, but CI will
-> break.
->
-> If you see failures in CI, you can either keep pushing changes to see if the
-> CI build passes, or you can enable the "synctest" package locally to see the
-> failures before pushing.
->
-> To enable the "synctest" package for testing, run the following command:
->
-> ```shell
-> GOEXPERIMENT=synctest go test ./...
-> ```
->
-> If you wish to enable synctest for all go commands, you can set the
-> `GOEXPERIMENT` environment variable in your shell profile or by using:
->
-> ```shell
-> go env -w GOEXPERIMENT=synctest
-> ```
->
-> Which will enable the "synctest" package for all go commands without needing
-> to set it for all shell sessions.
->
-> The synctest package is not required for production builds.
-
 ## Library detection
 
 Ollama looks for acceleration libraries in the following paths relative to the `ollama` executable:

docs/faq.md

@@ -20,7 +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 length to 8K, use: `OLLAMA_CONTEXT_LENGTH=8192 ollama serve`.
+By default, Ollama uses a context window size of 2048 tokens.
 
 To change this when using `ollama run`, use `/set parameter`:
 

docs/windows.md

@@ -81,11 +81,9 @@ help you keep up to date.
 
 If you'd like to install or integrate Ollama as a service, a standalone
 `ollama-windows-amd64.zip` zip file is available containing only the Ollama CLI
-and GPU library dependencies for Nvidia.  If you have an AMD GPU, also download
-and extract the additional ROCm package `ollama-windows-amd64-rocm.zip` into the
-same directory.  This allows for embedding Ollama in existing applications, or
-running it as a system service via `ollama serve` with tools such as
-[NSSM](https://nssm.cc/). 
+and GPU library dependencies for Nvidia and AMD. This allows for embedding
+Ollama in existing applications, or running it as a system service via `ollama
+serve` with tools such as [NSSM](https://nssm.cc/).
 
 > [!NOTE]  
 > If you are upgrading from a prior version, you should remove the old directories first.

fs/ggml/ggml.go

@@ -565,43 +565,6 @@ func (f GGML) GraphSize(context, batch uint64, kvCacheType string) (kv, partialO
 	return
 }
 
-func (llm GGML) VisionGraphSize() (weights, graphSize uint64) {
-	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
-
-		graphSize = 4 * (8 +
-			imageSize*imageSize*kv("num_channels")*maxNumTiles +
-			embeddingLength*numPatches*maxNumTiles +
-			9*embeddingLength*numPaddedPatches*maxNumTiles +
-			numPaddedPatches*maxNumTiles*numPaddedPatches*maxNumTiles*headCount)
-	}
-	return weights, graphSize
-}
-
 // SupportsKVCacheType checks if the requested cache type is supported
 func (f GGML) SupportsKVCacheType(cacheType string) bool {
 	return slices.Contains([]string{"f16", "q8_0", "q4_0"}, cacheType)

go.mod

@@ -24,7 +24,7 @@ require (
 	github.com/nlpodyssey/gopickle v0.3.0
 	github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c
 	golang.org/x/image v0.22.0
-	golang.org/x/tools v0.30.0
+	gonum.org/v1/gonum v0.15.0
 )
 
 require (
@@ -44,7 +44,6 @@ require (
 	github.com/xtgo/set v1.0.0 // indirect
 	go4.org/unsafe/assume-no-moving-gc v0.0.0-20231121144256-b99613f794b6 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
-	gonum.org/v1/gonum v0.15.0 // indirect
 	gorgonia.org/vecf32 v0.9.0 // indirect
 	gorgonia.org/vecf64 v0.9.0 // indirect
 )

go.sum

@@ -309,8 +309,6 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.4/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
-golang.org/x/tools v0.30.0 h1:BgcpHewrV5AUp2G9MebG4XPFI1E2W41zU1SaqVA9vJY=
-golang.org/x/tools v0.30.0/go.mod h1:c347cR/OJfw5TI+GfX7RUPNMdDRRbjvYTS0jPyvsVtY=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

kvcache/cache.go

@@ -4,7 +4,6 @@ import (
 	"errors"
 
 	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
 )
 
 var (
@@ -52,7 +51,7 @@ 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, opts input.Options) error
+	StartForward(ctx ml.Context, positions []int32, seqs []int) error
 
 	// CopyPrefix copies tokens in the range [0, len) from srcSeq to dstSeq
 	CopyPrefix(srcSeq, dstSeq int, len int32)

kvcache/causal.go

@@ -8,7 +8,6 @@ import (
 	"slices"
 
 	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
 )
 
 type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error)
@@ -21,7 +20,6 @@ type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, e
 type Causal struct {
 	DType      ml.DType
 	Capacity   int32
-	causal     bool
 	windowSize int32
 
 	// config controls mostly backend-specific optimizations
@@ -44,12 +42,6 @@ type Causal struct {
 	// locations in the cache that are needed for this batch
 	curCellRange cellRange
 
-	// curSequences is the sequences corresponding to this pass's entries in the cache
-	curSequences []int
-
-	// curPositions is the positions corresponding to this pass's entries in the cache
-	curPositions []int32
-
 	// ** cache metadata **
 
 	// for each possible location in the cache, stores the position and set of sequences
@@ -63,8 +55,8 @@ type Causal struct {
 
 	shiftFn      shiftFn
 	backend      ml.Backend
-	ctxs         map[int]ml.Context
-	keys, values map[int]ml.Tensor
+	cacheCtx     ml.Context
+	keys, values []ml.Tensor
 }
 
 type cacheCell struct {
@@ -78,25 +70,11 @@ type cellRange struct {
 }
 
 func NewCausalCache(shift shiftFn) *Causal {
-	return &Causal{
-		causal:     true,
-		windowSize: math.MaxInt32,
-		shiftFn:    shift,
-		ctxs:       make(map[int]ml.Context),
-		keys:       make(map[int]ml.Tensor),
-		values:     make(map[int]ml.Tensor),
-	}
+	return &Causal{windowSize: math.MaxInt32, shiftFn: shift}
 }
 
 func NewSWACache(windowSize int32, shift shiftFn) *Causal {
-	return &Causal{
-		causal:     true,
-		windowSize: windowSize,
-		shiftFn:    shift,
-		ctxs:       make(map[int]ml.Context),
-		keys:       make(map[int]ml.Tensor),
-		values:     make(map[int]ml.Tensor),
-	}
+	return &Causal{windowSize: windowSize, shiftFn: shift}
 }
 
 func (c *Causal) Init(backend ml.Backend, dtype ml.DType, capacity int32) {
@@ -125,6 +103,7 @@ func (c *Causal) Init(backend ml.Backend, dtype ml.DType, capacity int32) {
 	c.cells = make([]cacheCell, c.Capacity)
 	c.cellRanges = make(map[int]cellRange)
 	c.backend = backend
+	c.cacheCtx = backend.NewContext()
 }
 
 func (c *Causal) SetConfig(config ml.CacheConfig) {
@@ -136,15 +115,11 @@ func (c *Causal) SetConfig(config ml.CacheConfig) {
 }
 
 func (c *Causal) Close() {
-	for _, ctx := range c.ctxs {
-		ctx.Close()
-	}
+	c.cacheCtx.Close()
 }
 
-func (c *Causal) StartForward(ctx ml.Context, opts input.Options) error {
-	c.curBatchSize = len(opts.Positions)
-	c.curSequences = opts.Sequences
-	c.curPositions = opts.Positions
+func (c *Causal) StartForward(ctx ml.Context, positions []int32, seqs []int) error {
+	c.curBatchSize = len(positions)
 
 	var err error
 	c.curLoc, err = c.findStartLoc()
@@ -157,8 +132,8 @@ func (c *Causal) StartForward(ctx ml.Context, opts input.Options) error {
 	}
 
 	c.curCellRange = newRange()
-	for i, pos := range opts.Positions {
-		seq := opts.Sequences[i]
+	for i, pos := range positions {
+		seq := seqs[i]
 
 		c.cells[c.curLoc+i] = cacheCell{pos: pos, sequences: []int{seq}}
 
@@ -183,7 +158,7 @@ func (c *Causal) StartForward(ctx ml.Context, opts input.Options) error {
 		c.cellRanges[seq] = seqRange
 	}
 
-	c.curMask, err = c.buildMask(ctx)
+	c.curMask, err = c.buildMask(ctx, positions, seqs)
 
 	return err
 }
@@ -224,7 +199,7 @@ func roundUp(length, pad int) int {
 // Builds a mask of history x batch indicating whether for each token in the batch the
 // token in the history should apply. This is based on both the sequence and causality (the
 // position of the history is not ahead of the token in the batch).
-func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) {
+func (c *Causal) buildMask(ctx ml.Context, positions []int32, seqs []int) (ml.Tensor, error) {
 	// Align and pad the two dimensions as required by the backend
 	batchSize := roundUp(c.curBatchSize, c.config.MaskBatchPadding)
 
@@ -236,9 +211,8 @@ func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) {
 
 	for i := range c.curBatchSize {
 		for j := c.curCellRange.min; j <= c.curCellRange.max; j++ {
-			if !slices.Contains(c.cells[j].sequences, c.curSequences[i]) ||
-				(c.causal && c.cells[j].pos > c.curPositions[i]) ||
-				c.cells[j].pos < c.curPositions[i]-c.windowSize {
+			if !slices.Contains(c.cells[j].sequences, seqs[i]) || c.cells[j].pos > positions[i] ||
+				c.cells[j].pos < positions[i]-c.windowSize {
 				mask[i*length+(j-c.curCellRange.min)] = float32(math.Inf(-1))
 			}
 		}
@@ -250,13 +224,13 @@ func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) {
 		mask[i] = float32(math.Inf(-1))
 	}
 
-	maskTensor, err := ctx.Input().FromFloatSlice(mask, length, batchSize)
+	maskTensor, err := ctx.FromFloatSlice(mask, length, batchSize)
 	if err != nil {
 		return nil, err
 	}
 
 	if c.config.MaskDType != ml.DTypeF32 {
-		out := ctx.Input().Empty(c.config.MaskDType, maskTensor.Shape()...)
+		out := ctx.Empty(c.config.MaskDType, maskTensor.Shape()...)
 		ctx.Forward(maskTensor.Copy(ctx, out))
 		maskTensor = out
 	}
@@ -265,11 +239,13 @@ func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) {
 }
 
 func (c *Causal) moveCells(ctx ml.Context, src, dst, len int) {
-	for i, key := range c.keys {
-		if key == nil {
+	for i := range c.keys {
+		if c.keys[i] == nil {
 			continue
 		}
 
+		key := c.keys[i]
+
 		kHeadDim := key.Dim(0)
 		numKVHeads := key.Dim(1)
 		rowSize := key.Stride(2)
@@ -329,7 +305,7 @@ func (c *Causal) defrag() {
 		layers++
 	}
 
-	maxMoves := ctx.MaxGraphNodes() / (6 * layers)
+	maxMoves := ctx.MaxTensors() / (6 * layers)
 	moves := 0
 
 	var pendingSrc, pendingDst, pendingLen int
@@ -401,27 +377,12 @@ func (c *Causal) defrag() {
 }
 
 func (c *Causal) SetLayer(layer int) {
-	c.curLayer = layer
-}
-
-// SetCausal enables or disables causal mask generation for subsequent calls to Get.
-// This state carries over to future forward passes. The default value is true.
-//
-// ctx may be set to nil if this is called from outside of a forward pass, for
-// example, when initializing the cache.
-func (c *Causal) SetCausal(ctx ml.Context, causal bool) {
-	if c.causal != causal {
-		c.causal = causal
-
-		if ctx != nil {
-			var err error
-			c.curMask, err = c.buildMask(ctx)
-			if err != nil {
-				// This error should never occur because we have previously built a mask with the same shape
-				panic(fmt.Errorf("SetCausal: %w", err))
-			}
-		}
+	if layer >= len(c.keys) {
+		c.keys = append(c.keys, make([]ml.Tensor, layer-len(c.keys)+1)...)
+		c.values = append(c.values, make([]ml.Tensor, layer-len(c.values)+1)...)
 	}
+
+	c.curLayer = layer
 }
 
 func (c *Causal) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) {
@@ -472,19 +433,13 @@ func (c *Causal) Put(ctx ml.Context, key, value ml.Tensor) {
 		panic(fmt.Errorf("inconsistent batch sizes (layer: %v, batch size: %v layer batch size: %v)", c.curLayer, c.curBatchSize, batchSize))
 	}
 
-	if _, ok := c.ctxs[c.curLayer]; !ok {
-		c.ctxs[c.curLayer] = c.backend.NewContextSize(2).Layer(c.curLayer)
-	}
+	if c.keys[c.curLayer] == nil || c.values[c.curLayer] == nil {
+		c.keys[c.curLayer] = c.cacheCtx.Zeros(c.DType, kHeadDim, numKVHeads, int(c.Capacity))
 
-	if _, ok := c.keys[c.curLayer]; !ok {
-		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, int(c.Capacity), vHeadDim, numKVHeads)
+			c.values[c.curLayer] = c.cacheCtx.Zeros(c.DType, int(c.Capacity), vHeadDim, numKVHeads)
 		} else {
-			c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, vHeadDim, numKVHeads, int(c.Capacity))
+			c.values[c.curLayer] = c.cacheCtx.Zeros(c.DType, vHeadDim, numKVHeads, int(c.Capacity))
 		}
 	}
 
@@ -546,7 +501,7 @@ func (c *Causal) shift(seq int, beginIndex, offset int32) error {
 		}
 	}
 
-	kShift, err := ctx.Input().FromIntSlice(offsets, len(offsets))
+	kShift, err := ctx.FromIntSlice(offsets, len(offsets))
 	if err != nil {
 		return err
 	}

kvcache/causal_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 
 	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
 )
 
 type testCase struct {
@@ -270,7 +269,7 @@ func testCache(t *testing.T, backend ml.Backend, cache Cache, tests []testCase)
 			context := backend.NewContext()
 			defer context.Close()
 
-			err := cache.StartForward(context, input.Options{Positions: test.pos, Sequences: test.seqs})
+			err := cache.StartForward(context, test.pos, test.seqs)
 			if err != nil {
 				panic(err)
 			}
@@ -304,10 +303,6 @@ func (b *testBackend) NewContext() ml.Context {
 	return &testContext{}
 }
 
-func (b *testBackend) NewContextSize(int) ml.Context {
-	return &testContext{}
-}
-
 func (b *testBackend) SystemInfo() string {
 	return "not implemented"
 }
@@ -351,15 +346,11 @@ func (c *testContext) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
 	return out, nil
 }
 
-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 }
 
 func (c *testContext) Compute(...ml.Tensor) {}
 
-func (c *testContext) MaxGraphNodes() int {
+func (c *testContext) MaxTensors() int {
 	return 10
 }
 

kvcache/encoder.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 
 	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
 )
 
 // Encoder cache stores K and V tensors that are position independent
@@ -36,17 +35,13 @@ type EncoderCache struct {
 	encoderPos int32
 
 	// ** cache data storage **
-	backend      ml.Backend
-	ctxs         map[int]ml.Context
-	keys, values map[int]ml.Tensor
+
+	cacheCtx     ml.Context
+	keys, values []ml.Tensor
 }
 
 func NewEncoderCache() *EncoderCache {
-	return &EncoderCache{
-		ctxs:   make(map[int]ml.Context),
-		keys:   make(map[int]ml.Tensor),
-		values: make(map[int]ml.Tensor),
-	}
+	return &EncoderCache{}
 }
 
 func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, capacity int32) {
@@ -62,7 +57,7 @@ func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, capacity int32)
 		panic(fmt.Errorf("encoder cache is unable to enforce requested CachePadding (%v)", c.config.CachePadding))
 	}
 
-	c.backend = backend
+	c.cacheCtx = backend.NewContext()
 }
 
 func (c *EncoderCache) SetConfig(config ml.CacheConfig) {
@@ -74,21 +69,22 @@ func (c *EncoderCache) SetConfig(config ml.CacheConfig) {
 }
 
 func (c *EncoderCache) Close() {
-	for _, ctx := range c.ctxs {
-		ctx.Close()
-	}
+	c.cacheCtx.Close()
 }
 
-func (c *EncoderCache) StartForward(ctx ml.Context, opts input.Options) error {
-	// We work with the most recent image
-	if len(opts.Multimodal) > 0 {
-		c.curPos = opts.Positions[opts.Multimodal[len(opts.Multimodal)-1].Index]
-	}
+func (c *EncoderCache) StartForward(ctx ml.Context, positions []int32, seqs []int) error {
+	// The image is always in the first position
+	c.curPos = positions[0]
 
 	return nil
 }
 
 func (c *EncoderCache) SetLayer(layer int) {
+	if layer >= len(c.keys) {
+		c.keys = append(c.keys, make([]ml.Tensor, layer-len(c.keys)+1)...)
+		c.values = append(c.values, make([]ml.Tensor, layer-len(c.values)+1)...)
+	}
+
 	c.curLayer = layer
 }
 
@@ -108,16 +104,9 @@ func (c *EncoderCache) Put(ctx ml.Context, key, value ml.Tensor) {
 		value = value.Permute(ctx, 1, 2, 0, 3)
 	}
 
-	if _, ok := c.ctxs[c.curLayer]; !ok {
-		c.ctxs[c.curLayer] = c.backend.NewContextSize(2).Layer(c.curLayer)
-	}
-
-	if _, ok := c.keys[c.curLayer]; !ok {
-		c.keys[c.curLayer] = c.ctxs[c.curLayer].Empty(key.DType(), key.Shape()...)
-	}
-
-	if _, ok := c.values[c.curLayer]; !ok {
-		c.values[c.curLayer] = c.ctxs[c.curLayer].Empty(value.DType(), value.Shape()...)
+	if c.keys[c.curLayer] == nil || c.values[c.curLayer] == nil {
+		c.keys[c.curLayer] = c.cacheCtx.Empty(key.DType(), key.Shape()...)
+		c.values[c.curLayer] = c.cacheCtx.Empty(value.DType(), value.Shape()...)
 	}
 
 	ctx.Forward(

kvcache/wrapper.go

@@ -4,7 +4,6 @@ import (
 	"math"
 
 	"github.com/ollama/ollama/ml"
-	"github.com/ollama/ollama/model/input"
 )
 
 // Wrapper cache is a container for multiple types of caches,
@@ -41,14 +40,14 @@ func (c *WrapperCache) Close() {
 	}
 }
 
-func (c *WrapperCache) StartForward(ctx ml.Context, opts input.Options) error {
+func (c *WrapperCache) StartForward(ctx ml.Context, positions []int32, seqs []int) error {
 	for i, cache := range c.caches {
-		err := cache.StartForward(ctx, opts)
+		err := cache.StartForward(ctx, positions, seqs)
 		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 opts.Positions {
-					_ = c.caches[j].Remove(opts.Sequences[k], opts.Positions[k], math.MaxInt32)
+				for k := range positions {
+					_ = c.caches[j].Remove(seqs[k], positions[k], math.MaxInt32)
 				}
 			}
 			return err

llama/llama.cpp/src/llama-vocab.cpp

@@ -1443,7 +1443,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
 
             const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str());
             if (precompiled_charsmap_keyidx != -1) {
-                size_t n_precompiled_charsmap = gguf_get_arr_data_n(ctx, precompiled_charsmap_keyidx);
+                size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
                 const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
                 precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap);
 #ifdef IS_BIG_ENDIAN

llama/llama.go

@@ -21,6 +21,18 @@ package llama
 
 extern bool llamaProgressCallback(float progress, void *user_data);
 extern void llamaLog(int level, char* text, void* user_data);
+
+typedef enum {COMP_UNKNOWN,COMP_GCC,COMP_CLANG} COMPILER;
+COMPILER inline get_compiler() {
+#if defined(__clang__)
+	return COMP_CLANG;
+#elif defined(__GNUC__)
+	return COMP_GCC;
+#else
+	return UNKNOWN_COMPILER;
+#endif
+}
+
 */
 import "C"
 
@@ -60,6 +72,19 @@ func BackendInit() {
 	C.llama_backend_init()
 }
 
+func PrintSystemInfo() string {
+	var compiler string
+	switch C.get_compiler() {
+	case C.COMP_UNKNOWN:
+		compiler = "cgo(unknown_compiler)"
+	case C.COMP_GCC:
+		compiler = "cgo(gcc)"
+	case C.COMP_CLANG:
+		compiler = "cgo(clang)"
+	}
+	return C.GoString(C.llama_print_system_info()) + compiler
+}
+
 func GetModelArch(modelPath string) (string, error) {
 	mp := C.CString(modelPath)
 	defer C.free(unsafe.Pointer(mp))
@@ -245,20 +270,6 @@ func LoadModelFromFile(modelPath string, params ModelParams) (*Model, error) {
 	return &m, nil
 }
 
-func LoadVocabFromFile(path string) (*Vocab, error) {
-	mp := C.CString(path)
-	defer C.free(unsafe.Pointer(mp))
-	v := Vocab{c: C.llama_load_vocab_from_file(mp)}
-	if v.c == nil {
-		return nil, fmt.Errorf("unable to load vocab: %s", path)
-	}
-	return &v, nil
-}
-
-func FreeVocab(vocab *Vocab) {
-	C.llama_free_vocab(vocab.c)
-}
-
 func FreeModel(model *Model) {
 	C.llama_model_free(model.c)
 }
@@ -307,10 +318,6 @@ func (m *Model) ApplyLoraFromFile(context *Context, loraPath string, scale float
 	return nil
 }
 
-type Vocab struct {
-	c *C.struct_llama_vocab
-}
-
 func (m *Model) Vocab() *C.struct_llama_vocab {
 	return C.llama_model_get_vocab(m.c)
 }
@@ -687,53 +694,3 @@ func SchemaToGrammar(schema []byte) []byte {
 	}
 	return buf[:n]
 }
-
-type Sampler struct {
-	c *C.struct_llama_sampler
-}
-
-func NewGrammarSampler(vocab *Vocab, grammar string) *Sampler {
-	cGrammar := C.CString(grammar)
-	cRoot := C.CString("root")
-	defer C.free(unsafe.Pointer(cGrammar))
-	defer C.free(unsafe.Pointer(cRoot))
-
-	sampler := &Sampler{c: C.llama_sampler_init_grammar(vocab.c, cGrammar, cRoot)}
-
-	return sampler
-}
-
-func (s *Sampler) Accept(token int32) {
-	C.llama_sampler_accept(s.c, C.llama_token(token))
-}
-
-type TokenData struct {
-	Id    int32
-	Logit float32
-}
-
-func (s *Sampler) Apply(tokens []TokenData) {
-	tds := make([]C.struct_llama_token_data, len(tokens))
-	for i, token := range tokens {
-		tds[i] = C.struct_llama_token_data{
-			id:    C.int32_t(token.Id),
-			logit: C.float(token.Logit),
-			p:     C.float(0.0),
-		}
-	}
-	tda := &C.llama_token_data_array{
-		data:     (*C.struct_llama_token_data)(unsafe.Pointer(&tds[0])),
-		size:     C.size_t(len(tokens)),
-		selected: C.int64_t(-1),
-		sorted:   C.bool(false),
-	}
-
-	var pinner runtime.Pinner
-	pinner.Pin(&tds[0])
-	defer pinner.Unpin()
-
-	C.llama_sampler_apply(s.c, tda)
-	for i := range tokens {
-		tokens[i].Logit = float32(tds[i].logit)
-	}
-}

llama/patches/0015-try-catch-backend-load.patch

@@ -0,0 +1,69 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: Michael Yang <mxyng@pm.me>
+Date: Tue, 11 Feb 2025 14:06:36 -0800
+Subject: [PATCH] try/catch backend load
+
+---
+ ggml/src/ggml-backend-reg.cpp | 45 ++++++++++++++++++-----------------
+ 1 file changed, 23 insertions(+), 22 deletions(-)
+
+diff --git a/ggml/src/ggml-backend-reg.cpp b/ggml/src/ggml-backend-reg.cpp
+index 98d5e14d..1c19129a 100644
+--- a/ggml/src/ggml-backend-reg.cpp
++++ b/ggml/src/ggml-backend-reg.cpp
+@@ -512,32 +512,33 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
+         }
+         fs::directory_iterator dir_it(search_path, fs::directory_options::skip_permission_denied);
+         for (const auto & entry : dir_it) {
+-            if (entry.is_regular_file()) {
+-                std::wstring filename = entry.path().filename().wstring();
+-                std::wstring ext = entry.path().extension().wstring();
+-                if (filename.find(file_prefix) == 0 && ext == backend_filename_suffix()) {
+-                    dl_handle_ptr handle { dl_load_library(entry.path().wstring()) };
+-                    if (!handle && !silent) {
+-                        GGML_LOG_ERROR("%s: failed to load %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
+-                    }
+-                    if (handle) {
++            try {
++                if (entry.is_regular_file()) {
++                    std::wstring filename = entry.path().filename().wstring();
++                    std::wstring ext = entry.path().extension().wstring();
++                    if (filename.find(file_prefix) == 0 && ext == backend_filename_suffix()) {
++                        dl_handle_ptr handle { dl_load_library(entry.path().wstring()) };
++                        if (!handle) {
++                            GGML_LOG_ERROR("%s: failed to load %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
++                            continue;
++                        }
++
+                         auto score_fn = (ggml_backend_score_t) dl_get_sym(handle.get(), "ggml_backend_score");
+-                        if (score_fn) {
+-                            int s = score_fn();
+-#ifndef NDEBUG
+-                            GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str(), s);
+-#endif
+-                            if (s > best_score) {
+-                                best_score = s;
+-                                best_path = entry.path().wstring();
+-                            }
+-                        } else {
+-                            if (!silent) {
+-                                GGML_LOG_INFO("%s: failed to find ggml_backend_score in %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
+-                            }
++                        if (!score_fn) {
++                            GGML_LOG_DEBUG("%s: failed to find ggml_backend_score in %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
++                            continue;
++                        }
++
++                        int s = score_fn();
++                        GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str(), s);
++                        if (s > best_score) {
++                            best_score = s;
++                            best_path = entry.path().wstring();
+                         }
+                     }
+                 }
++            } catch (const std::exception & e) {
++                GGML_LOG_ERROR("%s: failed to load %s: %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str(), e.what());
+             }
+         }
+     }

llama/patches/0016-use-std-filesystem-path-instead-of-wstring.patch

@@ -4,11 +4,11 @@ Date: Sun, 16 Feb 2025 20:00:22 -0500
 Subject: [PATCH] use std::filesystem::path instead of wstring
 
 ---
- ggml/src/ggml-backend-reg.cpp | 199 +++++++++++++++-------------------
- 1 file changed, 88 insertions(+), 111 deletions(-)
+ ggml/src/ggml-backend-reg.cpp | 144 ++++++++++++++--------------------
+ 1 file changed, 58 insertions(+), 86 deletions(-)
 
 diff --git a/ggml/src/ggml-backend-reg.cpp b/ggml/src/ggml-backend-reg.cpp
-index 98d5e14d..799af5f3 100644
+index 1c19129a..c854e6bb 100644
 --- a/ggml/src/ggml-backend-reg.cpp
 +++ b/ggml/src/ggml-backend-reg.cpp
 @@ -66,26 +66,6 @@
@@ -264,55 +264,47 @@ index 98d5e14d..799af5f3 100644
      for (const auto & search_path : search_paths) {
          if (!fs::exists(search_path)) {
              continue;
-@@ -513,29 +485,26 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
-         fs::directory_iterator dir_it(search_path, fs::directory_options::skip_permission_denied);
+@@ -514,31 +486,31 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
          for (const auto & entry : dir_it) {
-             if (entry.is_regular_file()) {
--                std::wstring filename = entry.path().filename().wstring();
--                std::wstring ext = entry.path().extension().wstring();
-+                std::string filename = entry.path().filename().string();
-+                std::string ext = entry.path().extension().string();
-                 if (filename.find(file_prefix) == 0 && ext == backend_filename_suffix()) {
--                    dl_handle_ptr handle { dl_load_library(entry.path().wstring()) };
--                    if (!handle && !silent) {
--                        GGML_LOG_ERROR("%s: failed to load %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
-+                    dl_handle_ptr handle { dl_load_library(entry.path()) };
-+                    if (!handle) {
-+                        GGML_LOG_ERROR("%s: failed to load %s\n", __func__, path_to_string(entry.path()).c_str());
-+                        continue;
-                     }
--                    if (handle) {
--                        auto score_fn = (ggml_backend_score_t) dl_get_sym(handle.get(), "ggml_backend_score");
--                        if (score_fn) {
--                            int s = score_fn();
--#ifndef NDEBUG
--                            GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str(), s);
--#endif
--                            if (s > best_score) {
--                                best_score = s;
--                                best_path = entry.path().wstring();
--                            }
--                        } else {
--                            if (!silent) {
--                                GGML_LOG_INFO("%s: failed to find ggml_backend_score in %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
--                            }
--                        }
-+
-+                    auto score_fn = (ggml_backend_score_t) dl_get_sym(handle.get(), "ggml_backend_score");
-+                    if (!score_fn) {
-+                        GGML_LOG_DEBUG("%s: failed to find ggml_backend_score in %s\n", __func__, path_to_string(entry.path()).c_str());
-+                        continue;
-+                    }
-+
-+                    int s = score_fn();
-+                    GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, path_to_string(entry.path()).c_str(), s);
-+                    if (s > best_score) {
-+                        best_score = s;
-+                        best_path = entry.path();
+             try {
+                 if (entry.is_regular_file()) {
+-                    std::wstring filename = entry.path().filename().wstring();
+-                    std::wstring ext = entry.path().extension().wstring();
++                    std::string filename = entry.path().filename().string();
++                    std::string ext = entry.path().extension().string();
+                     if (filename.find(file_prefix) == 0 && ext == backend_filename_suffix()) {
+-                        dl_handle_ptr handle { dl_load_library(entry.path().wstring()) };
++                        dl_handle_ptr handle { dl_load_library(entry.path()) };
+                         if (!handle) {
+-                            GGML_LOG_ERROR("%s: failed to load %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
++                            GGML_LOG_ERROR("%s: failed to load %s\n", __func__, path_to_string(entry.path()).c_str());
+                             continue;
+                         }
+ 
+                         auto score_fn = (ggml_backend_score_t) dl_get_sym(handle.get(), "ggml_backend_score");
+                         if (!score_fn) {
+-                            GGML_LOG_DEBUG("%s: failed to find ggml_backend_score in %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str());
++                            GGML_LOG_DEBUG("%s: failed to find ggml_backend_score in %s\n", __func__, path_to_string(entry.path()).c_str());
+                             continue;
+                         }
+ 
+                         int s = score_fn();
+-                        GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str(), s);
++                        GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, path_to_string(entry.path()).c_str(), s);
+                         if (s > best_score) {
+                             best_score = s;
+-                            best_path = entry.path().wstring();
++                            best_path = entry.path();
+                         }
                      }
                  }
+             } catch (const std::exception & e) {
+-                GGML_LOG_ERROR("%s: failed to load %s: %s\n", __func__, utf16_to_utf8(entry.path().wstring()).c_str(), e.what());
++                GGML_LOG_ERROR("%s: failed to load %s: %s\n", __func__, path_to_string(entry.path()).c_str(), e.what());
              }
-@@ -545,7 +514,7 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
+         }
+     }
+@@ -546,7 +518,7 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
      if (best_score == 0) {
          // try to load the base backend
          for (const auto & search_path : search_paths) {
@@ -321,49 +313,3 @@ index 98d5e14d..799af5f3 100644
              if (fs::exists(path)) {
                  return get_reg().load_backend(path, silent);
              }
-@@ -560,6 +529,14 @@ void ggml_backend_load_all() {
-     ggml_backend_load_all_from_path(nullptr);
- }
- 
-+static void ggml_backend_try_load_best(const char * name, bool silent, const char * user_search_path) {
-+    try {
-+        ggml_backend_load_best(name, silent, user_search_path);
-+    } catch (const std::exception & e) {
-+        GGML_LOG_DEBUG("%s: failed to load %s: %s\n", __func__, name, e.what());
-+    }
-+}
-+
- void ggml_backend_load_all_from_path(const char * dir_path) {
- #ifdef NDEBUG
-     bool silent = true;
-@@ -567,18 +544,18 @@ void ggml_backend_load_all_from_path(const char * dir_path) {
-     bool silent = false;
- #endif
- 
--    ggml_backend_load_best("blas", silent, dir_path);
--    ggml_backend_load_best("cann", silent, dir_path);
--    ggml_backend_load_best("cuda", silent, dir_path);
--    ggml_backend_load_best("hip", silent, dir_path);
--    ggml_backend_load_best("kompute", silent, dir_path);
--    ggml_backend_load_best("metal", silent, dir_path);
--    ggml_backend_load_best("rpc", silent, dir_path);
--    ggml_backend_load_best("sycl", silent, dir_path);
--    ggml_backend_load_best("vulkan", silent, dir_path);
--    ggml_backend_load_best("opencl", silent, dir_path);
--    ggml_backend_load_best("musa", silent, dir_path);
--    ggml_backend_load_best("cpu", silent, dir_path);
-+    ggml_backend_try_load_best("blas", silent, dir_path);
-+    ggml_backend_try_load_best("cann", silent, dir_path);
-+    ggml_backend_try_load_best("cuda", silent, dir_path);
-+    ggml_backend_try_load_best("hip", silent, dir_path);
-+    ggml_backend_try_load_best("kompute", silent, dir_path);
-+    ggml_backend_try_load_best("metal", silent, dir_path);
-+    ggml_backend_try_load_best("rpc", silent, dir_path);
-+    ggml_backend_try_load_best("sycl", silent, dir_path);
-+    ggml_backend_try_load_best("vulkan", silent, dir_path);
-+    ggml_backend_try_load_best("opencl", silent, dir_path);
-+    ggml_backend_try_load_best("musa", silent, dir_path);
-+    ggml_backend_try_load_best("cpu", silent, dir_path);
-     // check the environment variable GGML_BACKEND_PATH to load an out-of-tree backend
-     const char * backend_path = std::getenv("GGML_BACKEND_PATH");
-     if (backend_path) {

llama/patches/0019-fix-string-arr-kv-loading.patch

@@ -1,64 +0,0 @@
-From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
-From: jmorganca <jmorganca@gmail.com>
-Date: Wed, 5 Mar 2025 17:41:07 -0800
-Subject: [PATCH] fix string arr kv loading
-
----
- ggml/include/gguf.h | 1 +
- ggml/src/gguf.cpp   | 7 +++++--
- src/llama-vocab.cpp | 2 +-
- 3 files changed, 7 insertions(+), 3 deletions(-)
-
-diff --git a/ggml/include/gguf.h b/ggml/include/gguf.h
-index 79ee2020..3efb22f0 100644
---- a/ggml/include/gguf.h
-+++ b/ggml/include/gguf.h
-@@ -114,6 +114,7 @@ extern "C" {
-     // get raw pointer to the first element of the array with the given key_id
-     // for bool arrays, note that they are always stored as int8 on all platforms (usually this makes no difference)
-     GGML_API const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id);
-+    GGML_API size_t       gguf_get_arr_data_n(const struct gguf_context * ctx, int64_t key_id);
- 
-     // get ith C string from array with given key_id
-     GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i);
-diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp
-index ab13669c..f75b923f 100644
---- a/ggml/src/gguf.cpp
-+++ b/ggml/src/gguf.cpp
-@@ -777,10 +777,14 @@ enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id
- 
- const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id) {
-     GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
--    GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING);
-     return ctx->kv[key_id].data.data();
- }
- 
-+size_t gguf_get_arr_data_n(const struct gguf_context * ctx, int64_t key_id) {
-+    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
-+    return ctx->kv[key_id].data.size();
-+}
-+
- const char * gguf_get_arr_str(const struct gguf_context * ctx, int64_t key_id, size_t i) {
-     GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
-     GGML_ASSERT(ctx->kv[key_id].get_type() == GGUF_TYPE_STRING);
-@@ -874,7 +878,6 @@ const char * gguf_get_val_str(const struct gguf_context * ctx, int64_t key_id) {
- const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id) {
-     GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
-     GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
--    GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING);
-     return ctx->kv[key_id].data.data();
- }
- 
-diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp
-index c7ff28be..7a185443 100644
---- a/src/llama-vocab.cpp
-+++ b/src/llama-vocab.cpp
-@@ -1443,7 +1443,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
- 
-             const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str());
-             if (precompiled_charsmap_keyidx != -1) {
--                size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
-+                size_t n_precompiled_charsmap = gguf_get_arr_data_n(ctx, precompiled_charsmap_keyidx);
-                 const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
-                 precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap);
- #ifdef IS_BIG_ENDIAN

llama/sampling_ext.cpp

@@ -2,9 +2,6 @@
 #include "sampling.h"
 #include "sampling_ext.h"
 #include "json-schema-to-grammar.h"
-#include "llama.h"
-#include "llama-model.h"
-#include "llama-model-loader.h"
 
 struct common_sampler *common_sampler_cinit(const struct llama_model *model, struct common_sampler_cparams *params) {
     try {
@@ -67,22 +64,3 @@ int schema_to_grammar(const char *json_schema, char *grammar, size_t max_len)
         return 0;
     }
 }
-
-struct llama_vocab * llama_load_vocab_from_file(const char * fname) {
-    llama_vocab * vocab = new llama_vocab();
-    try {
-        const auto kv = LLM_KV(LLM_ARCH_UNKNOWN);
-        std::vector<std::string> splits = {};
-        llama_model_loader ml(std::string(fname), splits, false, false, nullptr);
-        vocab->load(ml, kv);
-    } catch (const std::exception & err) {
-        LLAMA_LOG_ERROR("%s: error loading model: %s\n", __func__, err.what());
-        return nullptr;
-    }
-
-    return vocab;
-}
-
-void llama_free_vocab(struct llama_vocab * vocab) {
-    delete vocab;
-}

llama/sampling_ext.h

@@ -35,9 +35,6 @@ extern "C"
 
     int schema_to_grammar(const char *json_schema, char *grammar, size_t max_len);
 
-    struct llama_vocab * llama_load_vocab_from_file(const char * fname);
-    void llama_free_vocab(struct llama_vocab * vocab);
-
 #ifdef __cplusplus
 }
 #endif

llm/memory.go

@@ -115,9 +115,6 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 		// multimodal models require at least 2048 context
 		opts.NumCtx = max(opts.NumCtx, 2048)
 	}
-	if projectorWeights == 0 && projectorGraph == 0 {
-		projectorWeights, projectorGraph = f.VisionGraphSize()
-	}
 
 	layers := f.Tensors().GroupLayers()
 	// add one layer worth of memory as a buffer

llm/server.go

@@ -30,7 +30,6 @@ import (
 	"github.com/ollama/ollama/format"
 	"github.com/ollama/ollama/fs/ggml"
 	"github.com/ollama/ollama/llama"
-	"github.com/ollama/ollama/model"
 )
 
 type LlamaServer interface {
@@ -55,15 +54,8 @@ type llmServer struct {
 	options     api.Options
 	numParallel int
 	modelPath   string
-
-	// llamaModel is an instance of the cgo llama.cpp model definition
-	// nil if this server is running the new engine
-	llamaModel     *llama.Model
-	llamaModelLock sync.Mutex
-
-	// textProcessor handles text encoding/decoding for the model in the Ollama engine
-	// nil if this server is running the llama.cpp based engine
-	textProcessor model.TextProcessor
+	modelLock   sync.Mutex   // Temporary until we switch fully to Go server
+	model       *llama.Model // If non-nil, the runner is a new Go server
 
 	estimate    MemoryEstimate
 	totalLayers uint64
@@ -97,7 +89,7 @@ func LoadModel(model string, maxArraySize int) (*ggml.GGML, error) {
 
 // NewLlamaServer will run a server for the given GPUs
 // The gpu list must be a single family.
-func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, adapters, projectors []string, opts api.Options, numParallel int) (LlamaServer, error) {
+func NewLlamaServer(gpus discover.GpuInfoList, model string, f *ggml.GGML, adapters, projectors []string, opts api.Options, numParallel int) (LlamaServer, error) {
 	systemInfo := discover.GetSystemInfo()
 	systemTotalMemory := systemInfo.System.TotalMemory
 	systemFreeMemory := systemInfo.System.FreeMemory
@@ -138,7 +130,7 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 	slog.Info("offload", "", estimate)
 
 	params := []string{
-		"--model", modelPath,
+		"--model", model,
 		"--ctx-size", strconv.Itoa(opts.NumCtx),
 		"--batch-size", strconv.Itoa(opts.NumBatch),
 	}
@@ -161,6 +153,11 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 		}
 	}
 
+	if len(projectors) > 0 {
+		// TODO: applying multiple projectors is not supported by the llama.cpp server yet
+		params = append(params, "--mmproj", projectors[0])
+	}
+
 	defaultThreads := systemInfo.GetOptimalThreadCount()
 	if opts.NumThread > 0 {
 		params = append(params, "--threads", strconv.Itoa(opts.NumThread))
@@ -260,34 +257,6 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 		}
 	}
 	slog.Debug("compatible gpu libraries", "compatible", compatible)
-	exe, err := os.Executable()
-	if err != nil {
-		return nil, fmt.Errorf("unable to lookup executable path: %w", err)
-	}
-
-	if eval, err := filepath.EvalSymlinks(exe); err == nil {
-		exe = eval
-	}
-
-	var llamaModel *llama.Model
-	var textProcessor model.TextProcessor
-	if envconfig.NewEngine() {
-		textProcessor, err = model.NewTextProcessor(modelPath)
-		if err != nil {
-			// To prepare for opt-out mode, instead of treating this as an error, we fallback to the old runner
-			slog.Debug("model not yet supported by Ollama engine, switching to compatibility mode", "model", modelPath, "error", err)
-		}
-	}
-	if textProcessor == nil {
-		llamaModel, err = llama.LoadModelFromFile(modelPath, llama.ModelParams{VocabOnly: true})
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	if len(projectors) > 0 && llamaModel != nil {
-		params = append(params, "--mmproj", projectors[0])
-	}
 
 	// iterate through compatible GPU libraries such as 'cuda_v12', 'cuda_v11', 'rocm', etc.
 	// adding each library's respective path to the LD_LIBRARY_PATH, until finally running
@@ -306,9 +275,7 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 			port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
 		}
 		finalParams := []string{"runner"}
-		if textProcessor != nil {
-			// New engine
-			// TODO - if we have failure to load scenarios, add logic to retry with the old runner
+		if envconfig.NewEngine() {
 			finalParams = append(finalParams, "--ollama-engine")
 		}
 		finalParams = append(finalParams, params...)
@@ -348,20 +315,28 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 		// finally, add the root library path
 		libraryPaths = append(libraryPaths, discover.LibOllamaPath)
 
+		exe, err := os.Executable()
+		if err != nil {
+			return nil, fmt.Errorf("unable to lookup executable path: %w", err)
+		}
+
+		if eval, err := filepath.EvalSymlinks(exe); err == nil {
+			exe = eval
+		}
+
+		// TODO - once fully switched to the Go runner, load the model here for tokenize/detokenize cgo access
 		s := &llmServer{
-			port:          port,
-			cmd:           exec.Command(exe, finalParams...),
-			status:        NewStatusWriter(os.Stderr),
-			options:       opts,
-			modelPath:     modelPath,
-			llamaModel:    llamaModel,
-			textProcessor: textProcessor,
-			estimate:      estimate,
-			numParallel:   numParallel,
-			sem:           semaphore.NewWeighted(int64(numParallel)),
-			totalLayers:   f.KV().BlockCount() + 1,
-			gpus:          gpus,
-			done:          make(chan error, 1),
+			port:        port,
+			cmd:         exec.Command(exe, finalParams...),
+			status:      NewStatusWriter(os.Stderr),
+			options:     opts,
+			modelPath:   model,
+			estimate:    estimate,
+			numParallel: numParallel,
+			sem:         semaphore.NewWeighted(int64(numParallel)),
+			totalLayers: f.KV().BlockCount() + 1,
+			gpus:        gpus,
+			done:        make(chan error, 1),
 		}
 
 		s.cmd.Env = os.Environ()
@@ -430,9 +405,6 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a
 			}
 			err := fmt.Errorf("error starting runner: %v %s", err, msg)
 			if len(compatible) == 0 {
-				if llamaModel != nil {
-					llama.FreeModel(llamaModel)
-				}
 				return nil, err
 			}
 
@@ -961,25 +933,64 @@ type TokenizeResponse struct {
 }
 
 func (s *llmServer) Tokenize(ctx context.Context, content string) ([]int, error) {
-	s.llamaModelLock.Lock()
-	defer s.llamaModelLock.Unlock()
+	s.modelLock.Lock()
+	defer s.modelLock.Unlock()
+	if s.model != nil {
+		return s.model.Tokenize(content, false, true)
+	}
 
-	if s.llamaModel != nil {
-		return s.llamaModel.Tokenize(content, false, true)
+	// Make sure the server is ready
+	status, err := s.getServerStatus(ctx)
+	if err != nil {
+		return nil, err
+	} else if status != ServerStatusReady && status != ServerStatusNoSlotsAvailable {
+		return nil, fmt.Errorf("unexpected server status: %s", status.ToString())
 	}
-	if s.textProcessor != nil {
-		tokens, err := s.textProcessor.Encode(content, false)
-		if err != nil {
-			return nil, err
-		}
-		toks := make([]int, len(tokens))
-		for i, t := range tokens {
-			toks[i] = int(t)
-		}
-		return toks, nil
+
+	data, err := json.Marshal(TokenizeRequest{Content: content})
+	if err != nil {
+		return nil, fmt.Errorf("marshaling encode data: %w", err)
 	}
-	// not reached
-	return nil, fmt.Errorf("no tokenizer configured")
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/tokenize", s.port), bytes.NewBuffer(data))
+	if err != nil {
+		return nil, fmt.Errorf("encode request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("do encode request: %w", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode == http.StatusNotFound {
+		if s.model == nil {
+			slog.Debug("new runner detected, loading model for cgo tokenization")
+			m, err := llama.LoadModelFromFile(s.modelPath, llama.ModelParams{VocabOnly: true})
+			if err != nil {
+				return nil, err
+			}
+			s.model = m
+		}
+		return s.model.Tokenize(content, false, true)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("read encode request: %w", err)
+	}
+
+	if resp.StatusCode >= 400 {
+		log.Printf("llm encode error: %s", body)
+		return nil, fmt.Errorf("%s", body)
+	}
+
+	var encoded TokenizeResponse
+	if err := json.Unmarshal(body, &encoded); err != nil {
+		return nil, fmt.Errorf("unmarshal encode response: %w", err)
+	}
+
+	return encoded.Tokens, nil
 }
 
 type DetokenizeRequest struct {
@@ -991,38 +1002,80 @@ type DetokenizeResponse struct {
 }
 
 func (s *llmServer) Detokenize(ctx context.Context, tokens []int) (string, error) {
-	s.llamaModelLock.Lock()
-	defer s.llamaModelLock.Unlock()
-
-	if s.llamaModel != nil {
+	s.modelLock.Lock()
+	defer s.modelLock.Unlock()
+	if s.model != nil {
 		var resp string
 		for _, token := range tokens {
-			resp += s.llamaModel.TokenToPiece(token)
+			resp += s.model.TokenToPiece(token)
 		}
 		return resp, nil
 	}
-	if s.textProcessor != nil {
-		toks := make([]int32, len(tokens))
-		for i, t := range tokens {
-			toks[i] = int32(t)
-		}
-		content, err := s.textProcessor.Decode(toks)
-		if err != nil {
-			return "", err
-		}
-		return content, nil
+	// Make sure the server is ready
+	status, err := s.getServerStatus(ctx)
+	if err != nil {
+		return "", err
+	} else if status != ServerStatusReady && status != ServerStatusNoSlotsAvailable {
+		return "", fmt.Errorf("unexpected server status: %s", status.ToString())
 	}
-	// not reached
-	return "", fmt.Errorf("no tokenizer configured")
+
+	data, err := json.Marshal(DetokenizeRequest{Tokens: tokens})
+	if err != nil {
+		return "", fmt.Errorf("marshaling decode data: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/detokenize", s.port), bytes.NewBuffer(data))
+	if err != nil {
+		return "", fmt.Errorf("decode request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("do decode request: %w", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode == http.StatusNotFound {
+		if s.model == nil {
+			slog.Debug("new runner detected, loading model for cgo tokenization")
+			m, err := llama.LoadModelFromFile(s.modelPath, llama.ModelParams{VocabOnly: true})
+			if err != nil {
+				return "", err
+			}
+			s.model = m
+		}
+		var resp string
+		for _, token := range tokens {
+			resp += s.model.TokenToPiece(token)
+		}
+		return resp, nil
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", fmt.Errorf("read decode request: %w", err)
+	}
+
+	if resp.StatusCode >= 400 {
+		log.Printf("llm decode error: %s", body)
+		return "", fmt.Errorf("%s", body)
+	}
+
+	var decoded DetokenizeResponse
+	if err := json.Unmarshal(body, &decoded); err != nil {
+		return "", fmt.Errorf("unmarshal encode response: %w", err)
+	}
+
+	return decoded.Content, nil
 }
 
 func (s *llmServer) Close() error {
-	s.llamaModelLock.Lock()
-	if s.llamaModel != nil {
-		llama.FreeModel(s.llamaModel)
-		s.llamaModel = nil
+	s.modelLock.Lock()
+	if s.model != nil {
+		llama.FreeModel(s.model)
+		s.model = nil
 	}
-	s.llamaModelLock.Unlock()
+	s.modelLock.Unlock()
 
 	if s.cmd != nil {
 		slog.Debug("stopping llama server")

ml/backend.go

@@ -24,7 +24,7 @@ type Backend interface {
 	Config() Config
 	Get(name string) Tensor
 	NewContext() Context
-	NewContextSize(size int) Context
+	SystemInfo() string
 }
 
 // BackendCacheConfig should be implemented by backends that need special output
@@ -100,17 +100,8 @@ type Context interface {
 
 	Forward(...Tensor) Context
 	Compute(...Tensor)
-	MaxGraphNodes() int
+	MaxTensors() int
 	Close()
-
-	// 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
 }
 
 type Tensor interface {
@@ -215,7 +206,7 @@ func Dump(ctx Context, t Tensor, opts ...DumpOptions) string {
 		return dump[[]float32](ctx, t, opts[0].Items, func(f float32) string {
 			return strconv.FormatFloat(float64(f), 'f', opts[0].Precision, 32)
 		})
-	case DTypeF16, DTypeQ80, DTypeQ40:
+	case DTypeF16:
 		f32 := ctx.Empty(DTypeF32, t.Shape()...)
 		f32 = t.Copy(ctx, f32)
 		return dump[[]float32](ctx, f32, opts[0].Items, func(f float32) string {
@@ -283,7 +274,5 @@ const (
 	DTypeOther DType = iota
 	DTypeF32
 	DTypeF16
-	DTypeQ80
-	DTypeQ40
 	DTypeI32
 )

ml/backend/ggml/ggml.go

@@ -1,61 +1,91 @@
 package ggml
 
-// #cgo CPPFLAGS: -I${SRCDIR}/ggml/include
-// #include <stdlib.h>
-// #include <stdint.h>
-// #include "ggml.h"
-// #include "ggml-cpu.h"
-// #include "ggml-backend.h"
+/*
+#cgo CPPFLAGS: -I${SRCDIR}/ggml/include
+#include <stdlib.h>
+#include <stdint.h>
+#include "ggml.h"
+#include "ggml-cpu.h"
+#include "ggml-backend.h"
+static struct ggml_backend_feature * getBackendFeatures(void *fp, ggml_backend_reg_t reg) {return ((ggml_backend_get_features_t)(fp))(reg);}
+static struct ggml_backend_feature * getNextBackendFeatures(struct ggml_backend_feature * feature) { return &feature[1];}
+
+typedef enum {COMP_UNKNOWN,COMP_GCC,COMP_CLANG} COMPILER;
+COMPILER inline get_compiler() {
+#if defined(__clang__)
+	return COMP_CLANG;
+#elif defined(__GNUC__)
+	return COMP_GCC;
+#else
+	return UNKNOWN_COMPILER;
+#endif
+}
+
+*/
 import "C"
 
 import (
-	"errors"
 	"fmt"
 	"io"
 	"log/slog"
-	"maps"
 	"os"
-	"slices"
-	"strconv"
-	"strings"
-	"unicode"
+	"sync"
 	"unsafe"
 
 	"github.com/ollama/ollama/format"
 	fs "github.com/ollama/ollama/fs/ggml"
 	"github.com/ollama/ollama/ml"
-	ggml "github.com/ollama/ollama/ml/backend/ggml/ggml/src"
 	"golang.org/x/sync/errgroup"
+
+	ggml "github.com/ollama/ollama/ml/backend/ggml/ggml/src"
 )
 
-func devices() []*C.struct_ggml_backend_device {
-	ggml.OnceLoad()
-	ds := make([]*C.struct_ggml_backend_device, C.ggml_backend_dev_count())
-	for i := range ds {
-		ds[i] = C.ggml_backend_dev_get(C.size_t(i))
-	}
-
-	return ds
+type device struct {
+	d *C.struct_ggml_backend_device
 }
 
+func (d device) LogValue() slog.Value {
+	var free, total uint64
+	C.ggml_backend_dev_memory(d.d, (*C.size_t)(&free), (*C.size_t)(&total))
+
+	kind := "unknown"
+	switch C.ggml_backend_dev_type(d.d) {
+	case C.GGML_BACKEND_DEVICE_TYPE_CPU:
+		kind = "cpu"
+	case C.GGML_BACKEND_DEVICE_TYPE_GPU:
+		kind = "gpu"
+	case C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
+		kind = "accel"
+	}
+
+	return slog.GroupValue(
+		slog.String("name", C.GoString(C.ggml_backend_dev_name(d.d))),
+		slog.String("description", C.GoString(C.ggml_backend_dev_description(d.d))),
+		slog.String("kind", kind),
+		slog.String("free", format.HumanBytes2(free)),
+		slog.String("total", format.HumanBytes2(total)),
+	)
+}
+
+var devices = sync.OnceValue(func() []device {
+	ggml.OnceLoad()
+
+	s := make([]device, C.ggml_backend_dev_count())
+	for i := range s {
+		s[i] = device{C.ggml_backend_dev_get(C.size_t(i))}
+	}
+
+	return s
+})
+
 type Backend struct {
-	meta    *fs.GGML
-	sched   *C.struct_ggml_backend_sched
-	tensors map[string]*C.struct_ggml_tensor
-
-	// 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
-
 	flashAttention bool
 
-	// maxGraphNodes is the maximum allowed number of graph nodes in this scheduler
-	maxGraphNodes int
+	meta       *fs.GGML
+	cpus, gpus []Context
+	tensors    map[string]*Context
+
+	sched *C.struct_ggml_backend_sched
 }
 
 func New(r *os.File, params ml.BackendParams) (ml.Backend, error) {
@@ -74,310 +104,107 @@ func New(r *os.File, params ml.BackendParams) (ml.Backend, error) {
 		"num_key_values", len(meta.KV()),
 	)
 
-	type deviceBufferType struct {
-		d   *C.struct_ggml_backend_device
-		bts []*C.struct_ggml_backend_buffer_type
-	}
-
-	var cpus, accels, gpus []*C.struct_ggml_backend_device
+	var cpus, gpus []Context
 	for _, d := range devices() {
-		switch C.ggml_backend_dev_type(d) {
-		case C.GGML_BACKEND_DEVICE_TYPE_CPU:
-			if len(cpus) == 0 {
-				// only the first cpu device should be used
-				cpus = append(cpus, d)
-			}
-		case C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
-			accels = append(accels, d)
-		case C.GGML_BACKEND_DEVICE_TYPE_GPU:
-			gpus = append(gpus, d)
-		}
-	}
-
-	// create list of buffer types for the cpu
-	cpuDeviceBufferType := deviceBufferType{d: C.ggml_backend_dev_by_type(C.GGML_BACKEND_DEVICE_TYPE_CPU)}
-	for _, d := range append(accels, append(gpus, cpus...)...) {
-		switch C.ggml_backend_dev_type(d) {
+		switch C.ggml_backend_dev_type(d.d) {
 		case C.GGML_BACKEND_DEVICE_TYPE_CPU,
 			C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
-			cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, C.ggml_backend_dev_buffer_type(d))
-		}
-	}
-
-	// create list of buffer types for each gpu
-	var gpuDeviceBufferTypes []deviceBufferType
-	for _, d := range gpus {
-		bt := C.ggml_backend_dev_buffer_type(d)
-		gpuDeviceBufferTypes = append(gpuDeviceBufferTypes, deviceBufferType{
-			d:   d,
-			bts: append([]*C.struct_ggml_backend_buffer_type{bt}, cpuDeviceBufferType.bts...),
-		})
-	}
-
-	useDefaultSplit := true
-	for _, s := range params.TensorSplit {
-		if s != 0 {
-			useDefaultSplit = false
-			break
-		}
-	}
-
-	// calculate splits
-	splits := make([]float32, len(gpus))
-	if useDefaultSplit {
-		// default: split on free memory
-		for i := range splits {
-			var free, total C.size_t
-			C.ggml_backend_dev_memory(gpus[i], &free, &total)
-			splits[i] = float32(free)
-		}
-	} else {
-		splits = params.TensorSplit
-	}
-
-	var sum float32
-	// cumulative sum of all splits
-	for i := range splits {
-		sum += splits[i]
-		splits[i] = sum
-	}
-
-	// normalize splits
-	for i := range splits {
-		splits[i] /= sum
-	}
-
-	// inputs always use cpu
-	input := cpuDeviceBufferType
-
-	blocks := int(meta.KV().BlockCount())
-
-	// define a range of gpu layers. anything outside of this range is assigned to the cpu
-	gpuRangeStart := max(0, blocks-params.NumGPULayers)
-	gpuRangeStop := min(gpuRangeStart+params.NumGPULayers, blocks+1)
-	assignLayer := func(i int) deviceBufferType {
-		if i < gpuRangeStart || i >= gpuRangeStop {
-			return cpuDeviceBufferType
-		}
-
-		index := slices.IndexFunc(splits, func(f float32) bool { return float32(i-gpuRangeStart)/float32(gpuRangeStop-gpuRangeStart) < f })
-		if index < 0 || index >= len(gpuDeviceBufferTypes) {
-			return cpuDeviceBufferType
-		}
-
-		return gpuDeviceBufferTypes[index]
-	}
-
-	// repeating layers are assigned based on their index in reverse order, e.g. i / (block_count + 1)
-	layers := make([]deviceBufferType, blocks)
-	for i := range layers {
-		layers[i] = assignLayer(i)
-	}
-
-	// outputs are assigned iff allowed by splits and configured number of gpu layers
-	output := assignLayer(blocks)
-
-	maxTensors := len(meta.Tensors().Items())
-	maxTensors += 1
-	// each layer has at most 2 extra tensors for rope operations
-	maxTensors += blocks * 2
-
-	type tensor struct {
-		source *fs.Tensor
-		target string
-	}
-
-	// some tensors are mapped to different names so keep a list
-	targets := make(map[string][]string)
-
-	// contexts are shared by tensors of the same buffer type
-	ctxs := make(map[*C.struct_ggml_backend_buffer_type]*C.struct_ggml_context)
-	createTensor := func(t tensor, bts []*C.struct_ggml_backend_buffer_type) *C.struct_ggml_tensor {
-		for _, bt := range bts {
-			if _, ok := ctxs[bt]; !ok {
-				ctxs[bt] = C.ggml_init(C.struct_ggml_init_params{
-					mem_size: C.ggml_tensor_overhead() * C.size_t(maxTensors),
+			slog.Info("cpu", "device", d)
+			cpus = append(cpus, Context{
+				ctx: C.ggml_init(C.struct_ggml_init_params{
+					mem_size: C.size_t(int(C.ggml_tensor_overhead()) * (len(meta.Tensors().Items()) + 1 + int(meta.KV().BlockCount())*2)),
 					no_alloc: true,
-				})
-			}
-
-			targets[t.source.Name] = append(targets[t.source.Name], t.target)
-
-			name := t.source.Name
-			if t.target != "" {
-				name = t.target
-			}
-
-			cname := C.CString(name)
-			defer C.free(unsafe.Pointer(cname))
-			if tt := C.ggml_get_tensor(ctxs[bt], cname); tt != nil {
-				return tt
-			}
-
-			tt := C.ggml_new_tensor(ctxs[bt], t.source.Kind, C.int(len(t.source.Shape)), (*C.int64_t)(unsafe.Pointer(&t.source.Shape[0])))
-			C.ggml_set_name(tt, cname)
-
-			slog.Debug("created tensor", "name", name, "shape", t.source.Shape, "dtype", t.source.Kind, "buffer_type", C.GoString(C.ggml_backend_buft_name(bt)))
-			//nolint:staticcheck // TODO: check if buffer type supports this tensor
-			return tt
-		}
-
-		return nil
-	}
-
-	contains := func(s string, parts ...string) bool {
-		split := strings.Split(s, ".")
-		for _, part := range parts {
-			if slices.Contains(split, part) {
-				return true
-			}
-		}
-
-		return false
-	}
-
-	for _, t := range meta.Tensors().Items() {
-		switch {
-		case contains(t.Name, "position_embd", "token_embd", "token_norm_embd", "token_types"):
-			createTensor(tensor{source: t}, input.bts)
-		case contains(t.Name, "cls", "output", "output_norm"):
-			createTensor(tensor{source: t}, output.bts)
-		case strings.HasPrefix(t.Name, "v.") || strings.HasPrefix(t.Name, "mm."):
-			// TODO: assign vision tensors to the gpu if possible
-			createTensor(tensor{source: t}, input.bts)
-		default:
-			layerIndex := -1
-			if fields := strings.FieldsFunc(t.Name, func(r rune) bool { return !unicode.IsNumber(r) }); len(fields) > 0 {
-				if i, err := strconv.Atoi(fields[0]); err == nil {
-					layerIndex = i
-				}
-			}
-
-			if layerIndex >= 0 {
-				createTensor(tensor{source: t}, layers[layerIndex].bts)
-			} else {
-				// this is a repeating tensor that doesn't explicitly associated with a layer so
-				// duplicate it for each layer
-				for i, layer := range layers {
-					createTensor(tensor{
-						source: t,
-						target: "blk." + strconv.Itoa(i) + "." + t.Name,
-					}, layer.bts)
-				}
-			}
-		}
-	}
-
-	// allocate buffers for each context
-	bbs := make(map[*C.struct_ggml_context]*C.struct_ggml_backend_buffer, len(ctxs))
-	for bt, c := range ctxs {
-		if C.ggml_get_first_tensor(c) == nil {
-			continue
-		}
-
-		b := C.ggml_backend_alloc_ctx_tensors_from_buft(c, bt)
-		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
-		bbs[c] = b
-	}
-
-	for bs := range maps.Values(bbs) {
-		slog.Info("model weights", "buffer", C.GoString(C.ggml_backend_buffer_name(bs)), "size", format.HumanBytes2(uint64(C.ggml_backend_buffer_get_size(bs))))
-	}
-
-	// map tensor names to tensors for easy lookup later
-	tensors := make(map[string]*C.struct_ggml_tensor)
-	for _, c := range ctxs {
-		for t := C.ggml_get_first_tensor(c); t != nil; t = C.ggml_get_next_tensor(c, t) {
-			tensors[C.GoString(C.ggml_get_name(t))] = t
-		}
-	}
-
-	// 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() {
-		for _, target := range targets[t.Name] {
-			g.Go(func() error {
-				if target == "" {
-					target = t.Name
-				}
-
-				tt, ok := tensors[target]
-				if !ok {
-					return fmt.Errorf("unassigned tensor: %s", t.Name)
-				}
-
-				bts := make([]byte, t.Size())
-				n, err := io.ReadFull(io.NewSectionReader(sr, int64(t.Offset), int64(t.Size())), bts)
-				if err != nil {
-					return err
-				}
-
-				if n != len(bts) {
-					return errors.New("short read")
-				}
-
-				C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), 0, C.size_t(t.Size()))
-				return nil
+				}),
+				backend: C.ggml_backend_dev_init(d.d, nil),
+			})
+		case C.GGML_BACKEND_DEVICE_TYPE_GPU:
+			slog.Info("gpu", "device", d)
+			gpus = append(gpus, Context{
+				ctx: C.ggml_init(C.struct_ggml_init_params{
+					mem_size: C.size_t(int(C.ggml_tensor_overhead()) * (len(meta.Tensors().Items()) + 1 + int(meta.KV().BlockCount())*2)),
+					no_alloc: true,
+				}),
+				backend: C.ggml_backend_dev_init(d.d, nil),
 			})
 		}
 	}
 
-	if g.Wait() != nil {
+	ctxFunc := func(s []Context) (*Context, error) {
+		for _, e := range s {
+			return &e, nil
+		}
+
+		return nil, fmt.Errorf("no devices available")
+	}
+
+	tensors := make(map[*fs.Tensor]*Context, len(meta.Tensors().Items()))
+	for _, t := range meta.Tensors().Items() {
+		c, err := ctxFunc(append(gpus, cpus...))
+		if err != nil {
+			return nil, err
+		}
+
+		func() {
+			tt := C.ggml_new_tensor(c.ctx, t.Kind, C.int(len(t.Shape)), (*C.int64_t)(unsafe.Pointer(&t.Shape[0])))
+
+			cname := C.CString(t.Name)
+			defer C.free(unsafe.Pointer(cname))
+			C.ggml_set_name(tt, cname)
+
+			tensors[t] = c
+		}()
+	}
+
+	for _, b := range append(gpus, cpus...) {
+		C.ggml_backend_alloc_ctx_tensors(b.ctx, b.backend)
+	}
+
+	sr := io.NewSectionReader(r, int64(meta.Tensors().Offset), n-int64(meta.Tensors().Offset))
+
+	var g errgroup.Group
+	for t, c := range tensors {
+		g.Go(func() error {
+			bts := make([]byte, t.Size())
+			n, err := io.ReadFull(io.NewSectionReader(sr, int64(t.Offset), int64(t.Size())), bts)
+			if err != nil {
+				return err
+			}
+
+			if n != int(t.Size()) {
+				return fmt.Errorf("expected %d bytes, got %d", t.Size(), n)
+			}
+
+			cname := C.CString(t.Name)
+			defer C.free(unsafe.Pointer(cname))
+
+			C.ggml_backend_tensor_set(C.ggml_get_tensor(c.ctx, cname), unsafe.Pointer(&bts[0]), 0, C.size_t(n))
+			return nil
+		})
+	}
+
+	if err := g.Wait(); err != nil {
 		return nil, err
 	}
 
-	// map devices to backend buffer types so new tensors can be assigned to the correct device
-	deviceBufferTypes := make(map[*C.struct_ggml_backend_device]*C.struct_ggml_backend_buffer_type)
-
-	// create backends and buffer types used for the compute graph scheduler
-	var schedBackends []*C.struct_ggml_backend
-	var schedBufts []*C.struct_ggml_backend_buffer_type
-	for _, d := range append(gpus, append(accels, cpus...)...) {
-		b := C.ggml_backend_dev_init(d, nil)
-		bt := C.ggml_backend_get_default_buffer_type(b)
-		if d := C.ggml_backend_get_device(b); C.ggml_backend_dev_type(d) == C.GGML_BACKEND_DEVICE_TYPE_CPU && len(gpus) > 0 {
-			// use the first gpu host buffer type for gpu if possible
-			if hbt := C.ggml_backend_dev_host_buffer_type(gpus[0]); hbt != nil {
-				bt = hbt
-			}
-		}
-
-		deviceBufferTypes[d] = bt
-
-		schedBackends = append(schedBackends, b)
-		schedBufts = append(schedBufts, bt)
-
-		slog.Info("compute graph", "backend", C.GoString(C.ggml_backend_name(b)), "buffer_type", C.GoString(C.ggml_backend_buft_name(bt)))
-
-		if C.ggml_backend_is_cpu(b) {
-			// set number of threads for cpu backend
-			C.ggml_backend_cpu_set_n_threads(b, C.int(params.NumThreads))
-		}
+	backends := make([]*C.struct_ggml_backend, len(gpus)+len(cpus))
+	bufts := make([]*C.struct_ggml_backend_buffer_type, len(gpus)+len(cpus))
+	for i, c := range append(gpus, cpus...) {
+		backends[i] = c.backend
+		bufts[i] = C.ggml_backend_get_default_buffer_type(c.backend)
 	}
 
-	maxGraphNodes := max(8192, len(meta.Tensors().Items())*5)
 	return &Backend{
 		flashAttention: params.FlashAttention,
 		meta:           meta,
-		tensors:        tensors,
+		cpus:           cpus,
+		gpus:           gpus,
 		sched: C.ggml_backend_sched_new(
-			(*C.ggml_backend_t)(unsafe.Pointer(&schedBackends[0])),
-			(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&schedBufts[0])),
-			C.int(len(schedBackends)),
-			C.size_t(maxGraphNodes),
+			(*C.ggml_backend_t)(unsafe.Pointer(&backends[0])),
+			(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&bufts[0])),
+			C.int(len(backends)),
+			C.size_t(max(8192, len(meta.Tensors().Items())*5)),
 			true,
 		),
-		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 {
-				m[i] = deviceBufferTypes[layer.d]
-			}
-			return m
-		}(),
-		maxGraphNodes: maxGraphNodes,
 	}, nil
 }
 
@@ -390,29 +217,36 @@ func (b *Backend) Config() ml.Config {
 }
 
 func (b *Backend) Get(name string) ml.Tensor {
-	if t, ok := b.tensors[name]; ok {
-		return &Tensor{b: b, t: t}
+	cname := C.CString(name)
+	defer C.free(unsafe.Pointer(cname))
+
+	for _, c := range append(b.gpus, b.cpus...) {
+		if t := C.ggml_get_tensor(c.ctx, cname); t != nil {
+			return &Tensor{b: b, t: t}
+		}
 	}
 
 	return nil
 }
 
 func (b *Backend) NewContext() ml.Context {
-	return b.NewContextSize(b.maxGraphNodes)
-}
+	nodes := max(8192, len(b.meta.Tensors().Items())*5)
+	c := C.ggml_init(C.struct_ggml_init_params{
+		mem_buffer: nil,
+		mem_size:   C.size_t(nodes)*C.ggml_tensor_overhead() + C.ggml_graph_overhead_custom(C.size_t(nodes), false),
+		no_alloc:   true,
+	})
 
-func (b *Backend) NewContextSize(n int) ml.Context {
-	if n > b.maxGraphNodes {
-		panic(fmt.Errorf("requested number of graph nodes (%v) for new context exceeds maximum (%v)", n, b.maxGraphNodes))
+	backends := make([]*C.struct_ggml_backend, len(b.gpus)+len(b.cpus))
+	for i, c := range append(b.gpus, b.cpus...) {
+		backends[i] = c.backend
 	}
 
 	return &Context{
-		b:             b,
-		maxGraphNodes: n,
-		ctx: C.ggml_init(C.struct_ggml_init_params{
-			mem_size: C.size_t(n)*C.ggml_tensor_overhead() + C.ggml_graph_overhead_custom(C.size_t(n), false),
-			no_alloc: true,
-		}),
+		b:       b,
+		ctx:     c,
+		backend: backends[0],
+		nodes:   nodes,
 	}
 }
 
@@ -425,60 +259,17 @@ func (b *Backend) CacheConfig() ml.CacheConfig {
 }
 
 type Context struct {
-	b *Backend
+	b       *Backend
+	ctx     *C.struct_ggml_context
+	backend *C.struct_ggml_backend
 
-	ctx   *C.struct_ggml_context
 	graph *C.struct_ggml_cgraph
-
-	// buft is the buffer type used for new tensors
-	buft *C.struct_ggml_backend_buffer_type
-
-	// maxGraphNodes is the maximum allowed number of graph nodes in this context
-	maxGraphNodes int
-}
-
-func (c Context) Input() ml.Context {
-	if c.b.input != nil {
-		return &Context{
-			b:             c.b,
-			ctx:           c.ctx,
-			buft:          c.b.input,
-			maxGraphNodes: c.maxGraphNodes,
-		}
-	}
-
-	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{
-			b:             c.b,
-			ctx:           c.ctx,
-			buft:          buft,
-			maxGraphNodes: c.maxGraphNodes,
-		}
-	}
-
-	return &c
+	nodes int
 }
 
 func (c *Context) Forward(tensors ...ml.Tensor) ml.Context {
 	if c.graph == nil {
-		c.graph = C.ggml_new_graph_custom(c.ctx, C.size_t(c.maxGraphNodes), false)
+		c.graph = C.ggml_new_graph_custom(c.ctx, C.size_t(c.nodes), false)
 	}
 
 	for _, tensor := range tensors {
@@ -488,7 +279,7 @@ func (c *Context) Forward(tensors ...ml.Tensor) ml.Context {
 	return c
 }
 
-func (c Context) Compute(tensors ...ml.Tensor) {
+func (c *Context) Compute(tensors ...ml.Tensor) {
 	C.ggml_backend_sched_graph_compute_async(c.b.sched, c.graph)
 	C.ggml_backend_sched_reset(c.b.sched)
 
@@ -507,48 +298,21 @@ func (c Context) Compute(tensors ...ml.Tensor) {
 	}
 }
 
-func (c Context) MaxGraphNodes() int {
-	return c.maxGraphNodes
+func (c *Context) MaxTensors() int {
+	return c.nodes
 }
 
 func shapeToGGML(shape []int) *C.int64_t {
 	sh := make([]C.int64_t, len(shape))
 	for i, s := range shape {
-		sh[i] = C.int64_t(s)
+		sh[i] = (C.int64_t)(s)
 	}
 
 	return &sh[0]
 }
 
-func pad(length, pad C.size_t) C.size_t {
-	return ((length + pad - 1) / pad) * pad
-}
-
-func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
-	if c.buft == nil {
-		panic("set Input, Output, or Layer before creating tensors")
-	}
-
-	var cdtype uint32
-	switch dtype {
-	case ml.DTypeF32:
-		cdtype = C.GGML_TYPE_F32
-	case ml.DTypeF16:
-		cdtype = C.GGML_TYPE_F16
-	case ml.DTypeQ80:
-		cdtype = C.GGML_TYPE_Q8_0
-	case ml.DTypeQ40:
-		cdtype = C.GGML_TYPE_Q4_0
-	case ml.DTypeI32:
-		cdtype = C.GGML_TYPE_I32
-	default:
-		panic("unsupported dtype")
-	}
-
-	if len(shape) < 1 || shape[0] == 0 {
-		var shape C.int64_t = 0
-		return &Tensor{b: c.b, t: C.ggml_new_tensor(c.ctx, cdtype, 1, &shape)}
-	} else if len(shape) > 4 {
+func newTensor(ctx Context, dtype ml.DType, zero bool, shape []int) ml.Tensor {
+	if len(shape) < 1 || len(shape) > 4 {
 		panic("unsupported number of dimensions")
 	}
 
@@ -558,28 +322,41 @@ func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
 		}
 	}
 
-	t := C.ggml_new_tensor(c.ctx, cdtype, C.int(len(shape)), shapeToGGML(shape))
-	size := pad(C.ggml_backend_buft_get_alloc_size(c.buft, t), C.ggml_backend_buft_get_alignment(c.buft))
-	b := C.ggml_backend_buft_alloc_buffer(c.buft, size)
+	var t *C.struct_ggml_tensor
+	switch dtype {
+	case ml.DTypeF32:
+		t = C.ggml_new_tensor(ctx.ctx, C.GGML_TYPE_F32, C.int(len(shape)), shapeToGGML(shape))
+	case ml.DTypeF16:
+		t = C.ggml_new_tensor(ctx.ctx, C.GGML_TYPE_F16, C.int(len(shape)), shapeToGGML(shape))
+	case ml.DTypeI32:
+		t = C.ggml_new_tensor(ctx.ctx, C.GGML_TYPE_I32, C.int(len(shape)), shapeToGGML(shape))
+	default:
+		panic("unsupported dtype")
+	}
+
+	b := C.ggml_backend_alloc_buffer(ctx.backend, C.ggml_nbytes(t))
 	C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
-	return &Tensor{b: c.b, t: t}
+	if zero {
+		C.ggml_set_zero(t)
+	}
+	return &Tensor{b: ctx.b, t: t}
 }
 
 func (c Context) Empty(dtype ml.DType, shape ...int) ml.Tensor {
-	return c.newTensor(dtype, shape)
+	return newTensor(c, dtype, false, shape)
 }
 
 func (c Context) Zeros(dtype ml.DType, shape ...int) ml.Tensor {
-	t := c.newTensor(dtype, shape)
-	C.ggml_set_zero(t.(*Tensor).t)
-	return t
+	return newTensor(c, dtype, true, shape)
 }
 
-func checkShape[S ~[]E, E any](s S, shape ...int) error {
+func fromSlice[S ~[]E, E float32 | int32](ctx Context, s S, shape []int, dtype uint32) (ml.Tensor, error) {
 	n := len(s)
 
 	if n == 0 {
-		return nil
+		var shape C.int64_t = 0
+		t := C.ggml_new_tensor(ctx.ctx, dtype, 1, &shape)
+		return &Tensor{b: ctx.b, t: t}, nil
 	}
 
 	for _, v := range shape {
@@ -587,36 +364,22 @@ func checkShape[S ~[]E, E any](s S, shape ...int) error {
 	}
 
 	if n != 1 {
-		return fmt.Errorf("invalid shape: %v", shape)
+		return nil, fmt.Errorf("invalid shape %v for %d elements", shape, len(s))
 	}
 
-	return nil
+	t := C.ggml_new_tensor(ctx.ctx, dtype, C.int(len(shape)), shapeToGGML(shape))
+	b := C.ggml_backend_alloc_buffer(ctx.backend, C.ggml_nbytes(t))
+	C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
+	C.ggml_backend_tensor_set(t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t))
+	return &Tensor{b: ctx.b, t: t}, nil
 }
 
 func (c Context) FromFloatSlice(s []float32, shape ...int) (ml.Tensor, error) {
-	if err := checkShape(s, shape...); err != nil {
-		return nil, err
-	}
-
-	t := c.newTensor(ml.DTypeF32, shape)
-	if len(s) > 0 {
-		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
-	}
-
-	return t, nil
+	return fromSlice(c, s, shape, C.GGML_TYPE_F32)
 }
 
 func (c Context) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
-	if err := checkShape(s, shape...); err != nil {
-		return nil, err
-	}
-
-	t := c.newTensor(ml.DTypeI32, shape)
-	if len(s) > 0 {
-		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
-	}
-
-	return t, nil
+	return fromSlice(c, s, shape, C.GGML_TYPE_I32)
 }
 
 func (c *Context) Close() {
@@ -684,10 +447,6 @@ func (t *Tensor) DType() ml.DType {
 		return ml.DTypeF32
 	case C.GGML_TYPE_F16:
 		return ml.DTypeF16
-	case C.GGML_TYPE_Q8_0:
-		return ml.DTypeQ80
-	case C.GGML_TYPE_Q4_0:
-		return ml.DTypeQ40
 	case C.GGML_TYPE_I32:
 		return ml.DTypeI32
 	default:
@@ -970,3 +729,34 @@ func (t *Tensor) ScaledDotProductAttention(ctx ml.Context, key, value, mask ml.T
 		return kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
 	}
 }
+
+func (b *Backend) SystemInfo() string {
+	var compiler string
+	switch C.get_compiler() {
+	case C.COMP_UNKNOWN:
+		compiler = "cgo(unknown_compiler)"
+	case C.COMP_GCC:
+		compiler = "cgo(gcc)"
+	case C.COMP_CLANG:
+		compiler = "cgo(clang)"
+	}
+
+	var s string
+	for i := range C.ggml_backend_reg_count() {
+		reg := C.ggml_backend_reg_get(i)
+		fName := C.CString("ggml_backend_get_features")
+		defer C.free(unsafe.Pointer(fName))
+		get_features_fn := C.ggml_backend_reg_get_proc_address(reg, fName)
+		if get_features_fn != nil {
+			s += C.GoString(C.ggml_backend_reg_name(reg))
+			s += " : "
+			for features := C.getBackendFeatures(get_features_fn, reg); features.name != nil; features = C.getNextBackendFeatures(features) {
+				s += C.GoString(features.name)
+				s += " = "
+				s += C.GoString(features.value)
+				s += " | "
+			}
+		}
+	}
+	return s + compiler
+}

ml/backend/ggml/ggml/include/gguf.h

@@ -114,7 +114,6 @@ extern "C" {
     // get raw pointer to the first element of the array with the given key_id
     // for bool arrays, note that they are always stored as int8 on all platforms (usually this makes no difference)
     GGML_API const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id);
-    GGML_API size_t       gguf_get_arr_data_n(const struct gguf_context * ctx, int64_t key_id);
 
     // get ith C string from array with given key_id
     GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i);

ml/backend/ggml/ggml/src/ggml-backend-reg.cpp

@@ -484,29 +484,33 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
         }
         fs::directory_iterator dir_it(search_path, fs::directory_options::skip_permission_denied);
         for (const auto & entry : dir_it) {
-            if (entry.is_regular_file()) {
-                std::string filename = entry.path().filename().string();
-                std::string ext = entry.path().extension().string();
-                if (filename.find(file_prefix) == 0 && ext == backend_filename_suffix()) {
-                    dl_handle_ptr handle { dl_load_library(entry.path()) };
-                    if (!handle) {
-                        GGML_LOG_ERROR("%s: failed to load %s\n", __func__, path_to_string(entry.path()).c_str());
-                        continue;
-                    }
+            try {
+                if (entry.is_regular_file()) {
+                    std::string filename = entry.path().filename().string();
+                    std::string ext = entry.path().extension().string();
+                    if (filename.find(file_prefix) == 0 && ext == backend_filename_suffix()) {
+                        dl_handle_ptr handle { dl_load_library(entry.path()) };
+                        if (!handle) {
+                            GGML_LOG_ERROR("%s: failed to load %s\n", __func__, path_to_string(entry.path()).c_str());
+                            continue;
+                        }
 
-                    auto score_fn = (ggml_backend_score_t) dl_get_sym(handle.get(), "ggml_backend_score");
-                    if (!score_fn) {
-                        GGML_LOG_DEBUG("%s: failed to find ggml_backend_score in %s\n", __func__, path_to_string(entry.path()).c_str());
-                        continue;
-                    }
+                        auto score_fn = (ggml_backend_score_t) dl_get_sym(handle.get(), "ggml_backend_score");
+                        if (!score_fn) {
+                            GGML_LOG_DEBUG("%s: failed to find ggml_backend_score in %s\n", __func__, path_to_string(entry.path()).c_str());
+                            continue;
+                        }
 
-                    int s = score_fn();
-                    GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, path_to_string(entry.path()).c_str(), s);
-                    if (s > best_score) {
-                        best_score = s;
-                        best_path = entry.path();
+                        int s = score_fn();
+                        GGML_LOG_DEBUG("%s: %s score: %d\n", __func__, path_to_string(entry.path()).c_str(), s);
+                        if (s > best_score) {
+                            best_score = s;
+                            best_path = entry.path();
+                        }
                     }
                 }
+            } catch (const std::exception & e) {
+                GGML_LOG_ERROR("%s: failed to load %s: %s\n", __func__, path_to_string(entry.path()).c_str(), e.what());
             }
         }
     }
@@ -529,14 +533,6 @@ void ggml_backend_load_all() {
     ggml_backend_load_all_from_path(nullptr);
 }
 
-static void ggml_backend_try_load_best(const char * name, bool silent, const char * user_search_path) {
-    try {
-        ggml_backend_load_best(name, silent, user_search_path);
-    } catch (const std::exception & e) {
-        GGML_LOG_DEBUG("%s: failed to load %s: %s\n", __func__, name, e.what());
-    }
-}
-
 void ggml_backend_load_all_from_path(const char * dir_path) {
 #ifdef NDEBUG
     bool silent = true;
@@ -544,18 +540,18 @@ void ggml_backend_load_all_from_path(const char * dir_path) {
     bool silent = false;
 #endif
 
-    ggml_backend_try_load_best("blas", silent, dir_path);
-    ggml_backend_try_load_best("cann", silent, dir_path);
-    ggml_backend_try_load_best("cuda", silent, dir_path);
-    ggml_backend_try_load_best("hip", silent, dir_path);
-    ggml_backend_try_load_best("kompute", silent, dir_path);
-    ggml_backend_try_load_best("metal", silent, dir_path);
-    ggml_backend_try_load_best("rpc", silent, dir_path);
-    ggml_backend_try_load_best("sycl", silent, dir_path);
-    ggml_backend_try_load_best("vulkan", silent, dir_path);
-    ggml_backend_try_load_best("opencl", silent, dir_path);
-    ggml_backend_try_load_best("musa", silent, dir_path);
-    ggml_backend_try_load_best("cpu", silent, dir_path);
+    ggml_backend_load_best("blas", silent, dir_path);
+    ggml_backend_load_best("cann", silent, dir_path);
+    ggml_backend_load_best("cuda", silent, dir_path);
+    ggml_backend_load_best("hip", silent, dir_path);
+    ggml_backend_load_best("kompute", silent, dir_path);
+    ggml_backend_load_best("metal", silent, dir_path);
+    ggml_backend_load_best("rpc", silent, dir_path);
+    ggml_backend_load_best("sycl", silent, dir_path);
+    ggml_backend_load_best("vulkan", silent, dir_path);
+    ggml_backend_load_best("opencl", silent, dir_path);
+    ggml_backend_load_best("musa", silent, dir_path);
+    ggml_backend_load_best("cpu", silent, dir_path);
     // check the environment variable GGML_BACKEND_PATH to load an out-of-tree backend
     const char * backend_path = std::getenv("GGML_BACKEND_PATH");
     if (backend_path) {

ml/backend/ggml/ggml/src/ggml.go

@@ -7,20 +7,6 @@ package ggml
 // #include <stdlib.h>
 // #include "ggml-backend.h"
 // extern void sink(int level, char *text, void *user_data);
-// static struct ggml_backend_feature * first_feature(ggml_backend_get_features_t fp, ggml_backend_reg_t reg) { return fp(reg); }
-// static struct ggml_backend_feature * next_feature(struct ggml_backend_feature * feature) { return &feature[1]; }
-/*
-typedef enum { COMPILER_CLANG, COMPILER_GNUC, COMPILER_UNKNOWN } COMPILER;
-static COMPILER compiler_name(void) {
-#if defined(__clang__)
-	return COMPILER_CLANG;
-#elif defined(__GNUC__)
-	return COMPILER_GNUC;
-#else
-	return COMPILER_UNKNOWN;
-#endif
-}
-*/
 import "C"
 
 import (
@@ -30,7 +16,6 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
-	"strconv"
 	"strings"
 	"sync"
 	"unsafe"
@@ -105,43 +90,4 @@ var OnceLoad = sync.OnceFunc(func() {
 			visited[abspath] = struct{}{}
 		}
 	}
-
-	slog.Info("system", "", system{})
 })
-
-type system struct{}
-
-func (system) LogValue() slog.Value {
-	var attrs []slog.Attr
-	names := make(map[string]int)
-	for i := range C.ggml_backend_dev_count() {
-		r := C.ggml_backend_dev_backend_reg(C.ggml_backend_dev_get(i))
-
-		func() {
-			fName := C.CString("ggml_backend_get_features")
-			defer C.free(unsafe.Pointer(fName))
-
-			if fn := C.ggml_backend_reg_get_proc_address(r, fName); fn != nil {
-				var features []any
-				for f := C.first_feature(C.ggml_backend_get_features_t(fn), r); f.name != nil; f = C.next_feature(f) {
-					features = append(features, C.GoString(f.name), C.GoString(f.value))
-				}
-
-				name := C.GoString(C.ggml_backend_reg_name(r))
-				attrs = append(attrs, slog.Group(name+"."+strconv.Itoa(names[name]), features...))
-				names[name] += 1
-			}
-		}()
-	}
-
-	switch C.compiler_name() {
-	case C.COMPILER_CLANG:
-		attrs = append(attrs, slog.String("compiler", "cgo(clang)"))
-	case C.COMPILER_GNUC:
-		attrs = append(attrs, slog.String("compiler", "cgo(gcc)"))
-	default:
-		attrs = append(attrs, slog.String("compiler", "cgo(unknown)"))
-	}
-
-	return slog.GroupValue(attrs...)
-}

ml/backend/ggml/ggml/src/gguf.cpp

@@ -777,14 +777,10 @@ enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id
 
 const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id) {
     GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
+    GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING);
     return ctx->kv[key_id].data.data();
 }
 
-size_t gguf_get_arr_data_n(const struct gguf_context * ctx, int64_t key_id) {
-    GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
-    return ctx->kv[key_id].data.size();
-}
-
 const char * gguf_get_arr_str(const struct gguf_context * ctx, int64_t key_id, size_t i) {
     GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
     GGML_ASSERT(ctx->kv[key_id].get_type() == GGUF_TYPE_STRING);
@@ -878,6 +874,7 @@ const char * gguf_get_val_str(const struct gguf_context * ctx, int64_t key_id) {
 const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id) {
     GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx));
     GGML_ASSERT(ctx->kv[key_id].get_ne() == 1);
+    GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING);
     return ctx->kv[key_id].data.data();
 }
 

model/input/input.go

@@ -1,37 +0,0 @@
-package input
-
-// Input represents one token in the input stream
-type Input struct {
-	// Token is a single element of text.
-	Token int32
-
-	// Multimodal is opaque data representing a non-text
-	// element such as an image (or part of one if the image
-	// can be processed in pieces). It may be either together
-	// with Token or on its own.
-	Multimodal any
-
-	// MultimodalHash is a unique representation of the data
-	// stored in Multimodal, used for caching and comparing
-	// equality.
-	MultimodalHash uint64
-}
-
-// MultimodalIndex is a multimodal element (such as an image)
-// together with an index into the slice of Inputs with the
-// corresponding token. Note that the index is not the same
-// as the position - to find that use the index with the
-// Positions slice.
-type MultimodalIndex struct {
-	Index      int
-	Multimodal any
-}
-
-// Options contains the inputs for a model forward pass
-type Options struct {
-	Inputs     []int32
-	Multimodal []MultimodalIndex
-	Positions  []int32
-	Sequences  []int
-	Outputs    []int32
-}

model/model.go

@@ -3,6 +3,7 @@ package model
 import (
 	"errors"
 	"fmt"
+	"image"
 	_ "image/jpeg"
 	_ "image/png"
 	"log/slog"
@@ -15,50 +16,23 @@ import (
 	_ "golang.org/x/image/tiff"
 	_ "golang.org/x/image/webp"
 
-	fs "github.com/ollama/ollama/fs/ggml"
 	"github.com/ollama/ollama/kvcache"
 	"github.com/ollama/ollama/ml"
 	_ "github.com/ollama/ollama/ml/backend"
-	"github.com/ollama/ollama/model/input"
 )
 
-// Model implements a specific model architecture, defining the forward pass and any model-specific configuration
-type Model interface {
-	Forward(ml.Context, input.Options) (ml.Tensor, error)
+// Options contains the inputs for a model forward pass
+type Options struct {
+	Inputs    []int32
+	Positions []int32
+	Sequences []int
+	Outputs   []int32
 
-	Backend() ml.Backend
-	Config() config
+	Images []image.Image
 }
 
-// MultimodalProcessor must be implemented by multimodal models.
-type MultimodalProcessor interface {
-	// EncodeMultimodal processes a single input (such as an image) and
-	// generates an output (typically an embedding) that can be used by the model.
-	//
-	// The return value is most typically an ml.Tensor, however, different
-	// type are possible, such as an object containing a tensor plus
-	// additional metadata, a slice of tensors or even just the original input.
-	//
-	// The result may be cached by the runner.
-	EncodeMultimodal(ml.Context, []byte) (any, error)
-
-	// PostTokenize is called after tokenization to allow the model to edit the
-	// input stream to correctly arrange multimodal elements.
-	//
-	// The input is a slice of tokens with the results of EncodeMultimodal interleaved
-	// in the order that the user provided them. Each element of the slice will be
-	// either a single token or single multimodal object.
-	//
-	// The model must ensure that inputs are stored according to how they will be
-	// processed and stored in the cache. For example, Llava-style models should insert
-	// placeholder tokens equal to the feature size of the corresponding image with
-	// the image itself attached to and split across these tokens. When Forward is called
-	// a partial subset of these tokens may be submitted according to the batch size.
-	//
-	// 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(ml.Context, []input.Input) ([]input.Input, error)
+type config struct {
+	Cache kvcache.Cache
 }
 
 // Base implements the common fields and methods for all models
@@ -67,10 +41,6 @@ type Base struct {
 	config
 }
 
-type config struct {
-	Cache kvcache.Cache
-}
-
 // Backend returns the underlying backend that will run the model
 func (m *Base) Backend() ml.Backend {
 	return m.b
@@ -80,6 +50,14 @@ func (m *Base) Config() config {
 	return m.config
 }
 
+// Model implements a specific model architecture, defining the forward pass and any model-specific configuration
+type Model interface {
+	Forward(ml.Context, Options) (ml.Tensor, error)
+
+	Backend() ml.Backend
+	Config() config
+}
+
 var models = make(map[string]func(ml.Config) (Model, error))
 
 // Register registers a model constructor for the given architecture
@@ -122,36 +100,6 @@ func New(modelPath string, params ml.BackendParams) (Model, error) {
 	return m, nil
 }
 
-func NewTextProcessor(s string) (TextProcessor, error) {
-	r, err := os.Open(s)
-	if err != nil {
-		return nil, err
-	}
-	defer r.Close()
-	meta, _, err := fs.Decode(r, -1)
-	if err != nil {
-		return nil, err
-	}
-	return getTextProcessor(meta.KV())
-}
-
-func getTextProcessor(kv fs.KV) (TextProcessor, error) {
-	arch := kv.Architecture()
-	f, ok := models[arch]
-	if !ok {
-		return nil, fmt.Errorf("unsupported model architecture %q", arch)
-	}
-	m, err := f(kv)
-	if err != nil {
-		return nil, err
-	}
-	tp, ok := m.(TextProcessor)
-	if !ok {
-		return nil, fmt.Errorf("%v is not a TextProcessor", m)
-	}
-	return tp, nil
-}
-
 func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
 	t := v.Type()
 
@@ -278,7 +226,7 @@ func canNil(t reflect.Type) bool {
 		t.Kind() == reflect.Slice
 }
 
-func Forward(ctx ml.Context, m Model, opts input.Options) (ml.Tensor, error) {
+func Forward(ctx ml.Context, m Model, opts 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))
 	}
@@ -289,7 +237,7 @@ func Forward(ctx ml.Context, m Model, opts input.Options) (ml.Tensor, error) {
 
 	cache := m.Config().Cache
 	if cache != nil {
-		err := cache.StartForward(ctx, opts)
+		err := cache.StartForward(ctx, opts.Positions, opts.Sequences)
 		if err != nil {
 			return nil, err
 		}

model/model_test.go

@@ -3,15 +3,12 @@ package model
 import (
 	"reflect"
 	"slices"
-	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
-	fs "github.com/ollama/ollama/fs/ggml"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/ml/backend/ggml"
 	"github.com/ollama/ollama/ml/nn"
-	"github.com/ollama/ollama/model/input"
 )
 
 func TestParseTags(t *testing.T) {
@@ -137,40 +134,3 @@ func TestPopulateFieldsAlternateName(t *testing.T) {
 		t.Errorf("populateFields() set incorrect values (-want +got):\n%s", diff)
 	}
 }
-
-func TestGetTextProcessor(t *testing.T) {
-	tp, err := getTextProcessor(fs.KV{})
-	if err == nil {
-		t.Error("expected error")
-	} else if !strings.Contains(err.Error(), "unsupported model architecture") {
-		t.Errorf("unexpected error: %v", err)
-	} else if tp != nil {
-		t.Error("expected nil tp")
-	}
-
-	models["dummy"] = func(ml.Config) (Model, error) {
-		return notTextProcessorModel{}, nil
-	}
-	tp, err = getTextProcessor(fs.KV{"general.architecture": "dummy"})
-	if err == nil {
-		t.Error("expected error")
-	} else if !strings.Contains(err.Error(), "not a TextProcessor") {
-		t.Errorf("unexpected error: %v", err)
-	} else if tp != nil {
-		t.Error("expected nil tp")
-	}
-}
-
-type notTextProcessorModel struct{}
-
-func (notTextProcessorModel) Forward(ml.Context, input.Options) (ml.Tensor, error) {
-	panic("unimplemented")
-}
-
-func (notTextProcessorModel) Backend() ml.Backend {
-	panic("unimplemented")
-}
-
-func (notTextProcessorModel) Config() config {
-	panic("unimplemented")
-}

model/models/llama/model.go

@@ -1,18 +1,16 @@
 package llama
 
 import (
-	"fmt"
 	"math"
-	"strings"
 
 	"github.com/ollama/ollama/kvcache"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/ml/nn"
 	"github.com/ollama/ollama/model"
-	"github.com/ollama/ollama/model/input"
 )
 
 type Options struct {
+	RopeFactors                      ml.Tensor `gguf:"rope_freqs.weight"`
 	hiddenSize, numHeads, numKVHeads int
 	eps, ropeBase, ropeScale         float32
 	ropeDim                          uint32
@@ -31,10 +29,6 @@ type Model struct {
 }
 
 func New(c ml.Config) (model.Model, error) {
-	if !strings.EqualFold(c.String("tokenizer.ggml.model"), "gpt2") {
-		return nil, fmt.Errorf("tokenizer %s not yet supported", c.String("tokenizer.ggml.model"))
-	}
-
 	m := Model{
 		BytePairEncoding: model.NewBytePairEncoding(
 			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+`),
@@ -66,11 +60,10 @@ func New(c ml.Config) (model.Model, error) {
 }
 
 type SelfAttention struct {
-	Query       *nn.Linear `gguf:"attn_q"`
-	Key         *nn.Linear `gguf:"attn_k"`
-	Value       *nn.Linear `gguf:"attn_v"`
-	Output      *nn.Linear `gguf:"attn_output"`
-	RopeFactors ml.Tensor  `gguf:"rope_freqs.weight"`
+	Query  *nn.Linear `gguf:"attn_q"`
+	Key    *nn.Linear `gguf:"attn_k"`
+	Value  *nn.Linear `gguf:"attn_v"`
+	Output *nn.Linear `gguf:"attn_output"`
 }
 
 func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
@@ -79,11 +72,11 @@ func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Ten
 
 	q := sa.Query.Forward(ctx, hiddenState)
 	q = q.Reshape(ctx, headDim, opts.numHeads, batchSize)
-	q = q.RoPE(ctx, positionIDs, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
+	q = q.RoPE(ctx, positionIDs, opts.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
 
 	k := sa.Key.Forward(ctx, hiddenState)
 	k = k.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
-	k = k.RoPE(ctx, positionIDs, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
+	k = k.RoPE(ctx, positionIDs, opts.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
 
 	v := sa.Value.Forward(ctx, hiddenState)
 	v = v.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
@@ -96,7 +89,7 @@ func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Ten
 }
 
 func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
-	return key.RoPE(ctx, shift, m.Layers[layer].SelfAttention.RopeFactors, m.ropeDim, m.ropeBase, m.ropeScale), nil
+	return key.RoPE(ctx, shift, m.Options.RopeFactors, m.Options.ropeDim, m.Options.ropeBase, m.Options.ropeScale), nil
 }
 
 type MLP struct {
@@ -138,18 +131,18 @@ func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs, outputs ml.Ten
 	return hiddenState.Add(ctx, residual)
 }
 
-func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
-	inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs))
+func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
+	inputs, err := ctx.FromIntSlice(opts.Inputs, len(opts.Inputs))
 	if err != nil {
 		return nil, err
 	}
 
-	positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions))
+	positions, err := ctx.FromIntSlice(opts.Positions, len(opts.Positions))
 	if err != nil {
 		return nil, err
 	}
 
-	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}

model/models/mllama/model.go

@@ -1,18 +1,10 @@
 package mllama
 
 import (
-	"bytes"
-	"encoding/binary"
-	"fmt"
-	"hash/fnv"
-	"image"
-	"slices"
-
 	"github.com/ollama/ollama/kvcache"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/ml/nn"
 	"github.com/ollama/ollama/model"
-	"github.com/ollama/ollama/model/input"
 )
 
 type Model struct {
@@ -33,10 +25,6 @@ const (
 )
 
 func New(c ml.Config) (model.Model, error) {
-	// Verify unified config
-	if c.Uint("vision.block_count") == 0 {
-		return nil, fmt.Errorf("non-unified vision model not supported")
-	}
 	m := Model{
 		BytePairEncoding: model.NewBytePairEncoding(
 			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+`),
@@ -62,92 +50,54 @@ func New(c ml.Config) (model.Model, error) {
 	return &m, nil
 }
 
-func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, error) {
-	image, _, err := image.Decode(bytes.NewReader(multimodalData))
-	if err != nil {
-		return nil, err
-	}
-
-	f32s, aspectRatioID, err := m.ImageProcessor.ProcessImage(image)
-	if err != nil {
-		return nil, err
-	}
-
-	pixelValues, err := ctx.Input().FromFloatSlice(f32s,
-		m.ImageProcessor.imageSize,
-		m.ImageProcessor.imageSize,
-		m.ImageProcessor.numChannels,
-		m.ImageProcessor.maxNumTiles,
-	)
-	if err != nil {
-		return nil, err
-	}
-
-	aspectRatio, err := ctx.Input().FromIntSlice([]int32{int32(aspectRatioID)}, 1)
-	if err != nil {
-		return nil, err
-	}
-
-	positions := make([]int32, 1601)
-	for i := range positions {
-		positions[i] = int32(i)
-	}
-
-	positionIDs, err := ctx.Input().FromIntSlice(positions, len(positions))
-	if err != nil {
-		return nil, err
-	}
-
-	crossAttentionStates := m.VisionModel.Forward(ctx, pixelValues, positionIDs, aspectRatio)
-	return m.Projector.Forward(ctx, crossAttentionStates), nil
-}
-
-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 = images[0].Multimodal
-				inputs[i].MultimodalHash = images[0].MultimodalHash
-				for j := 1; j < len(images); j++ {
-					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)
-					inputs[i].MultimodalHash = fnvHash.Sum64()
-				}
-				images = nil
-			}
-		} else {
-			images = append(images, inputs[i])
-			inputs[i].Token = -1
-		}
-	}
-
-	inputs = slices.DeleteFunc(inputs, func(input input.Input) bool { return input.Token == -1 })
-
-	return inputs, nil
-}
-
-func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
+func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
 	var crossAttentionStates ml.Tensor
-	if len(opts.Multimodal) > 0 {
-		crossAttentionStates = opts.Multimodal[len(opts.Multimodal)-1].Multimodal.(ml.Tensor)
+	if opts.Images != nil {
+		f32s, aspectRatioID, err := m.ImageProcessor.ProcessImage(opts.Images[0])
+		if err != nil {
+			return nil, err
+		}
+
+		pixelValues, err := ctx.FromFloatSlice(f32s,
+			m.ImageProcessor.imageSize,
+			m.ImageProcessor.imageSize,
+			m.ImageProcessor.numChannels,
+			m.ImageProcessor.maxNumTiles,
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		aspectRatio, err := ctx.FromIntSlice([]int32{int32(aspectRatioID)}, 1)
+		if err != nil {
+			return nil, err
+		}
+
+		positions := make([]int32, 1601)
+		for i := range positions {
+			positions[i] = int32(i)
+		}
+
+		positionIDs, err := ctx.FromIntSlice(positions, len(positions))
+		if err != nil {
+			return nil, err
+		}
+
+		crossAttentionStates = m.VisionModel.Forward(ctx, pixelValues, positionIDs, aspectRatio)
+		crossAttentionStates = m.Projector.Forward(ctx, crossAttentionStates)
 	}
 
-	inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs))
+	inputs, err := ctx.FromIntSlice(opts.Inputs, len(opts.Inputs))
 	if err != nil {
 		return nil, err
 	}
 
-	positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions))
+	positions, err := ctx.FromIntSlice(opts.Positions, len(opts.Positions))
 	if err != nil {
 		return nil, err
 	}
 
-	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}

model/models/mllama/model_text.go

@@ -10,11 +10,10 @@ import (
 )
 
 type TextSelfAttention struct {
-	Query       *nn.Linear `gguf:"attn_q"`
-	Key         *nn.Linear `gguf:"attn_k"`
-	Value       *nn.Linear `gguf:"attn_v"`
-	Output      *nn.Linear `gguf:"attn_output"`
-	RopeFactors ml.Tensor  `gguf:"rope_freqs.weight"`
+	Query  *nn.Linear `gguf:"attn_q"`
+	Key    *nn.Linear `gguf:"attn_k"`
+	Value  *nn.Linear `gguf:"attn_v"`
+	Output *nn.Linear `gguf:"attn_output"`
 }
 
 func (sa *TextSelfAttention) Forward(ctx ml.Context, hiddenState, positions, _ ml.Tensor, cache *kvcache.WrapperCache, opts *TextModelOptions) ml.Tensor {
@@ -23,11 +22,11 @@ func (sa *TextSelfAttention) Forward(ctx ml.Context, hiddenState, positions, _ m
 
 	query := sa.Query.Forward(ctx, hiddenState)
 	query = query.Reshape(ctx, headDim, opts.numHeads, batchSize)
-	query = query.RoPE(ctx, positions, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
+	query = query.RoPE(ctx, positions, opts.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
 
 	key := sa.Key.Forward(ctx, hiddenState)
 	key = key.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
-	key = key.RoPE(ctx, positions, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
+	key = key.RoPE(ctx, positions, opts.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale)
 
 	value := sa.Value.Forward(ctx, hiddenState)
 	value = value.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
@@ -40,11 +39,8 @@ func (sa *TextSelfAttention) Forward(ctx ml.Context, hiddenState, positions, _ m
 }
 
 func (m *TextModel) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
-	if sa, ok := m.Transformer.Layers[layer].(*TextSelfAttentionDecoderLayer); ok {
-		return key.RoPE(ctx, shift, sa.SelfAttention.RopeFactors, m.ropeDim, m.ropeBase, m.ropeScale), nil
-	}
-
-	return key, nil
+	// This will only get called for layers in the causal cache, which are just the self attention layers
+	return key.RoPE(ctx, shift, m.RopeFactors, m.ropeDim, m.ropeBase, m.ropeScale), nil
 }
 
 type TextMLP struct {
@@ -195,6 +191,8 @@ func (d *TextDecoder) Forward(ctx ml.Context, hiddenState, positionIDs, outputs,
 }
 
 type TextModelOptions struct {
+	RopeFactors ml.Tensor `gguf:"rope_freqs.weight"`
+
 	hiddenSize, numHeads, numKVHeads int
 	eps, ropeBase, ropeScale         float32
 	ropeDim                          uint32

model/process_text.go

@@ -19,7 +19,7 @@ const (
 )
 
 type TextProcessor interface {
-	Encode(s string, addSpecial bool) ([]int32, error)
+	Encode(string) ([]int32, error)
 	Decode([]int32) (string, error)
 	Is(int32, Special) bool
 }
@@ -144,7 +144,7 @@ type merge struct {
 	runes []rune
 }
 
-func (bpe BytePairEncoding) Encode(s string, addSpecial bool) ([]int32, error) {
+func (bpe BytePairEncoding) Encode(s string) ([]int32, error) {
 	fragments := []fragment{{value: s}}
 	for _, special := range bpe.vocab.SpecialVocabulary() {
 		// TODO: process special tokens concurrently
@@ -177,6 +177,7 @@ func (bpe BytePairEncoding) Encode(s string, addSpecial bool) ([]int32, error) {
 	for _, frag := range fragments {
 		if len(frag.ids) > 0 {
 			ids = append(ids, frag.ids...)
+			slog.Debug("encoded", "text", frag.value, "ids", frag.ids, "special", true)
 			continue
 		}
 
@@ -200,6 +201,7 @@ func (bpe BytePairEncoding) Encode(s string, addSpecial bool) ([]int32, error) {
 			// short circuit if the fragment is in the vocabulary
 			if id := bpe.vocab.Encode(sb.String()); id >= 0 {
 				ids = append(ids, id)
+				slog.Debug("encoded", "text", sb.String(), "ids", []int32{id})
 				continue
 			}
 
@@ -273,13 +275,14 @@ func (bpe BytePairEncoding) Encode(s string, addSpecial bool) ([]int32, error) {
 					// TODO: handle the edge case where the rune isn't in the vocabulary
 					if id := bpe.vocab.Encode(string(merge.runes)); id >= 0 {
 						ids = append(ids, id)
+						slog.Debug("encoded", "text", string(merge.runes), "ids", []int32{id})
 					}
 				}
 			}
 		}
 	}
 
-	if addSpecial && len(ids) > 0 {
+	if len(ids) > 0 {
 		if bpe.vocab.AddBOS {
 			if ids[0] == bpe.vocab.BOS {
 				slog.Warn("adding bos token to prompt which already has it", "id", bpe.vocab.BOS)
@@ -326,5 +329,6 @@ func (bpe BytePairEncoding) Decode(ids []int32) (string, error) {
 		}
 	}
 
+	slog.Debug("decoded", "ids", ids, "text", sb.String())
 	return sb.String(), nil
 }

model/process_text_test.go

@@ -74,7 +74,7 @@ func TestLlama(t *testing.T) {
 	t.Run("simple", func(t *testing.T) {
 		t.Parallel()
 
-		ids, err := tokenizer.Encode("hello world", true)
+		ids, err := tokenizer.Encode("hello world")
 		if err != nil {
 			t.Error(err)
 		}
@@ -92,7 +92,7 @@ func TestLlama(t *testing.T) {
 			t.Errorf("got %q, want hello world", s)
 		}
 
-		ids, err = tokenizer.Encode("hello <|end_of_text|>", true)
+		ids, err = tokenizer.Encode("hello <|end_of_text|>")
 		if err != nil {
 			t.Error(err)
 		}
@@ -126,7 +126,7 @@ func TestLlama(t *testing.T) {
 		}
 
 		for s, want := range cases {
-			ids, err := tokenizer.Encode(s, true)
+			ids, err := tokenizer.Encode(s)
 			if err != nil {
 				t.Error(err)
 			}
@@ -152,7 +152,7 @@ func TestLlama(t *testing.T) {
 		}
 
 		for _, want := range cases {
-			ids, err := tokenizer.Encode(want, true)
+			ids, err := tokenizer.Encode(want)
 			if err != nil {
 				t.Error(err)
 			}
@@ -176,7 +176,7 @@ func TestLlama(t *testing.T) {
 		}
 
 		for s, want := range cases {
-			ids, err := tokenizer.Encode(s, true)
+			ids, err := tokenizer.Encode(s)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -222,7 +222,7 @@ func BenchmarkBytePairEncoding(b *testing.B) {
 		b.Run("encode"+strconv.Itoa(n), func(b *testing.B) {
 			b.ResetTimer()
 			for range b.N {
-				_, err := tokenizer.Encode(string(bts), true)
+				_, err := tokenizer.Encode(string(bts))
 				if err != nil {
 					b.Fatal(err)
 				}
@@ -230,7 +230,7 @@ func BenchmarkBytePairEncoding(b *testing.B) {
 		})
 
 		b.Run("decode"+strconv.Itoa(n), func(b *testing.B) {
-			ids, err := tokenizer.Encode(string(bts), true)
+			ids, err := tokenizer.Encode(string(bts))
 			if err != nil {
 				b.Fatal(err)
 			}

runner/llamarunner/runner.go

@@ -931,6 +931,7 @@ func Execute(args []string) error {
 	slog.Info("starting go runner")
 
 	llama.BackendInit()
+	slog.Info("system", "info", llama.PrintSystemInfo(), "threads", *threads)
 
 	server := &Server{
 		batchSize: *batchSize,

runner/ollamarunner/cache.go

@@ -5,12 +5,12 @@ import (
 	"fmt"
 	"log/slog"
 	"math"
+	"reflect"
 	"time"
 
 	"github.com/ollama/ollama/kvcache"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/model"
-	"github.com/ollama/ollama/model/input"
 )
 
 type InputCache struct {
@@ -39,7 +39,10 @@ func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots
 	slots := make([]InputCacheSlot, numSlots)
 
 	for i := range slots {
-		slots[i] = InputCacheSlot{Id: i}
+		slots[i] = InputCacheSlot{
+			Id:     i,
+			Inputs: make([]input, 0),
+		}
 	}
 
 	cache := model.Config().Cache
@@ -59,9 +62,9 @@ func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots
 func kvCacheTypeFromStr(s string) ml.DType {
 	switch s {
 	case "q8_0":
-		return ml.DTypeQ80
+		panic("kv cache quantization not yet implemented")
 	case "q4_0":
-		return ml.DTypeQ40
+		panic("kv cache quantization not yet implemented")
 	default:
 		return ml.DTypeF16
 	}
@@ -80,7 +83,7 @@ type InputCacheSlot struct {
 	Id int
 
 	// Inputs that are stored in the KV cache
-	Inputs []input.Input
+	Inputs []input
 
 	// is this cache actively being processed as part of a sequence?
 	InUse bool
@@ -89,7 +92,7 @@ type InputCacheSlot struct {
 	lastUsed time.Time
 }
 
-func (c *InputCache) LoadCacheSlot(prompt []input.Input, cachePrompt bool) (*InputCacheSlot, []input.Input, error) {
+func (c *InputCache) LoadCacheSlot(prompt []input, cachePrompt bool) (*InputCacheSlot, []input, error) {
 	var slot *InputCacheSlot
 	var numPast int32
 	var err error
@@ -140,7 +143,7 @@ func (c *InputCache) LoadCacheSlot(prompt []input.Input, cachePrompt bool) (*Inp
 	return slot, prompt, nil
 }
 
-func (c *InputCache) findLongestCacheSlot(prompt []input.Input) (*InputCacheSlot, int32, error) {
+func (c *InputCache) findLongestCacheSlot(prompt []input) (*InputCacheSlot, int32, error) {
 	longest := int32(-1)
 	var longestSlot *InputCacheSlot
 
@@ -163,7 +166,7 @@ func (c *InputCache) findLongestCacheSlot(prompt []input.Input) (*InputCacheSlot
 	return longestSlot, longest, nil
 }
 
-func (c *InputCache) findBestCacheSlot(prompt []input.Input) (*InputCacheSlot, int32, error) {
+func (c *InputCache) findBestCacheSlot(prompt []input) (*InputCacheSlot, int32, error) {
 	oldest := time.Now()
 	var oldestSlot *InputCacheSlot
 
@@ -199,7 +202,7 @@ func (c *InputCache) findBestCacheSlot(prompt []input.Input) (*InputCacheSlot, i
 	if longest > 0 && longestSlot != oldestSlot {
 		slog.Debug("forking cache slot", "src", longestSlot.Id, "dst", oldestSlot.Id, "inputs", longest, "total",
 			len(longestSlot.Inputs))
-		oldestSlot.Inputs = make([]input.Input, longest)
+		oldestSlot.Inputs = make([]input, longest)
 		copy(oldestSlot.Inputs, longestSlot.Inputs[:longest])
 		if c.cache != nil {
 			c.cache.CopyPrefix(longestSlot.Id, oldestSlot.Id, longest)
@@ -209,7 +212,7 @@ func (c *InputCache) findBestCacheSlot(prompt []input.Input) (*InputCacheSlot, i
 	return oldestSlot, longest, nil
 }
 
-func countCommonPrefix(a []input.Input, b []input.Input) int32 {
+func countCommonPrefix(a []input, b []input) int32 {
 	var count int32
 
 	for i := range a {
@@ -217,7 +220,7 @@ func countCommonPrefix(a []input.Input, b []input.Input) int32 {
 			break
 		}
 
-		if a[i].Token != b[i].Token || a[i].MultimodalHash != b[i].MultimodalHash {
+		if !reflect.DeepEqual(a[i], b[i]) {
 			break
 		}
 

runner/ollamarunner/cache_test.go

@@ -4,8 +4,6 @@ import (
 	"image"
 	"testing"
 	"time"
-
-	"github.com/ollama/ollama/model/input"
 )
 
 func TestCountCommon(t *testing.T) {
@@ -15,50 +13,44 @@ func TestCountCommon(t *testing.T) {
 
 	tests := []struct {
 		name     string
-		t1       []input.Input
-		t2       []input.Input
+		t1       []input
+		t2       []input
 		expected int32
 	}{
 		{
 			name:     "Equal",
-			t1:       []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
-			t2:       []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
+			t1:       []input{{token: 1}, {token: 2}, {token: 3}},
+			t2:       []input{{token: 1}, {token: 2}, {token: 3}},
 			expected: 3,
 		},
 		{
 			name:     "Prefix",
-			t1:       []input.Input{{Token: 1}},
-			t2:       []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
+			t1:       []input{{token: 1}},
+			t2:       []input{{token: 1}, {token: 2}, {token: 3}},
 			expected: 1,
 		},
 		{
 			name:     "Image Prefix",
-			t1:       []input.Input{{Multimodal: imgA, MultimodalHash: 1}},
-			t2:       []input.Input{{Multimodal: imgA, MultimodalHash: 1}, {Multimodal: imgB, MultimodalHash: 2}, {Multimodal: imgC, MultimodalHash: 3}},
+			t1:       []input{{image: imgA}},
+			t2:       []input{{image: imgA}, {image: imgB}, {image: imgC}},
 			expected: 1,
 		},
 		{
 			name:     "Mixed",
-			t1:       []input.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}},
-			t2:       []input.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}, {Token: 5}},
+			t1:       []input{{token: 1}, {image: imgA}},
+			t2:       []input{{token: 1}, {image: imgA}, {token: 5}},
 			expected: 2,
 		},
-		{
-			name:     "Mixed, Same Length",
-			t1:       []input.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}},
-			t2:       []input.Input{{Token: 1}, {Multimodal: imgB, MultimodalHash: 2}},
-			expected: 1,
-		},
 		{
 			name:     "Empty",
-			t1:       []input.Input{},
-			t2:       []input.Input{{Token: 1}, {Token: 2}, {Token: 3}},
+			t1:       []input{},
+			t2:       []input{{token: 1}, {token: 2}, {token: 3}},
 			expected: 0,
 		},
 		{
 			name:     "Both Empty",
-			t1:       []input.Input{},
-			t2:       []input.Input{},
+			t1:       []input{},
+			t2:       []input{},
 			expected: 0,
 		},
 	}
@@ -82,7 +74,7 @@ func TestFindCacheSlot(t *testing.T) {
 	tests := []struct {
 		name    string
 		cache   InputCache
-		prompt  []input.Input
+		prompt  []input
 		longest expected
 		best    expected
 	}{
@@ -91,18 +83,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input.Input{},
+					Inputs:   []input{},
 					InUse:    false,
 					lastUsed: time.Time{},
 				},
 				{
 					Id:       1,
-					Inputs:   []input.Input{},
+					Inputs:   []input{},
 					InUse:    false,
 					lastUsed: time.Time{},
 				},
 			}},
-			prompt:  []input.Input{{Token: 1}},
+			prompt:  []input{{token: 1}},
 			longest: expected{result: 0, len: 0},
 			best:    expected{result: 0, len: 0},
 		},
@@ -111,18 +103,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input.Input{{Token: 1}},
+					Inputs:   []input{{token: 1}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input.Input{{Token: 1}, {Token: 2}},
+					Inputs:   []input{{token: 1}, {token: 2}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-2 * time.Second),
 				},
 			}},
-			prompt:  []input.Input{{Token: 1}, {Token: 2}},
+			prompt:  []input{{token: 1}, {token: 2}},
 			longest: expected{result: 1, len: 2},
 			best:    expected{result: 1, len: 2},
 		},
@@ -131,18 +123,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input.Input{{Token: 1}, {Token: 2}},
+					Inputs:   []input{{token: 1}, {token: 2}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input.Input{},
+					Inputs:   []input{},
 					InUse:    false,
 					lastUsed: time.Time{},
 				},
 			}},
-			prompt:  []input.Input{{Token: 2}},
+			prompt:  []input{{token: 2}},
 			longest: expected{result: 0, len: 0},
 			best:    expected{result: 1, len: 0},
 		},
@@ -152,19 +144,19 @@ func TestFindCacheSlot(t *testing.T) {
 				slots: []InputCacheSlot{
 					{
 						Id:       0,
-						Inputs:   []input.Input{{Token: 1}, {Token: 2}},
+						Inputs:   []input{{token: 1}, {token: 2}},
 						InUse:    false,
 						lastUsed: time.Now().Add(-time.Second),
 					},
 					{
 						Id:       1,
-						Inputs:   []input.Input{},
+						Inputs:   []input{},
 						InUse:    false,
 						lastUsed: time.Time{},
 					},
 				},
 			},
-			prompt:  []input.Input{{Token: 1}},
+			prompt:  []input{{token: 1}},
 			longest: expected{result: 0, len: 1},
 			best:    expected{result: 1, len: 1},
 		},
@@ -173,18 +165,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input.Input{{Token: 1}},
+					Inputs:   []input{{token: 1}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input.Input{{Token: 1}, {Token: 2}},
+					Inputs:   []input{{token: 1}, {token: 2}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-2 * time.Second),
 				},
 			}},
-			prompt:  []input.Input{{Token: 2}, {Token: 3}},
+			prompt:  []input{{token: 2}, {token: 3}},
 			longest: expected{result: 0, len: 0},
 			best:    expected{result: 1, len: 0},
 		},
@@ -193,18 +185,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input.Input{{Token: 1}, {Token: 2}},
+					Inputs:   []input{{token: 1}, {token: 2}},
 					InUse:    true,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input.Input{{Token: 1}},
+					Inputs:   []input{{token: 1}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-2 * time.Second),
 				},
 			}},
-			prompt:  []input.Input{{Token: 1}, {Token: 2}},
+			prompt:  []input{{token: 1}, {token: 2}},
 			longest: expected{result: 1, len: 1},
 			best:    expected{result: 1, len: 2},
 		},

runner/ollamarunner/runner.go

@@ -1,12 +1,13 @@
 package ollamarunner
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
-	"hash/maphash"
+	"image"
 	"log"
 	"log/slog"
 	"net"
@@ -26,26 +27,28 @@ import (
 	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/model"
-	"github.com/ollama/ollama/model/input"
 	"github.com/ollama/ollama/runner/common"
 	"github.com/ollama/ollama/sample"
 
 	_ "github.com/ollama/ollama/model/models"
 )
 
-type Sequence struct {
-	// ctx for allocating tensors that last the lifetime of the sequence, such as
-	// multimodal embeddings
-	ctx ml.Context
+// input is an element of the prompt to process, either a token or an image
+type input struct {
+	token int32
 
+	image image.Image
+}
+
+type Sequence struct {
 	// batch index
 	iBatch int
 
 	// prompt inputs left to evaluate
-	inputs []input.Input
+	inputs []input
 
 	// inputs that have been added to a batch but not yet submitted to Forward
-	pendingInputs []input.Input
+	pendingInputs []input
 
 	// tokens that have been generated but not returned yet (e.g. for stop sequences)
 	pendingResponses []string
@@ -98,9 +101,8 @@ func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequen
 	s.ready.Wait()
 
 	startTime := time.Now()
-	ctx := s.model.Backend().NewContext()
 
-	inputs, err := s.inputs(ctx, prompt, images)
+	inputs, err := s.inputs(prompt, images)
 	if err != nil {
 		return nil, fmt.Errorf("failed to process inputs: %w", err)
 	} else if len(inputs) == 0 {
@@ -126,7 +128,6 @@ func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequen
 	// TODO(jessegross): Ingest cached history for grammar
 
 	return &Sequence{
-		ctx:                 ctx,
 		inputs:              inputs,
 		numPromptInputs:     len(inputs),
 		startProcessingTime: startTime,
@@ -145,31 +146,28 @@ func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequen
 // 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(ctx ml.Context, prompt string, images []ImageData) ([]input.Input, error) {
-	var inputs []input.Input
+func (s *Server) inputs(prompt string, images []ImageData) ([]input, error) {
+	var inputs []input
 	var parts []string
 	var matches [][]string
 
-	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
+	// TODO(jessegross): This can sometimes trigger for matching text in the
+	// user's prompt. We previously tried to avoid it by only looking for images
+	// on image models. We don't have a clear indication now but it would be better
+	// to properly escape it in any case.
+	re := regexp.MustCompile(`\[img-(\d+)\]`)
+	parts = re.Split(prompt, -1)
+	matches = re.FindAllStringSubmatch(prompt, -1)
 
-	if visionModel {
-		re := regexp.MustCompile(`\[img-(\d+)\]`)
-		parts = re.Split(prompt, -1)
-		matches = re.FindAllStringSubmatch(prompt, -1)
-	} else {
-		parts = []string{prompt}
-	}
-
-	postTokenize := false
 	for i, part := range parts {
 		// text - tokenize
-		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
+		tokens, err := s.model.(model.TextProcessor).Encode(part)
 		if err != nil {
 			return nil, err
 		}
 
 		for _, t := range tokens {
-			inputs = append(inputs, input.Input{Token: t})
+			inputs = append(inputs, input{token: t})
 		}
 
 		// image - decode and store
@@ -188,25 +186,12 @@ func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]in
 				return nil, fmt.Errorf("invalid image index: %d", n)
 			}
 
-			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
+			image, _, err := image.Decode(bytes.NewReader(images[imageIndex].Data))
 			if err != nil {
 				return nil, err
 			}
 
-			s.multimodalHash.Reset()
-			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
-			imageHash := s.multimodalHash.Sum64()
-
-			inputs = append(inputs, input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash})
-			postTokenize = true
-		}
-	}
-
-	if visionModel && postTokenize {
-		var err error
-		inputs, err = multimodalProcessor.PostTokenize(ctx, inputs)
-		if err != nil {
-			return nil, err
+			inputs = append(inputs, input{image: image})
 		}
 	}
 
@@ -251,15 +236,8 @@ type Server struct {
 	// KV cache
 	cache *InputCache
 
-	// multimodalHash generates hashes for comparing equality
-	// of non-text data
-	multimodalHash maphash.Hash
-
-	// vocab is a llama.cpp vocab required for gammar-based
-	// constrained generation (json mode, structured outputs)
-	// TODO: this is temporary until Ollama sampling supports
-	// constrained generation
-	vocab *sample.Vocab
+	// next sequence for prompt processing to avoid starvation
+	nextSeq int
 }
 
 func (s *Server) allNil() bool {
@@ -305,7 +283,6 @@ 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)
 }
@@ -333,25 +310,30 @@ func (s *Server) processBatch() error {
 	}
 	defer s.mu.Unlock()
 
-	var options input.Options
+	var options model.Options
+	imgSeq := -1
+
+	seqIdx := s.nextSeq - 1
+	for range s.seqs {
+		seqIdx = (seqIdx + 1) % len(s.seqs)
+		seq := s.seqs[seqIdx]
 
-	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(i, "limit")
+			s.removeSequence(seqIdx, "limit")
 			continue
 		}
 
 		if !s.cache.enabled {
 			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
-			seq.cache.Inputs = []input.Input{}
+			seq.cache.Inputs = []input{}
 		}
 
-		for j, inp := range seq.inputs {
+		for i, input := 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)
@@ -363,23 +345,37 @@ func (s *Server) processBatch() error {
 				}
 			}
 
-			if j >= s.batchSize {
+			if i >= s.batchSize {
 				break
 			}
 
-			options.Inputs = append(options.Inputs, inp.Token)
-			if inp.Multimodal != nil {
-				options.Multimodal = append(options.Multimodal, input.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: inp.Multimodal})
+			// TODO(jessegross): Image inputs need to be rethought - it's
+			// it doesn't work well for different types of models or multiple sequences
+			if input.image != nil {
+				if len(seq.pendingInputs) != len(options.Images) {
+					break
+				}
+
+				if imgSeq != seqIdx && imgSeq != -1 {
+					s.nextSeq = seqIdx
+					break
+				}
+
+				imgSeq = seqIdx
+				options.Images = append(options.Images, input.image)
+				seq.pendingInputs = append(seq.pendingInputs, input)
+				continue
 			}
 
+			options.Inputs = append(options.Inputs, input.token)
 			options.Positions = append(options.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs)))
 			options.Sequences = append(options.Sequences, seq.cache.Id)
 
 			seq.iBatch = len(options.Outputs)
-			if j+1 == len(seq.inputs) {
+			if i+1 == len(seq.inputs) {
 				options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1))
 			}
-			seq.pendingInputs = append(seq.pendingInputs, inp)
+			seq.pendingInputs = append(seq.pendingInputs, input)
 		}
 
 		seq.inputs = seq.inputs[len(seq.pendingInputs):]
@@ -407,7 +403,7 @@ func (s *Server) processBatch() error {
 		// After calling Forward, pending inputs are now in the cache
 		if len(seq.pendingInputs) > 0 {
 			seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...)
-			seq.pendingInputs = []input.Input{}
+			seq.pendingInputs = []input{}
 		}
 
 		// don't sample prompt processing
@@ -426,7 +422,6 @@ func (s *Server) processBatch() error {
 		// if done processing the prompt, generate an embedding and return
 		if seq.embeddingOnly {
 			// TODO(jessegross): Embedding support
-			slog.Warn("generation of embedding outputs not yet supported")
 			s.removeSequence(i, "")
 			continue
 		}
@@ -454,7 +449,7 @@ func (s *Server) processBatch() error {
 			return err
 		}
 
-		seq.inputs = []input.Input{{Token: token}}
+		seq.inputs = []input{{token: token}}
 
 		seq.pendingResponses = append(seq.pendingResponses, piece)
 		sequence := strings.Join(seq.pendingResponses, "")
@@ -580,30 +575,11 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	var grammar *sample.Grammar
-	var err error
-	if req.Grammar != "" {
-		grammar, err = sample.NewGrammar(s.vocab, req.Grammar)
-		if err != nil {
-			http.Error(w, "failed to load model vocabulary required for format", http.StatusInternalServerError)
-			return
-		}
-	}
-
-	sampler := sample.NewSampler(
-		req.Temperature,
-		req.TopK,
-		req.TopP,
-		req.MinP,
-		req.Seed,
-		grammar,
-	)
-
 	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
 		numPredict: req.NumPredict,
 		stop:       req.Stop,
 		numKeep:    int32(req.NumKeep),
-		sampler:    sampler,
+		sampler:    sample.Greedy(), // TODO: add support for different samplers when performance is optimized
 		embedding:  false,
 	})
 	if err != nil {
@@ -810,7 +786,7 @@ func (s *Server) loadModel(
 		panic(err)
 	}
 
-	s.vocab = sample.NewVocab(mpath)
+	slog.Info("system", "info", s.model.Backend().SystemInfo(), "threads", params.NumThreads)
 
 	// TODO(jessegross): LoRA loading
 	if lpath.String() != "" {

sample/samplers.go

@@ -3,225 +3,118 @@ package sample
 import (
 	"errors"
 	"math"
-	"math/rand/v2"
-	"slices"
-	"sync"
 
-	"github.com/ollama/ollama/llama"
+	"golang.org/x/exp/rand"
+	"gonum.org/v1/gonum/stat/sampleuv"
 )
 
-// token represents information about a single token during sampling
-type token struct {
-	id    int32   // The token's unique identifier
-	value float32 // The raw logit or probability from the model
+type Sampler interface {
+	Sample([]float32) (int32, error)
 }
 
-type Sampler struct {
-	rng         *rand.Rand
-	topK        int
-	topP        float32
-	minP        float32
-	temperature float32
-	grammar     *Grammar
+type weighted struct {
+	src        rand.Source
+	transforms []Transform
 }
 
-func (s *Sampler) Sample(logits []float32) (int32, error) {
-	tokens := make([]token, len(logits))
+// TODO(parthsareen): remove uv sample dependency https://github.com/ollama/ollama/issues/9279
+func Weighted(seed *uint64, transforms ...Transform) Sampler {
+	var src rand.Source
+	if seed != nil {
+		src = rand.NewSource(*seed)
+	}
+	return weighted{src: src, transforms: transforms}
+}
+
+func (s weighted) Sample(logits []float32) (int32, error) {
+	logits64 := make([]float64, len(logits))
+	for i, v := range logits {
+		logits64[i] = float64(v)
+	}
+
+	for _, t := range s.transforms {
+		logits64 = t.Apply(logits64)
+	}
+
+	logitsCopy := make([]float64, 0, len(logits))
+	indices := make([]int, 0, len(logits))
+	for i, logit := range logits64 {
+		if !math.IsInf(logit, -1) {
+			logitsCopy = append(logitsCopy, logit)
+			indices = append(indices, i)
+		}
+	}
+
+	if len(logitsCopy) == 0 {
+		return -1, errors.New("no valid logits found for weighed sampling")
+	}
+
+	probs := softmax(logitsCopy)
+	w := sampleuv.NewWeighted(probs, s.src)
+	if idx, ok := w.Take(); ok {
+		return int32(indices[idx]), nil
+	}
+	return -1, errors.New("weighted sampler failed, no valid token found")
+}
+
+type greedy struct{}
+
+func Greedy() Sampler {
+	return greedy{}
+}
+
+// Sample returns the index of the maximum value in logits.
+func (s greedy) Sample(logits []float32) (int32, error) {
+	if len(logits) == 0 {
+		return -1, errors.New("no logits provided for greedy sampling")
+	}
+
+	maxIdx := 0
 	for i := range logits {
-		tokens[i].id = int32(i)
-		tokens[i].value = logits[i]
-	}
-
-	t, err := s.sample(tokens)
-	if err != nil {
-		return -1, err
-	}
-
-	if s.grammar != nil {
-		// optimization: first check if the max logit is accepted by the grammar
-		// if the max logit is rejected, apply the grammar to all logits (slower)
-		top := []token{t}
-		s.grammar.Apply(top)
-		if !math.IsInf(float64(top[0].value), -1) {
-			s.grammar.Accept(top[0].id)
-			return top[0].id, nil
-		}
-
-		// since .sample has side effects of modifying the tokens
-		// we need to reset them before applying the grammar and
-		// sampling again
-		for i := range logits {
-			tokens[i].id = int32(i)
-			tokens[i].value = logits[i]
-		}
-		s.grammar.Apply(tokens)
-		t, err = s.sample(tokens)
-		if err != nil {
-			return -1, err
-		}
-		s.grammar.Accept(t.id)
-	}
-
-	return t.id, nil
-}
-
-// greedy returns the highest probability token from the tokens
-func greedy(tokens []token) token {
-	max := tokens[0]
-	for i := 1; i < len(tokens); i++ {
-		if tokens[i].value > max.value {
-			max = tokens[i]
+		if logits[i] > logits[maxIdx] {
+			maxIdx = i
 		}
 	}
 
-	return max
-}
-
-// sample returns the highest probability token from the tokens
-// given sampler parameters. It also has side effects of modifying the tokens
-func (s *Sampler) sample(tokens []token) (token, error) {
-	if s.temperature == 0 {
-		return greedy(tokens), nil
-	}
-
-	if s.topK > 0 {
-		tokens = topK(tokens, s.topK)
-	} else {
-		sortLogits(tokens)
-	}
-
-	tokens = temperature(tokens, s.temperature)
-	tokens = softmax(tokens)
-	tokens = topP(tokens, s.topP)
-	tokens = minP(tokens, s.minP)
-
-	// 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")
-	}
-
-	var r float32
-	if s.rng != nil {
-		r = s.rng.Float32()
-	} else {
-		r = rand.Float32()
-	}
-
-	// 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
+	return int32(maxIdx), 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))
-	}
-	if temperature < 0.0 {
-		temperature = 0.0
+func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int) (Sampler, error) {
+	if temperature == 0 {
+		return Greedy(), nil
 	}
 
-	if topP < 0.0 {
-		topP = 0.0
-	}
-	if topP >= 1.0 {
-		topP = 1.0
+	if temperature < 0 || temperature > 2 {
+		return nil, errors.New("temperature must be between 0 and 2")
 	}
 
-	if minP < 0.0 {
-		minP = 0.0
-	}
-	if minP >= 1.0 {
-		minP = 1.0
-	}
+	transforms := []Transform{Temperature(temperature)}
 
-	return Sampler{
-		rng:         rng,
-		topK:        topK,
-		topP:        topP,
-		minP:        minP,
-		temperature: temperature,
-		grammar:     grammar,
-	}
-}
-
-type Grammar struct {
-	vocab   *Vocab
-	grammar string
-	sampler *llama.Sampler
-}
-
-func NewGrammar(vocab *Vocab, grammar string) (*Grammar, error) {
-	v, err := vocab.Load()
-	if err != nil {
-		return nil, err
-	}
-
-	return &Grammar{
-		vocab:   vocab,
-		grammar: grammar,
-		sampler: llama.NewGrammarSampler(v, grammar),
-	}, nil
-}
-
-func (g *Grammar) Apply(tokens []token) {
-	tds := make([]llama.TokenData, len(tokens))
-	for i, token := range tokens {
-		tds[i].Id = token.id
-		tds[i].Logit = token.value
-	}
-
-	g.sampler.Apply(tds)
-
-	for i := range tokens {
-		tokens[i].value = tds[i].Logit
-	}
-}
-
-func (g *Grammar) Accept(token int32) {
-	g.sampler.Accept(token)
-}
-
-type Vocab struct {
-	once  sync.Once
-	vocab *llama.Vocab
-	err   error
-	path  string
-}
-
-func NewVocab(path string) *Vocab {
-	return &Vocab{path: path}
-}
-
-// Load returns the lazily-loaded vocabulary
-func (v *Vocab) Load() (*llama.Vocab, error) {
-	v.once.Do(func() {
-		vocab, err := llama.LoadVocabFromFile(v.path)
-		if err != nil {
-			v.err = err
-			return
+	if topK != 0 {
+		if topK <= 0 {
+			return nil, errors.New("topK must be greater than 0")
 		}
-		v.vocab = vocab
-	})
-	return v.vocab, v.err
+		transforms = append(transforms, TopK(topK))
+	}
+
+	if topP != 0 {
+		if topP < 0 || topP >= 1 {
+			return nil, errors.New("topP must be between 0 and 1")
+		}
+		transforms = append(transforms, TopP(topP))
+	}
+
+	if minP != 0 {
+		if minP < 0 || minP >= 1 {
+			return nil, errors.New("minP must be between 0 and 1")
+		}
+		transforms = append(transforms, MinP(minP))
+	}
+
+	if seed >= 0 {
+		seed64 := uint64(seed)
+		return Weighted(&seed64, transforms...), nil
+	}
+	return Weighted(nil, transforms...), nil
 }

sample/samplers_benchmark_test.go

@@ -1,92 +0,0 @@
-package sample
-
-import (
-	"fmt"
-	"math/rand"
-	"testing"
-)
-
-func BenchmarkWeightedSampler(b *testing.B) {
-	sizes := []int{10, 100, 1000, 10000}
-
-	for _, size := range sizes {
-		b.Run(fmt.Sprintf("Size %d", size), func(b *testing.B) {
-			logits := make([]float32, size)
-			for i := range logits {
-				logits[i] = float32(rand.Float64()*10 - 5)
-			}
-
-			sampler := NewSampler(0.8, 0, 0, 0, 42, nil)
-			b.ResetTimer()
-			for b.Loop() {
-				sampler.Sample(logits)
-			}
-		})
-	}
-
-	configs := []struct {
-		name        string
-		temperature float32
-		topK        int
-		topP        float32
-		minP        float32
-		seed        int
-	}{
-		{"Greedy", 0, -1, 0, 0, -1},
-		{"Temperature", 0.8, -1, 0, 0, -1},
-		{"TopK", 0.8, 50, 0, 0, -1},
-		{"TopP", 0.8, -1, 0.9, 0, -1},
-		{"MinP", 0.8, -1, 0, 0.05, -1},
-		{"WithSeed", 0.8, 50, 0, 0, 42},
-	}
-
-	// Fixed size for common vocab size
-	size := 128000
-	logits := make([]float32, size)
-	for i := range logits {
-		logits[i] = float32(rand.Float64()*10 - 5)
-	}
-
-	for _, tc := range configs {
-		b.Run("Config"+tc.name, func(b *testing.B) {
-			sampler := NewSampler(tc.temperature, tc.topK, tc.topP, tc.minP, tc.seed, nil)
-			sampler.Sample(logits)
-
-			b.ResetTimer()
-
-			for b.Loop() {
-				sampler.Sample(logits)
-			}
-		})
-	}
-
-	// Test with combined transforms separately - topK influences performance greatly
-	b.Run("TransformCombined", func(b *testing.B) {
-		sampler := NewSampler(0.8, 50, 0.9, 0.05, 42, nil)
-		b.ResetTimer()
-
-		for b.Loop() {
-			sampler.Sample(logits)
-		}
-	})
-}
-
-func BenchmarkGreedySampler(b *testing.B) {
-	sizes := []int{10, 100, 1000, 10000, 100000}
-
-	for _, size := range sizes {
-		b.Run(fmt.Sprintf("Size %d", size), func(b *testing.B) {
-			logits := make([]float32, size)
-			for i := range logits {
-				logits[i] = float32(rand.Float64()*10 - 5)
-			}
-
-			sampler := NewSampler(0, -1, 0, 0, -1, nil)
-			b.ResetTimer()
-
-			for b.Loop() {
-				sampler.Sample(logits)
-			}
-		})
-	}
-}

sample/samplers_test.go

@@ -1,14 +1,15 @@
 package sample
 
 import (
+	"math"
 	"math/rand/v2"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
 )
 
 func TestWeighted(t *testing.T) {
-	logits := []float32{-10, 3, -10, -10}
-	sampler := NewSampler(0, 0, 0, 0, 0, nil)
-	got, err := sampler.Sample(logits)
+	got, err := Weighted(nil).Sample([]float32{float32(math.Inf(-1)), 2, float32(math.Inf(-1)), float32(math.Inf(-1))})
 	if err != nil {
 		t.Error(err)
 		return
@@ -18,26 +19,194 @@ func TestWeighted(t *testing.T) {
 		t.Errorf("index mismatch: want %d, got %d", want, got)
 	}
 
-	logits = []float32{-100, -10, 0, 10}
-	sampler = NewSampler(0, 0, 0, 0, 0, nil)
-	got, err = sampler.Sample(logits)
+	got, err = Weighted(nil).Sample([]float32{float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1))})
+	if err == nil {
+		t.Error("expected error for no valid tokens, got index", got)
+	}
+
+	seed := uint64(42)
+	got, err = Weighted(&seed).Sample([]float32{1, 2, 3, 4})
 	if err != nil {
 		t.Error(err)
 		return
 	}
-	want = int32(3) // Should pick highest probability with this r value
+	// With seed 42, we expect a consistent sample
+	want = int32(3) // This will be deterministic due to the seed
 	if want != got {
 		t.Errorf("index mismatch: want %d, got %d", want, got)
 	}
 }
 
-func BenchmarkSample(b *testing.B) {
-	samplers := map[string]Sampler{
-		"Greedy":   NewSampler(0, 0, 0, 0, 0, nil), // Use NewSampler with temp=0 for greedy
-		"Weighted": NewSampler(0.5, 10, 0.9, 0.2, -1, nil),
+type testTransform struct {
+	id        int
+	callOrder *[]int
+}
+
+func (ts *testTransform) Apply(logits []float64) []float64 {
+	if ts.callOrder != nil {
+		*ts.callOrder = append(*ts.callOrder, ts.id)
+	}
+	return logits
+}
+
+func TestSample(t *testing.T) {
+	input := []float32{1, 2, 3, 4}
+
+	var callOrder []int
+	mock1 := &testTransform{
+		id:        1,
+		callOrder: &callOrder,
+	}
+	mock2 := &testTransform{
+		id:        2,
+		callOrder: &callOrder,
+	}
+	mock3 := &testTransform{
+		id:        3,
+		callOrder: &callOrder,
+	}
+
+	_, err := Weighted(nil, mock1, mock2, mock3).Sample(input)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	wantOrder := []int{1, 2, 3}
+	if diff := cmp.Diff(wantOrder, callOrder); diff != "" {
+		t.Errorf("call order mismatch (-want +got):\n%s", diff)
+	}
+}
+
+func TestNewSampler(t *testing.T) {
+	tests := []struct {
+		name        string
+		temperature float32
+		topK        int
+		topP        float32
+		minP        float32
+		seed        int
+		wantErr     bool
+	}{
+		{
+			name: "no transforms",
+			// temperature is 0, so greedy should be used
+			wantErr: false,
+		},
+		{
+			name:        "temperature",
+			temperature: 0.5,
+			wantErr:     false,
+		},
+		{
+			name:        "invalid temperature negative",
+			temperature: -1,
+			wantErr:     true,
+		},
+		{
+			name:        "invalid temperature too high",
+			temperature: 2.1,
+			wantErr:     true,
+		},
+		{
+			name:        "top k",
+			topK:        10,
+			temperature: 0.8,
+			wantErr:     false,
+		},
+		{
+			name:        "invalid top k negative",
+			topK:        -1,
+			temperature: 0.8,
+			wantErr:     true,
+		},
+		{
+			name:        "top p",
+			topP:        0.9,
+			temperature: 0.8,
+			wantErr:     false,
+		},
+		{
+			name:        "invalid top p negative",
+			topP:        -0.1,
+			temperature: 0.8,
+			wantErr:     true,
+		},
+		{
+			name:        "invalid top p one",
+			topP:        1.0,
+			temperature: 0.8,
+			wantErr:     true,
+		},
+		{
+			name:        "min p",
+			minP:        0.2,
+			temperature: 0.8,
+			wantErr:     false,
+		},
+		{
+			name:        "invalid min p negative",
+			minP:        -0.1,
+			temperature: 0.8,
+			wantErr:     true,
+		},
+		{
+			name:        "invalid min p one",
+			minP:        1.0,
+			temperature: 0.8,
+			wantErr:     true,
+		},
+		{
+			name:        "default values",
+			temperature: 0.8,
+			topK:        40,
+			topP:        0.9,
+			minP:        0.0,
+			seed:        0,
+			wantErr:     false,
+		},
+		{
+			name:        "all zeroes",
+			temperature: 0.0,
+			topK:        0,
+			topP:        0.0,
+			minP:        0.0,
+			seed:        0,
+			wantErr:     false, // all zeroes means no transforms
+		},
+		{
+			name:        "all transforms",
+			temperature: 0.8,
+			topK:        50,
+			topP:        0.95,
+			minP:        0.1,
+			seed:        42,
+			wantErr:     false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := NewSampler(tt.temperature, tt.topK, tt.topP, tt.minP, tt.seed)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewSampler() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}
+
+func BenchmarkSample(b *testing.B) {
+	transforms := []Transform{
+		Temperature(0.5),
+		TopK(10),
+		TopP(0.9),
+		MinP(0.2),
+	}
+
+	samplers := map[string]Sampler{
+		"Greedy":   Greedy(),
+		"Weighted": Weighted(nil, transforms...),
 	}
 
-	// Generate random logits for benchmarking
 	logits := make([]float32, 1<<16)
 	for i := range logits {
 		logits[i] = rand.Float32()
@@ -46,9 +215,9 @@ func BenchmarkSample(b *testing.B) {
 	for name, s := range samplers {
 		b.Run(name, func(b *testing.B) {
 			b.ResetTimer()
-			for b.Loop() {
+			for range b.N {
 				if _, err := s.Sample(logits); err != nil {
-					b.Fatalf("error sampling: %v", err)
+					b.Error(err)
 				}
 			}
 		})

sample/transforms.go

@@ -1,203 +1,120 @@
 package sample
 
 import (
+	"cmp"
 	"math"
 	"slices"
+
+	pq "github.com/emirpasic/gods/v2/queues/priorityqueue"
 )
 
-func softmax(ts []token) []token {
-	var sum float32
-	for i, v := range ts {
-		ts[i].value = float32(math.Exp(float64(v.value)))
-		sum += ts[i].value
-	}
-
-	for i := range ts {
-		ts[i].value /= sum
-	}
-
-	return ts
+type Transform interface {
+	Apply([]float64) []float64
 }
 
-func temperature(ti []token, t float32) []token {
-	if t == 1 {
-		return ti
+// TODO(parthsareen): potentially cache softmax values
+func softmax(logits []float64) []float64 {
+	var sum float64
+	probs := make([]float64, len(logits))
+	for i, v := range logits {
+		probs[i] = math.Exp(v)
+		sum += probs[i]
 	}
 
-	temp := max(t, 1e-7)
-	maxLogit := float32(math.Inf(-1))
-	for _, token := range ti {
-		if token.value > maxLogit {
-			maxLogit = token.value
-		}
+	for i := range probs {
+		probs[i] /= sum
 	}
 
+	return probs
+}
+
+type Temperature float64
+
+func (t Temperature) Apply(logits []float64) []float64 {
+	temp := math.Max(float64(t), 1e-7)
+
 	// subtracting max logit to avoid under/overflow
-	for i := range ti {
-		ti[i].value = (ti[i].value - maxLogit) / temp
+	maxLogit := slices.Max(logits)
+	for i := range logits {
+		logits[i] = (logits[i] - maxLogit) / temp
 	}
 
-	return ti
+	return logits
 }
 
-// siftDown maintains a min-heap property by recursively moving larger elements down the heap.
-//
-// The heap is represented as an array where for any node at index i:
-// - Left child is at index 2i + 1
-// - Right child is at index 2i + 2
-// - Parent is at index (i-1)/2
-//
-// The function compares a node with its children and:
-// 1. Finds the smallest value between the node and its children
-// 2. If the node is not the smallest, swaps it with its smallest child
-// 3. Continues this process down the affected path until the min-heap property is restored
-func siftDown(data []token, start, end int) {
-	root := start
-	for {
-		child := 2*root + 1
-		if child >= end {
+type logitMap struct {
+	index int
+	logit float64
+}
+
+type TopK int
+
+// TODO(parthsareen): avoid having to check all logits after this transform
+func (k TopK) Apply(logits []float64) []float64 {
+	if int(k) >= len(logits) {
+		return logits
+	}
+	q := pq.NewWith(func(a, b logitMap) int {
+		return -cmp.Compare(a.logit, b.logit)
+	})
+
+	for i, logit := range logits {
+		q.Enqueue(logitMap{index: i, logit: logit})
+	}
+
+	validLogits := make(map[int]float64)
+	for range k {
+		logitMap, _ := q.Dequeue()
+		validLogits[logitMap.index] = logitMap.logit
+	}
+
+	for i := range logits {
+		if _, ok := validLogits[i]; !ok {
+			logits[i] = math.Inf(-1)
+		}
+	}
+
+	return logits
+}
+
+type TopP float64
+
+func (p TopP) Apply(logits []float64) []float64 {
+	probs := softmax(logits)
+	indices := make([]int, len(probs))
+	for i := range indices {
+		indices[i] = i
+	}
+
+	// sort in descending order
+	slices.SortFunc(indices, func(i, j int) int {
+		return cmp.Compare(probs[j], probs[i])
+	})
+
+	var sum float64
+	for i, idx := range indices {
+		sum += probs[idx]
+		if sum > float64(p) {
+			for _, idx := range indices[i+1:] {
+				logits[idx] = math.Inf(-1)
+			}
 			break
 		}
-		// Find smaller child (we want min heap)
-		if child+1 < end && data[child+1].value < data[child].value {
-			child++
-		}
-		// Exit if root is already smaller than children
-		if data[root].value <= data[child].value {
-			break
-		}
-		// Swap with smaller child and continue
-		data[root], data[child] = data[child], data[root]
-		root = child
 	}
+	return logits
 }
 
-// topK limits the number of tokens considered to the k highest logits
-func topK(ts []token, k int) []token {
-	if k >= len(ts) {
-		return ts
-	}
-	// Heapify + siftDown - O(nlog(k))
-	// Build min-heap of first k elements
-	heap := ts[:k]
-	for i := k/2 - 1; i >= 0; i-- {
-		siftDown(heap, i, k)
-	}
+type MinP float64
 
-	// Process remaining elements - if larger than heap root, replace root
-	for i := k; i < len(ts); i++ {
-		if ts[i].value > heap[0].value {
-			heap[0] = ts[i]
-			siftDown(heap, 0, k)
+func (p MinP) Apply(logits []float64) []float64 {
+	probs := softmax(logits)
+	threshold := slices.Max(probs) * float64(p)
+
+	for i, prob := range probs {
+		if prob < threshold {
+			logits[i] = math.Inf(-1)
 		}
 	}
 
-	slices.Reverse(heap)
-
-	ts = heap
-	return ts
-}
-
-// topP limits tokens to those with cumulative probability p
-func topP(ts []token, p float32) []token {
-	if p == 1.0 {
-		return ts
-	}
-
-	// Find cutoff index where cumulative sum exceeds p
-	var sum float32
-	for i, t := range ts {
-		sum += t.value
-		if sum > float32(p) {
-			ts = ts[:i+1]
-			return ts
-		}
-	}
-
-	return ts
-}
-
-// minP limits tokens to those with cumulative probability p
-func minP(ts []token, p float32) []token {
-	if p == 1.0 {
-		return ts
-	}
-
-	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
-}
-
-// TODO(parthsareen): possibly replace with simpler implementation https://github.com/ollama/ollama/issues/9584
-// Conting sort implementation to sort tokens by logits
-func sortLogits(tokens []token) {
-	if len(tokens) <= 1 {
-		return
-	}
-
-	// Find max/min in a single pass
-	minLogit, maxLogit := tokens[0].value, tokens[0].value
-	for _, t := range tokens[1:] {
-		if t.value < minLogit {
-			minLogit = t.value
-		} else if t.value > maxLogit {
-			maxLogit = t.value
-		}
-	}
-
-	// Calculate scaling to map to uint32 range
-	logitRange := maxLogit - minLogit
-	if logitRange < 1e-6 {
-		return // All values effectively equal
-	}
-
-	// Count frequencies directly from tokens
-	const maxInt = (1 << 24) - 1 // Use 24 bits for good granularity
-	var counts [256]int          // For first byte
-
-	// First pass: count frequencies
-	for _, t := range tokens {
-		// Map to [0, maxInt] range
-		score := min(uint32((t.value-minLogit)*float32(maxInt)/logitRange), maxInt)
-		counts[score>>16]++
-	}
-
-	// Calculate offsets
-	var offset int
-	for i := range counts {
-		count := counts[i]
-		counts[i] = offset
-		offset += count
-	}
-
-	// Second pass: place elements in correct position
-	output := make([]token, len(tokens))
-	// Track current positions
-	countsCopy := counts
-
-	for i, t := range tokens {
-		score := min(uint32((t.value-minLogit)*float32(maxInt)/logitRange), maxInt)
-
-		pos := countsCopy[score>>16]
-		countsCopy[score>>16]++
-		output[len(tokens)-1-pos] = tokens[i]
-	}
-
-	copy(tokens, output)
+	return logits
 }

sample/transforms_test.go

@@ -4,182 +4,77 @@ import (
 	"math"
 	"math/rand/v2"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
 )
 
-// Helper to convert float64 slice to logit slice
-func toTokens(values []float64) []token {
-	tokens := make([]token, len(values))
-	for i, v := range values {
-		tokens[i] = token{
-			id:    int32(i),
-			value: float32(v),
-		}
-	}
-	return tokens
-}
-
-// Helper to compare logit slices
-func compareLogits(t *testing.T, name string, want []float64, got []token) {
-	t.Helper()
-	if len(want) != len(got) {
-		t.Errorf("%s: length mismatch: want %d, got %d", name, len(want), len(got))
-		return
-	}
-	for i := range want {
-		if math.Abs(float64(got[i].value)-want[i]) > 1e-6 {
-			t.Errorf("%s: index %d: want %f, got %f", name, i, want[i], got[i].value)
-		}
-	}
-}
-
 func TestTemperature(t *testing.T) {
-	input := []float64{2, -1, 4, -3, 1, -2, 0}
-	want := []float64{-4, -10, 0, -14, -6, -12, -8} // (logit - max logit) / temp
-
-	got := temperature(toTokens(input), 0.5)
-	compareLogits(t, "Temperature", want, got)
+	got := Temperature(0.5).Apply([]float64{2, -1, 4, -3, 1, -2, 0})
+	want := []float64{-4, -10, 0, -14, -6, -12, -8}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("logits mismatch (-want +got):\n%s", diff)
+	}
 }
 
 func TestSoftmax(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4}
-	got := softmax(toTokens(input))
+	got := softmax([]float64{-3, -2, -1, 0, 1, 2, 4})
 
-	// 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)
-	}
-
-	// Check relative ordering is preserved
-	for i := 1; i < len(got); i++ {
-		if got[i].value < got[i-1].value {
-			t.Errorf("probability ordering not preserved at index %d", i)
-		}
+	want := []float64{0.000751406628089903, 0.0020425349829204676, 0.005552185728064613, 0.015092405572827691, 0.04102541181635154, 0.11151863144543739, 0.8240174238263085}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("probs mismatch (-want +got):\n%s", diff)
 	}
 }
 
 func TestTopK(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4}
-
-	// Test k=3
-	got := topK(toTokens(input), 3)
-	if len(got) != 3 {
-		t.Errorf("topK(3): wrong length: want 3, got %d", len(got))
+	got := TopK(3).Apply([]float64{-3, -2, -1, 0, 1, 2, 4})
+	want := []float64{math.Inf(-1), math.Inf(-1), math.Inf(-1), math.Inf(-1), 1, 2, 4}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("logits mismatch (-want +got):\n%s", diff)
 	}
-	// Should keep highest 3 values: 4, 2, 1
-	want := []float64{4, 2, 1}
-	compareLogits(t, "topK(3)", want, got)
 
-	// Test k > len
-	got = topK(toTokens(input), 10)
-	compareLogits(t, "topK(10)", input, got)
+	got = TopK(10).Apply([]float64{-3, -2, -1, 0, 1, 2, 4})
+
+	want = []float64{-3, -2, -1, 0, 1, 2, 4}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("logits mismatch (-want +got):\n%s", diff)
+	}
 }
 
 func TestTopP(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4}
-	tokens := toTokens(input)
-
-	// First apply temperature and softmax to get probabilities
-	tokens = temperature(tokens, 1)
-	tokens = softmax(tokens)
-	sortLogits(tokens)
-
-	// 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(got))
-		t.Logf("got: %v", got)
+	got := TopP(0.9).Apply([]float64{-3, -2, -1, 0, 1, 2, 4})
+	want := []float64{math.Inf(-1), math.Inf(-1), math.Inf(-1), math.Inf(-1), math.Inf(-1), 2, 4}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("logits mismatch (-want +got):\n%s", diff)
 	}
 }
 
 func TestMinP(t *testing.T) {
-	input := []float64{-3, -2, -1, 0, 1, 2, 4, 3}
-	tokens := toTokens(input)
-
-	// First apply temperature and softmax
-	tokens = temperature(tokens, 1)
-	tokens = softmax(tokens)
-
-	// Then apply minP
-	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))
+	got := MinP(0.2).Apply([]float64{-3, -2, -1, 0, 1, 2, 4, 3})
+	want := []float64{math.Inf(-1), math.Inf(-1), math.Inf(-1), math.Inf(-1), math.Inf(-1), math.Inf(-1), 4, 3}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("logits mismatch (-want +got):\n%s", diff)
 	}
 }
 
-func TestSortLogits(t *testing.T) {
-	input := []float64{3, 1, 4, 2, -1, 0, -2}
-	tokens := toTokens(input)
-
-	sortLogits(tokens)
-
-	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)
-		}
+func BenchmarkTransform(b *testing.B) {
+	transforms := map[string]Transform{
+		"Temperature": Temperature(0.5),
+		"TopK":        TopK(10),
+		"TopP":        TopP(0.9),
+		"MinP":        MinP(0.2),
 	}
 
-	want := []float64{4, 3, 2, 1, 0, -1, -2}
-	compareLogits(t, "sortLogits", want, tokens)
-}
-
-func BenchmarkTransforms(b *testing.B) {
-	// Generate random logits
-	tokens := make([]token, 1<<16)
-	for i := range tokens {
-		tokens[i] = token{
-			id:    int32(i),
-			value: rand.Float32(),
-		}
+	logits := make([]float64, 1<<16)
+	for i := range logits {
+		logits[i] = rand.Float64()
 	}
 
-	tokensCopy := make([]token, len(tokens))
-
-	b.Run("Temperature", func(b *testing.B) {
-		b.ResetTimer()
-		for b.Loop() {
-			copy(tokensCopy, tokens)
-			temperature(tokensCopy, 0.5)
-		}
-	})
-
-	b.Run("TopK", func(b *testing.B) {
-		b.ResetTimer()
-		for b.Loop() {
-			copy(tokensCopy, tokens)
-			topK(tokensCopy, 10)
-		}
-	})
-
-	b.Run("TopP", func(b *testing.B) {
-		b.ResetTimer()
-		for b.Loop() {
-			copy(tokensCopy, tokens)
-			topP(tokensCopy, 0.9)
-		}
-	})
-
-	b.Run("MinP", func(b *testing.B) {
-		b.ResetTimer()
-		for b.Loop() {
-			copy(tokensCopy, tokens)
-			minP(tokensCopy, 0.2)
-		}
-	})
-
-	b.Run("SortTokens", func(b *testing.B) {
-		b.ResetTimer()
-		for b.Loop() {
-			copy(tokensCopy, tokens)
-			sortLogits(tokensCopy)
-		}
-	})
+	for name, transform := range transforms {
+		b.Run(name, func(b *testing.B) {
+			b.ResetTimer()
+			for range b.N {
+				transform.Apply(logits)
+			}
+		})
+	}
 }

scripts/install.sh

@@ -77,12 +77,11 @@ if [ -d "$OLLAMA_INSTALL_DIR/lib/ollama" ] ; then
 fi
 status "Installing ollama to $OLLAMA_INSTALL_DIR"
 $SUDO install -o0 -g0 -m755 -d $BINDIR
-$SUDO install -o0 -g0 -m755 -d "$OLLAMA_INSTALL_DIR/lib/ollama"
+$SUDO install -o0 -g0 -m755 -d "$OLLAMA_INSTALL_DIR"
 status "Downloading Linux ${ARCH} bundle"
 curl --fail --show-error --location --progress-bar \
     "https://ollama.com/download/ollama-linux-${ARCH}.tgz${VER_PARAM}" | \
     $SUDO tar -xzf - -C "$OLLAMA_INSTALL_DIR"
-
 if [ "$OLLAMA_INSTALL_DIR/bin/ollama" != "$BINDIR/ollama" ] ; then
     status "Making ollama accessible in the PATH in $BINDIR"
     $SUDO ln -sf "$OLLAMA_INSTALL_DIR/ollama" "$BINDIR/ollama"

server/internal/client/ollama/registry.go

@@ -45,9 +45,9 @@ import (
 
 // Errors
 var (
-	// ErrModelNotFound is returned when a manifest is not found in the
+	// ErrManifestNotFound is returned when a manifest is not found in the
 	// cache or registry.
-	ErrModelNotFound = errors.New("model not found")
+	ErrManifestNotFound = errors.New("manifest not found")
 
 	// ErrManifestInvalid is returned when a manifest found in a local or
 	// remote cache is invalid.
@@ -114,18 +114,7 @@ type Error struct {
 }
 
 func (e *Error) Error() string {
-	var b strings.Builder
-	b.WriteString("registry responded with status ")
-	b.WriteString(strconv.Itoa(e.Status))
-	if e.Code != "" {
-		b.WriteString(": code ")
-		b.WriteString(e.Code)
-	}
-	if e.Message != "" {
-		b.WriteString(": ")
-		b.WriteString(e.Message)
-	}
-	return b.String()
+	return fmt.Sprintf("registry responded with status %d: %s %s", e.Status, e.Code, e.Message)
 }
 
 func (e *Error) LogValue() slog.Value {
@@ -266,6 +255,10 @@ func DefaultRegistry() (*Registry, error) {
 	}
 
 	var rc Registry
+	rc.Cache, err = defaultCache()
+	if err != nil {
+		return nil, err
+	}
 	rc.Key, err = ssh.ParseRawPrivateKey(keyPEM)
 	if err != nil {
 		return nil, err
@@ -366,7 +359,7 @@ func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error {
 				n.Model(),
 				l.Digest,
 			)
-			res, err := r.send(ctx, "POST", startURL, nil)
+			res, err := r.doOK(ctx, "POST", startURL, nil)
 			if err != nil {
 				return err
 			}
@@ -390,7 +383,7 @@ func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error {
 			}
 			req.ContentLength = l.Size
 
-			res, err = sendRequest(r.client(), req)
+			res, err = doOK(r.client(), req)
 			if err == nil {
 				res.Body.Close()
 			}
@@ -410,7 +403,7 @@ func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error {
 		n.Model(),
 		n.Tag(),
 	)
-	res, err := r.send(ctx, "PUT", path, bytes.NewReader(m.Data))
+	res, err := r.doOK(ctx, "PUT", path, bytes.NewReader(m.Data))
 	if err == nil {
 		res.Body.Close()
 	}
@@ -459,15 +452,10 @@ func (r *Registry) Pull(ctx context.Context, name string) error {
 
 	t := traceFromContext(ctx)
 
-	g, ctx := errgroup.WithContext(ctx)
+	var g errgroup.Group
 	g.SetLimit(r.maxStreams())
 
-	layers := m.Layers
-	if m.Config != nil && m.Config.Digest.IsValid() {
-		layers = append(layers, m.Config)
-	}
-
-	for _, l := range layers {
+	for _, l := range m.Layers {
 		if exists(l) {
 			t.update(l, l.Size, ErrCached)
 			continue
@@ -484,9 +472,7 @@ func (r *Registry) Pull(ctx context.Context, name string) error {
 
 		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)
+				res, err := doOK(r.client(), req)
 				if err != nil {
 					return err
 				}
@@ -512,21 +498,19 @@ func (r *Registry) Pull(ctx context.Context, name string) error {
 			// 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)
+			res, err := doOK(r.client(), req)
 			if err != nil {
 				return err
 			}
 			res.Body.Close()
 			req = res.Request.WithContext(req.Context())
 
-			wp := writerPool{size: r.maxChunkSize()}
-
+			streamNo := 0
+			tws := make([]*bufio.Writer, r.maxStreams()-1)
 			for chunk := range chunks.Of(l.Size, r.maxChunkSize()) {
-				if ctx.Err() != nil {
-					break
-				}
-
 				ticket := q.Take()
+				bufIdx := streamNo % len(tws)
+				streamNo++
 				g.Go(func() (err error) {
 					defer func() {
 						if err != nil {
@@ -540,18 +524,23 @@ func (r *Registry) Pull(ctx context.Context, name string) error {
 						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)
+							res, err := doOK(r.client(), req)
 							if err != nil {
 								return err
 							}
 							defer res.Body.Close()
 
-							tw := wp.get()
+							tw := tws[bufIdx]
+							if tw == nil {
+								tw = bufio.NewWriterSize(nil, int(r.maxChunkSize()))
+								tws[bufIdx] = tw
+							}
 							tw.Reset(ticket)
-							defer wp.put(tw)
+							defer tw.Reset(nil) // release ticket
 
 							_, err = io.CopyN(tw, res.Body, chunk.Size())
 							if err != nil {
@@ -610,9 +599,6 @@ type Manifest struct {
 	Name   string   `json:"-"` // the canonical name of the model
 	Data   []byte   `json:"-"` // the raw data of the manifest
 	Layers []*Layer `json:"layers"`
-
-	// For legacy reasons, we still have to download the config layer.
-	Config *Layer `json:"config"`
 }
 
 var emptyDigest, _ = blob.ParseDigest("sha256:0000000000000000000000000000000000000000000000000000000000000000")
@@ -696,7 +682,7 @@ func (r *Registry) ResolveLocal(name string) (*Manifest, error) {
 	data, err := os.ReadFile(c.GetFile(d))
 	if err != nil {
 		if errors.Is(err, fs.ErrNotExist) {
-			return nil, fmt.Errorf("%w: %s", ErrModelNotFound, name)
+			return nil, fmt.Errorf("%w: %s", ErrManifestNotFound, name)
 		}
 		return nil, err
 	}
@@ -719,7 +705,7 @@ func (r *Registry) Resolve(ctx context.Context, name string) (*Manifest, error)
 		manifestURL = fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", scheme, n.Host(), n.Namespace(), n.Model(), d)
 	}
 
-	res, err := r.send(ctx, "GET", manifestURL, nil)
+	res, err := r.doOK(ctx, "GET", manifestURL, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -744,7 +730,7 @@ func (r *Registry) client() *http.Client {
 }
 
 // newRequest constructs a new request, ready to use, with the given method,
-// url, and body, pre-signed with client [Key] and [UserAgent].
+// url, and body, presigned with client Key and UserAgent.
 func (r *Registry) newRequest(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) {
 	req, err := http.NewRequestWithContext(ctx, method, url, body)
 	if err != nil {
@@ -763,17 +749,11 @@ func (r *Registry) newRequest(ctx context.Context, method, url string, body io.R
 	return req, nil
 }
 
-// sendRequest makes a request with the given client and request, and returns the
+// doOK makes a request with the given client and request, and returns the
 // response if the status code is 200. If the status code is not 200, an Error
 // is parsed from the response body and returned. If any other error occurs, it
 // is returned.
-func sendRequest(c *http.Client, r *http.Request) (_ *http.Response, err error) {
-	defer func() {
-		if err != nil {
-			err = fmt.Errorf("request error %s: %w", r.URL, err)
-		}
-	}()
-
+func doOK(c *http.Client, r *http.Request) (*http.Response, error) {
 	if r.URL.Scheme == "https+insecure" {
 		// TODO(bmizerany): clone client.Transport, set
 		// InsecureSkipVerify, etc.
@@ -816,26 +796,20 @@ func sendRequest(c *http.Client, r *http.Request) (_ *http.Response, err error)
 			// Use the raw body if we can't parse it as an error object.
 			re.Message = string(out)
 		}
-
-		// coerce MANIFEST_UNKNOWN to ErrManifestNotFound
-		if strings.EqualFold(re.Code, "MANIFEST_UNKNOWN") {
-			return nil, ErrModelNotFound
-		}
-
 		re.Status = res.StatusCode
 		return nil, &re
 	}
 	return res, nil
 }
 
-// send is a convenience method for making a request with newRequest and
-// passing it to send with r.client().
-func (r *Registry) send(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
+// doOK is a convenience method for making a request with newRequest and
+// passing it to doOK with r.client().
+func (r *Registry) doOK(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
 	req, err := r.newRequest(ctx, method, path, body)
 	if err != nil {
 		return nil, err
 	}
-	return sendRequest(r.client(), req)
+	return doOK(r.client(), req)
 }
 
 // makeAuthToken creates an Ollama auth token for the given private key.
@@ -990,28 +964,3 @@ func splitExtended(s string) (scheme, name, digest string) {
 	}
 	return scheme, s, digest
 }
-
-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

@@ -608,7 +608,7 @@ func TestInsecureSkipVerify(t *testing.T) {
 	url := fmt.Sprintf("https://%s/%s", s.Listener.Addr(), name)
 	_, err := rc.Resolve(t.Context(), url)
 	if err == nil || !strings.Contains(err.Error(), "failed to verify") {
-		t.Errorf("err = %v; want cert verification failure", err)
+		t.Errorf("err = %v; want cert verifiction failure", err)
 	}
 
 	url = fmt.Sprintf("https+insecure://%s/%s", s.Listener.Addr(), name)

server/internal/client/ollama/trace.go

@@ -13,13 +13,9 @@ type Trace struct {
 	// Update is called during [Registry.Push] and [Registry.Pull] to
 	// report the progress of blob uploads and downloads.
 	//
-	// The n argument is the number of bytes transferred so far, and err is
-	// any error that has occurred. If n == 0, and err is nil, the download
-	// or upload has just started. If err is [ErrCached], the download or
-	// upload has been skipped because the blob is already present in the
-	// local cache or remote registry, respectively. Otherwise, if err is
-	// non-nil, the download or upload has failed. When l.Size == n, and
-	// err is nil, the download or upload has completed.
+	// It is called once at the beginning of the download with a zero n and
+	// then once per read operation with the number of bytes read so far,
+	// and an error if any.
 	//
 	// A function assigned must be safe for concurrent use. The function is
 	// called synchronously and so should not block or take long to run.

server/internal/internal/backoff/backoff_test.go

@@ -1,5 +1,3 @@
-//go:build goexperiment.synctest
-
 package backoff
 
 import (

server/internal/internal/syncs/line_test.go

@@ -1,5 +1,3 @@
-//go:build goexperiment.synctest
-
 package syncs
 
 import (

server/internal/registry/server.go

@@ -7,14 +7,10 @@ import (
 	"cmp"
 	"encoding/json"
 	"errors"
-	"fmt"
 	"io"
 	"log/slog"
 	"net/http"
-	"sync"
-	"time"
 
-	"github.com/ollama/ollama/server/internal/cache/blob"
 	"github.com/ollama/ollama/server/internal/client/ollama"
 )
 
@@ -35,10 +31,6 @@ type Local struct {
 	// Fallback, if set, is used to handle requests that are not handled by
 	// this handler.
 	Fallback http.Handler
-
-	// Prune, if set, is called to prune the local disk cache after a model
-	// is deleted.
-	Prune func() error // optional
 }
 
 // serverError is like ollama.Error, but with a Status field for the HTTP
@@ -113,8 +105,6 @@ func (s *Local) serveHTTP(rec *statusCodeRecorder, r *http.Request) {
 		switch r.URL.Path {
 		case "/api/delete":
 			return false, s.handleDelete(rec, r)
-		case "/api/pull":
-			return false, s.handlePull(rec, r)
 		default:
 			if s.Fallback != nil {
 				s.Fallback.ServeHTTP(rec, r)
@@ -214,100 +204,6 @@ func (s *Local) handleDelete(_ http.ResponseWriter, r *http.Request) error {
 	if !ok {
 		return &serverError{404, "not_found", "model not found"}
 	}
-	if s.Prune == nil {
-		return nil
-	}
-	return s.Prune()
-}
-
-type progressUpdateJSON struct {
-	Status    string      `json:"status"`
-	Digest    blob.Digest `json:"digest,omitempty,omitzero"`
-	Total     int64       `json:"total,omitempty,omitzero"`
-	Completed int64       `json:"completed,omitempty,omitzero"`
-}
-
-func (s *Local) handlePull(w http.ResponseWriter, r *http.Request) error {
-	if r.Method != "POST" {
-		return errMethodNotAllowed
-	}
-
-	p, err := decodeUserJSON[*params](r.Body)
-	if err != nil {
-		return err
-	}
-
-	maybeFlush := func() {
-		fl, _ := w.(http.Flusher)
-		if fl != nil {
-			fl.Flush()
-		}
-	}
-	defer maybeFlush()
-
-	var mu sync.Mutex
-	enc := json.NewEncoder(w)
-	enc.Encode(progressUpdateJSON{Status: "pulling manifest"})
-
-	ctx := ollama.WithTrace(r.Context(), &ollama.Trace{
-		Update: func(l *ollama.Layer, n int64, err error) {
-			mu.Lock()
-			defer mu.Unlock()
-
-			// TODO(bmizerany): coalesce these updates; writing per
-			// update is expensive
-			enc.Encode(progressUpdateJSON{
-				Digest:    l.Digest,
-				Status:    "pulling",
-				Total:     l.Size,
-				Completed: n,
-			})
-		},
-	})
-
-	done := make(chan error, 1)
-	go func() {
-		// TODO(bmizerany): continue to support non-streaming responses
-		done <- s.Client.Pull(ctx, p.model())
-	}()
-
-	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
-				}
-
-				// 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
 }
 

server/internal/registry/server_test.go

@@ -1,27 +1,17 @@
 package registry
 
 import (
-	"bytes"
-	"context"
 	"encoding/json"
-	"fmt"
-	"io"
-	"io/fs"
-	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"regexp"
 	"strings"
-	"sync"
 	"testing"
 
 	"github.com/ollama/ollama/server/internal/cache/blob"
 	"github.com/ollama/ollama/server/internal/client/ollama"
 	"github.com/ollama/ollama/server/internal/testutil"
-	"golang.org/x/tools/txtar"
-
-	_ "embed"
 )
 
 type panicTransport struct{}
@@ -40,7 +30,7 @@ type bytesResetter interface {
 	Reset()
 }
 
-func newTestServer(t *testing.T, upstreamRegistry http.HandlerFunc) *Local {
+func newTestServer(t *testing.T) *Local {
 	t.Helper()
 	dir := t.TempDir()
 	err := os.CopyFS(dir, os.DirFS("testdata/models"))
@@ -51,25 +41,10 @@ func newTestServer(t *testing.T, upstreamRegistry http.HandlerFunc) *Local {
 	if err != nil {
 		t.Fatal(err)
 	}
-
-	client := panicOnRoundTrip
-	if upstreamRegistry != nil {
-		s := httptest.NewTLSServer(upstreamRegistry)
-		t.Cleanup(s.Close)
-		tr := s.Client().Transport.(*http.Transport).Clone()
-		tr.DialContext = func(ctx context.Context, _, _ string) (net.Conn, error) {
-			var d net.Dialer
-			return d.DialContext(ctx, "tcp", s.Listener.Addr().String())
-		}
-		client = &http.Client{Transport: tr}
-	}
-
 	rc := &ollama.Registry{
 		Cache:      c,
-		HTTPClient: client,
-		Mask:       "example.com/library/_:latest",
+		HTTPClient: panicOnRoundTrip,
 	}
-
 	l := &Local{
 		Client: rc,
 		Logger: testutil.Slogger(t),
@@ -110,7 +85,7 @@ func captureLogs(t *testing.T, s *Local) (*Local, bytesResetter) {
 func TestServerDelete(t *testing.T) {
 	check := testutil.Checker(t)
 
-	s := newTestServer(t, nil)
+	s := newTestServer(t)
 
 	_, err := s.Client.ResolveLocal("smol")
 	check(err)
@@ -152,105 +127,8 @@ func TestServerDelete(t *testing.T) {
 	}
 }
 
-//go:embed testdata/registry.txt
-var registryTXT []byte
-
-var registryFS = sync.OnceValue(func() fs.FS {
-	// Txtar gets hung up on \r\n line endings, so we need to convert them
-	// 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)
-	}
-	return fsys
-})
-
-func TestServerPull(t *testing.T) {
-	modelsHandler := http.FileServerFS(registryFS())
-	s := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
-		switch r.URL.Path {
-		case "/v2/library/BOOM/manifests/latest":
-			w.WriteHeader(999)
-			io.WriteString(w, `{"error": "boom"}`)
-		case "/v2/library/unknown/manifests/latest":
-			w.WriteHeader(404)
-			io.WriteString(w, `{"errors": [{"code": "MANIFEST_UNKNOWN", "message": "manifest unknown"}]}`)
-		default:
-			t.Logf("serving file: %s", r.URL.Path)
-			modelsHandler.ServeHTTP(w, r)
-		}
-	})
-
-	checkResponse := func(got *httptest.ResponseRecorder, wantlines string) {
-		t.Helper()
-
-		if got.Code != 200 {
-			t.Fatalf("Code = %d; want 200", got.Code)
-		}
-		gotlines := got.Body.String()
-		t.Logf("got:\n%s", gotlines)
-		for want := range strings.Lines(wantlines) {
-			want = strings.TrimSpace(want)
-			want, unwanted := strings.CutPrefix(want, "!")
-			want = strings.TrimSpace(want)
-			if !unwanted && !strings.Contains(gotlines, want) {
-				t.Fatalf("! missing %q in body", want)
-			}
-			if unwanted && strings.Contains(gotlines, 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, `
-		{"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"}
-	`)
-
-	got = s.send(t, "DELETE", "/api/pull", `{"model": "smol"}`)
-	checkErrorResponse(t, got, 405, "method_not_allowed", "method not allowed")
-
-	got = s.send(t, "POST", "/api/pull", `!`)
-	checkErrorResponse(t, got, 400, "bad_request", "invalid character '!' looking for beginning of value")
-
-	got = s.send(t, "POST", "/api/pull", ``)
-	checkErrorResponse(t, got, 400, "bad_request", "empty request body")
-
-	got = s.send(t, "POST", "/api/pull", `{"model": "://"}`)
-	checkResponse(got, `
-		{"status":"pulling manifest"}
-		{"status":"error: invalid or missing name: \"\""}
-
-		!verifying
-		!writing
-		!success
-	`)
-}
-
 func TestServerUnknownPath(t *testing.T) {
-	s := newTestServer(t, nil)
+	s := newTestServer(t)
 	got := s.send(t, "DELETE", "/api/unknown", `{}`)
 	checkErrorResponse(t, got, 404, "not_found", "not found")
 }

server/internal/registry/testdata/registry.txt

@@ -1,22 +0,0 @@
--- v2/library/smol/manifests/latest --
-{
-  "schemaVersion": 2,
-  "mediaType": "application/vnd.docker.distribution.manifest.v2+json",
-  "config": {
-    "mediaType": "application/vnd.docker.container.image.v1+json",
-    "digest": "sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356",
-    "size": 3
-  },
-  "layers": [
-    {
-      "mediaType": "application/vnd.ollama.image.model",
-      "digest": "sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312",
-      "size": 5
-    }
-  ]
-}
-
--- v2/library/smol/blobs/sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312 --
-GGUF
--- v2/library/smol/blobs/sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356 --
-{}

server/prompt.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/llm"
 	"github.com/ollama/ollama/model/models/mllama"
 	"github.com/ollama/ollama/template"
@@ -92,7 +93,7 @@ func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api.
 			var imgData llm.ImageData
 
 			if isMllama {
-				if len(m.ProjectorPaths) == 0 {
+				if envconfig.NewEngine() {
 					imgData = llm.ImageData{
 						ID:   len(images),
 						Data: i,

server/routes.go

@@ -42,12 +42,6 @@ import (
 	"github.com/ollama/ollama/version"
 )
 
-func experimentEnabled(name string) bool {
-	return slices.Contains(strings.Split(os.Getenv("OLLAMA_EXPERIMENT"), ","), name)
-}
-
-var useClient2 = experimentEnabled("client2")
-
 var mode string = gin.DebugMode
 
 type Server struct {
@@ -211,7 +205,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 
 	images := make([]llm.ImageData, len(req.Images))
 	for i := range req.Images {
-		if isMllama && len(model.ProjectorPaths) > 0 {
+		if isMllama && !envconfig.NewEngine() {
 			data, opts, err := mllama.Preprocess(bytes.NewReader(req.Images[i]))
 			if err != nil {
 				c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "error processing image"})
@@ -1179,7 +1173,6 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) {
 	r.HEAD("/api/tags", s.ListHandler)
 	r.GET("/api/tags", s.ListHandler)
 	r.POST("/api/show", s.ShowHandler)
-	r.DELETE("/api/delete", s.DeleteHandler)
 
 	// Create
 	r.POST("/api/create", s.CreateHandler)
@@ -1201,19 +1194,14 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) {
 	r.GET("/v1/models", openai.ListMiddleware(), s.ListHandler)
 	r.GET("/v1/models/:model", openai.RetrieveMiddleware(), s.ShowHandler)
 
-	if rc != nil {
-		// wrap old with new
-		rs := &registry.Local{
-			Client:   rc,
-			Logger:   slog.Default(), // TODO(bmizerany): Take a logger, do not use slog.Default()
-			Fallback: r,
-
-			Prune: PruneLayers,
-		}
-		return rs, nil
+	// wrap old with new
+	rs := &registry.Local{
+		Client:   rc,
+		Logger:   slog.Default(), // TODO(bmizerany): Take a logger, do not use slog.Default()
+		Fallback: r,
 	}
 
-	return r, nil
+	return rs, nil
 }
 
 func Serve(ln net.Listener) error {
@@ -1268,20 +1256,15 @@ func Serve(ln net.Listener) error {
 
 	s := &Server{addr: ln.Addr()}
 
-	var rc *ollama.Registry
-	if useClient2 {
-		var err error
-		rc, err = ollama.DefaultRegistry()
-		if err != nil {
-			return err
-		}
+	rc, err := ollama.DefaultRegistry()
+	if err != nil {
+		return err
 	}
 
 	h, err := s.GenerateRoutes(rc)
 	if err != nil {
 		return err
 	}
-
 	http.Handle("/", h)
 
 	ctx, done := context.WithCancel(context.Background())

server/sched.go

@@ -194,6 +194,11 @@ func (s *Scheduler) processPending(ctx context.Context) {
 						break
 					}
 
+					// Embedding models should always be loaded with parallel=1
+					if pending.model.CheckCapabilities(CapabilityCompletion) != nil {
+						numParallel = 1
+					}
+
 					// Evaluate if the model will fit in the available system memory, or if we should unload a model first
 					if len(gpus) == 1 && gpus[0].Library == "cpu" {
 						// simplifying assumption of defaultParallel when in CPU mode