Repository: ollama/ollama
Pull request: runner.go: Support for vision models

In addition to bringing feature parity with the C++ runner, this also incorporates several improvements:
 - Cache prompting works with images, avoiding the need to re-decode embeddings for every message in a conversation
 - Parallelism is supported, avoiding the need to restrict to one sequence at a time. (Though for now Ollama will not schedule them while we might need to fall back to the old runner.)

Co-authored-by: jmorganca <jmorganca@gmail.com>

Changed files:
filename: llama/runner/cache.go
patch: @@ -2,6 +2,7 @@ package main
 
 import (
 	"errors"
+	"hash/maphash"
 	"log/slog"
 	"reflect"
 	"time"
@@ -13,9 +14,13 @@ type InputCache struct {
 	// context window size (per slot)
 	numCtx int
 
-	// individual caches
+	// individual KV caches
 	slots []InputCacheSlot
 
+	// cache of images to embeddings
+	images    []imageCache
+	imageHash maphash.Hash
+
 	lc *llama.Context
 }
 
@@ -32,6 +37,7 @@ func NewInputCache(lc *llama.Context, kvSize int, numSlots int) *InputCache {
 	return &InputCache{
 		numCtx: kvSize / numSlots,
 		slots:  slots,
+		images: make([]imageCache, numSlots),
 		lc:     lc,
 	}
 }
@@ -160,3 +166,55 @@ func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int, numDiscar
 	}
 	slot.Inputs = slot.Inputs[:len(slot.Inputs)-numDiscard]
 }
+
+// Locking: Lookup and store operations on imageCache require a lock
+// to be held that serializes these with each other. Hash does not
+// require a lock nor they need to be serialized with InputCacheSlot.
+
+type imageCache struct {
+	key      uint64
+	val      [][]float32
+	lastUsed time.Time
+}
+
+func (c *InputCache) HashImage(image []byte) uint64 {
+	c.imageHash.Reset()
+	_, _ = c.imageHash.Write(image)
+	return c.imageHash.Sum64()
+}
+
+var ErrImageNotFound = errors.New("image not found in cache")
+
+func (c *InputCache) FindImage(hash uint64) ([][]float32, error) {
+	for i := range c.images {
+		if c.images[i].key == hash {
+			slog.Debug("loading image embeddings from cache", "entry", i)
+			c.images[i].lastUsed = time.Now()
+			return c.images[i].val, nil
+		}
+	}
+
+	return nil, ErrImageNotFound
+}
+
+func (c *InputCache) AddImage(hash uint64, embed [][]float32) {
+	best := time.Now()
+	var bestImage int
+
+	for i := range c.images {
+		if c.images[i].key == hash {
+			bestImage = i
+			break
+		}
+
+		if c.images[i].lastUsed.Compare(best) < 0 {
+			best = c.images[i].lastUsed
+			bestImage = i
+		}
+	}
+
+	slog.Debug("storing image embeddings in cache", "entry", bestImage, "used", c.images[bestImage].lastUsed)
+	c.images[bestImage].key = hash
+	c.images[bestImage].val = embed
+	c.images[bestImage].lastUsed = time.Now()
+}

filename: llama/runner/cache_test.go
patch: @@ -1,6 +1,7 @@
 package main
 
 import (
+	"reflect"
 	"testing"
 	"time"
 )
@@ -210,3 +211,77 @@ func TestFindCacheSlot(t *testing.T) {
 		})
 	}
 }
+
+func TestImageCache(t *testing.T) {
+	cache := NewInputCache(nil, 2048, 4)
+
+	valA := [][]float32{{0.1, 0.2}, {0.3}}
+	valB := [][]float32{{0.4}, {0.5}, {0.6}}
+	valC := [][]float32{{0.7}}
+	valD := [][]float32{{0.8}}
+	valE := [][]float32{{0.9}}
+
+	// Empty cache
+	result, err := cache.FindImage(0x5adb61d31933a946)
+	if err != ErrImageNotFound {
+		t.Errorf("found result in empty cache: result %v, err %v", result, err)
+	}
+
+	// Insert A
+	cache.AddImage(0x5adb61d31933a946, valA)
+
+	result, err = cache.FindImage(0x5adb61d31933a946)
+	if !reflect.DeepEqual(result, valA) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+
+	// Insert B
+	cache.AddImage(0x011551369a34a901, valB)
+
+	result, err = cache.FindImage(0x5adb61d31933a946)
+	if !reflect.DeepEqual(result, valA) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+	result, err = cache.FindImage(0x011551369a34a901)
+	if !reflect.DeepEqual(result, valB) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+
+	// Replace B with C
+	cache.AddImage(0x011551369a34a901, valC)
+
+	result, err = cache.FindImage(0x5adb61d31933a946)
+	if !reflect.DeepEqual(result, valA) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+	result, err = cache.FindImage(0x011551369a34a901)
+	if !reflect.DeepEqual(result, valC) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+
+	// Evict A
+	cache.AddImage(0x756b218a517e7353, valB)
+	cache.AddImage(0x75e5e8d35d7e3967, valD)
+	cache.AddImage(0xd96f7f268ca0646e, valE)
+
+	result, err = cache.FindImage(0x5adb61d31933a946)
+	if reflect.DeepEqual(result, valA) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+	result, err = cache.FindImage(0x756b218a517e7353)
+	if !reflect.DeepEqual(result, valB) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+	result, err = cache.FindImage(0x011551369a34a901)
+	if !reflect.DeepEqual(result, valC) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+	result, err = cache.FindImage(0x75e5e8d35d7e3967)
+	if !reflect.DeepEqual(result, valD) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+	result, err = cache.FindImage(0xd96f7f268ca0646e)
+	if !reflect.DeepEqual(result, valE) {
+		t.Errorf("failed to find expected value: result %v, err %v", result, err)
+	}
+}

filename: llama/runner/runner.go
patch: @@ -190,9 +190,15 @@ func (s *Server) inputs(prompt string, images []ImageData) ([]input, error) {
 				return nil, fmt.Errorf("invalid image index: %d", n)
 			}
 
+			hash := s.cache.HashImage(images[imageIndex].Data)
+
 			// Vision models cannot be accessed concurrently
 			s.clip.mu.Lock()
-			embed := llama.NewLlavaImageEmbed(s.lc, s.clip.cc, images[imageIndex].Data)
+			embed, err := s.cache.FindImage(hash)
+			if err != nil {
+				embed = llama.NewLlavaImageEmbed(s.lc, s.clip.cc, images[imageIndex].Data)
+				s.cache.AddImage(hash, embed)
+			}
 			s.clip.mu.Unlock()
 
 			for _, e := range embed {