cache.go 7.6 KB
Newer Older
Jesse Gross's avatar
Jesse Gross committed
1
package ollamarunner
2
3
4

import (
	"errors"
5
	"fmt"
6
	"log/slog"
Jesse Gross's avatar
Jesse Gross committed
7
	"math"
8
9
	"time"

Jesse Gross's avatar
Jesse Gross committed
10
11
12
	"github.com/ollama/ollama/kvcache"
	"github.com/ollama/ollama/ml"
	"github.com/ollama/ollama/model"
13
	"github.com/ollama/ollama/model/input"
14
15
16
17
)

type InputCache struct {
	// context window size (per slot)
Jesse Gross's avatar
Jesse Gross committed
18
19
20
21
22
23
	numCtx int32

	// does the cache store data or do we need to always send the full input?
	// note that when enabled is false the underlying cache may either be nil
	// or a non-nil dummy that doesn't actually store anything
	enabled bool
24
25
26
27
28
29
30

	// individual KV caches
	slots []InputCacheSlot

	// optimize cache eviction for multiple users
	multiUserCache bool

Jesse Gross's avatar
Jesse Gross committed
31
	cache kvcache.Cache
32
33
}

34
35
36
37
func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots int, batchSize int, multiUserCache bool) (*InputCache, error) {
	numCtx := kvSize / int32(numSlots)

	if numCtx < 1 {
38
39
40
		return nil, fmt.Errorf("must have at least one kv cache entry per parallel sequence (kv: %v parallel: %v)", kvSize, numSlots)
	}

41
42
43
	slots := make([]InputCacheSlot, numSlots)

	for i := range slots {
44
		slots[i] = InputCacheSlot{Id: i}
45
46
	}

Jesse Gross's avatar
Jesse Gross committed
47
48
	cache := model.Config().Cache
	if cache != nil {
49
		cache.Init(model.Backend(), kvCacheTypeFromStr(kvCacheType), numSlots, int(numCtx), batchSize)
Jesse Gross's avatar
Jesse Gross committed
50
51
	}

52
	return &InputCache{
53
		numCtx:         numCtx,
Jesse Gross's avatar
Jesse Gross committed
54
		enabled:        cache != nil,
55
56
		slots:          slots,
		multiUserCache: multiUserCache,
Jesse Gross's avatar
Jesse Gross committed
57
		cache:          cache,
58
	}, nil
59
60
}

Jesse Gross's avatar
Jesse Gross committed
61
62
63
func kvCacheTypeFromStr(s string) ml.DType {
	switch s {
	case "q8_0":
64
		return ml.DTypeQ80
Jesse Gross's avatar
Jesse Gross committed
65
	case "q4_0":
66
		return ml.DTypeQ40
Jesse Gross's avatar
Jesse Gross committed
67
68
69
70
71
72
	default:
		return ml.DTypeF16
	}
}

func (c *InputCache) Close() {
Jesse Gross's avatar
Jesse Gross committed
73
74
75
76
	if c == nil {
		return
	}

Jesse Gross's avatar
Jesse Gross committed
77
78
79
	c.cache.Close()
}

80
// Locking: Operations on InputCacheSlot (including finding one
81
// through LoadCacheSlot) require a lock to be held that serializes
Jesse Gross's avatar
Jesse Gross committed
82
// these operations with each other and processBatch
83
84
85
86
87
88

type InputCacheSlot struct {
	// Index in the KV cache
	Id int

	// Inputs that are stored in the KV cache
89
	Inputs []*input.Input
90
91
92
93
94
95
96
97

	// 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
}

98
func (c *InputCache) LoadCacheSlot(prompt []*input.Input) (*InputCacheSlot, []*input.Input, error) {
99
	var slot *InputCacheSlot
Jesse Gross's avatar
Jesse Gross committed
100
	var numPast int32
101
102
103
	var err error

	// In single-user scenarios, the longest cache slot works fine for getting good input
Jesse Gross's avatar
Jesse Gross committed
104
	// cache hit rates and it keeps the footprint of the cache small, which improves throughput.
105
	// For multiple users, the "best" cache slot produces better input cache hit rates
Jesse Gross's avatar
Jesse Gross committed
106
	// at the cost of worse performance when we miss the input cache.
107
108
109
110
111
112
	if !c.multiUserCache {
		slot, numPast, err = c.findLongestCacheSlot(prompt)
	} else {
		slot, numPast, err = c.findBestCacheSlot(prompt)
	}
	if err != nil {
113
		return nil, nil, err
114
115
116
117
118
	}

	slot.InUse = true
	slot.lastUsed = time.Now()

Jesse Gross's avatar
Jesse Gross committed
119
	if numPast == int32(len(prompt)) {
120
121
122
123
		// Leave one input to sample so we can get a response
		numPast--
	}

Jesse Gross's avatar
Jesse Gross committed
124
	if c.cache != nil {
125
126
127
128
		if numPast > 0 && !c.cache.CanResume(slot.Id, numPast) {
			numPast = 0
		}

Jesse Gross's avatar
Jesse Gross committed
129
130
131
132
133
134
135
136
137
		err = c.cache.Remove(slot.Id, numPast, math.MaxInt32)
		if err != nil {
			// Some models don't support partial erasure
			err = c.cache.Remove(slot.Id, 0, math.MaxInt32)
			if err != nil {
				return nil, nil, err
			}
			numPast = 0
		}
138
139
140
	}

	slog.Debug("loading cache slot", "id", slot.Id, "cache", len(slot.Inputs), "prompt", len(prompt),
Jesse Gross's avatar
Jesse Gross committed
141
		"used", numPast, "remaining", int32(len(prompt))-numPast)
142

143
	slot.Inputs = prompt[:numPast]
144
145
	prompt = prompt[numPast:]

146
	return slot, prompt, nil
147
148
}

149
func (c *InputCache) findLongestCacheSlot(prompt []*input.Input) (*InputCacheSlot, int32, error) {
Jesse Gross's avatar
Jesse Gross committed
150
	longest := int32(-1)
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
	var longestSlot *InputCacheSlot

	for i, s := range c.slots {
		if s.InUse {
			continue
		}

		count := countCommonPrefix(s.Inputs, prompt)
		if count > longest {
			longest = count
			longestSlot = &c.slots[i]
		}
	}

	if longestSlot == nil {
		return nil, 0, errors.New("no available cache slots")
	}

	return longestSlot, longest, nil
}

172
func (c *InputCache) findBestCacheSlot(prompt []*input.Input) (*InputCacheSlot, int32, error) {
173
174
175
	oldest := time.Now()
	var oldestSlot *InputCacheSlot

Jesse Gross's avatar
Jesse Gross committed
176
	longest := int32(-1)
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
	var longestSlot *InputCacheSlot

	for i, s := range c.slots {
		count := countCommonPrefix(s.Inputs, prompt)
		if count > longest {
			longest = count
			longestSlot = &c.slots[i]
		}

		if s.lastUsed.Compare(oldest) < 0 && !s.InUse {
			oldest = s.lastUsed
			oldestSlot = &c.slots[i]
		}
	}

Jesse Gross's avatar
Jesse Gross committed
192
	if longest == int32(len(longestSlot.Inputs)) && !longestSlot.InUse {
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
		return longestSlot, longest, nil
	}

	if oldestSlot.InUse {
		return nil, 0, errors.New("no available cache slots")
	}

	if len(oldestSlot.Inputs) != 0 {
		slog.Debug("evicting cache slot", "id", oldestSlot.Id, "inputs", len(oldestSlot.Inputs),
			"used", oldestSlot.lastUsed)
	}

	if longest > 0 && longestSlot != oldestSlot {
		slog.Debug("forking cache slot", "src", longestSlot.Id, "dst", oldestSlot.Id, "inputs", longest, "total",
			len(longestSlot.Inputs))
208
		oldestSlot.Inputs = make([]*input.Input, longest)
209
		copy(oldestSlot.Inputs, longestSlot.Inputs[:longest])
Jesse Gross's avatar
Jesse Gross committed
210
211
		if c.cache != nil {
			c.cache.CopyPrefix(longestSlot.Id, oldestSlot.Id, longest)
212
213
214
215
216
217
		}
	}

	return oldestSlot, longest, nil
}

218
func countCommonPrefix(a []*input.Input, b []*input.Input) int32 {
Jesse Gross's avatar
Jesse Gross committed
219
	var count int32
220
221
222
223
224
225

	for i := range a {
		if i >= len(b) {
			break
		}

226
		if a[i].Token != b[i].Token || a[i].MultimodalHash != b[i].MultimodalHash {
227
228
229
230
231
232
233
234
235
			break
		}

		count++
	}

	return count
}

236
237
// TODO(jessegross): If we need to reprocess the inputs we should ensure that
// we don't split up a SameBatch
Jesse Gross's avatar
Jesse Gross committed
238
func (c *InputCache) ShiftDiscard(inputLen int32, numKeep int32) int32 {
239
240
241
242
243
244
245
246
247
248
249
250
251
	targetFree := (c.numCtx - numKeep) / 2
	targetFree = max(targetFree, 1)

	currentFree := c.numCtx - inputLen
	discard := targetFree - currentFree

	if discard < 0 {
		discard = 0
	}

	return discard
}

252
type ErrReprocessInputs struct {
253
	Inputs []*input.Input
254
255
256
257
258
259
}

func (e *ErrReprocessInputs) Error() string {
	return fmt.Sprintf("kv cache shift not supported, inputs need reprocessing (input count: %v)", len(e.Inputs))
}

260
261
262
263
// Frees up space in the KV cache by deleting the oldest half of history and shifting
// the newest half into that space (saving numKeep inputs at the beginning).
//
// Assumes that at least 1 entry can be freed up by shifting (i.e. numKeep < numCtx)
Jesse Gross's avatar
Jesse Gross committed
264
func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int32) error {
265
266
267
268
	if numKeep >= c.numCtx {
		return fmt.Errorf("unable to shift context - keep exceeds context (keep: %v context: %v)", numKeep, c.numCtx)
	}

Jesse Gross's avatar
Jesse Gross committed
269
270
	inputLen := int32(len(slot.Inputs))
	discard := c.ShiftDiscard(inputLen, numKeep)
271
272

	if discard <= 0 {
273
		return nil
274
275
	}

276
	slog.Debug("context limit hit - shifting", "id", slot.Id, "limit", c.numCtx, "input", len(slot.Inputs),
277
278
		"keep", numKeep, "discard", discard)

Jesse Gross's avatar
Jesse Gross committed
279
280
281
	if c.cache != nil {
		err := c.cache.Remove(slot.Id, numKeep, numKeep+discard)
		if err != nil {
282
283
284
285
			slog.Debug("kv cache removal unsupported, clearing cache and returning inputs for reprocessing",
				"id", slot.Id, "error", err)

			// Create new input slice with preserved tokens (numKeep + remaining tokens after discard)
286
			newInputs := make([]*input.Input, numKeep+inputLen-(numKeep+discard))
287
288
289
290
			copy(newInputs[:numKeep], slot.Inputs[:numKeep])
			copy(newInputs[numKeep:], slot.Inputs[numKeep+discard:])

			// Reset the cache
291
			_ = c.cache.Remove(slot.Id, 0, math.MaxInt32)
292
			slot.Inputs = []*input.Input{}
293
294
295

			// Return error with inputs that need to be reprocessed
			return &ErrReprocessInputs{Inputs: newInputs}
Jesse Gross's avatar
Jesse Gross committed
296
		}
297
	}
298

Jesse Gross's avatar
Jesse Gross committed
299
	for i := numKeep + discard; i < inputLen; i++ {
300
		slot.Inputs[i-discard] = slot.Inputs[i]
301
	}
Jesse Gross's avatar
Jesse Gross committed
302
	slot.Inputs = slot.Inputs[:inputLen-discard]
303
304

	return nil
305
}