llama: add patch to avoid errors on arrays of strigns

README.md

@@ -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`            |
@@ -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)
@@ -434,7 +432,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,7 +509,6 @@ 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)
@@ -564,7 +560,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
 

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`:
 

go.mod

@@ -25,6 +25,7 @@ require (
 	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 +45,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
 )

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

@@ -245,20 +245,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 +293,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 +669,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/0019-allow-arrays-of-strings.patch

@@ -0,0 +1,21 @@
+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] allow arrays of strings
+
+---
+ ggml/src/gguf.cpp | 1 -
+ 1 file changed, 1 deletion(-)
+
+diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp
+index ab13669c..9dfefd43 100644
+--- a/ggml/src/gguf.cpp
++++ b/ggml/src/gguf.cpp
+@@ -777,7 +777,6 @@ 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();
+ }
+ 

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/server.go

@@ -729,24 +729,29 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
 	}
 
 	if len(req.Format) > 0 {
-		switch string(req.Format) {
-		case `null`, `""`:
-			// Field was set, but "missing" a value. We accept
-			// these as "not set".
-			break
-		case `"json"`:
-			request["grammar"] = grammarJSON
-		default:
-			if req.Format[0] != '{' {
-				return fmt.Errorf("invalid format: %q; expected \"json\" or a valid JSON Schema object", req.Format)
-			}
+		format := string(req.Format)
+		if format != `null` && format != `""` {
+			if s.textProcessor != nil {
+				// New engine handles this on the backend
+				request["format"] = req.Format
+			} else {
+				// old engine
+				switch format {
+				case `"json"`:
+					request["grammar"] = grammarJSON
+				default:
+					if req.Format[0] != '{' {
+						return fmt.Errorf("invalid format: %q; expected \"json\" or a valid JSON Schema object", req.Format)
+					}
 
-			// User provided a JSON schema
-			g := llama.SchemaToGrammar(req.Format)
-			if g == nil {
-				return fmt.Errorf("invalid JSON schema in format")
+					// User provided a JSON schema
+					g := llama.SchemaToGrammar(req.Format)
+					if g == nil {
+						return fmt.Errorf("invalid JSON schema in format")
+					}
+					request["grammar"] = string(g)
+				}
 			}
-			request["grammar"] = string(g)
 		}
 	}
 
@@ -968,7 +973,7 @@ func (s *llmServer) Tokenize(ctx context.Context, content string) ([]int, error)
 		return s.llamaModel.Tokenize(content, false, true)
 	}
 	if s.textProcessor != nil {
-		tokens, err := s.textProcessor.Encode(content, false)
+		tokens, err := s.textProcessor.Encode(content)
 		if err != nil {
 			return nil, err
 		}

ml/backend.go

@@ -24,7 +24,6 @@ type Backend interface {
 	Config() Config
 	Get(name string) Tensor
 	NewContext() Context
-	NewContextSize(size int) Context
 }
 
 // BackendCacheConfig should be implemented by backends that need special output
@@ -100,17 +99,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 +205,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 +273,5 @@ const (
 	DTypeOther DType = iota
 	DTypeF32
 	DTypeF16
-	DTypeQ80
-	DTypeQ40
 	DTypeI32
 )

ml/backend/ggml/ggml.go

@@ -9,53 +9,67 @@ package ggml
 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 +88,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 +201,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 +243,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 +263,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 +282,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 +306,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 +348,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 +431,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:

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/gguf.cpp

@@ -780,11 +780,6 @@ const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id)
     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 +873,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"
@@ -19,46 +20,20 @@ import (
 	"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 +42,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 +51,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
@@ -278,7 +257,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 +268,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

@@ -11,7 +11,6 @@ import (
 	"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) {
@@ -163,7 +162,7 @@ func TestGetTextProcessor(t *testing.T) {
 
 type notTextProcessorModel struct{}
 
-func (notTextProcessorModel) Forward(ml.Context, input.Options) (ml.Tensor, error) {
+func (notTextProcessorModel) Forward(ml.Context, Options) (ml.Tensor, error) {
 	panic("unimplemented")
 }
 

model/models/llama/model.go

@@ -9,10 +9,10 @@ import (
 	"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
@@ -66,11 +66,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 +78,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 +95,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 +137,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,12 @@
 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 {
@@ -62,92 +56,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/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,8 +786,6 @@ func (s *Server) loadModel(
 		panic(err)
 	}
 
-	s.vocab = sample.NewVocab(mpath)
-
 	// TODO(jessegross): LoRA loading
 	if lpath.String() != "" {
 		panic("loras are not yet implemented")

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)
+			}
+		})
+	}
 }

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