Repository: ollama/ollama

Issue:
Fix seeded sampling in Ollama engine


Discussion:
@ParthSareen: closed by #9374 

Solution patch:
diff --git a/README.md b/README.md
index 5aa4801e..1162f891 100644
--- a/README.md
+++ b/README.md
@@ -54,6 +54,7 @@ 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`            |
diff --git a/go.mod b/go.mod
index c45c9892..cc578900 100644
--- a/go.mod
+++ b/go.mod
@@ -25,7 +25,6 @@ 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 (
@@ -45,6 +44,7 @@ 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
 )
diff --git a/llama/llama.cpp/src/llama-vocab.cpp b/llama/llama.cpp/src/llama-vocab.cpp
index c7ff28be..7a185443 100644
--- a/llama/llama.cpp/src/llama-vocab.cpp
+++ b/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_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
diff --git a/llama/patches/0019-fix-string-arr-kv-loading.patch b/llama/patches/0019-fix-string-arr-kv-loading.patch
new file mode 100644
index 00000000..aa7b4d3c
--- /dev/null
+++ b/llama/patches/0019-fix-string-arr-kv-loading.patch
@@ -0,0 +1,64 @@
+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
diff --git a/llm/server.go b/llm/server.go
index 09690a5f..9553ba8f 100644
--- a/llm/server.go
+++ b/llm/server.go
@@ -973,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)
+		tokens, err := s.textProcessor.Encode(content, false)
 		if err != nil {
 			return nil, err
 		}
diff --git a/ml/backend/ggml/ggml/include/gguf.h b/ml/backend/ggml/ggml/include/gguf.h
index 79ee2020..3efb22f0 100644
--- a/ml/backend/ggml/ggml/include/gguf.h
+++ b/ml/backend/ggml/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/ml/backend/ggml/ggml/src/gguf.cpp b/ml/backend/ggml/ggml/src/gguf.cpp
index ab13669c..f75b923f 100644
--- a/ml/backend/ggml/ggml/src/gguf.cpp
+++ b/ml/backend/ggml/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/model/model.go b/model/model.go
index f8ed8741..75b7f639 100644
--- a/model/model.go
+++ b/model/model.go
@@ -3,7 +3,6 @@ package model
 import (
 	"errors"
 	"fmt"
-	"image"
 	_ "image/jpeg"
 	_ "image/png"
 	"log/slog"
@@ -22,14 +21,40 @@ import (
 	_ "github.com/ollama/ollama/ml/backend"
 )
 
+// 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
-	Positions []int32
-	Sequences []int
-	Outputs   []int32
-
-	Images []image.Image
+	Inputs     []int32
+	Multimodal []MultimodalIndex
+	Positions  []int32
+	Sequences  []int
+	Outputs    []int32
 }
 
 type config struct {
@@ -59,6 +84,37 @@ type Model interface {
 	Config() config
 }
 
+// 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, error)
+}
+
 var models = make(map[string]func(ml.Config) (Model, error))
 
 // Register registers a model constructor for the given architecture
diff --git a/model/models/mllama/model.go b/model/models/mllama/model.go
index 8fee0cdb..945c7295 100644
--- a/model/models/mllama/model.go
+++ b/model/models/mllama/model.go
@@ -1,7 +1,12 @@
 package mllama
 
 import (
+	"bytes"
+	"encoding/binary"
 	"fmt"
+	"hash/fnv"
+	"image"
+	"slices"
 
 	"github.com/ollama/ollama/kvcache"
 	"github.com/ollama/ollama/ml"
@@ -56,41 +61,79 @@ func New(c ml.Config) (model.Model, error) {
 	return &m, nil
 }
 
-func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
-	var crossAttentionStates ml.Tensor
-	if opts.Images != nil {
-		f32s, aspectRatioID, err := m.ImageProcessor.ProcessImage(opts.Images[0])
-		if err != nil {
-			return nil, err
-		}
+func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, error) {
+	image, _, err := image.Decode(bytes.NewReader(multimodalData))
+	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
-		}
+	f32s, aspectRatioID, err := m.ImageProcessor.ProcessImage(image)
+	if err != nil {
+		return nil, err
+	}
 
-		aspectRatio, err := ctx.FromIntSlice([]int32{int32(aspectRatioID)}, 1)
-		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
+	}
 
-		positions := make([]int32, 1601)
-		for i := range positions {
-			positions[i] = int32(i)
-		}
+	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
+	}
 
-		positionIDs, err := ctx.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 []model.Input) ([]model.Input, error) {
+	var images []model.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 model.Input) bool { return input.Token == -1 })
 
-		crossAttentionStates = m.VisionModel.Forward(ctx, pixelValues, positionIDs, aspectRatio)
-		crossAttentionStates = m.Projector.Forward(ctx, crossAttentionStates)
+	return inputs, nil
+}
+
+func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
+	var crossAttentionStates ml.Tensor
+	if opts.Multimodal != nil {
+		crossAttentionStates = opts.Multimodal[0].Multimodal.(ml.Tensor)
 	}
 
 	inputs, err := ctx.FromIntSlice(opts.Inputs, len(opts.Inputs))
diff --git a/model/process_text.go b/model/process_text.go
index 7083f36f..bfb0a5f2 100644
--- a/model/process_text.go
+++ b/model/process_text.go
@@ -19,7 +19,7 @@ const (
 )
 
 type TextProcessor interface {
-	Encode(string) ([]int32, error)
+	Encode(s string, addSpecial bool) ([]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) ([]int32, error) {
+func (bpe BytePairEncoding) Encode(s string, addSpecial bool) ([]int32, error) {
 	fragments := []fragment{{value: s}}
 	for _, special := range bpe.vocab.SpecialVocabulary() {
 		// TODO: process special tokens concurrently
@@ -282,7 +282,7 @@ func (bpe BytePairEncoding) Encode(s string) ([]int32, error) {
 		}
 	}
 
-	if len(ids) > 0 {
+	if addSpecial && 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)
diff --git a/runner/ollamarunner/cache.go b/runner/ollamarunner/cache.go
index e1fa98b1..2fd060a1 100644
--- a/runner/ollamarunner/cache.go
+++ b/runner/ollamarunner/cache.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"log/slog"
 	"math"
-	"reflect"
 	"time"
 
 	"github.com/ollama/ollama/kvcache"
@@ -39,10 +38,7 @@ func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots
 	slots := make([]InputCacheSlot, numSlots)
 
 	for i := range slots {
-		slots[i] = InputCacheSlot{
-			Id:     i,
-			Inputs: make([]input, 0),
-		}
+		slots[i] = InputCacheSlot{Id: i}
 	}
 
 	cache := model.Config().Cache
@@ -83,7 +79,7 @@ type InputCacheSlot struct {
 	Id int
 
 	// Inputs that are stored in the KV cache
-	Inputs []input
+	Inputs []model.Input
 
 	// is this cache actively being processed as part of a sequence?
 	InUse bool
@@ -92,7 +88,7 @@ type InputCacheSlot struct {
 	lastUsed time.Time
 }
 
-func (c *InputCache) LoadCacheSlot(prompt []input, cachePrompt bool) (*InputCacheSlot, []input, error) {
+func (c *InputCache) LoadCacheSlot(prompt []model.Input, cachePrompt bool) (*InputCacheSlot, []model.Input, error) {
 	var slot *InputCacheSlot
 	var numPast int32
 	var err error
@@ -143,7 +139,7 @@ func (c *InputCache) LoadCacheSlot(prompt []input, cachePrompt bool) (*InputCach
 	return slot, prompt, nil
 }
 
-func (c *InputCache) findLongestCacheSlot(prompt []input) (*InputCacheSlot, int32, error) {
+func (c *InputCache) findLongestCacheSlot(prompt []model.Input) (*InputCacheSlot, int32, error) {
 	longest := int32(-1)
 	var longestSlot *InputCacheSlot
 
@@ -166,7 +162,7 @@ func (c *InputCache) findLongestCacheSlot(prompt []input) (*InputCacheSlot, int3
 	return longestSlot, longest, nil
 }
 
-func (c *InputCache) findBestCacheSlot(prompt []input) (*InputCacheSlot, int32, error) {
+func (c *InputCache) findBestCacheSlot(prompt []model.Input) (*InputCacheSlot, int32, error) {
 	oldest := time.Now()
 	var oldestSlot *InputCacheSlot
 
@@ -202,7 +198,7 @@ func (c *InputCache) findBestCacheSlot(prompt []input) (*InputCacheSlot, int32,
 	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, longest)
+		oldestSlot.Inputs = make([]model.Input, longest)
 		copy(oldestSlot.Inputs, longestSlot.Inputs[:longest])
 		if c.cache != nil {
 			c.cache.CopyPrefix(longestSlot.Id, oldestSlot.Id, longest)
@@ -212,7 +208,7 @@ func (c *InputCache) findBestCacheSlot(prompt []input) (*InputCacheSlot, int32,
 	return oldestSlot, longest, nil
 }
 
-func countCommonPrefix(a []input, b []input) int32 {
+func countCommonPrefix(a []model.Input, b []model.Input) int32 {
 	var count int32
 
 	for i := range a {
@@ -220,7 +216,7 @@ func countCommonPrefix(a []input, b []input) int32 {
 			break
 		}
 
-		if !reflect.DeepEqual(a[i], b[i]) {
+		if a[i].Token != b[i].Token || a[i].MultimodalHash != b[i].MultimodalHash {
 			break
 		}
 
diff --git a/runner/ollamarunner/runner.go b/runner/ollamarunner/runner.go
index 1a4bbf19..81e06562 100644
--- a/runner/ollamarunner/runner.go
+++ b/runner/ollamarunner/runner.go
@@ -1,13 +1,12 @@
 package ollamarunner
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
-	"image"
+	"hash/maphash"
 	"log"
 	"log/slog"
 	"net"
@@ -33,22 +32,19 @@ import (
 	_ "github.com/ollama/ollama/model/models"
 )
 
-// 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 {
+	// ctx for allocating tensors that last the lifetime of the sequence, such as
+	// multimodal embeddings
+	ctx ml.Context
+
 	// batch index
 	iBatch int
 
 	// prompt inputs left to evaluate
-	inputs []input
+	inputs []model.Input
 
 	// inputs that have been added to a batch but not yet submitted to Forward
-	pendingInputs []input
+	pendingInputs []model.Input
 
 	// tokens that have been generated but not returned yet (e.g. for stop sequences)
 	pendingResponses []string
@@ -101,8 +97,9 @@ 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(prompt, images)
+	inputs, err := s.inputs(ctx, prompt, images)
 	if err != nil {
 		return nil, fmt.Errorf("failed to process inputs: %w", err)
 	} else if len(inputs) == 0 {
@@ -128,6 +125,7 @@ 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,
@@ -146,28 +144,31 @@ 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(prompt string, images []ImageData) ([]input, error) {
-	var inputs []input
+func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]model.Input, error) {
+	var inputs []model.Input
 	var parts []string
 	var matches [][]string
 
-	// 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)
+	multimodalProcessor, visionModel := s.model.(model.MultimodalProcessor)
 
+	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)
+		tokens, err := s.model.(model.TextProcessor).Encode(part, i == 0)
 		if err != nil {
 			return nil, err
 		}
 
 		for _, t := range tokens {
-			inputs = append(inputs, input{token: t})
+			inputs = append(inputs, model.Input{Token: t})
 		}
 
 		// image - decode and store
@@ -186,12 +187,25 @@ func (s *Server) inputs(prompt string, images []ImageData) ([]input, error) {
 				return nil, fmt.Errorf("invalid image index: %d", n)
 			}
 
-			image, _, err := image.Decode(bytes.NewReader(images[imageIndex].Data))
+			imageEmbeddings, err := multimodalProcessor.EncodeMultimodal(ctx, images[imageIndex].Data)
 			if err != nil {
 				return nil, err
 			}
 
-			inputs = append(inputs, input{image: image})
+			s.multimodalHash.Reset()
+			_, _ = s.multimodalHash.Write(images[imageIndex].Data)
+			imageHash := s.multimodalHash.Sum64()
+
+			inputs = append(inputs, model.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
 		}
 	}
 
@@ -238,6 +252,10 @@ type Server struct {
 
 	// next sequence for prompt processing to avoid starvation
 	nextSeq int
+
+	// multimodalHash generates hashes for comparing equality
+	// of non-text data
+	multimodalHash maphash.Hash
 }
 
 func (s *Server) allNil() bool {
@@ -283,6 +301,7 @@ func (s *Server) removeSequence(seqIndex int, reason string) {
 	close(seq.responses)
 	close(seq.embedding)
 	seq.cache.InUse = false
+	seq.ctx.Close()
 	s.seqs[seqIndex] = nil
 	s.seqsSem.Release(1)
 }
@@ -311,7 +330,6 @@ func (s *Server) processBatch() error {
 	defer s.mu.Unlock()
 
 	var options model.Options
-	imgSeq := -1
 
 	seqIdx := s.nextSeq - 1
 	for range s.seqs {
@@ -330,7 +348,7 @@ func (s *Server) processBatch() error {
 
 		if !s.cache.enabled {
 			seq.inputs = append(seq.cache.Inputs, seq.inputs...)
-			seq.cache.Inputs = []input{}
+			seq.cache.Inputs = []model.Input{}
 		}
 
 		for i, input := range seq.inputs {
@@ -349,25 +367,21 @@ func (s *Server) processBatch() error {
 				break
 			}
 
-			// 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
-				}
+			// TODO(jessegross): This is a workaround for generating an attention mask and also providing a hint
+			// to the encoder cache.
+			//
+			// Break the batch when switching from text to images so that images are always at the beginning.
+			if input.Multimodal != nil && !(len(seq.pendingInputs) == 0 ||
+				(len(options.Multimodal) > 0 && options.Multimodal[len(options.Multimodal)-1].Index == len(options.Inputs)-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)
+			if input.Multimodal != nil {
+				options.Multimodal = append(options.Multimodal, model.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: input.Multimodal})
 			}
 
-			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)
 
@@ -403,7 +417,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{}
+			seq.pendingInputs = []model.Input{}
 		}
 
 		// don't sample prompt processing
@@ -449,7 +463,7 @@ func (s *Server) processBatch() error {
 			return err
 		}
 
-		seq.inputs = []input{{token: token}}
+		seq.inputs = []model.Input{{Token: token}}
 
 		seq.pendingResponses = append(seq.pendingResponses, piece)
 		sequence := strings.Join(seq.pendingResponses, "")
@@ -575,11 +589,19 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	sampler := sample.NewSampler(
+		req.Temperature,
+		req.TopK,
+		req.TopP,
+		req.MinP,
+		req.Seed,
+	)
+
 	seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{
 		numPredict: req.NumPredict,
 		stop:       req.Stop,
 		numKeep:    int32(req.NumKeep),
-		sampler:    sample.Greedy(), // TODO: add support for different samplers when performance is optimized
+		sampler:    sampler,
 		embedding:  false,
 	})
 	if err != nil {
diff --git a/sample/samplers.go b/sample/samplers.go
index 1b8a5edd..a5a0507c 100644
--- a/sample/samplers.go
+++ b/sample/samplers.go
@@ -2,76 +2,103 @@ package sample
 
 import (
 	"errors"
-	"math"
-
-	"golang.org/x/exp/rand"
-	"gonum.org/v1/gonum/stat/sampleuv"
+	"math/rand/v2"
+	"slices"
 )
 
+// Sampler is not thread-safe. Each goroutine should have its own instance
 type Sampler interface {
 	Sample([]float32) (int32, error)
 }
 
+// logit represents information about a single token during sampling
+type logit struct {
+	id    int32   // The token's unique identifier
+	value float32 // The raw logit or probability from the model
+}
+
 type weighted struct {
-	src        rand.Source
-	transforms []Transform
+	rng         *rand.Rand
+	tokens      []logit
+	topK        int
+	topP        float32
+	minP        float32
+	temperature float32
 }
 
-// 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)
+func (s *weighted) Sample(logits []float32) (int32, error) {
+	if len(s.tokens) < len(logits) {
+		s.tokens = make([]logit, len(logits))
 	}
-	return weighted{src: src, transforms: transforms}
-}
 
-func (s weighted) Sample(logits []float32) (int32, error) {
-	logits64 := make([]float64, len(logits))
+	tokens := s.tokens[:len(logits)]
+
 	for i, v := range logits {
-		logits64[i] = float64(v)
+		tokens[i].id = int32(i)
+		tokens[i].value = v
 	}
 
-	for _, t := range s.transforms {
-		logits64 = t.Apply(logits64)
+	// Tokens are sorted by logits in TopK or SortTokens
+	if s.topK > 0 {
+		tokens = topK(tokens, s.topK)
+	} else {
+		sortLogits(tokens)
 	}
 
-	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)
-		}
+	tokens = temperature(tokens, s.temperature)
+	tokens = softmax(tokens)
+
+	tokens = topP(tokens, s.topP)
+	tokens = minP(tokens, s.minP)
+
+	if len(tokens) == 0 {
+		return -1, errors.New("no valid logits found for weighted sampling")
 	}
 
-	if len(logitsCopy) == 0 {
-		return -1, errors.New("no valid logits found for weighed sampling")
+	var r float32
+	if s.rng != nil {
+		r = s.rng.Float32()
+	} else {
+		r = rand.Float32()
 	}
 
-	probs := softmax(logitsCopy)
-	w := sampleuv.NewWeighted(probs, s.src)
-	if idx, ok := w.Take(); ok {
-		return int32(indices[idx]), nil
+	// Calculate cumulative sum of probabilities
+	var sum float32
+	for i := range tokens {
+		sum += tokens[i].value
+		tokens[i].value = sum
 	}
-	return -1, errors.New("weighted sampler failed, no valid token found")
-}
+	r *= tokens[len(tokens)-1].value
 
-type greedy struct{}
+	idx, _ := slices.BinarySearchFunc(tokens, r, func(token logit, target float32) int {
+		// Compare cumulative probabilities
+		if token.value < target {
+			return -1
+		}
+		// First token that exceeds target
+		return 1
+	})
 
-func Greedy() Sampler {
-	return greedy{}
+	if idx >= len(tokens) {
+		idx = len(tokens) - 1
+	}
+
+	return tokens[idx].id, nil
 }
 
-// Sample returns the index of the maximum value in logits.
+type greedy struct{}
+
+// 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 {
-		if logits[i] > logits[maxIdx] {
+	maxVal := logits[0]
+	for i := 1; i < len(logits); i++ {
+		if logits[i] > maxVal {
+			maxVal = logits[i]
 			maxIdx = i
 		}
 	}
@@ -80,41 +107,40 @@ func (s greedy) Sample(logits []float32) (int32, error) {
 }
 
 // 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) (Sampler, error) {
+func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int) Sampler {
 	if temperature == 0 {
-		return Greedy(), nil
+		return &greedy{}
 	}
 
-	if temperature < 0 || temperature > 2 {
-		return nil, errors.New("temperature must be between 0 and 2")
+	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))
 	}
+	temperature = max(temperature, 1)
 
-	transforms := []Transform{Temperature(temperature)}
-
-	if topK != 0 {
-		if topK <= 0 {
-			return nil, errors.New("topK must be greater than 0")
-		}
-		transforms = append(transforms, TopK(topK))
+	if topP < 0.0 {
+		topP = 0.0
 	}
-
-	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 topP >= 1.0 {
+		topP = 1.0
 	}
 
-	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 minP < 0.0 {
+		minP = 0.0
+	}
+	if minP >= 1.0 {
+		minP = 1.0
 	}
 
-	if seed >= 0 {
-		seed64 := uint64(seed)
-		return Weighted(&seed64, transforms...), nil
+	return &weighted{
+		rng:         rng,
+		topK:        topK,
+		topP:        topP,
+		minP:        minP,
+		temperature: temperature,
 	}
-	return Weighted(nil, transforms...), nil
 }
diff --git a/sample/transforms.go b/sample/transforms.go
index 2dc6ebae..f1f4f3b1 100644
--- a/sample/transforms.go
+++ b/sample/transforms.go
@@ -1,120 +1,203 @@
 package sample
 
 import (
-	"cmp"
 	"math"
 	"slices"
-
-	pq "github.com/emirpasic/gods/v2/queues/priorityqueue"
 )
 
-type Transform interface {
-	Apply([]float64) []float64
-}
-
-// 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]
+func softmax(ts []logit) []logit {
+	var sum float32
+	for i, v := range ts {
+		ts[i].value = float32(math.Exp(float64(v.value)))
+		sum += ts[i].value
 	}
 
-	for i := range probs {
-		probs[i] /= sum
+	for i := range ts {
+		ts[i].value /= sum
 	}
 
-	return probs
+	return ts
 }
 
-type Temperature float64
+func temperature(ti []logit, t float32) []logit {
+	if t == 1 {
+		return ti
+	}
 
-func (t Temperature) Apply(logits []float64) []float64 {
-	temp := math.Max(float64(t), 1e-7)
+	temp := max(t, 1e-7)
+	maxLogit := float32(math.Inf(-1))
+	for _, token := range ti {
+		if token.value > maxLogit {
+			maxLogit = token.value
+		}
+	}
 
 	// subtracting max logit to avoid under/overflow
-	maxLogit := slices.Max(logits)
-	for i := range logits {
-		logits[i] = (logits[i] - maxLogit) / temp
+	for i := range ti {
+		ti[i].value = (ti[i].value - maxLogit) / temp
 	}
 
-	return logits
+	return ti
 }
 
-type logitMap struct {
-	index int
-	logit float64
+// 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 []logit, start, end int) {
+	root := start
+	for {
+		child := 2*root + 1
+		if child >= end {
+			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
+	}
 }
 
-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
+// topK limits the number of tokens considered to the k highest logits
+func topK(ts []logit, k int) []logit {
+	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)
 	}
-	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})
+	// 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)
+		}
 	}
 
-	validLogits := make(map[int]float64)
-	for range k {
-		logitMap, _ := q.Dequeue()
-		validLogits[logitMap.index] = logitMap.logit
+	slices.Reverse(heap)
+
+	ts = heap
+	return ts
+}
+
+// topP limits tokens to those with cumulative probability p
+func topP(ts []logit, p float32) []logit {
+	if p == 1.0 {
+		return ts
 	}
 
-	for i := range logits {
-		if _, ok := validLogits[i]; !ok {
-			logits[i] = math.Inf(-1)
+	// 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 logits
+	return ts
 }
 
-type TopP float64
+// minP limits tokens to those with cumulative probability p
+func minP(ts []logit, p float32) []logit {
+	if p == 1.0 {
+		return ts
+	}
 
-func (p TopP) Apply(logits []float64) []float64 {
-	probs := softmax(logits)
-	indices := make([]int, len(probs))
-	for i := range indices {
-		indices[i] = i
+	maxProb := float32(math.Inf(-1))
+	for _, token := range ts {
+		if token.value > maxProb {
+			maxProb = token.value
+		}
 	}
 
-	// sort in descending order
-	slices.SortFunc(indices, func(i, j int) int {
-		return cmp.Compare(probs[j], probs[i])
-	})
+	threshold := maxProb * float32(p)
 
-	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
+	// Filter tokens in-place
+	validTokens := ts[:0]
+	for i, token := range ts {
+		if token.value >= threshold {
+			validTokens = append(validTokens, ts[i])
 		}
 	}
-	return logits
-}
 
-type MinP float64
+	ts = validTokens
+	return ts
+}
 
-func (p MinP) Apply(logits []float64) []float64 {
-	probs := softmax(logits)
-	threshold := slices.Max(probs) * float64(p)
+// 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 []logit) {
+	if len(tokens) <= 1 {
+		return
+	}
 
-	for i, prob := range probs {
-		if prob < threshold {
-			logits[i] = math.Inf(-1)
+	// 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
 		}
 	}
 
-	return logits
+	// 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([]logit, 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)
 }



Test patch:
diff --git a/model/process_text_test.go b/model/process_text_test.go
index cad1f94f..f4830321 100644
--- a/model/process_text_test.go
+++ b/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")
+		ids, err := tokenizer.Encode("hello world", true)
 		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 
")
+		ids, err = tokenizer.Encode("hello 
", true)
 		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)
+			ids, err := tokenizer.Encode(s, true)
 			if err != nil {
 				t.Error(err)
 			}
@@ -152,7 +152,7 @@ func TestLlama(t *testing.T) {
 		}
 
 		for _, want := range cases {
-			ids, err := tokenizer.Encode(want)
+			ids, err := tokenizer.Encode(want, true)
 			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)
+			ids, err := tokenizer.Encode(s, true)
 			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))
+				_, err := tokenizer.Encode(string(bts), true)
 				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))
+			ids, err := tokenizer.Encode(string(bts), true)
 			if err != nil {
 				b.Fatal(err)
 			}
diff --git a/runner/ollamarunner/cache_test.go b/runner/ollamarunner/cache_test.go
index 99e67b4f..9ce03b73 100644
--- a/runner/ollamarunner/cache_test.go
+++ b/runner/ollamarunner/cache_test.go
@@ -4,6 +4,8 @@ import (
 	"image"
 	"testing"
 	"time"
+
+	"github.com/ollama/ollama/model"
 )
 
 func TestCountCommon(t *testing.T) {
@@ -13,44 +15,50 @@ func TestCountCommon(t *testing.T) {
 
 	tests := []struct {
 		name     string
-		t1       []input
-		t2       []input
+		t1       []model.Input
+		t2       []model.Input
 		expected int32
 	}{
 		{
 			name:     "Equal",
-			t1:       []input{{token: 1}, {token: 2}, {token: 3}},
-			t2:       []input{{token: 1}, {token: 2}, {token: 3}},
+			t1:       []model.Input{{Token: 1}, {Token: 2}, {Token: 3}},
+			t2:       []model.Input{{Token: 1}, {Token: 2}, {Token: 3}},
 			expected: 3,
 		},
 		{
 			name:     "Prefix",
-			t1:       []input{{token: 1}},
-			t2:       []input{{token: 1}, {token: 2}, {token: 3}},
+			t1:       []model.Input{{Token: 1}},
+			t2:       []model.Input{{Token: 1}, {Token: 2}, {Token: 3}},
 			expected: 1,
 		},
 		{
 			name:     "Image Prefix",
-			t1:       []input{{image: imgA}},
-			t2:       []input{{image: imgA}, {image: imgB}, {image: imgC}},
+			t1:       []model.Input{{Multimodal: imgA, MultimodalHash: 1}},
+			t2:       []model.Input{{Multimodal: imgA, MultimodalHash: 1}, {Multimodal: imgB, MultimodalHash: 2}, {Multimodal: imgC, MultimodalHash: 3}},
 			expected: 1,
 		},
 		{
 			name:     "Mixed",
-			t1:       []input{{token: 1}, {image: imgA}},
-			t2:       []input{{token: 1}, {image: imgA}, {token: 5}},
+			t1:       []model.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}},
+			t2:       []model.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}, {Token: 5}},
 			expected: 2,
 		},
+		{
+			name:     "Mixed, Same Length",
+			t1:       []model.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}},
+			t2:       []model.Input{{Token: 1}, {Multimodal: imgB, MultimodalHash: 2}},
+			expected: 1,
+		},
 		{
 			name:     "Empty",
-			t1:       []input{},
-			t2:       []input{{token: 1}, {token: 2}, {token: 3}},
+			t1:       []model.Input{},
+			t2:       []model.Input{{Token: 1}, {Token: 2}, {Token: 3}},
 			expected: 0,
 		},
 		{
 			name:     "Both Empty",
-			t1:       []input{},
-			t2:       []input{},
+			t1:       []model.Input{},
+			t2:       []model.Input{},
 			expected: 0,
 		},
 	}
@@ -74,7 +82,7 @@ func TestFindCacheSlot(t *testing.T) {
 	tests := []struct {
 		name    string
 		cache   InputCache
-		prompt  []input
+		prompt  []model.Input
 		longest expected
 		best    expected
 	}{
@@ -83,18 +91,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input{},
+					Inputs:   []model.Input{},
 					InUse:    false,
 					lastUsed: time.Time{},
 				},
 				{
 					Id:       1,
-					Inputs:   []input{},
+					Inputs:   []model.Input{},
 					InUse:    false,
 					lastUsed: time.Time{},
 				},
 			}},
-			prompt:  []input{{token: 1}},
+			prompt:  []model.Input{{Token: 1}},
 			longest: expected{result: 0, len: 0},
 			best:    expected{result: 0, len: 0},
 		},
@@ -103,18 +111,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input{{token: 1}},
+					Inputs:   []model.Input{{Token: 1}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input{{token: 1}, {token: 2}},
+					Inputs:   []model.Input{{Token: 1}, {Token: 2}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-2 * time.Second),
 				},
 			}},
-			prompt:  []input{{token: 1}, {token: 2}},
+			prompt:  []model.Input{{Token: 1}, {Token: 2}},
 			longest: expected{result: 1, len: 2},
 			best:    expected{result: 1, len: 2},
 		},
@@ -123,18 +131,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input{{token: 1}, {token: 2}},
+					Inputs:   []model.Input{{Token: 1}, {Token: 2}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input{},
+					Inputs:   []model.Input{},
 					InUse:    false,
 					lastUsed: time.Time{},
 				},
 			}},
-			prompt:  []input{{token: 2}},
+			prompt:  []model.Input{{Token: 2}},
 			longest: expected{result: 0, len: 0},
 			best:    expected{result: 1, len: 0},
 		},
@@ -144,19 +152,19 @@ func TestFindCacheSlot(t *testing.T) {
 				slots: []InputCacheSlot{
 					{
 						Id:       0,
-						Inputs:   []input{{token: 1}, {token: 2}},
+						Inputs:   []model.Input{{Token: 1}, {Token: 2}},
 						InUse:    false,
 						lastUsed: time.Now().Add(-time.Second),
 					},
 					{
 						Id:       1,
-						Inputs:   []input{},
+						Inputs:   []model.Input{},
 						InUse:    false,
 						lastUsed: time.Time{},
 					},
 				},
 			},
-			prompt:  []input{{token: 1}},
+			prompt:  []model.Input{{Token: 1}},
 			longest: expected{result: 0, len: 1},
 			best:    expected{result: 1, len: 1},
 		},
@@ -165,18 +173,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input{{token: 1}},
+					Inputs:   []model.Input{{Token: 1}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input{{token: 1}, {token: 2}},
+					Inputs:   []model.Input{{Token: 1}, {Token: 2}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-2 * time.Second),
 				},
 			}},
-			prompt:  []input{{token: 2}, {token: 3}},
+			prompt:  []model.Input{{Token: 2}, {Token: 3}},
 			longest: expected{result: 0, len: 0},
 			best:    expected{result: 1, len: 0},
 		},
@@ -185,18 +193,18 @@ func TestFindCacheSlot(t *testing.T) {
 			cache: InputCache{slots: []InputCacheSlot{
 				{
 					Id:       0,
-					Inputs:   []input{{token: 1}, {token: 2}},
+					Inputs:   []model.Input{{Token: 1}, {Token: 2}},
 					InUse:    true,
 					lastUsed: time.Now().Add(-time.Second),
 				},
 				{
 					Id:       1,
-					Inputs:   []input{{token: 1}},
+					Inputs:   []model.Input{{Token: 1}},
 					InUse:    false,
 					lastUsed: time.Now().Add(-2 * time.Second),
 				},
 			}},
-			prompt:  []input{{token: 1}, {token: 2}},
+			prompt:  []model.Input{{Token: 1}, {Token: 2}},
 			longest: expected{result: 1, len: 1},
 			best:    expected{result: 1, len: 2},
 		},
diff --git a/sample/samplers_benchmark_test.go b/sample/samplers_benchmark_test.go
new file mode 100644
index 00000000..41c0b487
--- /dev/null
+++ b/sample/samplers_benchmark_test.go
@@ -0,0 +1,104 @@
+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)
+			b.ResetTimer()
+			for b.Loop() {
+				_, err := sampler.Sample(logits)
+				if err != nil {
+					b.Fatalf("Sampling failed: %v", err)
+				}
+			}
+		})
+	}
+
+	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)
+			sampler.Sample(logits)
+
+			b.ResetTimer()
+
+			for b.Loop() {
+				_, err := sampler.Sample(logits)
+				if err != nil {
+					b.Fatalf("Sampling failed: %v", err)
+				}
+			}
+		})
+	}
+
+	// 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)
+		b.ResetTimer()
+
+		for b.Loop() {
+			_, err := sampler.Sample(logits)
+			if err != nil {
+				b.Fatalf("Sampling failed: %v", err)
+			}
+		}
+	})
+}
+
+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)
+			b.ResetTimer()
+
+			for b.Loop() {
+				_, err := sampler.Sample(logits)
+				if err != nil {
+					b.Fatalf("Sampling failed: %v", err)
+				}
+			}
+		})
+	}
+}
diff --git a/sample/samplers_test.go b/sample/samplers_test.go
index 32364a3b..dbbee17b 100644
--- a/sample/samplers_test.go
+++ b/sample/samplers_test.go
@@ -1,15 +1,14 @@
 package sample
 
 import (
-	"math"
 	"math/rand/v2"
 	"testing"
-
-	"github.com/google/go-cmp/cmp"
 )
 
 func TestWeighted(t *testing.T) {
-	got, err := Weighted(nil).Sample([]float32{float32(math.Inf(-1)), 2, float32(math.Inf(-1)), float32(math.Inf(-1))})
+	logits := []float32{-10, 3, -10, -10}
+	sampler := NewSampler(0, 0, 0, 0, 0)
+	got, err := sampler.Sample(logits)
 	if err != nil {
 		t.Error(err)
 		return
@@ -19,64 +18,19 @@ func TestWeighted(t *testing.T) {
 		t.Errorf("index mismatch: want %d, got %d", want, got)
 	}
 
-	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})
+	logits = []float32{-100, -10, 0, 10}
+	sampler = NewSampler(0, 0, 0, 0, 0)
+	got, err = sampler.Sample(logits)
 	if err != nil {
 		t.Error(err)
 		return
 	}
-	// With seed 42, we expect a consistent sample
-	want = int32(3) // This will be deterministic due to the seed
+	want = int32(3) // Should pick highest probability with this r value
 	if want != got {
 		t.Errorf("index mismatch: want %d, got %d", want, got)
 	}
 }
 
-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
@@ -85,75 +39,41 @@ func TestNewSampler(t *testing.T) {
 		topP        float32
 		minP        float32
 		seed        int
-		wantErr     bool
+		wantGreedy  bool // Instead of wantErr, check if we get greedy sampler
 	}{
-		{
-			name: "no transforms",
-			// temperature is 0, so greedy should be used
-			wantErr: false,
-		},
 		{
 			name:        "temperature",
 			temperature: 0.5,
-			wantErr:     false,
+			wantGreedy:  false,
 		},
 		{
-			name:        "invalid temperature negative",
-			temperature: -1,
-			wantErr:     true,
-		},
-		{
-			name:        "invalid temperature too high",
-			temperature: 2.1,
-			wantErr:     true,
+			name:        "zero temperature - greedy",
+			temperature: 0,
+			wantGreedy:  true,
 		},
 		{
 			name:        "top k",
+			temperature: 0.1,
 			topK:        10,
-			temperature: 0.8,
-			wantErr:     false,
-		},
-		{
-			name:        "invalid top k negative",
-			topK:        -1,
-			temperature: 0.8,
-			wantErr:     true,
+			wantGreedy:  false,
 		},
 		{
 			name:        "top p",
+			temperature: 0.1,
 			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,
+			wantGreedy:  false,
 		},
 		{
 			name:        "min p",
+			temperature: 0.1,
 			minP:        0.2,
-			temperature: 0.8,
-			wantErr:     false,
-		},
-		{
-			name:        "invalid min p negative",
-			minP:        -0.1,
-			temperature: 0.8,
-			wantErr:     true,
+			wantGreedy:  false,
 		},
 		{
-			name:        "invalid min p one",
-			minP:        1.0,
-			temperature: 0.8,
-			wantErr:     true,
+			name:        "seed - weighted",
+			temperature: 0.1,
+			seed:        42,
+			wantGreedy:  false,
 		},
 		{
 			name:        "default values",
@@ -162,16 +82,16 @@ func TestNewSampler(t *testing.T) {
 			topP:        0.9,
 			minP:        0.0,
 			seed:        0,
-			wantErr:     false,
+			wantGreedy:  false,
 		},
 		{
-			name:        "all zeroes",
+			name:        "all zeroes - greedy",
 			temperature: 0.0,
 			topK:        0,
 			topP:        0.0,
 			minP:        0.0,
 			seed:        0,
-			wantErr:     false, // all zeroes means no transforms
+			wantGreedy:  true,
 		},
 		{
 			name:        "all transforms",
@@ -180,33 +100,28 @@ func TestNewSampler(t *testing.T) {
 			topP:        0.95,
 			minP:        0.1,
 			seed:        42,
-			wantErr:     false,
+			wantGreedy:  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)
+			sampler := NewSampler(tt.temperature, tt.topK, tt.topP, tt.minP, tt.seed)
+			_, isGreedy := sampler.(*greedy)
+			if isGreedy != tt.wantGreedy {
+				t.Errorf("NewSampler() got greedy = %v, want %v", isGreedy, tt.wantGreedy)
 			}
 		})
 	}
 }
 
 func BenchmarkSample(b *testing.B) {
-	transforms := []Transform{
-		Temperature(0.5),
-		TopK(10),
-		TopP(0.9),
-		MinP(0.2),
-	}
-
+	weighted := NewSampler(0.5, 10, 0.9, 0.2, -1)
 	samplers := map[string]Sampler{
-		"Greedy":   Greedy(),
-		"Weighted": Weighted(nil, transforms...),
+		"Greedy":   NewSampler(0, 0, 0, 0, 0), // Use NewSampler with temp=0 for greedy
+		"Weighted": weighted,
 	}
 
+	// Generate random logits for benchmarking
 	logits := make([]float32, 1<<16)
 	for i := range logits {
 		logits[i] = rand.Float32()
@@ -215,7 +130,7 @@ func BenchmarkSample(b *testing.B) {
 	for name, s := range samplers {
 		b.Run(name, func(b *testing.B) {
 			b.ResetTimer()
-			for range b.N {
+			for b.Loop() {
 				if _, err := s.Sample(logits); err != nil {
 					b.Error(err)
 				}
diff --git a/sample/transforms_test.go b/sample/transforms_test.go
index 05f76a27..950d79b3 100644
--- a/sample/transforms_test.go
+++ b/sample/transforms_test.go
@@ -4,77 +4,182 @@ import (
 	"math"
 	"math/rand/v2"
 	"testing"
-
-	"github.com/google/go-cmp/cmp"
 )
 
-func TestTemperature(t *testing.T) {
-	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)
+// Helper to convert float64 slice to logit slice
+func toLogits(values []float64) []logit {
+	tokens := make([]logit, len(values))
+	for i, v := range values {
+		tokens[i] = logit{
+			id:    int32(i),
+			value: float32(v),
+		}
+	}
+	return tokens
+}
+
+// Helper to compare logit slices
+func compareLogits(t *testing.T, name string, want []float64, got []logit) {
+	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(toLogits(input), 0.5)
+	compareLogits(t, "Temperature", want, got)
+}
+
 func TestSoftmax(t *testing.T) {
-	got := softmax([]float64{-3, -2, -1, 0, 1, 2, 4})
+	input := []float64{-3, -2, -1, 0, 1, 2, 4}
+	got := softmax(toLogits(input))
 
-	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)
+	// 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)
 	}
-}
 
-func TestTopK(t *testing.T) {
-	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)
+	// 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)
+		}
 	}
+}
 
-	got = TopK(10).Apply([]float64{-3, -2, -1, 0, 1, 2, 4})
+func TestTopK(t *testing.T) {
+	input := []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)
+	// Test k=3
+	got := topK(toLogits(input), 3)
+	if len(got) != 3 {
+		t.Errorf("topK(3): wrong length: want 3, got %d", len(got))
 	}
+	// Should keep highest 3 values: 4, 2, 1
+	want := []float64{4, 2, 1}
+	compareLogits(t, "topK(3)", want, got)
+
+	// Test k > len
+	got = topK(toLogits(input), 10)
+	compareLogits(t, "topK(10)", input, got)
 }
 
 func TestTopP(t *testing.T) {
-	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)
+	input := []float64{-3, -2, -1, 0, 1, 2, 4}
+	tokens := toLogits(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)
 	}
 }
 
 func TestMinP(t *testing.T) {
-	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)
+	input := []float64{-3, -2, -1, 0, 1, 2, 4, 3}
+	tokens := toLogits(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))
 	}
 }
 
-func BenchmarkTransform(b *testing.B) {
-	transforms := map[string]Transform{
-		"Temperature": Temperature(0.5),
-		"TopK":        TopK(10),
-		"TopP":        TopP(0.9),
-		"MinP":        MinP(0.2),
-	}
+func TestSortLogits(t *testing.T) {
+	input := []float64{3, 1, 4, 2, -1, 0, -2}
+	tokens := toLogits(input)
+
+	sortLogits(tokens)
 
-	logits := make([]float64, 1<<16)
-	for i := range logits {
-		logits[i] = rand.Float64()
+	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)
+		}
 	}
 
-	for name, transform := range transforms {
-		b.Run(name, func(b *testing.B) {
-			b.ResetTimer()
-			for range b.N {
-				transform.Apply(logits)
-			}
-		})
+	want := []float64{4, 3, 2, 1, 0, -1, -2}
+	compareLogits(t, "sortLogits", want, tokens)
+}
+
+func BenchmarkTransforms(b *testing.B) {
+	// Generate random logits
+	tokens := make([]logit, 1<<16)
+	for i := range tokens {
+		tokens[i] = logit{
+			id:    int32(i),
+			value: rand.Float32(),
+		}
 	}
+
+	tokensCopy := make([]logit, 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)
+		}
+	})
 }