Repository navigation
Add native Apple Silicon GPU inference with Metal and MPS - #2
Conversation
There was a problem hiding this comment.
🔵 Needs a closer look
It introduces a large new hand-written Metal/Objective-C++ GPU backend with subtle fp32 numerics that cannot be compiled or executed in this review environment, so correctness depends on hardware validation that warrants human sign-off.
No open issues 🎉
What changed in this PR
This PR adds an optional native Apple Silicon GPU backend (--device metal --precision fp32) that runs the anemoi GraphTransformer forecaster on a unified-memory Metal GPU, using MPSMatrixMultiplication for dense GEMMs and hand-written Metal compute shaders for attention, normalisation, GELU, residual and bounding. It builds on the Apple Silicon CPU support (#1), depends on no PyTorch, and has no silent CPU fallback. The new backend mirrors the CPU fp32 oracle's operation ordering (pairwise dot accumulation, separate mul/add, destination-row chunking) and is verified against the existing synthetic fp32 tolerances rather than a bitwise CPU/CUDA match.
Changes:
- New Metal/MPS backend:
kernels.metal,backend_metal.mm,metal_backend.h, embedded shader source, and thevarda-metal-infoprobe app, wired throughDevice::Metal,make_metal_backend, andmodel.cpp. - Build integration:
VARDA_METALCMake option (defaults to${APPLE}, Objective‑C++/ARC, Foundation/Metal/MPS frameworks, configure-time shader embedding) andVARDA_WITH_METALguards. - CLI/test/docs/CI:
metalaccepted byvarda,varda-run/rollout/parity/forecastandcommon.h; synthetic suite +test_metal.pyintegration tests; macOS CI step; newdocs/metal.mdand doc updates.
| File | Description |
|---|---|
| src/varda/metal/kernels.metal | Metal compute kernels for attention, norm, GELU, residual, bounding, (de)normalise |
| src/varda/metal/backend_metal.mm | Metal backend: buffer arena, MPS GEMMs, dispatch, chunked forward pass |
| src/varda/metal/metal_backend.h | Device availability/info API |
| src/varda/metal/kernels_source.h.in | Configure-time shader-source embedding template |
| src/varda/model.cpp / varda.h / backend.h | Device::Metal, to_string, backend selection, factory decl |
| apps/metal/varda-metal-info.cpp | GPU probe (exit 77 when no GPU) |
| apps/common.h / varda-run/rollout/parity/io/varda-forecast | metal device parsing and usage strings |
| CMakeLists.txt | VARDA_METAL option, OBJCXX/ARC, frameworks, shader embedding, app wiring |
| varda / tests/run_synthetic.sh / python/tests/test_metal.py / .github/workflows/cpu-portability.yml | Setup/forecast plumbing, synthetic Metal runs, integration tests, macOS CI |
| README.md / docs/metal.md / docs/apple-silicon.md / CLAUDE.md | Documentation of the Metal backend and build |
I traced the Metal kernels against the CPU oracle (attention scores/softmax/values, edge projection, assemble, residual, bounding stages, (de)normalise) and the chunk-local vs. absolute buffer indexing, and found them consistent, including empty-graph/zero-in-degree handling, the compensated fp32 LayerNorm, and the 32-bit-index/working-set guards. I did not identify a concrete functional defect to comment on.
Review effort: Balanced
There was a problem hiding this comment.
🔵 Needs a closer look
It adds a large new Metal/MPS GPU backend whose shaders, Objective-C++ and build wiring cannot be compiled or executed in this environment, so its correctness depends on GPU-only validation that requires human review.
No open issues 🎉
Review effort: Balanced
Adds `--device metal --precision fp16` for inference, host rollout and the forecast wrapper. Metal previously rejected FP16. This PR is stacked on #2 and targets `codex/apple-silicon-metal`; it contains only the additional mixed-precision support. The precision policy follows the CUDA/autocast approach: FP16 weights and activation storage, FP32 GEMM results/accumulation and bias addition, FP32 LayerNorm statistics/affine and attention reductions/softmax, and FP32 physical output operations. Each dense result is rounded once to FP16 after bias. Edge projection uses half operands, FP32 sums and an explicit half projection rounding. This is not a bitwise CUDA implementation. MPS never receives an all-half GEMM. Padded weights and reusable input packing handle odd widths and unaligned chunks; compatible half activations are used directly. FP32 result scratch and embedding chunks keep temporary storage bounded. Public input/output remain FP32, and hooks upcast the same stored production values after synchronization. FP32 retains its existing path; CPU stays FP32-only and CUDA is unchanged. Validation on Apple M1, without changing any tolerance files: - 24 standard GPU parity runs: six models, FP32/FP16, default/37-row chunks. Worst FP16 activation is `proc.3.out` in e: rel_absmax 1.395e-3, max_abs 2.488e-2 against the existing 3e-3 tolerance. Input stages are bitwise exact, and hooks are inert in both modes. - 20 additional independent model checks cover 7/35/96 channels, one hidden node, chunks down to one row, and 8,201 nodes crossing the default boundary. Both precisions pass existing tolerances. - A real-backend dense cancellation regression (default/37-row chunks) checks products of +/-65536 summing to exactly zero; an all-half accumulation would overflow. - Seven integration tests pass, including repeated inference, four-step host rollout with computed forcings compared against CPU at existing FP32/FP16 tolerances, empty graphs, and invalid chunk / CPU FP16 rejection. - Direct production LayerNorm shaders pass 768 cases (384 per precision), including near-constant rows, optional bias and output guards. Worst FP16 rel_absmax 4.81e-4, max_abs 1.95e-3 at 3e-3 tolerance. - 9/9 Metal-enabled unit tests, 8/8 Metal-disabled unit tests and 18/18 CPU synthetic checks pass. All six CPU and all six Metal FP32 outputs are bitwise identical to #2. Ruff, actionlint and ShellCheck pass. The full local GPU suite passes with Metal API and shader validation enabled. Hosted CI retains API validation for MPS and shader validation for the direct kernel tests, following #2's documented Paravirtual-GPU instrumentation limitation. Two alternating 20-call FP32/FP16 benchmark rounds on M1 (first three discarded) measured warm mean times of 27.1 -> 23.0 ms for synthetic e and 32.6 -> 27.8 ms for d, about 1.18x faster. Repeats are bitwise stable. These small-model timings do not establish full-Varda performance. Half storage reduces selected buffers, not the entire footprint. Full-checkpoint and real-data forecast/rollout qualification still require a larger-memory Mac and real reference assets. CUDA execution was not retested. No Metal graph capture or GPU-resident rollout window is added. [Linux x86_64 and macOS arm64 CI pass at 5358459](https://github.com/frazane/varda.cpp/actions/runs/40018539850). The macOS job executed FP32/FP16 GPU parity and all seven integration tests on an Apple Paravirtual device; GPU checks were not skipped.
Adds native Apple Silicon GPU inference through
--device metal --precision fp32, using MPS matrix multiplication and Metal compute shaders for graph attention, normalization, GELU, residuals and bounding. Builds on the Apple Silicon CPU support merged in #1 and targetsmain.The backend uses shared reusable buffers and destination-row chunking. LayerNorm cooperates across a SIMD group, queries its width at runtime, and computes centered, compensated FP32 moments to avoid cancellation in nearly constant rows with large offsets. Command buffers batch up to 64 encoded operations, preserving completion waits before host reads and propagating GPU errors. Shader source is embedded in the executable; no PyTorch dependency or CPU compute fallback.
Host rollout, the GRIB forecast CLI and the setup wrapper accept Metal.
VARDA_METAL=OFFremoves Apple GPU framework dependencies; Linux defaults to OFF and never enables Objective-C++. Existing CPU/CUDA numerical kernels and tolerance files are unchanged. FP16 remains explicitly unsupported.Validation on Apple M1:
proc.1.outin model d:rel_absmax2.72e-6,max_abs4.18e-5, relative tolerance 1e-5. Hooked and hook-free outputs are bitwise equal.rel_absmax3.39e-7,max_abs1.62e-6 at tolerance 1e-5. The original shader fails 90 cases; the revised shader passes all.Two alternating 30-call FP32 microbenchmark rounds on M1, discarding the first three calls, measured warm mean forward times of 40.8 -> 25.9 ms (synthetic e) and 53.2 -> 31.1 ms (synthetic d) after parallel normalization and bounded batching. Repeats are bitwise stable. These synthetic timings do not establish full-Varda speed.
Limits: FP32 only, host-driven rollout, no Metal graph capture. The full checkpoint and real-data GRIB forecasts still need qualification on a Mac with more unified memory and real reference assets; the 16 GB test Mac cannot hold the full model. CUDA execution was not retested. GPU-less CI runners explicitly report GPU execution as skipped.
Linux x86_64 and macOS arm64 CI passed at 317c63e. The macOS job executed GPU checks on an Apple Paravirtual device.