Add MLX backend for Apple Silicon neural network inference

[ChinChangYang]

Summary

This PR adds a new MLX backend for KataGo, enabling neural network inference on Apple Silicon using Apple's MLX framework. The implementation includes full model support with several performance optimizations.

Key Features

• Complete neural network support: Implements all required layers (convolution, batch normalization, residual blocks, policy/value heads, SGF metadata encoding)
• Graph compilation: Uses mx::compile() to JIT-compile inference functions into fused Metal kernels, with thread-safe caching per configuration
• Performance optimizations:
  • Skip mask operations when requireExactNNLen=true (all boards are exactly nnXLen × nnYLen)
  • Use logaddexp for Mish activation (reduces 7 ops to 2 ops)
  • Use addmm for fused matmul+bias operations (reduces 2 ops to 1 op)
• FP16 support: Optional half-precision mode (defaults to FP32 as FP16 doesn't improve performance on Apple Silicon MLX)
• Batch normalization: Computes merged scale/bias parameters for compatibility with layer tests

Build Instructions

cd cpp
cmake -G Ninja -DUSE_BACKEND=MLX
ninja

Requires CMake 3.27+

Configuration

Enable FP16 (optional, not recommended):

mlxUseFP16 = true

Testing

• Passes ./katago runtests (unit tests)
• Passes ./katago runnnlayertests (neural network layer tests)
• Cross-backend validation via ./katago testgpuerror with Eigen reference
• Benchmark testing shows functional parity with other backends

Changes

• cpp/neuralnet/mlxbackend.cpp: New 1518-line implementation
• cpp/CMakeLists.txt: MLX backend build configuration
• cpp/configs/gtp_example.cfg: MLX configuration options
• cpp/main.cpp, cpp/command/benchmark.cpp, cpp/program/setup.cpp: Backend registration

Performance Notes

• Graph compilation reduces dispatch overhead and enables operation fusion
• Compiled functions are cached per (batchSize, nnXLen, nnYLen, useMask, hasMeta)
• FP16 mode available but not recommended (no performance benefit observed)

🤖 Generated with Claude Code

[ChinChangYang]

Here’s what we learned making the MLX backend faster, written as general advice — not about our code, just things you can apply in your own fp32 backend.
The starting problem was that the plain fp32 path was noticeably slower than the MPSGraph backend, and the cost was almost all in the convolutions. Everything useful came from attacking that.
- Replace the generic convolution with a Winograd convolution for the main trunk. That is where the speed comes back. It is the single biggest lever for a fp32 network like this.
- Don’t hand-tune the convolution kernel’s launch settings. Parameterize them and let an autotuner search, ideally copying the method from a mature tuner you already trust. When we did this, the only setting that actually mattered was the data layout (looping so channels are the fast axis); the rest were noise. A hand-guess would almost certainly have tuned the wrong knob.
- Once fp32 is competitive, fp16 is the next big jump, but expect it to break first. In a deep network, fused low-precision accumulation can overflow the fp16 range and produce non-finite outputs, even though no single operation looks dangerous on its own. The fix is cheap: at the accumulation step that overflows (for us, BatchNorm), accumulate in fp32 and store back as fp16. Keep everything else fp16. Done right, the fp32 path is bit-for-bit unchanged and fp16 ends up much faster — for us, roughly a quarter to over a third faster than fp32, at the same accuracy.
- Be very careful how you measure. A thread sweep prints one number per thread count, and a careless parser will grab the slowest one and hand you a fake speedup; this once gave us a bogus “+20%”. Always pin the thread count, run the two versions back to back, and check the raw output. Measured honestly, our fp32 path only ties MPSGraph — the real win is fp16.
So the order is: Winograd to catch up, autotuning to find the one setting that matters, then fp16 (expect an overflow, fix it with fp32 accumulation at the bad step) to pull clearly ahead — and don’t trust any benchmark number whose parser you haven’t read.

---

Drafted by Claude Opus 4.7. Reviewed, experimented, confirmed, and edited by me.

[ChinChangYang]

Update: MLX tuner adaptive scoring landed

Pushed 29 commits to feature/mlx-backend (51c9db60..c656f840). Two related features completed in this batch:

MLX tuner shape diagnostic (predecessor, commits 49e6dccb..67fe303a)

Added two pure-observability log lines so a single tuner run produces empirical data on the model's 3×3 conv shape distribution. No scoring logic change.

MLX tuner adaptive scoring (3cb2df5c..c656f840)

Replaced the hardcoded {trunk, mid, max} three-slot scoring rotation with a model-adaptive rotation driven by planShapeRotation over the loaded model's actual 3×3 conv channel distribution, work-weighted by count × channels.

Motivation: On b18c384nbt (the recommended model), the dominant 3×3 conv runs at C=192 (mid, 72 of 78 calls) but the old scoring weighted it ~33% while the rarely-used C=384 trunk shape received ~67%. The new scoring derives weights from the actual distribution.

Selection rule: top-3 shapes, drop below 3% of total work, 3-rep floor per included shape, proportional remainder.

Validation on this host:

• runnnlayertests + runtests: all pass
• Gated tests (LOG_FORMAT, SWEEP, PER_SHAPE): all pass
• testgpuerror vs Eigen reference: bit-identical to the pre-rewrite MLX FP32 (winrateError max 0.00065%)
• Benchmark MLX FP16 on b18c384nbt: 491.78 v/s (mean of 2 re-tuned samples) vs pre-rewrite 499.45 — within run-to-run noise

Empirical: the new scoring picked a meaningfully different winner: input tg0=8 tg1=25 wpt=1 vw=4 (vs pre-rewrite tg0=64 tg1=16 wpt=1 vw=2), output tg0=32 tg1=2 wpt=1 (vs tg0=192 tg1=2 wpt=2). Throughput-equivalent on this model.

Specs and plans included in the diff:

• docs/superpowers/specs/2026-05-21-mlx-tuner-shape-diagnostic-design.md
• docs/superpowers/plans/2026-05-21-mlx-tuner-shape-diagnostic.md
• docs/superpowers/specs/2026-05-21-mlx-tuner-adaptive-scoring-design.md
• docs/superpowers/plans/2026-05-21-mlx-tuner-adaptive-scoring.md

🤖 Generated with Claude Code

[ChinChangYang]

Moved to lightvector#1199.