mirror of
https://github.com/ollama/ollama.git
synced 2026-02-20 00:05:06 -05:00
44bdd9a2ef
This change adds a new MLX based runner which includes: * Method-based MLX bindings * Subprocess-based MLX runner (x/mlxrunner) * KV cache with tree management * A basic sampler The GLM4-MoE-Lite model has been ported to use the new bindings. --------- Co-authored-by: Michael Yang <git@mxy.ng>
88 lines
1.5 KiB
Go
88 lines
1.5 KiB
Go
//go:build mlx
|
|
|
|
package mlx
|
|
|
|
// #include "generated.h"
|
|
import "C"
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
"strconv"
|
|
)
|
|
|
|
func (b Byte) String() string {
|
|
return strconv.FormatInt(int64(b), 10) + " B"
|
|
}
|
|
|
|
func (b KibiByte) String() string {
|
|
return strconv.FormatFloat(float64(b)/(1<<10), 'f', 2, 64) + " KiB"
|
|
}
|
|
|
|
func (b MebiByte) String() string {
|
|
return strconv.FormatFloat(float64(b)/(1<<(2*10)), 'f', 2, 64) + " MiB"
|
|
}
|
|
|
|
func (b GibiByte) String() string {
|
|
return strconv.FormatFloat(float64(b)/(1<<(3*10)), 'f', 2, 64) + " GiB"
|
|
}
|
|
|
|
func (b TebiByte) String() string {
|
|
return strconv.FormatFloat(float64(b)/(1<<(4*10)), 'f', 2, 64) + " TiB"
|
|
}
|
|
|
|
func PrettyBytes(n int) fmt.Stringer {
|
|
switch {
|
|
case n < 1<<10:
|
|
return Byte(n)
|
|
case n < 1<<(2*10):
|
|
return KibiByte(n)
|
|
case n < 1<<(3*10):
|
|
return MebiByte(n)
|
|
case n < 1<<(4*10):
|
|
return GibiByte(n)
|
|
default:
|
|
return TebiByte(n)
|
|
}
|
|
}
|
|
|
|
func ActiveMemory() int {
|
|
var active C.size_t
|
|
C.mlx_get_active_memory(&active)
|
|
return int(active)
|
|
}
|
|
|
|
func CacheMemory() int {
|
|
var cache C.size_t
|
|
C.mlx_get_cache_memory(&cache)
|
|
return int(cache)
|
|
}
|
|
|
|
func PeakMemory() int {
|
|
var peak C.size_t
|
|
C.mlx_get_peak_memory(&peak)
|
|
return int(peak)
|
|
}
|
|
|
|
type Memory struct{}
|
|
|
|
func (Memory) LogValue() slog.Value {
|
|
return slog.GroupValue(
|
|
slog.Any("active", PrettyBytes(ActiveMemory())),
|
|
slog.Any("cache", PrettyBytes(CacheMemory())),
|
|
slog.Any("peak", PrettyBytes(PeakMemory())),
|
|
)
|
|
}
|
|
|
|
type (
|
|
Byte int
|
|
KibiByte int
|
|
MebiByte int
|
|
GibiByte int
|
|
TebiByte int
|
|
)
|
|
|
|
func ClearCache() {
|
|
C.mlx_clear_cache()
|
|
}
|