Repository: ollama/ollama
Pull request: runner.go: Prompt caching

Currently, KV cache entries from a sequence are discarded at the end of each processing run. In a typical chat conversation, this results in each message taking longer and longer to process as the entire history of the conversation needs to be replayed.

Prompt caching retains the KV entries as long as possible so that we only need to process the newest message in the converation, at least until there are too many simultaneous conversations and something needs to be evicted.

Changed files:
filename: llama/llama.go
patch: @@ -139,6 +139,10 @@ func (c *Context) KvCacheSeqRm(seqId int, p0 int, p1 int) bool {
 	return bool(C.llama_kv_cache_seq_rm(c.c, C.int(seqId), C.int(p0), C.int(p1)))
 }
 
+func (c *Context) KvCacheSeqCp(srcSeqId int, dstSeqId int, p0 int, p1 int) {
+	C.llama_kv_cache_seq_cp(c.c, C.int(srcSeqId), C.int(dstSeqId), C.int(p0), C.int(p1))
+}
+
 // Get the embeddings for a sequence id
 func (c *Context) GetEmbeddingsSeq(seqId int) []float32 {
 	embeddings := unsafe.Pointer(C.llama_get_embeddings_seq(c.c, C.int(seqId)))

filename: llama/runner/cache.go
patch: @@ -0,0 +1,156 @@
+package main
+
+import (
+	"log/slog"
+	"time"
+
+	"github.com/ollama/ollama/llama"
+)
+
+type TokenCache struct {
+	// context window size (per slot)
+	numCtx int
+
+	// individual caches
+	slots []TokenCacheSlot
+
+	lc *llama.Context
+}
+
+func NewTokenCache(lc *llama.Context, kvSize int, numSlots int) *TokenCache {
+	slots := make([]TokenCacheSlot, numSlots)
+
+	for i := range slots {
+		slots[i] = TokenCacheSlot{
+			id:     i,
+			tokens: make([]int, 0),
+		}
+	}
+
+	return &TokenCache{
+		numCtx: kvSize / numSlots,
+		slots:  slots,
+		lc:     lc,
+	}
+}
+
+// Locking: Operations on TokenCacheSlot (including finding one
+// through LoadCacheSlot) require a lock to be be held that serializes
+// these operations with each other and llama.Decode
+
+type TokenCacheSlot struct {
+	// Index in the KV cache
+	id int
+
+	// tokens that are stored in the KV cache
+	tokens []int
+
+	// is this cache actively being processed as part of a sequence?
+	inUse bool
+
+	// last time this cache was used (as of start of processing)
+	lastUsed time.Time
+}
+
+func (t *TokenCache) LoadCacheSlot(prompt []int) (*TokenCacheSlot, []int, int) {
+	slot, numPast := t.findCacheSlot(prompt)
+
+	slot.inUse = true
+	slot.lastUsed = time.Now()
+
+	if numPast == len(prompt) {
+		// Leave one token to sample so we can get a response
+		numPast--
+	}
+
+	if !t.lc.KvCacheSeqRm(slot.id, numPast, -1) {
+		// Some models don't support partial erasure
+		t.lc.KvCacheSeqRm(slot.id, 0, -1)
+		numPast = 0
+	}
+
+	slog.Debug("loading cache slot", "id", slot.id, "cache", len(slot.tokens), "prompt", len(prompt),
+		"used", numPast, "remaining", len(prompt)-numPast)
+
+	prompt = prompt[numPast:]
+	slot.tokens = slot.tokens[:numPast]
+
+	return slot, prompt, numPast
+}
+
+func (t *TokenCache) findCacheSlot(prompt []int) (*TokenCacheSlot, int) {
+	oldest := time.Now()
+	var oldestSlot *TokenCacheSlot
+
+	longest := -1
+	var longestSlot *TokenCacheSlot
+
+	for i, s := range t.slots {
+		count := countCommonPrefix(s.tokens, prompt)
+		if count > longest {
+			longest = count
+			longestSlot = &t.slots[i]
+		}
+
+		if s.lastUsed.Compare(oldest) < 0 && !s.inUse {
+			oldest = s.lastUsed
+			oldestSlot = &t.slots[i]
+		}
+	}
+
+	if longest == len(longestSlot.tokens) && !longestSlot.inUse {
+		return longestSlot, longest
+	}
+
+	if oldestSlot.inUse {
+		panic("no available cache slots")
+	}
+
+	if len(oldestSlot.tokens) != 0 {
+		slog.Debug("evicting cache slot", "id", oldestSlot.id, "tokens", len(oldestSlot.tokens),
+			"used", oldestSlot.lastUsed)
+	}
+
+	if longest > 0 && longestSlot != oldestSlot {
+		slog.Debug("forking cache slot", "src", longestSlot.id, "dst", oldestSlot.id, "tokens", longest, "total",
+			len(longestSlot.tokens))
+		oldestSlot.tokens = make([]int, longest)
+		copy(oldestSlot.tokens, longestSlot.tokens[:longest])
+		// This is only nil for unit tests
+		if t.lc != nil {
+			t.lc.KvCacheSeqCp(longestSlot.id, oldestSlot.id, 0, longest)
+		}
+	}
+
+	return oldestSlot, longest
+}
+
+func countCommonPrefix(a []int, b []int) int {
+	var count int
+
+	for i := range a {
+		if i >= len(b) {
+			break
+		}
+
+		if a[i] != b[i] {
+			break
+		}
+
+		count++
+	}
+
+	return count
+}
+
+func (t *TokenCache) ShiftCacheSlot(slot *TokenCacheSlot, numKeep int, numDiscard int, numPast int) {
+	// TODO (jessegross): KV cache removal can fail for certain types of models
+	// server.cpp doesn't handle this, though we can be more graceful
+	t.lc.KvCacheSeqRm(slot.id, numKeep, numKeep+numDiscard)
+	t.lc.KvCacheSeqAdd(slot.id, numKeep+numDiscard, numPast, -numDiscard)
+
+	for i := numKeep + numDiscard; i < len(slot.tokens); i++ {
+		slot.tokens[i-numDiscard] = slot.tokens[i]
+	}
+	slot.tokens = slot.tokens[:len(slot.tokens)-numDiscard]
+}

filename: llama/runner/cache_test.go
patch: @@ -0,0 +1,192 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+func TestCountCommon(t *testing.T) {
+	tests := []struct {
+		name     string
+		t1       []int
+		t2       []int
+		expected int
+	}{
+		{
+			name:     "Equal",
+			t1:       []int{1, 2, 3},
+			t2:       []int{1, 2, 3},
+			expected: 3,
+		},
+		{
+			name:     "Prefix",
+			t1:       []int{1},
+			t2:       []int{1, 2, 3},
+			expected: 1,
+		},
+		{
+			name:     "Empty",
+			t1:       []int{},
+			t2:       []int{1, 2, 3},
+			expected: 0,
+		},
+		{
+			name:     "Both Empty",
+			t1:       []int{},
+			t2:       []int{},
+			expected: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := countCommonPrefix(tt.t1, tt.t2)
+			if result != tt.expected {
+				t.Errorf("countCommonPrefix(%v, %v): have %v; want %v", tt.t1, tt.t2, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestFindCacheSlot(t *testing.T) {
+	tests := []struct {
+		name        string
+		cache       TokenCache
+		prompt      []int
+		expected    int
+		expectedLen int
+	}{
+		{
+			name: "Empty",
+			cache: TokenCache{slots: []TokenCacheSlot{
+				{
+					id:       0,
+					tokens:   []int{},
+					inUse:    false,
+					lastUsed: time.Time{},
+				},
+				{
+					id:       1,
+					tokens:   []int{},
+					inUse:    false,
+					lastUsed: time.Time{},
+				},
+			}},
+			prompt:      []int{1},
+			expected:    0,
+			expectedLen: 0,
+		},
+		{
+			name: "Extend",
+			cache: TokenCache{slots: []TokenCacheSlot{
+				{
+					id:       0,
+					tokens:   []int{1},
+					inUse:    false,
+					lastUsed: time.Now().Add(-time.Second),
+				},
+				{
+					id:       1,
+					tokens:   []int{1, 2},
+					inUse:    false,
+					lastUsed: time.Now().Add(-2 * time.Second),
+				},
+			}},
+			prompt:      []int{1, 2},
+			expected:    1,
+			expectedLen: 2,
+		},
+		{
+			name: "New",
+			cache: TokenCache{slots: []TokenCacheSlot{
+				{
+					id:       0,
+					tokens:   []int{1, 2},
+					inUse:    false,
+					lastUsed: time.Now().Add(-time.Second),
+				},
+				{
+					id:       1,
+					tokens:   []int{},
+					inUse:    false,
+					lastUsed: time.Time{},
+				},
+			}},
+			prompt:      []int{2},
+			expected:    1,
+			expectedLen: 0,
+		},
+		{
+			name: "Fork",
+			cache: TokenCache{
+				slots: []TokenCacheSlot{
+					{
+						id:       0,
+						tokens:   []int{1, 2},
+						inUse:    false,
+						lastUsed: time.Now().Add(-time.Second),
+					},
+					{
+						id:       1,
+						tokens:   []int{},
+						inUse:    false,
+						lastUsed: time.Time{},
+					},
+				},
+			},
+			prompt:      []int{1},
+			expected:    1,
+			expectedLen: 1,
+		},
+		{
+			name: "Evict",
+			cache: TokenCache{slots: []TokenCacheSlot{
+				{
+					id:       0,
+					tokens:   []int{1},
+					inUse:    false,
+					lastUsed: time.Now().Add(-time.Second),
+				},
+				{
+					id:       1,
+					tokens:   []int{1, 2},
+					inUse:    false,
+					lastUsed: time.Now().Add(-2 * time.Second),
+				},
+			}},
+			prompt:      []int{2, 3},
+			expected:    1,
+			expectedLen: 0,
+		},
+		{
+			name: "In use",
+			cache: TokenCache{slots: []TokenCacheSlot{
+				{
+					id:       0,
+					tokens:   []int{1},
+					inUse:    false,
+					lastUsed: time.Now().Add(-time.Second),
+				},
+				{
+					id:       1,
+					tokens:   []int{1, 2},
+					inUse:    true,
+					lastUsed: time.Now().Add(-2 * time.Second),
+				},
+			}},
+			prompt:      []int{1, 2},
+			expected:    0,
+			expectedLen: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, resultLen := tt.cache.findCacheSlot(tt.prompt)
+			if result.id != tt.expected || resultLen != tt.expectedLen {
+				t.Errorf("findCacheSlot: slot have %v, want %v len have %v, want %v",
+					result.id, tt.expected, resultLen, tt.expectedLen)
+			}
+		})
+	}
+}

filename: llama/runner/runner.go
patch: @@ -25,7 +25,7 @@ import (
 
 type Sequence struct {
 	// number of tokens evaluated
-	nPast int
+	numPast int
 
 	// batch index
 	iBatch int
@@ -40,6 +40,9 @@ type Sequence struct {
 	// TODO (jmorganca): simplify this
 	pendingResponses []string
 
+	// token cache being used by this sequence
+	cache *TokenCacheSlot
+
 	// channel to send responses over
 	responses chan string
 
@@ -94,18 +97,18 @@ func (s *Server) NewSequence(prompt string, params NewSequenceParams) *Sequence
 
 	if !params.embedding {
 		// Subtracting 4 ensures that at least 1 token can be discarded during shift
-		params.numKeep = min(params.numKeep, s.numCtx-4)
+		params.numKeep = min(params.numKeep, s.cache.numCtx-4)
 		params.numKeep += s.bosToken
 	} else {
 		// Embeddings are 1 shot - just truncate to the context window, without ever shifting
-		params.numKeep = min(params.numKeep, s.numCtx)
+		params.numKeep = min(params.numKeep, s.cache.numCtx)
 	}
 
 	// truncate to fit in context window
-	if len(tokens) > s.numCtx {
-		slog.Warn("truncating input prompt", "limit", s.numCtx, "prompt", len(tokens), "numKeep", params.numKeep)
+	if len(tokens) > s.cache.numCtx {
+		slog.Warn("truncating input prompt", "limit", s.cache.numCtx, "prompt", len(tokens), "numKeep", params.numKeep)
 		newTokens := tokens[:params.numKeep]
-		newTokens = append(newTokens, tokens[len(tokens)-s.numCtx+params.numKeep:]...)
+		newTokens = append(newTokens, tokens[len(tokens)-s.cache.numCtx+params.numKeep:]...)
 		tokens = newTokens
 	}
 
@@ -146,8 +149,8 @@ type Server struct {
 	// TODO (jmorganca): this can probably be moved into run()
 	seqs []*Sequence
 
-	// context window size
-	numCtx int
+	// KV cache
+	cache *TokenCache
 
 	// does this model require a beginning of sequence token?
 	bosToken int
@@ -173,21 +176,16 @@ func (s *Server) allNil() bool {
 	return true
 }
 
-func (s *Server) shiftContext(seqIndex int) {
-	seq := s.seqs[seqIndex]
-
-	numLeft := seq.nPast - seq.numKeep
+func (s *Server) shiftContext(seq *Sequence) {
+	numLeft := seq.numPast - seq.numKeep
 	numDiscard := numLeft / 2
 
-	slog.Debug("context limit hit - shifting", "limit", s.numCtx, "nPast", seq.nPast,
+	slog.Debug("context limit hit - shifting", "limit", s.cache.numCtx, "numPast", seq.numPast,
 		"numKeep", seq.numKeep, "numLeft", numLeft, "numDiscard", numDiscard)
 
-	// TODO (jessegross): KV cache removal can fail for certain types of models
-	// server.cpp doesn't handle this, though we can be more graceful
-	s.lc.KvCacheSeqRm(seqIndex, seq.numKeep, seq.numKeep+numDiscard)
-	s.lc.KvCacheSeqAdd(seqIndex, seq.numKeep+numDiscard, seq.nPast, -numDiscard)
+	s.cache.ShiftCacheSlot(seq.cache, seq.numKeep, numDiscard, seq.numPast)
 
-	seq.nPast -= numDiscard
+	seq.numPast -= numDiscard
 }
 
 func incompleteUnicode(token string) bool {
@@ -227,8 +225,8 @@ func (s *Server) removeSequence(seqIndex int, reason string) {
 	close(seq.responses)
 	close(seq.embedding)
 	seq.pendingResponses = []string{}
+	seq.cache.inUse = false
 	seq.samplingCtx.Free()
-	s.lc.KvCacheSeqRm(seqIndex, 0, -1)
 	s.seqs[seqIndex] = nil
 }
 
@@ -266,8 +264,8 @@ func (s *Server) processBatch() {
 			continue
 		}
 
-		if seq.nPast+len(seq.tokens) > s.numCtx {
-			s.shiftContext(i)
+		if seq.numPast+len(seq.tokens) > s.cache.numCtx {
+			s.shiftContext(seq)
 		}
 
 		if seq.startProcessingTime.IsZero() {
@@ -280,10 +278,11 @@ func (s *Server) processBatch() {
 			if j >= s.batchSize {
 				break
 			}
-			batch.Add(t, seq.nPast, []int{i}, numTokensProcessed+1 == len(seq.tokens))
-			seq.nPast++
+			batch.Add(t, seq.numPast, []int{seq.cache.id}, numTokensProcessed+1 == len(seq.tokens))
+			seq.numPast++
 			numTokensProcessed++
 		}
+		seq.cache.tokens = append(seq.cache.tokens, seq.tokens[:numTokensProcessed]...)
 		seq.tokens = seq.tokens[numTokensProcessed:]
 		seq.iBatch = batch.NumTokens() - 1
 	}
@@ -486,6 +485,7 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	s.mu.Lock()
 	for i, sq := range s.seqs {
 		if sq == nil {
+			seq.cache, seq.tokens, seq.numPast = s.cache.LoadCacheSlot(seq.tokens)
 			s.seqs[i] = seq
 			s.cond.Signal()
 			break
@@ -561,6 +561,7 @@ func (s *Server) embeddings(w http.ResponseWriter, r *http.Request) {
 	s.mu.Lock()
 	for i, sq := range s.seqs {
 		if sq == nil {
+			seq.cache, seq.tokens, seq.numPast = s.cache.LoadCacheSlot(seq.tokens)
 			s.seqs[i] = seq
 			s.cond.Signal()
 			break
@@ -646,6 +647,8 @@ func (s *Server) loadModel(
 		panic("Vision models are not yet supported")
 	}
 
+	s.cache = NewTokenCache(s.lc, kvSize, s.parallel)
+
 	s.status = ServerStatusReady
 	s.ready.Done()
 }
@@ -691,7 +694,6 @@ func main() {
 	slog.SetDefault(slog.New(handler))
 
 	server := &Server{
-		numCtx:    *kvSize / *parallel,
 		batchSize: *batchSize,
 		parallel:  *parallel,
 		seqs:      make([]*Sequence, *parallel),