backend.go 10.5 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
package ml

import (
	"bytes"
5
	"context"
Michael Yang's avatar
Michael Yang committed
6
7
	"encoding/binary"
	"fmt"
8
	"math"
Michael Yang's avatar
Michael Yang committed
9
	"slices"
Michael Yang's avatar
Michael Yang committed
10
11
12
	"strconv"
	"strings"

13
14
	"github.com/ollama/ollama/fs"
)
Michael Yang's avatar
Michael Yang committed
15
16

type Backend interface {
Jesse Gross's avatar
Jesse Gross committed
17
18
19
	// Close frees all memory associated with this backend
	Close()

20
	Load(ctx context.Context, progress func(float32)) error
21
22
23
24

	// BackendMemory returns the memory allocations that were made for this model
	BackendMemory() BackendMemory

25
	Config() fs.Config
Michael Yang's avatar
Michael Yang committed
26
27
	Get(name string) Tensor
	NewContext() Context
28
	NewContextSize(size int) Context
29
30
31

	// Enumerate the devices available for inference via this backend
	BackendDevices() []DeviceInfo
Michael Yang's avatar
Michael Yang committed
32
33
}

34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
// BackendCacheConfig should be implemented by backends that need special output
// from the cache to meet specific requirements. It is frequently implemented in
// conjunction with ScaledDotProductAttention.
type BackendCacheConfig interface {
	CacheConfig() CacheConfig
}

// CacheConfig controls optimizations (mostly backend-specific) that may transform
// the output the cache to work better with specific kernels.
type CacheConfig struct {
	// CachePadding specifies the multiple for the number of tokens of cache history
	// that will be returned from cache Get for k, v and mask. The capacity of the
	// cache itself will also be increased to a multiple of this size if needed.
	CachePadding int

	// PermutedV performs Permute(ctx, 1, 2, 0, 3) on v tensors stored via Put
	// and return the permuted version via Get. This uses the cache copy operation
	// to avoid a Contiguous call on the permuted tensor.
	PermutedV bool
53
54
55
56
57
58
59
60

	// MaskDType specifies the data type for generating the mask. If unset it will
	// default to DTypeF32.
	MaskDType DType

	// MaskBatchPadding specifies the multiple for the batch size dimension in the mask.
	// Any position that does not correspond to an actual token will be filled with -Inf.
	MaskBatchPadding int
61
62
}

63
64
// BackendParams controls how the backend loads and executes models
type BackendParams struct {
Jesse Gross's avatar
Jesse Gross committed
65
66
67
68
69
	// AllocMemory causes the backend to allocate memory for the model. If
	// false, this is only being used for discovering the required amount of
	// memory and cannot load the model for running.
	AllocMemory bool

70
71
	// NumThreads sets the number of threads to use if running on the CPU
	NumThreads int
Michael Yang's avatar
Michael Yang committed
72

Jesse Gross's avatar
Jesse Gross committed
73
74
	// GPULayers is the set of layers to offload to GPUs
	GPULayers GPULayersList
75
76
77

	// FlashAttention indicates that we should use a fused flash attention kernel
	FlashAttention bool
78
79
}

80
var backends = make(map[string]func(string, BackendParams) (Backend, error))
81

82
func RegisterBackend(name string, f func(string, BackendParams) (Backend, error)) {
Michael Yang's avatar
Michael Yang committed
83
84
85
86
87
88
89
	if _, ok := backends[name]; ok {
		panic("backend: backend already registered")
	}

	backends[name] = f
}

90
func NewBackend(modelPath string, params BackendParams) (Backend, error) {
Michael Yang's avatar
Michael Yang committed
91
	if backend, ok := backends["ggml"]; ok {
92
		return backend(modelPath, params)
Michael Yang's avatar
Michael Yang committed
93
94
95
96
97
98
	}

	return nil, fmt.Errorf("unsupported backend")
}

type Context interface {
99
	Empty(dtype DType, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
100
	Zeros(dtype DType, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
101
102
103
	FromBytes(dtype DType, s []byte, shape ...int) Tensor
	FromFloats(s []float32, shape ...int) Tensor
	FromInts(s []int32, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
104

Michael Yang's avatar
arange  
Michael Yang committed
105
106
107
	// Arange creates a 1D tensor with values within an interval (start, stop] increased by step.
	Arange(start, stop, step float32, dtype DType) Tensor

108
	Forward(...Tensor) Context
109
110
111
112
113

	// SetBatchSize provides a hint on the batch size to optimize processing
	// Uses heuristics if not set
	SetBatchSize(int)

114
	Compute(...Tensor)
115
	ComputeWithNotify(func(), ...Tensor) // notify callback once compute has begun
116
117
118
119
120

	// Reserve is analogous to Compute but rather than executing a
	// graph, simply preallocates memory. Typically called with a
	// worst case graph to ensure all resources are available for
	// for future inference.
121
	Reserve()
122

123
	MaxGraphNodes() int
124
	Close()
125

126
127
	// Input returns a context appropriate for creating tensors that are
	// inputs to the model (which includes things like output locations)
128
129
130
131
	Input() Context

	// Layer returns a context appropriate for creating intermediate tensors
	Layer(int) Context
Michael Yang's avatar
Michael Yang committed
132
133
134
}

type Tensor interface {
135
136
	Dim(n int) int
	Stride(n int) int
Michael Yang's avatar
Michael Yang committed
137

138
	Shape() []int
Michael Yang's avatar
Michael Yang committed
139
	DType() DType
140
	Cast(ctx Context, dtype DType) Tensor
Michael Yang's avatar
Michael Yang committed
141
142
143
144

	Bytes() []byte
	Floats() []float32

Michael Yang's avatar
Michael Yang committed
145
146
147
	FromBytes([]byte)
	FromFloats([]float32)
	FromInts([]int32)
148

Michael Yang's avatar
Michael Yang committed
149
	Add(ctx Context, t2 Tensor) Tensor
Michael Yang's avatar
Michael Yang committed
150
	Sub(ctx Context, t2 Tensor) Tensor
Michael Yang's avatar
Michael Yang committed
151
	Mul(ctx Context, t2 Tensor) Tensor
152
153
	Div(ctx Context, t2 Tensor) Tensor

Michael Yang's avatar
Michael Yang committed
154
	Mulmat(ctx Context, t2 Tensor) Tensor
155
	MulmatFullPrec(ctx Context, t2 Tensor) Tensor
Michael Yang's avatar
llama4  
Michael Yang committed
156
	MulmatID(ctx Context, t2, ids Tensor) Tensor
157
	AddID(ctx Context, t2, ids Tensor) Tensor
Michael Yang's avatar
Michael Yang committed
158
159

	Softmax(ctx Context) Tensor
160
	L2Norm(ctx Context, eps float32) Tensor
Michael Yang's avatar
Michael Yang committed
161
162
163
	LayerNorm(ctx Context, weight, bias Tensor, eps float32) Tensor
	RMSNorm(ctx Context, weight Tensor, eps float32) Tensor
	Scale(ctx Context, s float64) Tensor
164
	SumRows(ctx Context) Tensor
Michael Yang's avatar
Michael Yang committed
165

Michael Yang's avatar
Michael Yang committed
166
	AvgPool2D(ctx Context, k, s int, p float32) Tensor
Michael Yang's avatar
Michael Yang committed
167
	Conv2D(ctx Context, weight Tensor, s0, s1, p0, p1, d0, d1 int) Tensor
168
	Conv3D(ctx Context, weight Tensor, c, s0, s1, s2, p0, p1, p2, d0, d1, d2 int) Tensor
Michael Yang's avatar
Michael Yang committed
169

170
	IM2Col(ctx Context, weight Tensor, s0, s1, p0, p1, d0, d1 int) Tensor
Michael Yang's avatar
Michael Yang committed
171

172
173
	Sin(ctx Context) Tensor
	Cos(ctx Context) Tensor
Michael Yang's avatar
Michael Yang committed
174
	Tanh(ctx Context) Tensor
175
	GELU(ctx Context, up ...Tensor) Tensor
Michael Yang's avatar
Michael Yang committed
176
	QuickGELU(ctx Context, up ...Tensor) Tensor
177
178
	SILU(ctx Context, up ...Tensor) Tensor
	RELU(ctx Context, up ...Tensor) Tensor
Michael Yang's avatar
llama4  
Michael Yang committed
179
	Sigmoid(ctx Context) Tensor
180
181
182

	// AlphaLimitSILU is a variant of SILU that clamps the input to the range [-limit, limit]
	SILUAlphaLimit(ctx Context, up Tensor, alpha, limit float32) Tensor
Michael Yang's avatar
Michael Yang committed
183

184
	Reshape(ctx Context, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
185
186
	View(ctx Context, offset int, shape ...int) Tensor
	Permute(ctx Context, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
187
	Contiguous(ctx Context, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
188

189
	Pad(ctx Context, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
190
191

	Stack(ctx Context, dim int, s ...Tensor) Tensor
192
193
194

	// Repeat repeats the tensor n times along dimension dim
	Repeat(ctx Context, dim, n int) Tensor
Michael Yang's avatar
Michael Yang committed
195
196
197
	Concat(ctx Context, t2 Tensor, dim int) Tensor
	Rows(ctx Context, t2 Tensor) Tensor
	Copy(ctx Context, t2 Tensor) Tensor
198
	Duplicate(ctx Context) Tensor
Michael Yang's avatar
llama4  
Michael Yang committed
199

200
201
202
203
	Slice(ctx Context, dim, low, high, step int) Tensor
	Chunk(ctx Context, dim int, size int) []Tensor
	ChunkSections(ctx Context, dim int, sections ...int) []Tensor

Michael Yang's avatar
llama4  
Michael Yang committed
204
	TopK(ctx Context, k int) Tensor
205
	Argsort(ctx Context) Tensor
Michael Yang's avatar
Michael Yang committed
206
207
208
209
210
	Mean(ctx Context) Tensor
	Variance(ctx Context) Tensor
	Stddev(ctx Context) Tensor
	Sqr(ctx Context) Tensor
	Sqrt(ctx Context) Tensor
Michael Yang's avatar
Michael Yang committed
211
212

	Interpolate(ctx Context, dims [4]int, samplingMode SamplingMode) Tensor
Michael Yang's avatar
Michael Yang committed
213
214
}

215
216
217
218
// ScaledDotProductAttention implements a fused attention
// operation equivalent to following code on a tensor named
// query:
//
219
220
221
222
// query = query.Permute(ctx, 0, 2, 1, 3)
// key = key.Permute(ctx, 0, 2, 1, 3)
// value = value.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx)
//
223
224
225
226
227
228
229
230
231
232
233
234
235
// kq := key.MulmatFullPrec(ctx, query)
//
// kq = kq.Scale(ctx, scale)
//
//	if mask != nil {
//		kq = kq.Add(ctx, mask)
//	}
//
// kq = kq.Softmax(ctx)
//
// kqv := value.Mulmat(ctx, kq)
// return kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
type ScaledDotProductAttention interface {
Grace's avatar
Grace committed
236
	ScaledDotProductAttention(ctx Context, key, value, mask, sinks Tensor, vmla Tensor, scale float64) Tensor
237
238
}

Michael Yang's avatar
Michael Yang committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
type number interface {
	~int | ~int8 | ~int16 | ~int32 | ~int64 |
		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 |
		~float32 | ~float64 |
		~complex64 | ~complex128
}

func mul[T number](s ...T) T {
	p := T(1)
	for _, v := range s {
		p *= v
	}

	return p
}

255
type DumpOptions func(*dumpOptions)
Michael Yang's avatar
Michael Yang committed
256

257
258
259
260
261
// DumpWithPrecision sets the number of decimal places to print. Applies to float32 and float64.
func DumpWithPrecision(n int) DumpOptions {
	return func(opts *dumpOptions) {
		opts.Precision = n
	}
Michael Yang's avatar
Michael Yang committed
262
263
}

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
// DumpWithThreshold sets the threshold for printing the entire tensor. If the number of elements
// is less than or equal to this value, the entire tensor will be printed. Otherwise, only the
// beginning and end of each dimension will be printed.
func DumpWithThreshold(n int) DumpOptions {
	return func(opts *dumpOptions) {
		opts.Threshold = n
	}
}

// DumpWithEdgeItems sets the number of elements to print at the beginning and end of each dimension.
func DumpWithEdgeItems(n int) DumpOptions {
	return func(opts *dumpOptions) {
		opts.EdgeItems = n
	}
}

type dumpOptions struct {
	Precision, Threshold, EdgeItems int
}

func Dump(ctx Context, t Tensor, optsFuncs ...DumpOptions) string {
	opts := dumpOptions{Precision: 4, Threshold: 1000, EdgeItems: 3}
	for _, optsFunc := range optsFuncs {
		optsFunc(&opts)
	}

	if mul(t.Shape()...) <= opts.Threshold {
		opts.EdgeItems = math.MaxInt
Michael Yang's avatar
Michael Yang committed
292
293
294
295
	}

	switch t.DType() {
	case DTypeF32:
296
297
		return dump[[]float32](ctx, t, opts.EdgeItems, func(f float32) string {
			return strconv.FormatFloat(float64(f), 'f', opts.Precision, 32)
Jesse Gross's avatar
Jesse Gross committed
298
		})
299
	case DTypeF16, DTypeQ80, DTypeQ40:
300
		f32 := ctx.Input().Empty(DTypeF32, t.Shape()...)
Jesse Gross's avatar
Jesse Gross committed
301
		f32 = t.Copy(ctx, f32)
302
303
		return dump[[]float32](ctx, f32, opts.EdgeItems, func(f float32) string {
			return strconv.FormatFloat(float64(f), 'f', opts.Precision, 32)
Michael Yang's avatar
Michael Yang committed
304
305
		})
	case DTypeI32:
306
		return dump[[]int32](ctx, t, opts.EdgeItems, func(i int32) string {
Michael Yang's avatar
Michael Yang committed
307
308
309
310
311
312
313
			return strconv.FormatInt(int64(i), 10)
		})
	default:
		return "<unsupported>"
	}
}

Jesse Gross's avatar
Jesse Gross committed
314
315
func dump[S ~[]E, E number](ctx Context, t Tensor, items int, fn func(E) string) string {
	if t.Bytes() == nil {
316
		ctx.Forward(t).Compute(t)
Michael Yang's avatar
Michael Yang committed
317
318
319
320
321
322
323
324
	}

	s := make(S, mul(t.Shape()...))
	if err := binary.Read(bytes.NewBuffer(t.Bytes()), binary.LittleEndian, &s); err != nil {
		panic(err)
	}

	shape := t.Shape()
Michael Yang's avatar
Michael Yang committed
325
	slices.Reverse(shape)
Michael Yang's avatar
Michael Yang committed
326
327

	var sb strings.Builder
328
329
	var f func([]int, int)
	f = func(dims []int, stride int) {
Michael Yang's avatar
Michael Yang committed
330
		prefix := strings.Repeat(" ", len(shape)-len(dims)+1)
Michael Yang's avatar
Michael Yang committed
331
332
		sb.WriteString("[")
		defer func() { sb.WriteString("]") }()
333
		for i := 0; i < dims[0]; i++ {
Michael Yang's avatar
Michael Yang committed
334
			if i >= items && i < dims[0]-items {
Michael Yang's avatar
Michael Yang committed
335
				sb.WriteString("..., ")
Michael Yang's avatar
Michael Yang committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349
				// skip to next printable element
				skip := dims[0] - 2*items
				if len(dims) > 1 {
					stride += mul(append(dims[1:], skip)...)
					fmt.Fprint(&sb, strings.Repeat("\n", len(dims)-1), prefix)
				}
				i += skip - 1
			} else if len(dims) > 1 {
				f(dims[1:], stride)
				stride += mul(dims[1:]...)
				if i < dims[0]-1 {
					fmt.Fprint(&sb, ",", strings.Repeat("\n", len(dims)-1), prefix)
				}
			} else {
Michael Yang's avatar
Michael Yang committed
350
351
352
353
354
355
				text := fn(s[stride+i])
				if len(text) > 0 && text[0] != '-' {
					sb.WriteString(" ")
				}

				sb.WriteString(text)
Michael Yang's avatar
Michael Yang committed
356
				if i < dims[0]-1 {
Michael Yang's avatar
Michael Yang committed
357
					sb.WriteString(", ")
Michael Yang's avatar
Michael Yang committed
358
359
360
361
362
363
364
365
366
367
368
369
				}
			}
		}
	}
	f(shape, 0)

	return sb.String()
}

type DType int

const (
Jesse Gross's avatar
Jesse Gross committed
370
371
372
	DTypeOther DType = iota
	DTypeF32
	DTypeF16
373
374
	DTypeQ80
	DTypeQ40
Michael Yang's avatar
Michael Yang committed
375
	DTypeI32
Michael Yang's avatar
Michael Yang committed
376
	DTypeMXFP4
Michael Yang's avatar
Michael Yang committed
377
)
Michael Yang's avatar
Michael Yang committed
378
379
380
381
382
383
384

type SamplingMode int

const (
	SamplingModeNearest SamplingMode = iota
	SamplingModeBilinear
)