Add MLX backend for Apple Silicon (GPU + ANE) - #1199
Open
ChinChangYang wants to merge 59 commits into
Open
Conversation
ChinChangYang
force-pushed
the
mlx-backend-squash
branch
from
May 23, 2026 04:32
b544c66 to
dcf296a
Compare
Introduces a new neural-net backend (USE_BACKEND=MLX) targeting Apple
Silicon via Apple's MLX framework. The backend implements the full
nninterface contract (model load, batched evaluation, FP16/FP32 paths)
and ships with a Winograd 3x3 convolution path plus an adaptive
per-shape tuner that picks the fastest implementation for each
conv-3x3 shape at model load.
Backend
- cpp/neuralnet/mlxbackend.cpp: backend implementation. Supports
variable board sizes via input masking (same nnXLen/nnYLen
contract as other backends; the global COMPILE_MAX_BOARD_LEN
bound still applies). FP16/FP32 selected by the mlxUseFP16 config
(default auto -> fp16); same input feature layout as the other
backends. Mish activation runs FP16-safe (asserts on
ACTIVATION_MISH_SCALE8 so out-of-range variants are caught
explicitly rather than silently truncated).
- cpp/neuralnet/mlxwinograd.h: F(4x4, 3x3) Winograd transform with
fused activation + residual add.
- cpp/neuralnet/mlxwinotuner.{cpp,h}: per-shape Winograd tuner with
adaptive scoring (rotates the candidate set per shape, scores by
median-time delta against a baked-default baseline). Logs the
conv-3x3 shape distribution at model load.
- cpp/neuralnet/mlxtests.cpp: unit tests for the Winograd path
and tuner numeric-consistency, gated under runnnlayertests.
Build / wiring
- cpp/CMakeLists.txt: USE_BACKEND=MLX target. MLX requires CMake
3.27 (cmake_minimum_required stays at 3.18.2 so other backends
keep building on older CMake). Links Homebrew's prebuilt
libmlx.dylib; OSX deployment target intentionally not pinned so
the executable's minos matches the dylib it was linked against.
- cpp/main.cpp, cpp/program/setup.cpp, cpp/command/benchmark.cpp:
wire MLX into backend selection / benchmark.
- cpp/configs/{gtp,analysis,match,contribute}_example.cfg: document
mlxUseFP16 (auto / true / false), default auto -> fp16.
- Compiling.md: build instructions for the MLX backend.
Validation
- Cross-backend validation against an Eigen reference (testgpuerror)
for b18c384nbt, b40v8, and humanv0 nets shows FP32 max winrate
error 0.00095% and FP16 max 2.63%, well within the existing
backend tolerances.
This is the squash of 130 commits from feature/mlx-backend.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ChinChangYang
force-pushed
the
mlx-backend-squash
branch
from
May 23, 2026 08:06
dcf296a to
81b00db
Compare
…arity smoke test (#26)
ChinChangYang
force-pushed
the
mlx-backend-squash
branch
from
May 26, 2026 23:38
31eec63 to
628e377
Compare
ChinChangYang
changed the title
# Conflicts: # cpp/rungpuerrortest.sh
master consolidated createComputeContext's trailing params (openCLTunerFile, openCLReTunePerBoardSize, useNHWCMode) into a single ConfigParser& cfg. The Metal backend was already updated; update the MLX backend to match so it compiles against the merged interface. NHWC is still enforced per-handle via inputsUseNHWC in createComputeHandle. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The MLX backend implemented only the convnet path; transformer nets crashed (empty-(0)-tensor broadcast on rmsnorm/SiLU tips) or produced garbage (GQA). Implement the transformer trunk/tip path in mlxbackend.cpp, mirroring eigenbackend.cpp: - ACTIVATION_SILU (x * sigmoid(x)) - TransformerRMSNormLayer (spatial rmsnorm tip) + TransformerTrunkRMSNormLayer (pre-LN) - GQA TransformerAttentionBlock + SwiGLU FFNBlock - branch the trunk tip on trunkNormKind; wire the new block kinds into the block-variant and nested-bottleneck loops; thread nnX/nnY through Verified via testgpuerror against fresh Eigen references (boardsize 19): fp32 winrateError max — rope 0.00094%, silu 0.00046%, gqa 0.00029% (bar 0.10%); convnet g170-b6c96 unregressed (0.00036%); runtests + runnnlayertests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The MLX backend's ANE mux path (mlxDeviceToUseThread0=100) drives inference through the shared external/katagocoreml converter -- the same library the Metal backend uses. Bring PR#1205's CoreML/ANE transformer work into that converter so the MLX ANE path supports the v15+ transformer trunk: - Transformer MIL support: attention (incl. grouped-query attention), learnable RoPE, SiLU, RMSNorm/batchnorm tips, SwiGLU FFN. - FP16 accuracy precision tiers, gated on actual transformer-block presence (blocksContainTransformer, recursing into nested-bottleneck blocks): narrow trunks (<256ch) build fully FP32; wider ones escalate non-spatial matmuls + global pooling to FP32; very wide (>=320ch) also escalate convs; RMSNorm reductions FP32 in FP16 mode. Plain convnets stay pure FP16 on the ANE (the d052d2a regression-fix behavior is preserved). The converter's public API is unchanged, so the MLX call site (CoreMLConversion::convertModelToTemp) needs no edits. The Metal-GPU/MPSGraph portions of PR#1205 (metalbackend.cpp, metallayers.swift) are intentionally not ported -- the MLX backend's native GPU path already has transformer support. Verified on the MLX ANE mux (testgpuerror vs fresh Eigen FP32 references): all 3 transformer test nets pass FP16 thresholds across board sizes/buffer configs (7 configs), a plain convnet stays pure FP16 (non-regression), and runtests + runnnlayertests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChinChangYang
changed the title
applyGlobalPooling / applyValueHeadPooling summed in fp16 but produced an fp32 mean (division by the fp32 maskSum), which also leaked fp32 into the downstream gpool-bias and value-v2 head matmuls. Cast maskSum to the input dtype so the whole pooling and the heads stay in the compute dtype (fp16 when useFP16), maximizing fp16 utilization rather than escalating to fp32 for negligible accuracy gain. The masked-max keeps its 1e9 constant in fp32 (1e9 overflows fp16 -> inf -> 0*inf=NaN), then casts the max result back to the compute dtype. The fp32 path is unaffected (the astype casts are no-ops in fp32). Verified via testgpuerror vs fresh Eigen fp32 references on all 3 transformer nets (7 board-size/buffer configs): fp16 winrate error max <= 2.07% (within tolerance, winrate unchanged vs baseline), fp32 path byte-identical, ownership output bit-identical, runnnlayertests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…nt warmup
M2: A candidate whose threadgroup exceeds the pipeline's register-pressure-
dependent maxTotalThreadsPerThreadgroup (can be < 1024), or that hits a
transient GPU error, throws out of mx::eval during the flat sweep. Previously
this propagated out of loadOrAutoTune and aborted model load with no fallback.
Now each candidate's scoring is wrapped in try/catch: a throw is counted and
skipped (mirroring the OpenCL tuner's mark-bad-and-continue), and best/bestTime
are seeded with the baked default so even a fully-failing sweep returns a valid
result. A separate "flatSweep{Input,Output} skipped=N" log line is emitted only
when skips occur; it intentionally omits the colon after the function name so it
cannot collide with the regex-tested "flatSweepInput: considered" log line.
M3: timeOneInputTransform/timeOneOutputUntransform ran an untimed warmup eval on
every call, but the scoring functions already warmed up once before the measured
loop -- so every measured rep paid an extra full warmup (~doubling tuning cost).
Add a doWarmup parameter gating the internal warmup; the scoring functions drop
their explicit warmup and pass (r == 0), so each shape warms exactly once on its
first measured rep.
Verified by triggering autotuning: gated flat-sweep tests (convergence,
log-format, baseline-consistency, per-shape) pass; an end-to-end re-tune via
loadOrAutoTune runs a fresh sweep and saves valid fp16/fp32 caches; testgpuerror
output is unchanged (tuner params are numerically inert); runtests passes.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The METAL and MLX backend branches each ran set(CMAKE_OSX_DEPLOYMENT_TARGET 13.0) *after* project(), where it is a silent no-op: for this Swift project the deployment target is fixed during project()/enable_language, so a later set() never affects the produced binary. Both shipped binaries already carry minos 26.0 (the build host / libmlx's floor), not 13.0, confirming the pins were inert dead code that contradicted the pre-project comment explaining why the deployment target is deliberately not pinned. Delete both pins so code, comment, and reality agree; the comment becomes literally true. Add a guard note documenting that a post-project pin is a no-op so it is not reintroduced. No behavior change: binaries still build at minos 26.0, matching libmlx's minos and MLX's macOS >= 14 requirement. Verified: MLX reconfigure+build clean; METAL branch configures clean; binary minos unchanged (26.0); runtests pass; testgpuerror unchanged (fp32 max 0.00036%, fp16 max 0.863%). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The katagocoreml parser read the q/k/v/out projection matmuls of a
transformer attention block without checking their declared dimensions
against the head geometry or trunk width. The CoreML graph builder
(MILBuilder) then reshapes each flat projection into a [seq, heads,
headDim] grid and reshapes the out-projection result back into the trunk,
so a mismatched dimension would either read past the weight buffer or
build a graph that compiles but computes nonsense - the exact failure
mode the existing checkBlockChannels() guard was added to prevent for
conv blocks.
Thread trunk_num_channels into parseTransformerAttentionBlock (mirroring
parseNestedBottleneckBlock) and add a checkAttentionProjDim() helper in
the style of checkBlockChannels(), then validate all four projections:
qProj.outChannels == numHeads * qHeadDim (master desc.cpp:1129)
kProj.outChannels == numKVHeads * qHeadDim (master desc.cpp:1131)
vProj.outChannels == numKVHeads * vHeadDim (master desc.cpp:1133)
outProj.inChannels == numHeads * vHeadDim (master desc.cpp:1135)
qProj.inChannels == trunkNumChannels (master desc.cpp:1430)
outProj.outChannels== trunkNumChannels (master desc.cpp:1437)
k/vProj.inChannels == trunkNumChannels (gap master leaves
implicit to the backend)
Six checks mirror master desc.cpp's transformer attention consistency
checks exactly; the k/v inChannels checks additionally close a gap master
leaves to the backend (all three QKV projections consume the same
normed-trunk input, so their inChannels must equal the trunk width). K
pairs with Q in the QK^T dot product, so kProj uses qHeadDim; only V
carries vHeadDim.
Purely additive: throws std::runtime_error on a malformed model, no-op on
valid ones, no numerics touched. Verified: MLX build clean, runtests
pass, and ANE-path testgpuerror on all three transformer nets (incl. the
GQA net with 6 heads/3 KV heads, qk=32/v=16) loads and converts with zero
false-positive throws and unchanged numerics.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Cuts memory during the on-device KataGo -> CoreML conversion and while running the ANE/CoreML path, with byte-identical converter output: - The converter's weight tensors become non-owning views into the parsed model instead of owning extra FP32 copies; derived/transposed tensors keep an owned buffer. This drops redundant resident weight copies during conversion. CoreML model serialization is made deterministic (SetSerializationDeterministic) so the output is byte-stable. - The KataGo model parser streams the gzip through a bounded ~1 MB refill buffer instead of decompressing the whole file into memory, while preserving the existing NaN/Inf weight validation. - ModelDesc gains releaseWeights(), which frees the in-memory weight arrays (keeping scalar shape metadata). The Metal backend calls it on the ANE (CoreML) path after converting from the model file on disk, gated by a new ComputeContext::aneOnly flag so it only fires when every configured device is ANE -- the GPU/MPSGraph path keeps its weights. The call is serialized under computeHandleMutex and only scalar dims are read afterward. Measured on b18c384nbt (19x19) over the ANE path: idle steady-state RSS 0.59 GB -> 0.19 GB; peak (load+convert) 0.87 GB -> 0.48 GB. Cross-backend parity vs an Eigen reference is unchanged on both the GPU and ANE paths. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit b05f559)
WeightEntry stores a non-owning view (const float*, count) into the live KataGoModelDesc, so the backing std::vector must outlive serialization. addConstOp/registerWeight took the data by const& and silently stored a pointer to it; a caller passing a temporary would bind to that const& and leave the view dangling, read much later during serialization. Delete the rvalue overloads of both so any such call fails to compile, forcing temporaries through addOwnedConstOp/registerOwnedWeight (which take ownership). Named lvalues (the model-member call sites) still bind to the const& overload, so no existing caller changes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 971fa9d)
Own the gzFile with a custom-deleter unique_ptr so it closes on every exit path (normal return, exception, bad_alloc); removes the manual try/catch+gzclose in parse() and the ordering caveat on buffer allocation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit eeefc97)
Introduce a KataGo-local non-owning FloatView for WeightEntry::data instead of a raw const float*/size_t pair; convert to MILBlob::Util::Span only inside WeightSerializer, keeping the MILBlob dependency out of Operations.hpp. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 6bfa617)
The ComputeHandle member-order comment claimed that declaring mpsGraphOnlyHandle before coremlOnlyHandle is what prevents a GPU handle from reading freed weights. That overstates the ordering's role: within a single ComputeHandle exactly one handle is built (mutually exclusive on gpuIdx, enforced by the ctor's exactly-one check), and releaseWeights() only fires on an aneOnly context where no MPSGraph handle is ever built. Reframe the declaration order as belt-and-suspenders and point at ComputeContext::aneOnly as the actual invariant. Comment-only change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 4159930)
Replace the file-local releaseXXX free functions in desc.cpp (which reached into each desc struct's internals from outside) with releaseWeights() member methods on each weight-bearing struct, matching the existing OO convention used by applyScale8ToReduceActivations() and iterConvLayers(). Each container delegates to its members; type-erased block dispatch is inlined with the same cast pattern those methods use. Behavior-preserving: same set of freed vectors, same block recursion, same metaEncoderVersion guard. ModelDesc::releaseWeights() keeps its signature, so the metalbackend.cpp call site is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 44342a3)
Move the 11 leaf/container releaseWeights() definitions in desc.cpp out of the bottom cluster (inherited from the old free-function layout) and place each immediately after its struct's last existing method, matching the file's per-struct grouping convention used by every other method. ModelDesc::releaseWeights() stays put, already adjacent to its siblings. Pure relocation: function bodies and desc.h are unchanged; only two stray double-blank lines were normalized to single. Verified clean Metal build, testgpuerror vs Eigen reference (g170-b6c96) at <0.0004% winrate error, and runtests all pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 98b17eb)
…contract The transformer attention builder emits four function-local std::vector<float> tensors: RoPE cos/sin tables, the rotation matrix R, and per-head out-projection weight slices. After merging the transformer support onto the FloatView branch, these needed two fixes: 1. Dangling view. lightvector#1202 made WeightEntry::data a non-owning FloatView, so addConstOp registers a view whose backing buffer must outlive serialization. These locals were passed to addConstOp and would dangle once the build function returns (serialization runs afterwards). Route them through addOwnedConstOp so KataGoOps owns the buffer until serialization. (Under lightvector#1205's owning WeightEntry they were copied, so this only surfaces post-merge.) 2. dtype mismatch. emitConstOp declares each const's dtype as m_weight_dtype, but addOwnedConstOp / registerOwnedWeight stored at the global mode (is_fp32 hardcoded false). In an FP16 model these derived consts land in the attention / value-head FP32 sub-region (m_weight_dtype == FLOAT32), so they were declared FP32 but stored FP16. CoreML/ANE then rejects the model at load ("Metadata data type does not match requested type", BNNS error -14), which SIGABRT'd every FP16 ANE transformer. Thread is_fp32 through registerOwnedWeight and have addOwnedConstOp pass is_fp32 = (m_weight_dtype == FLOAT32), mirroring addConstOp so the stored dtype always matches the declared dtype. This also fixes the same latent mismatch for addLinearOp's transposed value-head weights. Verified with testgpuerror against fresh Eigen FP32 references: b7c96h3tfrs and b7c96h6gqa, which previously SIGABRT'd on the FP16 ANE path, now load and match to <0.0005% winrate; convnet ANE output is byte-identical and the Metal GPU path is unchanged. katago runtests and runnnlayertests also pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 8481a94)
The cherry-picked per-struct releaseWeights() refactor (44342a3/98b17ebb) predates this branch's MLX transformer port, so it only added releaseWeights() to the non-transformer descriptors. Extend the coverage to the transformer descriptors present on this branch (RMSNormLayerDesc, TransformerRMSNormDesc, TransformerAttentionDesc incl. ropeFreqs, TransformerFFNDesc) and handle TRANSFORMER_ATTENTION_BLOCK_KIND / TRANSFORMER_FFN_BLOCK_KIND plus trunkTipRMSNorm in the trunk release walk. Without this, releasing weights on a transformer model would hit ASSERT_UNREACHABLE. This makes desc.cpp/desc.h byte-identical to the lightvector#1202 feature branch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port the lightvector#1202 ANE steady-state memory lever to the MLX backend. Add ComputeContext::aneOnly, set in createComputeContext when every configured device index is MLX_MUX_ANE, and call ModelDesc::releaseWeights() in convertAndCreateCoreMLOnlyHandleMLX after the model has been converted to CoreML on disk. Safe because: the ANE path re-reads the model from modelPath (not the in-memory weight arrays); the ComputeHandle ctor takes the MLX_MUX_ANE early-return before building any MLX/GPU model (the only weight-array consumer); only scalar dims are read afterward, which releaseWeights() preserves; and it runs under computeHandleMutex. Mirrors the Metal backend's aneOnly release. GPU path unaffected (aneOnly is false whenever any thread uses MLX_MUX_GPU). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
parseTransformerFFNBlock only checked num_channels/ffn_channels > 0, while the attention block validates all of its projection dimensions. Thread trunk_num_channels through and add the mirror checks: num_channels must equal the trunk width (the block adds its output back into the trunk residually) and the linear layers must chain numChannels -> ffnChannels -> numChannels (with the SwiGLU gate also numChannels -> ffnChannels). A malformed FFN block now fails at parse time instead of producing an opaque CoreML compile error or silently-wrong activations. Reuses the existing checkAttentionProjDim helper. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Six conv/matmul/RMSNorm/FFN sites hand-rolled save/restore of m_weight_dtype around their FP32 escalation windows. An exception thrown inside a window would leave m_weight_dtype stuck at FLOAT32, causing later FP16 consts to be tagged FP32 -> the BNNS "Metadata data type does not match" SIGABRT on the FP16 ANE. Give ScopedFp32 an active flag (so a conditional window needs no construction-time branch) and an idempotent restore() (to end the window before a trailing cast-down while keeping the dtor's exception-safe restore), then route all six sites through it. The guard is constructed exactly where the manual flip was and restore() called exactly where the manual restore was, so op-emission order -- and thus the serialized converter output -- is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The raw-output forward path had no callers: production inference goes through getOutput() -> Model::applyCompiled(). It also duplicated applyCompiled's input setup and output copy. Delete it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… test comment CMakeLists.txt: uppercase USE_BACKEND into USE_BACKEND_NORMALIZED before the pre-project() MLX version guard and the Swift language selection, mirroring the post-project() string(TOUPPER). Previously a lowercase -DUSE_BACKEND=mlx skipped the CMake 3.27 guard and Swift enablement, then still tried to build the Swift sources later, producing a confusing failure instead of a clear message. rungpuerrortest.sh: the gpu/ane modes drive whichever backend the binary was built with (backend-agnostic deviceToUseThread0), so reword the usage comment from "the Metal backend" to "the active backend (Metal or MLX)". Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… param, transformer test Address four follow-up items from the PR#1199 review (test/tuner-only; no inference forward-pass changes): 1. Atomic tuner-cache save. MLXWinogradTuneParams::save now writes to a per-process temp path (filename + ".tmp.<pid>") and FileUtils::rename's it onto the final path, so two processes that cache-miss and tune the same model concurrently can no longer tear the shared cache file. 2. Independent Winograd oracle. The GPU and FP16 Winograd metal_kernel tests previously asserted only against cpuConv2d3x3, itself a Winograd F(2,3) impl sharing the kernel's B/G/A transform matrices -- a shared sign/transpose error would cancel and pass. They now also assert against the independent naive direct-conv oracle. 3. Remove the dead seedOverride parameter from MLXWinogradTuner:: loadOrAutoTune (declaration, definition, and both call sites). It was documented "reserved ... currently ignored" and always passed nullptr. 4. Transformer-layer numeric test (runMLXTransformerLayerFP16Test): the transformer path (RMSNorm / attention / RoPE) had no layer-level coverage -- only end-to-end via testgpuerror. Adds RMSNorm fp32-vs-CPU correctness, attention fp16 output-dtype preservation (the regression guard for the just-fixed fp16->fp32 attention-scale promotion), fp16/ fp32 closeness, and a zero-outProj residual-identity anchor; covers fixed-RoPE on/off and mask on/off. Verified: build clean; runtests and runnnlayertests pass; the three transformer nets pass testgpuerror on the MLX GPU path within thresholds. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…delete, help text, fp16 accum test - mlxwinotuner.cpp: non-finite median in the two scoring paths now maps to +infinity instead of 0.0 — the tuner minimizes time, so a NaN/inf from a failed kernel run was making that candidate win selection. Diagnostic-only per-shape guards left as-is. - Operations.hpp: add symmetric 4-arg rvalue registerWeight(..., bool) = delete to close the arity hole; registerWeight(name, std::move(vec), shape, true) previously bound a temporary to the const& view overload, leaving a dangling FloatView. No-default form avoids lvalue-overload ambiguity. - main.cpp: tuner help text "(OpenCL only)" -> "(OpenCL and MLX)". - mlxtests.cpp: add a self-calibrating fp16 Winograd accumulation guard. Measures the scale-invariant normalized error (maxAbsErr/outMagMax) at Cin=8 and Cin=384 and asserts it stays small AND flat in Cin (ratio < 3). fp32 accumulation keeps it flat (~1.0); an fp16-accum regression would grow with the term count. Hardware-independent (no absolute-magnitude tuning). Verified: runtests, runnnlayertests (accum guard ratio 1.02), and testgpuerror vs eigen_reference.json (fp32 max winErr 0.00036%, fp16 0.55%) all pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Mirrors the build-macos-metal job: installs the mlx Homebrew formula alongside the shared CoreML deps (protobuf/abseil), configures with -DUSE_BACKEND=MLX under Ninja, builds, and runs runtests on the Apple Silicon (arm64) macos-latest runner that MLX requires. Carries over the metal job's dependency-version cache keying: the CMake build bakes in version-pinned Homebrew Cellar paths (protobuf/abseil/mlx), so the installed versions are folded into the cache key to force a fresh configure when a formula is bumped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
MLX's GPU streams have no per-stream worker thread, so gpu::eval runs inline on the calling thread and every ComputeHandle shares MLX's single global default GPU stream. With more than one thread driving inference concurrently -- numNNServerThreadsPerModel > 1, a second loaded model (e.g. a human SL net), or the multi-threaded analysis engine -- two threads open two compute command encoders on the same MTLCommandBuffer, which aborts with the Metal assertion "A command encoder is already encoding to this command buffer". Guard the whole MLX graph-build + eval + result read in Model::applyCompiled with a file-scope mlxGpuEvalMutex. Input prep in getOutput stays outside the lock so it still overlaps; one Apple GPU serializes the actual work anyway and KataGo's batching remains the throughput lever. No-op for the default single-model, single-server- thread GTP config. Cherry-picked (cpp portion) from ios-dev 8d07258. That commit's other half -- a JIT threadgroup-size assertion patch -- lives in the iOS app's vendored mlx-swift and does not apply to this Homebrew-MLX CLI build. Co-Authored-By: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ss-net memo) Squashes the model-load autotuner overhaul developed on ios-dev (c4f150f, ed7415a, 771756d, b12c736, 469cea8, 69d1fac, 9483fdd, fb422d2), restricted to the C++ backend (the iOS app/UI companion changes are dropped). What changes: - Fast coarse model-load tune. planShapeRotation takes a `full` flag: the per-load coarse tune now uses a 7-rep / 2-rep-floor budget and a trimmed 240-config grid (was 19/3 and 360), ~3.75x fewer GPU dispatches on a cache miss. `tuner -full` keeps the precise 19/3, wide grid. The documented broad plateau (geometry moves end-to-end <=1.5%) justifies the coarse budget. - Greedy coordinate-descent on the coarse path (useGreedy), backed by a new header-only GreedySearch::coordinateDescent core (neuralnet/greedysearch.h) with a standalone unit test (greedysearch_test.cpp, not wired into the katago build). A self-test gate asserts greedy stays within 5% of coarse-exhaustive. - Cross-context, session-scoped tune memo so the main and human SL nets (identical b18c384 3x3-conv shapes) tune once per session, not twice; cleared when the last ComputeContext is freed. - Separate cache files per mode: defaultFileName gains a `_full` suffix; the coarse "fast" tune keeps the legacy name so existing caches hit. - mlxTunerFull / mlxReTune are now read from -override-config in createComputeContext and fed to loadOrAutoTune (were hardcoded full=false, reTune=false). Cache format VERSION unchanged (3). - Diagnostics: tuner candidate-count line + MLX_TUNE_STUDY per-candidate dump; dropped the always-on [MLX-TUNE] sweep stderr line. Deliberately excluded (iOS-app only, no CLI value): the CoreML cache bridge wiring (686f823), the iOS extern-C katagocoreml shims (aebf3e0), and the iOS/visionOS set_cache_limit cap (97d3453). Co-Authored-By: Chin-Chang Yang <2770271+ChinChangYang@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Pure, behavior-preserving refactor of the model-load path the autotuner
and cross-net memo landed in. No change to inference, tuning results, or
the eval hot path.
- Fold the three file-scope memo globals (g_winoTuneMemoMutex,
g_winoTuneMemo, g_liveComputeContexts) into one WinogradTuneMemo unit
with tryGet/put/retain/release. The locking obligation and the
"clear on last context release" session-scope invariant now live in
one place instead of being open-coded across createComputeContext,
freeComputeContext, and the ComputeHandle ctor.
- Extract the ~70-line tuner orchestration block out of the
ComputeHandle ctor into a free resolveTuneParams(context, loadedModel,
useFP16): shape-key build, memo lookup, loadOrAutoTune, memo store. It
reads only ComputeContext + LoadedModel state, so it stands alone.
- Collapse the verbatim-duplicated "exactly one inference path" invariant
check (ANE early-return and GPU end) into a checkExactlyOnePath()
member.
Validation (Apple M3 Max, b18 uec vs eigen_reference_b18.json):
runtests + runnnlayertests pass; FP32 winrate max 0.00065% (bit-identical
to baseline); 2 GPU server threads reuse the memo across handles
("Reusing MLX Winograd tuning for shape ..._fast") with no deadlock and
no "encoder is already encoding".
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Pure, behavior-preserving refactor of the autotuner's sweep logging. No change to candidate enumeration, winner selection, tuning results, or the eval hot path -- only how the already-computed diagnostic line is built. flatSweepInput and flatSweepOutput each carried a ~40-line copy of the same tail: an optional skipped-count line, the per-shape median suffix loop, and the considered/best/baseline/delta_pct summary. The two copies differed only in the sweep label, the per-shape scorer, and whether the best-config body prints vw/gridOrder. - Extract renderPerShapeMs() for the " shape_ms=c<C>:<ms>,..." suffix. - Extract logFlatSweep() to own the skipped + summary lines, including the %+.1f delta_pct (sign-forced for the gated regex) and the best=none / delta_pct=nan degenerate branch. The regex-pinned log format (mlxtests.cpp) now has a single source of truth. - Each sweep keeps only its own best-config field rendering (input adds vw/gridOrder) and hands the rest off. perShapeStr is built on the same (best && baseline>=1e-9) condition logFlatSweep uses for delta_pct, so the nan branch stays in lockstep. Net -69/+67 lines. Validation (Apple M3 Max, b18 uec vs eigen_reference_b18.json): runtests + runnnlayertests pass, including the gated flatSweepInput / flatSweepOutput log-format regex checks (KATAGO_MLX_WINOTUNER_RUN_LOG_FORMAT_TEST=1). A forced retune (mlxReTune=true) emits byte-compatible sweep lines and FP32 winrate max stays 0.00065% (bit-identical to baseline). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tor#1208) The windows-latest runner image now ships Visual Studio 2026 only (actions/runner-images#14017), so configuring with -G "Visual Studio 17 2022" fails with "could not find any instance of Visual Studio". Drop the explicit generator and let CMake auto-detect the installed Visual Studio, keeping -A x64. https://claude.ai/code/session_018eaSTE1PhvV7SsiNyJrq76 Co-authored-by: Claude <noreply@anthropic.com>
Narrow mlxGpuEvalMutex to graph construction + async_eval (which encodes synchronously); move the completion wait and result readback outside the lock. With numNNServerThreadsPerModel=2, one server thread now encodes batch N+1 while the other waits on batch N, eliminating the ~1.4ms/batch GPU idle bubble. Single-server-thread configs are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Fold midBN+activation (after regularConv) and the residual add (after finalConv/postConv) into the Winograd output-untransform kernel for useMask=false, eliminating per-conv elementwise kernel launches and full-tensor round-trips. Residual fuses in T arithmetic (bit-identical); BN+act consumes the rounded-to-T conv value in fp32. Non-Winograd convs, masked configs, gpool blocks, and heads keep the unfused path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Apply fusedConvResidual to GlobalPoolingResidualBlock's input+finalOut residual, mirroring ResidualBlock/NestedBottleneckResidualBlock. Closes the one gpool special-case left unfused in 63446b6; every Winograd-conv->residual in the trunk now fuses uniformly. Bit-identical (T+T residual); FP32 stays 0.00065% vs Eigen. The gpool bias-add (broadcast) path stays unfused. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…nsform Add a BiasBNAct epilogue mode (BNAct kernel + one broadcast-bias-add line) and fuse GlobalPoolingResidualBlock's regularConv -> (+bias) -> midBN chain into the regularConv untransform, for useMask=false. The bias add is T+T (bit-identical to regularOut + bias); round-to-T then BN+act in fp32 matches midBN. Completes the gpool block: no elementwise kernel left between regularConv and finalConv. FP32 stays 0.00065% vs Eigen. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- applyCompiled result memcpys: cast the leading operand to size_t so the byte-count product is computed in size_t (the int sub-product could overflow on a large board with a big batch before the sizeof promotion). - MLXWinogradTuneParams::isValid(): bound each threadgroup dim to <=1024 before multiplying (a corrupt cache pair could overflow the int product and slip past the >1024 gate), and reject gridOrder values outside the defined enum so a corrupt cache re-tunes instead of running an unintended geometry. Both surface only on a hand-corrupted local tuner cache. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- KataGoParser: reject non-positive matmul in/out channels, qHeadDim/vHeadDim < 1, and ropeTheta <= 0, matching master desc.cpp. The ropeTheta check matters most: a non-positive theta yields NaN in the builder-derived cos/sin tables, which (being derived, not parsed) bypass the readFloats NaN/Inf gate and would otherwise produce a valid-but-garbage model. - Converter/MILBuilder: derive the serialized IO dtype from the builder's effective use_fp16_io (post narrow-transformer FP32 downgrade) via a new getUseFp16Io() getter, instead of the raw request. Fixes a spec/program mismatch where a narrow transformer with use_fp16_io=true emitted FP32 IO tensors while the model spec still declared FP16 IO. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…oud) - Session tune-memo key: include a signature of the 3x3-conv distribution (the histogram that actually drives planShapeRotation), so two nets with the same trunk width but different conv shapes no longer share a tune. modelVersion stays omitted so same-shape/different-version nets still share. - MLXWinogradTuneParams::isValid(): upper-bound wpt (<=8) and vw (<=4) from a corrupt cache, matching the existing tg<=1024 caps. - Greedy seed: guard the hard-coded seed indices (input + output sweeps) with a runtime check that they decode to the baked default, so a future reorder of the coarse value sets fails loudly instead of silently degrading tuning. - Policy-optimism postprocessor: throw on an unsupported numPolicyChannels instead of a release-elided assert that would silently mis-stride. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Both parseTransformerRMSNorm and parseRMSNormLayer validated only num_channels; master desc.cpp rejects epsilon <= 0 || > 1.0f. Add the same check so a malformed model fails loudly instead of building a valid-but-garbage model (rsqrt(x+eps) with a degenerate eps). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Match master desc.cpp: reject nested-bottleneck numBlocks < 1 (before parseBlockStack reserves on the count — a crafted model could otherwise force a multi-GB reserve from a few bytes) and matbias numChannels <= 0. Both fail loud on a malformed/corrupt model instead of throwing a confusing length_error or attempting a huge allocation. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Extract each backend's body into a file-static helper (runOpenCLTuner / runMLXTuner) gated by its own backend macro, leaving MainCmds::tuner as a short #if/#elif/#else dispatcher. Pure relocation — the bodies are byte-identical, no args/defaults/messages/control-flow change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
Author
|
Verified and validated by a real app on Apple devices. This PR will be set to ready when #1202 is merged. |
Upstream has merged PR lightvector#1202 (CoreML conversion memory levers) and PR lightvector#1205 (Metal GPU + CoreML/ANE transformer support), plus later polish. This branch carried its own cherry-picked copies of that work, so all 7 conflicts are in cpp/external/katagocoreml/ and are resolved by taking master's side wholesale: master is a strict superset there (+195/-561 in this branch's direction), and every branch-side validation has an equal-or-stronger, NaN-safe counterpart upstream (branch `x <= 0` vs master `!(x > 0) || !isfinite(x)`). The branch's getUseFp16Io() is master's getEffectiveUseFp16IO() under a different name, and the branch's ScopedFp32 generalization is ergonomics-only. The MLX backend uses the converter through its unchanged public API, so it needs nothing from these files. Conflicts resolved (all take-master): external/katagocoreml/src/Converter.cpp external/katagocoreml/src/builder/MILBuilder.{cpp,hpp} external/katagocoreml/src/builder/Operations.hpp external/katagocoreml/src/parser/KataGoParser.{cpp,hpp} external/katagocoreml/src/serializer/CoreMLSerializer.cpp Two defects the auto-merge produced silently, with no conflict markers: - neuralnet/desc.cpp: git kept BOTH this branch's and master's placement of BatchNormLayerDesc::releaseWeights() and MatMulLayerDesc::releaseWeights(), giving duplicate definitions that break every backend build. Took master's file (master already covers all transformer descriptors, incl. ropeFreqs). - CMakeLists.txt: master never touched this file since the merge base, so this branch's version was taken wholesale - including its deletion of set(CMAKE_OSX_DEPLOYMENT_TARGET 13.0) from the METAL branch. Restored the pin; leaving MLX unpinned stays deliberate (it must match libmlx.dylib). rungpuerrortest.sh keeps the gpu|ane mode + ${SUFFIX} convention, now extended to the 5 commented-out large-transformer entries master appended. Behavior inherited from master: the converter's FULL_FP32_MAX_TRUNK_CHANNELS moved 256 -> 320, so transformer trunks 256-319 channels wide now build fully FP32 on the ANE; attention mask constant 1e4 -> 3e4; ATTENTION_CORE_FP32 gate. Verified on this exact tree: MLX and METAL build clean; runtests and runnnlayertests pass; testgpuerror vs freshly regenerated Eigen FP32 references passes GPU 37/37 and ANE 31/31 supported configs (6 pre-v8 unsupported by the CoreML converter by design).
b10c384h6nbttflrs, b10c512h8nbt3tflrs-fson-silu-rsnh and
b11c768h12nbt3tflrs-fson-silu are published in the v1.17.1 release, so the
rationale for keeping these entries commented out ("not publicly
downloadable, place the model files in models/ manually") no longer holds.
Uncomment them and add the corresponding wget lines alongside the other
models the script fetches.
MODEL14 named b15c512h8nbttflrs-fson-silu, which is not in the release;
point it at b10c512h8nbt3tflrs-fson-silu-rsnh instead.
This takes the matrix from 37 to 42 configs. All five new entries keep the
-quick position subsample, which must match between reference generation
and checking.
Verified on an M3 Max at this commit, against Eigen FP32 references
generated from the same tree: GPU 5/5 pass (worst fp32 margin 0.00266x of
limit), ANE 5/5 pass (worst FP16 margin 0.51x of limit, on
b10c512h8nbt3tflrs-fson-silu-rsnh).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GyzjzGx3gXWH4gY6RNFgMd
Upstream v1.17.2 -> v1.18.1 (71 commits): new ROCm/HIP and ONNX Runtime
backends, major CUDA refactor, benchmark rework, and the rules-version-3
excludeTerritoryAdjacentToAtari feature.
Four textual conflicts:
* cpp/CMakeLists.txt - kept our pre-project MLX block (USE_BACKEND_NORMALIZED,
the CMake >= 3.27 guard, the deliberately-unpinned deployment target) alongside
upstream's ROCm helper functions, and added the MLX arm back to the backend
STRINGS property, the source-selection chain and the link chain next to
upstream's new ONNX and ROCM arms. The METAL arm keeps its
CMAKE_OSX_DEPLOYMENT_TARGET 13.0 pin.
* cpp/main.cpp, cpp/program/setup.cpp - MLX and ONNX arms both kept in the
backend ladders; version is upstream's 1.18.1, which ~90 test expectations embed.
* cpp/rungpuerrortest.sh - resolved as take-upstream, dropping the fork-only
gpu|ane mode, the ${SUFFIX} on result files, and the b10c512h8 entry. This file
conflicted on nearly every past merge and once silently lost ${SUFFIX}; carrying
no delta here removes that hazard for good. ANE runs are still available as
./rungpuerrortest.sh "deviceToUseThread0=100", which builds the identical
override; only the _ane result filenames are gone.
Also ported the MLX backend to upstream's new NeuralNet interface
(getRuntimeBackendDetail / getBatchPolicy / getNumEffectiveDevices). Without
this the MLX build fails to link, and nothing in the merge state flags it,
since mlxbackend.cpp exists only on this branch.
The existing GPU-error reference files remain valid and were not regenerated:
preferExcludeTerritoryAdjacentToAtari claims the old must-be-zero option-C slot
(no stream desync), and its featurization path is unreachable for the test
corpus, which is area-scored with TAX_NONE. Confirmed by regenerating three
references on both sides of the modelVersion >= 15 gate byte-identically.
Verified: MLX, METAL and EIGEN all build and link clean; runtests and
runnnlayertests pass on MLX and METAL; the full GPU sweep passes every entry.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014WzZtpDg9i5kGwMf5kPErb
…tion
Upstream v1.18.1 added its own pre-project `string(TOUPPER "${USE_BACKEND}"
USE_BACKEND)` (for the ROCm pre-project setup) immediately above the
`project()` dispatch, which is exactly what the fork's USE_BACKEND_NORMALIZED
existed to provide. Move the MLX version guard below that normalization and
dispatch on USE_BACKEND directly, so the fork carries no extra variable.
This is behavior-preserving. Everything between the guard's old position and
upstream's normalization is `function()` definitions and comments - no
executed statements - and all eight helpers defined there
(katago_default_hip_archs, katago_find_rocm_prefix, katago_win_*) are called
only under `USE_BACKEND STREQUAL "ROCM"`, so the relocated
`cmake_policy(VERSION 3.27)` cannot affect any reachable call. Verified by
configuring MLX, mlx (lowercase), METAL and EIGEN: lowercase still enables
Swift and trips the 3.27 guard, METAL keeps -mmacosx-version-min=13.0, MLX
stays deliberately unpinned, and all three real build dirs relink unchanged.
The point is merge surface: cpp/CMakeLists.txt lines 1..331 are now
byte-identical to upstream, so the fork's pre-project delta collapses from two
hunks - one of them pinned to line 1, where upstream inserted 462 lines last
cycle - into a single contiguous hunk adjacent to upstream's own
normalization comment.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014WzZtpDg9i5kGwMf5kPErb
greedysearch.h was a whole new header in neuralnet/ for an algorithm with exactly one inbound edge (mlxwinotuner.cpp), and greedysearch_test.cpp had zero: it was absent from CMakeLists.txt, unreachable by any glob (the source listing is fully explicit), and defined its own main(), so it could not have linked into katago even if listed. Its only documented build was a hand-typed clang++ line in its own header comment, i.e. it never ran. Move coordinateDescent into mlxwinotuner.cpp as a file-local static in the existing anonymous namespace, directly above its two callers, matching how the file already places pure helpers next to their use (renderPerShapeMs, logFlatSweep, buildConv3x3HistogramsFromConvs). The body is an exact move; only the result type name changes. Axis orders, seeds, drift guards and the sweep logic are untouched. GreedyResult is declared in mlxwinotuner.h so the test wrapper can name it, mirroring the existing ShapePlan/planShapeRotation split. Its `evaluated` field comment is corrected while moving: it is not just instrumentation, it becomes `considered` and feeds logFlatSweep and the consideredOut that the MLX_TUNE_STUDY greedy-vs-coarse-exhaustive accept check reads. Port the four unit cases into runMLXWinotunerTests() ungated. They are pure and sub-millisecond, so unlike the sweep tests they need no GPU and no env gate. This is a net coverage gain: the descent previously had no automatic coverage at all, since the only other exercise of it sits behind KATAGO_MLX_WINOTUNER_RUN_SWEEP_TEST. The trade-off is that the tests now run only in an MLX build, where before they were portable but never built. mlxwinotuner.cpp was relying on greedysearch.h for <cassert> (12 assert call sites) and <functional> transitively; both are now included explicitly. No CMakeLists change: headers are never listed, and the test file never was. Verified: MLX/Metal/Eigen all build clean; runtests and runnnlayertests pass; mutation-testing the descent loop makes the new tests fail as expected; the gated end-to-end sweep converges; katago tuner still reports considered=23, confirming the greedy path rather than a fallback to enumeration; testgpuerror on b10c384h6nbttflrs (19x19 and rect) passes at worst 0.193x of limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014WzZtpDg9i5kGwMf5kPErb
Compiling.md already covers building with -DUSE_BACKEND=MLX, but the docs that enumerate backends still omitted MLX. Adds it in the places where every other backend already appears. README.md follows the Metal precedent rather than the ROCm/ONNX one: both are macOS-only, and Metal is deliberately excluded from the "six backends" count, the table of contents, the section heading and the two "which backend should I use" bullet lists, appearing only in the macOS parenthetical. MLX is Apple-Silicon-only, so it is documented the same way. The tuner line is updated because this branch extends `katago tuner` to MLX (cpp/command/tune.cpp), which cpp/main.cpp's help text already reflects. cpp/README.md's source-tree listing and docs/ONNX_Model_Files.md's backend enumeration both already name Metal, so their omission of MLX was purely an oversight. While correcting the latter, ROCm is added too - it reads .bin.gz models only and was missing from that sentence as well. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014WzZtpDg9i5kGwMf5kPErb
ChinChangYang
changed the title
The Metal, OpenCL, ROCm and Eigen blocks each announce which build the user is running before the FP16 config line. The MLX block printed only mlxUseFP16, so a benchmark log gave no direct evidence of the backend. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014WzZtpDg9i5kGwMf5kPErb
ChinChangYang
marked this pull request as ready for review
August 28, 2026 01:27
Contributor
Author
|
#1202 is merged, so this is ready for review. #1205 and #1202 both landed upstream since this branch opened, so what remains here is only the MLX backend. The description has been rewritten from scratch and carries no branch history, so two notes for anyone who read the old one: validation was re-run on the current head at 48 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
USE_BACKEND=MLX, a neural-net backend for Apple Silicon on Apple's MLX framework, with two dispatch paths: GPU (MLX/Metal, F(2×2, 3×3) Winograd, per-shape tuner, FP16 by default) and ANE (CoreML on CPU + Neural Engine atgpuIdx = 100, sharing the Metal backend's converter) — usable separately or muxed on one model. Both run the v15+ transformer trunk (GQA attention, learnable RoPE, SiLU, RMSNorm tips, SwiGLU FFN) as well as convnets.19 files, purely additive, nothing under
cpp/external/katagocoreml/,desc.*ormetalbackend.*. New:mlxbackend.cpp,mlxwinograd.h,mlxwinotuner.{cpp,h},mlxtests.cpp;./katago tunergains MLX; the rest is build, config, doc and CI wiring.Build
cd cpp cmake -G Ninja -DUSE_BACKEND=MLX ninjaRequires CMake ≥ 3.27 and
brew install mlx.Validation
testgpuerroragainst an Eigen FP32 reference through upstream's unmodifiedcpp/rungpuerrortest.sh— 48 entries, 15 nets, batched and unbatched. MLX FP16 vs Eigen FP32, winrate error per net (99%-ile / max):GPU 48/48, ANE 42/42 supported configs; worst 0.52× (GPU) and 0.48× (ANE) of the enforced limit, FP32-vs-Eigen ≤ 0.0017%.
runtestsandrunnnlayertestspass.The five
numNNServerThreadsPerModel=2entries are GPU+ANE mux — onlydeviceToUseThread0is overridden, so thread 1 lands on the GPU — hence the ANE column is the 37 pure-ANE configs.ᵗ These three sit below the converter's
FULL_FP32_MAX_TRUNK_CHANNELS= 320, so their ANE graph builds fully FP32 and FP16 == FP32 — hence the ~1e-4% figures. The three wide nets (384 / 512 / 768 channels) are on the partial-FP32 tier that real transformer nets will use.ᶜ
b11c768h12nbt3tflrsnever ran on the ANE: CoreML fails to compile it for the Neural Engine and falls back to CPU, so that row is the CoreML CPU path. Metal's ANE path fails identically.The worst GPU FP16 max (2.49%,
b7c96h3tfrs) is intrinsic FP16 rounding, not an MLX defect — KataGo's OpenCL FP16 path agrees within 0.04 points on the same net, config and reference file:b7c96h3tfrs, 19×19, vs Eigen FP32)GPU throughput: MLX vs Metal
Apple M3 Max, both backends GPU-only and built from this commit: MLX's FP16 Winograd path against Metal's FP32-only MPSGraph path.
benchmark -half-batch-size -t 8,16,32 -v 800 -n <per-net> -boardsize 19,-nraised per net for ≥ 30 s cells, interleaved, median of 3.visits/s @ t=8 / 16 / 32
visits/s @ t=8 / 16 / 32
Faster on 14 of the 15 nets, 1.03×–1.46×, the margin growing with trunk depth;
b5c192nbt-v16testat 0.98× is within run-to-run overlap.Status
Ready for review — feedback wanted on the backend's structure, the tuner approach, the GPU/ANE dispatch wiring and the native MLX transformer trunk.