"docs/en/datasets/waymo_det.md" did not exist on "4fb544126e07e797e9120885abcba94a84523000"
cache.go 6.72 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
}

Jesse Gross's avatar
Jesse Gross committed
34
35
func NewInputCache(model model.Model, kvCacheType string, kvSize int32, numSlots int, multiUserCache bool) (*InputCache, error) {
	if kvSize/int32(numSlots) < 1 {
36
37
38
		return nil, fmt.Errorf("must have at least one kv cache entry per parallel sequence (kv: %v parallel: %v)", kvSize, numSlots)
	}

39
40
41
	slots := make([]InputCacheSlot, numSlots)

	for i := range slots {
42
		slots[i] = InputCacheSlot{Id: i}
43
44
	}

Jesse Gross's avatar
Jesse Gross committed
45
46
47
48
49
	cache := model.Config().Cache
	if cache != nil {
		cache.Init(model.Backend(), kvCacheTypeFromStr(kvCacheType), kvSize)
	}

50
	return &InputCache{
Jesse Gross's avatar
Jesse Gross committed
51
52
		numCtx:         kvSize / int32(numSlots),
		enabled:        cache != nil,
53
54
		slots:          slots,
		multiUserCache: multiUserCache,
Jesse Gross's avatar
Jesse Gross committed
55
		cache:          cache,
56
	}, nil
57
58
}

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

func (c *InputCache) Close() {
	c.cache.Close()
}

74
75
// Locking: Operations on InputCacheSlot (including finding one
// through LoadCacheSlot) require a lock to be be held that serializes
Jesse Gross's avatar
Jesse Gross committed
76
// these operations with each other and processBatch
77
78
79
80
81
82

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

	// Inputs that are stored in the KV cache
83
	Inputs []input.Input
84
85
86
87
88
89
90
91

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

92
func (c *InputCache) LoadCacheSlot(prompt []input.Input) (*InputCacheSlot, []input.Input, error) {
93
	var slot *InputCacheSlot
Jesse Gross's avatar
Jesse Gross committed
94
	var numPast int32
95
96
97
	var err error

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

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

Jesse Gross's avatar
Jesse Gross committed
113
	if numPast == int32(len(prompt)) {
114
115
116
117
		// Leave one input to sample so we can get a response
		numPast--
	}

Jesse Gross's avatar
Jesse Gross committed
118
119
120
121
122
123
124
125
126
127
	if c.cache != nil {
		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
		}
128
129
130
	}

	slog.Debug("loading cache slot", "id", slot.Id, "cache", len(slot.Inputs), "prompt", len(prompt),
Jesse Gross's avatar
Jesse Gross committed
131
		"used", numPast, "remaining", int32(len(prompt))-numPast)
132
133
134
135

	prompt = prompt[numPast:]
	slot.Inputs = slot.Inputs[:numPast]

136
	return slot, prompt, nil
137
138
}

139
func (c *InputCache) findLongestCacheSlot(prompt []input.Input) (*InputCacheSlot, int32, error) {
Jesse Gross's avatar
Jesse Gross committed
140
	longest := int32(-1)
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
	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
}

162
func (c *InputCache) findBestCacheSlot(prompt []input.Input) (*InputCacheSlot, int32, error) {
163
164
165
	oldest := time.Now()
	var oldestSlot *InputCacheSlot

Jesse Gross's avatar
Jesse Gross committed
166
	longest := int32(-1)
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
	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
182
	if longest == int32(len(longestSlot.Inputs)) && !longestSlot.InUse {
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
		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))
198
		oldestSlot.Inputs = make([]input.Input, longest)
199
		copy(oldestSlot.Inputs, longestSlot.Inputs[:longest])
Jesse Gross's avatar
Jesse Gross committed
200
201
		if c.cache != nil {
			c.cache.CopyPrefix(longestSlot.Id, oldestSlot.Id, longest)
202
203
204
205
206
207
		}
	}

	return oldestSlot, longest, nil
}

208
func countCommonPrefix(a []input.Input, b []input.Input) int32 {
Jesse Gross's avatar
Jesse Gross committed
209
	var count int32
210
211
212
213
214
215

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

216
		if a[i].Token != b[i].Token || a[i].MultimodalHash != b[i].MultimodalHash {
217
218
219
220
221
222
223
224
225
			break
		}

		count++
	}

	return count
}

Jesse Gross's avatar
Jesse Gross committed
226
func (c *InputCache) ShiftDiscard(inputLen int32, numKeep int32) int32 {
227
228
229
230
231
232
233
234
235
236
237
238
239
	targetFree := (c.numCtx - numKeep) / 2
	targetFree = max(targetFree, 1)

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

	if discard < 0 {
		discard = 0
	}

	return discard
}

240
241
242
243
// 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
244
func (c *InputCache) ShiftCacheSlot(slot *InputCacheSlot, numKeep int32) error {
245
246
247
248
	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
249
250
	inputLen := int32(len(slot.Inputs))
	discard := c.ShiftDiscard(inputLen, numKeep)
251
252

	if discard <= 0 {
253
		return nil
254
255
	}

256
	slog.Debug("context limit hit - shifting", "id", slot.Id, "limit", c.numCtx, "input", len(slot.Inputs),
257
258
		"keep", numKeep, "discard", discard)

259
	// TODO (jessegross): KV cache removal can fail for certain types of models
Jesse Gross's avatar
Jesse Gross committed
260
261
262
263
264
	if c.cache != nil {
		err := c.cache.Remove(slot.Id, numKeep, numKeep+discard)
		if err != nil {
			return fmt.Errorf("unable to remove old kv cache entries (id: %v, keep: %v discard: %v): %w", slot.Id, numKeep, discard, err)
		}
265
	}
266

Jesse Gross's avatar
Jesse Gross committed
267
	for i := numKeep + discard; i < inputLen; i++ {
268
		slot.Inputs[i-discard] = slot.Inputs[i]
269
	}
Jesse Gross's avatar
Jesse Gross committed
270
	slot.Inputs = slot.Inputs[:inputLen-discard]
271
272

	return nil
273
}