Skip to content

Add mixed-precision FP16 inference to the Metal backend - #3

Merged
frazane merged 1 commit into
mainfrom
codex/apple-silicon-metal-fp16
Oct 24, 2026
Merged

frazane merged 1 commit into
mainfrom
codex/apple-silicon-metal-fp16

Conversation

@ofuhrer

@ofuhrer ofuhrer commented Oct 23, 2026 •

Copy link
Copy Markdown
Collaborator

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 Add native Apple Silicon GPU inference with Metal and MPS #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. The macOS job executed FP32/FP16 GPU parity and all seven integration tests on an Apple Paravirtual device; GPU checks were not skipped.

@ofuhrer
ofuhrer requested review from frazane and a balanced review from Copilot October 23, 2026 20:23

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Mixed-precision GPU numerics rely on MPS FP16 GEMM accumulation behavior, buffer alignment, and hazard tracking that cannot be fully verified by static review, and full-checkpoint/real-data qualification is explicitly deferred.

No open issues 🎉

What changed in this PR

This PR adds a mixed-precision FP16 path to the existing Apple Silicon Metal/MPS backend, enabling --device metal --precision fp16 across inference, host rollout, and the ./varda forecast wrapper. It is stacked on #2 (which introduced the FP32 Metal backend) and carries only the incremental FP16 support. The precision policy mirrors CUDA/autocast: FP16 storage for weights and activations, but FP32 for GEMM results/accumulation, bias addition, LayerNorm statistics/affine, attention reductions/softmax, and all physical output operations. A key design constraint is that MPS is never asked for an all-half GEMM (its intermediate range can overflow before cancellation); instead, FP16 operands feed a GEMM that writes an FP32 scratch result, then linear_store_half rounds once after bias.

Changes:

  • Remove the Metal FP16 rejection (in model.cpp and the varda wrapper) and wire precision selection through the backend, selecting half vs. float storage via an Activation typedef and a per-tensor element_size.
  • Add FP16 machinery: padded half weight rows, reusable packing/scratch buffers, pack_half_from_float/pack_half_from_half/linear_store_half kernels, half-aware arenas, chunked FP16 GEMMs, and hook upcasting; all kernels gain FP32-no-op casts so the FP32 path stays bitwise-identical to #2.
  • Extend tests (synthetic FP16 runs, per-precision parity/chunk loops, a dense-cancellation regression, CPU-FP16 rejection, FP16 LayerNorm shader cases) and update docs/CI accordingly.
File Description
src/​varda/​metal/​backend_metal.mm Core FP16 support: element_size-aware Tensor, half_data, half arenas, padded-row linear loads, FP16 GEMM path with packing/FP32 scratch/store, hook upcasting.
src/​varda/​metal/​kernels.metal Activation typedef, pack/store-half kernels, and FP32-no-op casts across elementwise/attention/LayerNorm/residual kernels.
src/​varda/​model.cpp Removes the Metal FP32-only guard so the Metal backend accepts FP16.
varda Drops the wrapper's Metal FP16 rejection; updates the setup note.
python/​tests/​test_metal.py Adds FP16 to repeated inference/parity/rollout/empty-graph tests, a dense-cancellation regression, and switches the rejection test to CPU FP16.
tests/​metal/​test_metal_kernels.mm Runs the direct LayerNorm shader test for both FP32 and FP16 with a half-aware tolerance/guard.
tests/​run_synthetic.sh Adds Metal FP16 default/chunk-37 synthetic runs.
.github/​workflows/​cpu-portability.yml Globs all a.metal-fp*.log files on failure.
docs/​metal.md, docs/​apple-silicon.md, README.md, CMakeLists.txt Document and describe the new mixed-FP16 capability.

Review effort: Balanced

@frazane
frazane force-pushed the codex/apple-silicon-metal-fp16 branch from 5358459 to 55b1eef Compare October 24, 2026 11:03
@frazane
frazane changed the base branch from codex/apple-silicon-metal to main October 24, 2026 11:03
@frazane
frazane merged commit f4749c5 into main Oct 24, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants