go-mlx/docs/development.md

# Development Guide

Module: `forge.lthn.ai/core/go-mlx`

---

## Prerequisites

### Platform

**macOS on Apple Silicon only.** All CGO source files carry `//go:build darwin && arm64`. The package will not build for native Metal inference on any other platform; a stub (`mlx_stub.go`) provides `MetalAvailable() bool` returning false elsewhere.

### Required Tools

| Tool | Version | Purpose |
|------|---------|---------|
| Go | 1.25.5+ | Module toolchain |
| CMake | 3.24+ | Builds mlx-c from source |
| AppleClang | 17.0+ | C/C++ compiler for mlx-c |
| macOS SDK | 26.2+ | Metal framework headers |
| Xcode Command Line Tools | Current | Provides `xcrun`, frameworks |

Install CMake if absent:

```bash
brew install cmake
```

### Go Workspace

go-mlx participates in a Go workspace alongside go-inference. The `go.mod` uses a `replace` directive for local development:

```
replace forge.lthn.ai/core/go-inference => ../go-inference
```

After adding modules or changing dependencies: `go work sync`

---

## Build

### Step 1: Build mlx-c

Run from the module root:

```bash
go generate ./...
```

This executes the `//go:generate` directives in `mlx.go`:

```
cmake -S . -B build -DCMAKE_INSTALL_PREFIX=dist -DCMAKE_BUILD_TYPE=Release
cmake --build build --parallel
cmake --install build
```

CMake fetches mlx-c v0.4.1 from GitHub, builds it with:
- `MLX_BUILD_SAFETENSORS=ON` (model loading)
- `MLX_BUILD_GGUF=OFF`
- `BUILD_SHARED_LIBS=ON`
- macOS deployment target: 13.3 (minimum required by MLX)

The built library installs to `dist/include/` and `dist/lib/`. Build time is approximately 2 minutes on M3 Ultra.

The `dist/` directory is gitignored and must be rebuilt on each fresh checkout.

### Step 2: Run Tests

```bash
go test ./...
```

Tests require a working mlx-c build. Integration tests that load model files are skipped automatically when model paths are absent (`/Volumes/Data/lem/safetensors/...`).

---

## CGO Flags

The `#cgo` directives in `internal/metal/metal.go` set all required flags automatically when building on darwin/arm64:

```c
#cgo CXXFLAGS: -std=c++17
#cgo CFLAGS: -mmacosx-version-min=26.0
#cgo CPPFLAGS: -I${SRCDIR}/../../dist/include
#cgo LDFLAGS: -L${SRCDIR}/../../dist/lib -lmlxc -lmlx
#cgo darwin LDFLAGS: -framework Foundation -framework Metal -framework Accelerate
#cgo darwin LDFLAGS: -Wl,-rpath,${SRCDIR}/../../dist/lib
```

`${SRCDIR}` is the directory containing `metal.go` at build time (`internal/metal/`), so the `../../dist/` path resolves to the module root `dist/`.

No manual environment variables are needed for `go build` or `go test`.

---

## Test Patterns

Tests use the `_Good`, `_Bad`, `_Ugly` suffix convention:

| Suffix | Meaning |
|--------|---------|
| `_Good` | Happy path; expected to succeed |
| `_Bad` | Expected error conditions |
| `_Ugly` | Panic / edge cases |

Example:

```go
func TestMatmul_Good(t *testing.T) { ... }
func TestMatmul_Bad(t *testing.T) { ... }
```

Tests that require model files on disk use `t.Skip()` when the path is absent:

```go
const modelPath = "/Volumes/Data/lem/safetensors/gemma-3/"
if _, err := os.Stat(modelPath); err != nil {
    t.Skip("model not available:", modelPath)
}
```

All 180+ tests in `internal/metal/` are unit or integration tests that exercise the CGO layer directly. The 11 tests in the root package (`mlx_test.go`) exercise the public API via go-inference.

### Running a Single Test

```bash
go test -run TestRMSNorm_Good ./internal/metal/
```

### Running with Race Detector

```bash
go test -race ./...
```

---

## Benchmarks

29 benchmarks in `internal/metal/bench_test.go`. Run with:

```bash
go test -bench=. -benchtime=2s ./internal/metal/
```

Key benchmarks:

| Benchmark group | What it measures |
|----------------|-----------------|
| `BenchmarkMatmul_*` | Matrix multiply at 128² through 4096², plus token projection |
| `BenchmarkSoftmax_*` | Softmax at 1K through 128K vocab |
| `BenchmarkElementWise_*` | Add, Mul, SiLU at 1M elements |
| `BenchmarkRMSNorm_*` | Fused RMSNorm at decode and prefill shapes |
| `BenchmarkRoPE_*` | RoPE at single-token and 512-token shapes |
| `BenchmarkSDPA_*` | Scaled dot-product attention at 1, 32, 512 sequence lengths |
| `BenchmarkLinear_*` | Linear layer forward at decode and prefill shapes |
| `BenchmarkSampler_*` | Greedy, TopK, TopP, and full chain on 32K vocab |

Model-level benchmarks (`model.Forward`, tokenizer) require model files on disk and are not included in the automated suite.

---

## Code Structure

### Adding a New Operation

1. Add the C binding to the appropriate file in `internal/metal/`:
   - `ops.go` — element-wise, reduction, matrix, shape operations
   - `fast.go` — fused Metal kernel wrappers
   - `slice.go` — slicing and scatter operations
2. Follow the `newArray("OP_NAME", inputs...)` pattern for tracking
3. Add tests in the corresponding `_test.go` file using `_Good`/`_Bad` suffixes
4. Add a benchmark in `bench_test.go` for any operation on the hot path

### Adding a New Model Architecture

1. Read `config.json` `model_type` and add a case in `model.go`:`loadModel`
2. Create `architecture.go` in `internal/metal/` implementing `InternalModel`
3. Add `ApplyLoRA` to the new model
4. Add a `close*` helper in `close.go` for deterministic resource cleanup
5. Add `formatXyzChat` in `generate.go` for the chat template
6. Add tokeniser BOS/EOS detection in `tokenizer.go`:`LoadTokenizer`
7. Write tests: config parsing, missing weights, end-to-end inference

---

## Coding Standards

### Language

UK English throughout: colour, organisation, centre, initialise, behaviour. Never American spellings.

### Go Style

- `declare(strict_types=1)` equivalent: all parameters and return types must be explicitly typed
- PSR-12 equivalent: `gofmt` + `goimports`; run before committing
- `go test ./...` must pass before every commit; no red tests in main

### Licence Header

Every new source file must carry the EUPL-1.2 licence identifier:

```go
// SPDX-Licence-Identifier: EUPL-1.2
```

### Conventional Commits

Format: `type(scope): description`

Types:
- `feat` — new capability
- `fix` — bug fix
- `test` — test additions or changes
- `bench` — benchmark additions or changes
- `refactor` — code restructuring without behaviour change
- `docs` — documentation only
- `chore` — maintenance (gitignore, go.mod, CMake)

Scopes: `metal`, `api`, `mlxlm`, `cpp`, `docs`

Examples:
```
feat(metal): add TopP nucleus sampling
fix(metal): auto-contiguous data access for non-contiguous arrays
test(metal): add model loading robustness tests
bench(metal): add 29 benchmarks baselined on M3 Ultra
```

### Co-Author

All commits must include:

```
Co-Authored-By: Virgil <virgil@lethean.io>
```

### Build Tags

- All CGO files: `//go:build darwin && arm64`
- Stub file: `//go:build !darwin || !arm64`
- mlxlm opt-out: `//go:build !nomlxlm`

---

## CMake Configuration

`CMakeLists.txt` at the module root. Key settings:

```cmake
set(MLX_BUILD_SAFETENSORS ON)   # Required for model loading
set(MLX_BUILD_GGUF OFF)         # GGUF not supported
set(BUILD_SHARED_LIBS ON)       # Shared .dylib for rpath loading
set(CMAKE_OSX_DEPLOYMENT_TARGET 13.3)  # MLX minimum
```

To force a clean rebuild:

```bash
rm -rf build dist
go generate ./...
```

---

## mlxlm Backend Development

The `mlxlm/` package has no CGO dependency and tests run on any platform where Python 3 is available. Tests use `testdata/mock_bridge.py` instead of the real `bridge.py`, so no `mlx-lm` installation is required.

Run mlxlm tests:

```bash
go test ./mlxlm/
```

The mock bridge responds to all commands with fixed fake data, enabling full subprocess protocol testing without GPU or Python ML dependencies.

To opt out of building the mlxlm backend:

```bash
go build -tags nomlxlm ./...
```

---

## Dependency Graph

```
go-mlx
├── forge.lthn.ai/core/go-inference  (shared interfaces, zero dependencies)
└── mlx-c v0.4.1                     (CMake, fetched from GitHub at generate time)
    └── Apple MLX (Metal GPU compute)
        └── Foundation, Metal, Accelerate frameworks
```

The root package and `mlxlm/` have no CGO dependency. Only `internal/metal/` links against mlx-c.