Skip to content

A Posterior on Any GPU

Latest
Latest

Choose a tag to compare

View all tags
@borodark borodark released this 23 Jul 12:35
· 1 commit to main since this release
v0.3.0
e511720

A Posterior on Any GPU

eXMC 0.3.0 · Read the full story

This release lands the Vulkan (nx_vulkan) f64 custom-synthesis GPU backend and
brings eXMC onto Nx 0.13, with EMLX dropped. NUTS now runs on the GPU via
nx_vulkan's VulkanoBackend (fused f64 leapfrog chain shaders) in addition to EXLA —
select with EXMC_COMPILER=vulkan. Works anywhere Vulkan does, including FreeBSD/MoltenVK
where CUDA and Metal don't exist.

Highlights

  • Vulkan f64 chain-shader sampling. Multi-RV models synthesize a fused f64 leapfrog
    chain shader (CustomSynth) and dispatch on the GPU; non-synthesizable models fall
    back to per-op or the Plan B′ guard.
  • Nx 0.13. Upgraded nx/exla to 0.13; nx_vulkan pinned to a Nx-0.13-compatible
    commit so a fresh mix deps.get resolves.
  • EMLX dropped. The f32-only Apple Metal path is retired; the default precision is now
    f64 end-to-end.

Fixes

  • Gamma/Beta priors now route through the synth chain path (added
    Push.prior_param_floats/1 encoders) instead of erroring.
  • Measurable-matmul under Vulkan: jit_solve/jit_determinant pinned to
    BinaryBackend so the LU host-fallback's index tensors no longer leak into
    Nx.BinaryBackend.slice/to_binary and crash under nx 0.13.
  • push_too_large graceful fallback: models whose priors exceed the 128-byte f64
    push-constants block now degrade to per-op sampling (with a warning) instead of crashing
    at dispatch.
  • f64 @data_sentinel: the no-observation sentinel is f64, matching the post-EMLX
    default precision.

Validation

  • posteriordb: 33/33 PASS on EXLA-GPU (external Stan gold-standard reference
    posteriors) — no inference regression.
  • Vulkan push_too_large fallback verified on FreeBSD/MoltenVK (mac-247):
    eight_schools_noncentered goes from a hard crash to a graceful per-op run.
  • Full vulkan suite on the FreeBSD-NVIDIA reference: 755 tests, residual failures confined
    to known synthesis-coverage / guard buckets.
  • Clean-room mix deps.get (fresh clone, no path overrides) resolves nx 0.13 + nx_vulkan.

Dependency

Requires nx_vulkan main (Nx-0.13-compatible). Consumers using EXMC_COMPILER=vulkan
need a Vulkan driver + glslangValidator.

Assets 2
Loading