mirror of
https://github.com/ollama/ollama.git
synced 2026-04-19 08:54:21 +02:00
5daf59cc66
The previous approach tracked array lifecycles through reference counting, where each array recorded its inputs and a reference count that was decremented as dependents were freed. This is not really necessary as MLX tracks references internally. It is also error prone as it is easy to create new arrays and forget to free them when the Go variable goes out of scope. Instead, we can pin just the arrays we want (typically outputs and specific intermediates, like the cache). All other arrays are freed by default when we run sweep. This avoids most causes of memory leaks while still giving the freedom to save what we want.
75 lines
1.5 KiB
Go
75 lines
1.5 KiB
Go
//go:build mlx
|
|
|
|
package mlx
|
|
|
|
// #include "generated.h"
|
|
import "C"
|
|
|
|
import (
|
|
"unsafe"
|
|
)
|
|
|
|
func ScaledDotProductAttention(query, key, value, mask *Array, scale float32) *Array {
|
|
if mask == nil {
|
|
mask = New("")
|
|
}
|
|
|
|
sinks := New("")
|
|
|
|
mode := "causal"
|
|
cMode := C.CString(mode)
|
|
defer C.free(unsafe.Pointer(cMode))
|
|
|
|
out := New("FAST_SDPA")
|
|
C.mlx_fast_scaled_dot_product_attention(&out.ctx, query.ctx, key.ctx, value.ctx, C.float(scale), cMode, mask.ctx, sinks.ctx, DefaultStream().ctx)
|
|
return out
|
|
}
|
|
|
|
type LayerNorm struct {
|
|
Weight Array `weight:"weight"`
|
|
Bias Array `weight:"bias"`
|
|
}
|
|
|
|
func (r *LayerNorm) Forward(x *Array, eps float32) *Array {
|
|
out := New("FAST_LAYERNORM")
|
|
C.mlx_fast_layer_norm(&out.ctx, x.ctx, r.Weight.ctx, r.Bias.ctx, C.float(eps), DefaultStream().ctx)
|
|
return out
|
|
}
|
|
|
|
type RMSNorm struct {
|
|
Weight Array `weight:"weight"`
|
|
}
|
|
|
|
func (r RMSNorm) Forward(x *Array, eps float32) *Array {
|
|
out := New("FAST_RMSNORM")
|
|
C.mlx_fast_rms_norm(&out.ctx, x.ctx, r.Weight.ctx, C.float(eps), DefaultStream().ctx)
|
|
return out
|
|
}
|
|
|
|
type RoPE struct {
|
|
Dims int
|
|
Traditional bool
|
|
Base float32 `json:"rope_theta"`
|
|
Scale float32
|
|
}
|
|
|
|
func (r RoPE) Forward(t *Array, offset int) *Array {
|
|
freqs := New("")
|
|
out := New("FAST_ROPE")
|
|
C.mlx_fast_rope(
|
|
&out.ctx,
|
|
t.ctx,
|
|
C.int(r.Dims),
|
|
C._Bool(r.Traditional),
|
|
C.mlx_optional_float{
|
|
value: C.float(r.Base),
|
|
has_value: C._Bool(func() bool { return r.Base != 0 }()),
|
|
},
|
|
C.float(r.Scale),
|
|
C.int(offset),
|
|
freqs.ctx,
|
|
DefaultStream().ctx,
|
|
)
|
|
return out
|
|
}
|