mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 13:38:14 -05:00
Daniel Hiltgen
43f4eda808
* glimmer: implement the Muse Glimmer model MLX model (language + vision encoder) with DFlash draft wiring, llama-server DFlash support and rope-interleave fix, renderer and parser, tokenizer fixes, and the import quantization policy. * mlxrunner: report committed prefill chunks after the sweep and eval The drafter's flush evaluates its report, and an eval that runs while the chunk's construction handles are still live cannot free any intermediate buffer. On media chunks that retention keeps the whole vision tower resident and grinds the Metal allocator at its limit until the request dies. Pin the report's inputs across the sweep, report after the chunk materializes, and release media items after the report so a drafter can still capture the rows its deferred flush embeds. * ci: retry CUDA pre-release download
100 lines
2.0 KiB
Go
100 lines
2.0 KiB
Go
package mlx
|
|
|
|
import (
|
|
"math"
|
|
"testing"
|
|
|
|
"github.com/ollama/ollama/x/internal/mlxthread"
|
|
)
|
|
|
|
func TestGELUCompiledMatchesEager(t *testing.T) {
|
|
values := []float32{-6, -2, -0.5, 0, 0.5, 2, 6}
|
|
tests := []struct {
|
|
name string
|
|
dtype DType
|
|
tolerance float32
|
|
}{
|
|
{name: "float32", dtype: DTypeFloat32, tolerance: 1e-6},
|
|
{name: "bfloat16", dtype: DTypeBFloat16, tolerance: 1e-2},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
withMLXThread(t, func() {
|
|
EnableCompile()
|
|
input := FromValues(values, len(values)).AsType(tt.dtype)
|
|
Pin(input)
|
|
|
|
want := gelu(input)
|
|
got := GELU(input)
|
|
wantF32 := want.AsType(DTypeFloat32)
|
|
gotF32 := got.AsType(DTypeFloat32)
|
|
Eval(wantF32, gotF32)
|
|
|
|
wantValues := wantF32.Floats()
|
|
gotValues := gotF32.Floats()
|
|
for i := range wantValues {
|
|
if delta := float32(math.Abs(float64(gotValues[i] - wantValues[i]))); delta > tt.tolerance {
|
|
t.Fatalf("%s GELU[%d] = %v, want %v (delta %v)", tt.name, i, gotValues[i], wantValues[i], delta)
|
|
}
|
|
}
|
|
Unpin(input)
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
func BenchmarkGELUEager(b *testing.B) {
|
|
benchmarkGELU(b, gelu)
|
|
}
|
|
|
|
func BenchmarkGELUCompiled(b *testing.B) {
|
|
benchmarkGELU(b, GELU)
|
|
}
|
|
|
|
func benchmarkGELU(b *testing.B, fn func(*Array) *Array) {
|
|
thread, err := mlxthread.Start("mlx-gelu-benchmark", func() error {
|
|
if err := CheckInit(); err != nil {
|
|
return err
|
|
}
|
|
if GPUIsAvailable() {
|
|
SetDefaultDeviceGPU()
|
|
}
|
|
EnableCompile()
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
b.Skipf("MLX not available: %v", err)
|
|
}
|
|
defer func() {
|
|
if err := thread.Stop(b.Context(), func() {
|
|
Sweep()
|
|
ClearCache()
|
|
resetDefaultStreamCache()
|
|
}); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}()
|
|
|
|
if err := thread.Do(b.Context(), func() error {
|
|
input := AddScalar(Zeros(DTypeBFloat16, 1, 4096, 8192), 1)
|
|
Eval(input)
|
|
Pin(input)
|
|
defer Unpin(input)
|
|
|
|
warmup := fn(input)
|
|
Eval(warmup)
|
|
Sweep()
|
|
|
|
b.ResetTimer()
|
|
for range b.N {
|
|
output := fn(input)
|
|
Eval(output)
|
|
Sweep()
|
|
}
|
|
return nil
|
|
}); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|