go-mlx¶

forge.lthn.ai/core/go-mlx provides native Apple Metal GPU inference and LoRA fine-tuning for Go. It wraps Apple's MLX framework through the mlx-c C API, implementing the inference.Backend interface from forge.lthn.ai/core/go-inference.

Platform: darwin/arm64 only (Apple Silicon M1-M4). A stub provides MetalAvailable() bool returning false on all other platforms.

Quick Start¶

import (
    "context"
    "fmt"

    "forge.lthn.ai/core/go-inference"
    _ "forge.lthn.ai/core/go-mlx" // registers "metal" backend via init()
)

func main() {
    m, err := inference.LoadModel("/path/to/safetensors/model/")
    if err != nil {
        panic(err)
    }
    defer m.Close()

    ctx := context.Background()
    for tok := range m.Generate(ctx, "What is 2+2?", inference.WithMaxTokens(128)) {
        fmt.Print(tok.Text)
    }
    if err := m.Err(); err != nil {
        panic(err)
    }
}

The blank import (_ "forge.lthn.ai/core/go-mlx") auto-registers the Metal backend. All interaction goes through the go-inference interfaces -- go-mlx itself exports only Metal-specific memory controls.

Features¶

• Streaming inference -- token-by-token generation via iter.Seq[Token] (range-over-func)
• Multi-turn chat -- native chat templates for Gemma 3, Qwen 2/3, and Llama 3
• Batch inference -- Classify (prefill-only) and BatchGenerate (autoregressive) for multiple prompts
• LoRA fine-tuning -- low-rank adaptation with AdamW optimiser and gradient checkpointing
• Quantisation -- transparent support for 4-bit and 8-bit quantised models via QuantizedMatmul
• Attention inspection -- extract post-RoPE K vectors from the KV cache for analysis
• Performance metrics -- prefill/decode tokens per second, GPU memory usage

Supported Models¶

Models must be in HuggingFace safetensors format (not GGUF). Architecture is auto-detected from config.json:

Package Layout¶

Metal Memory Controls¶

These control the Metal allocator directly, not individual models:

import mlx "forge.lthn.ai/core/go-mlx"

mlx.SetCacheLimit(4 << 30)   // 4 GB cache limit
mlx.SetMemoryLimit(32 << 30) // 32 GB hard limit
mlx.ClearCache()              // release cached memory between chat turns

fmt.Printf("active: %d MB, peak: %d MB\n",
    mlx.GetActiveMemory()/1024/1024,
    mlx.GetPeakMemory()/1024/1024)

Performance Baseline¶

Measured on M3 Ultra (60-core GPU, 96 GB unified memory):

Documentation¶

• Architecture -- CGO binding layer, lazy evaluation, memory model, attention, KV cache
• Models -- model loading, supported architectures, tokenisation, chat templates
• Training -- LoRA fine-tuning, gradient computation, AdamW optimiser, loss functions
• Build Guide -- prerequisites, CMake setup, build tags, testing

Downstream Consumers¶

Licence¶

EUPL-1.2