Skip to content

index.md

Latest commit

 

History

History
197 lines (152 loc) · 7.17 KB

index.md

File metadata and controls

197 lines (152 loc) · 7.17 KB
Raw
title go-mlx
description Native Metal GPU inference and training for Go on Apple Silicon.

go-mlx

dappco.re/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 dappco.re/go/inference and an RFC-style direct root-package API.

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

Quick Start

import (
    "context"
    "fmt"

    "dappco.re/go/inference"
    _ "dappco.re/go/mlx" // registers "metal" backend via init()
)

func main() {
    m, err := inference.LoadModel("/path/to/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 (_ "dappco.re/go/mlx") auto-registers the Metal backend. You can use either the go-inference interfaces or the direct root API:

import (
    "fmt"

    mlx "dappco.re/go/mlx"
)

model, err := mlx.LoadModel("/path/to/model/",
    mlx.WithContextLength(262144), // opt into larger Qwen-class contexts
    mlx.WithParallelSlots(1),      // one foreground local runner by default
)
if err != nil {
    panic(err)
}
defer model.Close()

if err := model.WarmPromptCache(stableSystemAndToolsPrefix); err != nil {
    panic(err)
}

text, err := model.Generate("What is 2+2?", mlx.WithMaxTokens(64))
if err != nil {
    panic(err)
}
fmt.Println(text)

Features

  • Streaming inference -- token-by-token generation via iter.Seq[Token] (range-over-func)
  • Multi-turn chat -- native chat templates for Gemma 3/4, Qwen 2/3, and Llama 3
  • Batch inference -- Classify (prefill-only) and BatchGenerate (autoregressive) for multiple prompts
  • Frame compute sessions -- non-LLM pixel-buffer pipelines with explicit per-frame lifecycle, scaling, swizzling, palette expansion, and format conversion
  • 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
  • Restorable model state -- capture KV, logits, token offsets, and generated-token history into reloadable sessions
  • State bundles -- strict JSON artifacts that bind model identity, tokenizer/chat-template metadata, prompt hash, sampler settings, LoRA identity, KV hash, SAMI/probe data, and optional memvid refs
  • Performance metrics -- prefill/decode tokens per second, GPU memory usage
  • Local-runner defaults -- GPU, 131k bounded context, one native slot, and exact token-prefix prompt cache enabled by default
  • Non-HTTP sidecar -- Violet serves native generation over a local Unix socket for harnesses that do not need an OpenAI-compatible HTTP layer

Supported Models

Models may be loaded from HuggingFace safetensors shards or GGUF checkpoints. Architecture is auto-detected from config.json:

Architecture model_type values Tested sizes
Gemma 3 gemma3, gemma3_text, gemma2 1B, 4B, 27B
Gemma 4 gemma4, gemma4_text E2B, E4B, 26B MoE, 31B
Qwen 3 qwen3, qwen2 8B+
Llama 3 llama 8B+

Package Layout

Package Purpose
Root (mlx) Public API: backend registration, direct model API, memory controls, training type exports
internal/metal/ All CGO code: array ops, model loaders, generation, training primitives
mlxlm/ Alternative subprocess backend via Python's mlx-lm (no CGO required)
pkg/daemon/ and cmd/violet Unix-socket sidecar for local native generation without HTTP

Violet Native Route

Violet is the direct local route for CoreAgent-style harnesses that already own tool execution and do not need an OpenAI-compatible server. Configure one or more model paths, run the daemon, then send one JSON frame per line over the Unix socket:

# violet.toml
[models]
default = "/path/to/mlx/model"
violet --config violet.toml --socket /tmp/violet.sock

Prompt generation:

{"action":"generate","prompt":"What is 2+2?","max_tokens":64}

Chat generation:

{"action":"generate","messages":[{"role":"system","content":"Be direct."},{"role":"user","content":"What is 2+2?"}],"max_tokens":64}

The native route uses the same mlx.LoadModel defaults as the direct API: GPU execution, 131k bounded context, one active native slot, and exact token-prefix prompt caching. Models are loaded on first use and kept resident until the daemon exits.

Metal Memory Controls

These control the Metal allocator directly, not individual models:

import mlx "dappco.re/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)
Function Purpose
SetCacheLimit(bytes) Soft limit on the allocator cache
SetMemoryLimit(bytes) Hard ceiling on Metal memory
SetWiredLimit(bytes) Wired memory limit
GetActiveMemory() Current live allocations in bytes
GetPeakMemory() High-water mark since last reset
GetCacheMemory() Cached (not yet freed) memory
ClearCache() Release cached memory to the OS
ResetPeakMemory() Reset the high-water mark
GetDeviceInfo() Metal GPU hardware information

Performance Baseline

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

Operation Throughput
Gemma3-1B 4-bit prefill 246 tok/s
Gemma3-1B 4-bit decode 82 tok/s
Gemma3-1B 4-bit classify (4 prompts) 152 prompts/s
DeepSeek R1 7B 4-bit decode 27 tok/s
Llama 3.1 8B 4-bit decode 30 tok/s

Documentation

  • Compute Guide -- frame-oriented Metal compute sessions, pixel buffers, kernels, metrics
  • 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
  • Model State Roadmap -- native session restore, state bundles, probes, training runner, model packs, memory planning, benchmarks
  • Build Guide -- prerequisites, CMake setup, build tags, testing

Downstream Consumers

Package Role
dappco.re/go/core/ml Imports go-inference + go-mlx for the Metal backend training loop
dappco.re/go/core/i18n Gemma3-1B domain classification (Phase 2a)
dappco.re/go/core/rocm Sibling AMD GPU backend, same go-inference interfaces

Licence

EUPL-1.2