backend.go 7.56 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
7
package ml

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
Michael Yang's avatar
Michael Yang committed
8
	"slices"
Michael Yang's avatar
Michael Yang committed
9
10
11
12
13
14
15
16
17
	"strconv"
	"strings"
)

type Config interface {
	Architecture() string
	String(string, ...string) string
	Uint(string, ...uint32) uint32
	Float(string, ...float32) float32
18
	Bool(string, ...bool) bool
Michael Yang's avatar
Michael Yang committed
19
20
21
22
23
24
25
26
27

	Strings(string, ...[]string) []string
	Uints(string, ...[]uint32) []uint32
}

type Backend interface {
	Config() Config
	Get(name string) Tensor
	NewContext() Context
28
	NewContextSize(size int) Context
Michael Yang's avatar
Michael Yang committed
29
30
}

31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
// 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
50
51
52
53
54
55
56
57

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

60
61
62
63
// BackendParams controls how the backend loads and executes models
type BackendParams struct {
	// NumThreads sets the number of threads to use if running on the CPU
	NumThreads int
Michael Yang's avatar
Michael Yang committed
64

65
66
67
68
69
70
71
72
	// MainGPU is the index of the primary GPU to use
	MainGPU int

	// NumGPULayers is the number of layers to offload to GPUs
	NumGPULayers int

	// TensorSplit is the fraction of the model to offload to each GPU
	TensorSplit []float32
73
74
75

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

var backends = make(map[string]func(*os.File, BackendParams) (Backend, error))

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

	backends[name] = f
}

88
func NewBackend(f *os.File, params BackendParams) (Backend, error) {
Michael Yang's avatar
Michael Yang committed
89
	if backend, ok := backends["ggml"]; ok {
90
		return backend(f, params)
Michael Yang's avatar
Michael Yang committed
91
92
93
94
95
96
	}

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

type Context interface {
97
	Empty(dtype DType, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
98
99
100
101
	Zeros(dtype DType, shape ...int) Tensor
	FromFloatSlice(s []float32, shape ...int) (Tensor, error)
	FromIntSlice(s []int32, shape ...int) (Tensor, error)

102
	Forward(...Tensor) Context
103
	Compute(...Tensor)
104
	MaxGraphNodes() int
105
	Close()
106
107
108
109
110
111
112
113
114

	// Input returns a context appropriate for creating input tensors
	Input() Context

	// Output returns a context appropriate for creating output tensors
	Output() Context

	// Layer returns a context appropriate for creating intermediate tensors
	Layer(int) Context
Michael Yang's avatar
Michael Yang committed
115
116
117
}

type Tensor interface {
118
119
	Dim(n int) int
	Stride(n int) int
Michael Yang's avatar
Michael Yang committed
120

121
	Shape() []int
Michael Yang's avatar
Michael Yang committed
122
123
124
125
126
127
128
129
	DType() DType

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

	Add(ctx Context, t2 Tensor) Tensor
	Mul(ctx Context, t2 Tensor) Tensor
	Mulmat(ctx Context, t2 Tensor) Tensor
130
	MulmatFullPrec(ctx Context, t2 Tensor) Tensor
Michael Yang's avatar
Michael Yang committed
131
132
133
134
135
136
137
138
139
140
141
142
143

	Softmax(ctx Context) Tensor
	LayerNorm(ctx Context, weight, bias Tensor, eps float32) Tensor
	RMSNorm(ctx Context, weight Tensor, eps float32) Tensor
	Scale(ctx Context, s float64) Tensor

	Conv2D(ctx Context, weight Tensor, s0, s1, p0, p1, d0, d1 int) Tensor
	RoPE(ctx Context, positionIDs, ropeFactors Tensor, dim uint32, base, scale float32) Tensor

	Tanh(ctx Context) Tensor
	GELU(ctx Context) Tensor
	SILU(ctx Context) Tensor

144
	Reshape(ctx Context, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
145
146
147
148
	View(ctx Context, offset int, shape ...int) Tensor
	Permute(ctx Context, shape ...int) Tensor
	Contiguous(ctx Context) Tensor

149
150
	Pad(ctx Context, shape ...int) Tensor
	Unpad(ctx Context, shape ...int) Tensor
Michael Yang's avatar
Michael Yang committed
151
152
153
154
155
156
157

	Stack(ctx Context, dim int, s ...Tensor) Tensor
	Concat(ctx Context, t2 Tensor, dim int) Tensor
	Rows(ctx Context, t2 Tensor) Tensor
	Copy(ctx Context, t2 Tensor) Tensor
}

158
159
160
161
// ScaledDotProductAttention implements a fused attention
// operation equivalent to following code on a tensor named
// query:
//
162
163
164
165
// 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)
//
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
// 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 {
	ScaledDotProductAttention(ctx Context, key, value, mask Tensor, scale float64) Tensor
}

Michael Yang's avatar
Michael Yang committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
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
}

type DumpOptions struct {
	// Items is the number of elements to print at the beginning and end of each dimension.
200
	Items int
Michael Yang's avatar
Michael Yang committed
201
202
203
204
205

	// Precision is the number of decimal places to print. Applies to float32 and float64.
	Precision int
}

Jesse Gross's avatar
Jesse Gross committed
206
func Dump(ctx Context, t Tensor, opts ...DumpOptions) string {
Michael Yang's avatar
Michael Yang committed
207
208
209
210
211
212
213
214
215
	if len(opts) < 1 {
		opts = append(opts, DumpOptions{
			Items:     3,
			Precision: 4,
		})
	}

	switch t.DType() {
	case DTypeF32:
Jesse Gross's avatar
Jesse Gross committed
216
217
218
		return dump[[]float32](ctx, t, opts[0].Items, func(f float32) string {
			return strconv.FormatFloat(float64(f), 'f', opts[0].Precision, 32)
		})
219
	case DTypeF16, DTypeQ80, DTypeQ40:
220
		f32 := ctx.Empty(DTypeF32, t.Shape()...)
Jesse Gross's avatar
Jesse Gross committed
221
222
		f32 = t.Copy(ctx, f32)
		return dump[[]float32](ctx, f32, opts[0].Items, func(f float32) string {
Michael Yang's avatar
Michael Yang committed
223
224
225
			return strconv.FormatFloat(float64(f), 'f', opts[0].Precision, 32)
		})
	case DTypeI32:
Jesse Gross's avatar
Jesse Gross committed
226
		return dump[[]int32](ctx, t, opts[0].Items, func(i int32) string {
Michael Yang's avatar
Michael Yang committed
227
228
229
230
231
232
233
			return strconv.FormatInt(int64(i), 10)
		})
	default:
		return "<unsupported>"
	}
}

Jesse Gross's avatar
Jesse Gross committed
234
235
func dump[S ~[]E, E number](ctx Context, t Tensor, items int, fn func(E) string) string {
	if t.Bytes() == nil {
236
		ctx.Forward(t).Compute(t)
Michael Yang's avatar
Michael Yang committed
237
238
239
240
241
242
243
244
	}

	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
245
	slices.Reverse(shape)
Michael Yang's avatar
Michael Yang committed
246
247

	var sb strings.Builder
248
249
	var f func([]int, int)
	f = func(dims []int, stride int) {
Michael Yang's avatar
Michael Yang committed
250
		prefix := strings.Repeat(" ", len(shape)-len(dims)+1)
Michael Yang's avatar
Michael Yang committed
251
252
		sb.WriteString("[")
		defer func() { sb.WriteString("]") }()
253
		for i := 0; i < dims[0]; i++ {
Michael Yang's avatar
Michael Yang committed
254
			if i >= items && i < dims[0]-items {
Michael Yang's avatar
Michael Yang committed
255
				sb.WriteString("..., ")
Michael Yang's avatar
Michael Yang committed
256
257
258
259
260
261
262
263
264
265
266
267
268
269
				// 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
270
271
272
273
274
275
				text := fn(s[stride+i])
				if len(text) > 0 && text[0] != '-' {
					sb.WriteString(" ")
				}

				sb.WriteString(text)
Michael Yang's avatar
Michael Yang committed
276
				if i < dims[0]-1 {
Michael Yang's avatar
Michael Yang committed
277
					sb.WriteString(", ")
Michael Yang's avatar
Michael Yang committed
278
279
280
281
282
283
284
285
286
287
288
289
				}
			}
		}
	}
	f(shape, 0)

	return sb.String()
}

type DType int

const (
Jesse Gross's avatar
Jesse Gross committed
290
291
292
	DTypeOther DType = iota
	DTypeF32
	DTypeF16
293
294
	DTypeQ80
	DTypeQ40
Michael Yang's avatar
Michael Yang committed
295
296
	DTypeI32
)