Skip to content

VMEX 0.8.1 — the cold-start performance release

Choose a tag to compare

@rogeriojorge rogeriojorge released this 02 Sep 04:11
· 529 commits to main since this release
9508c7a

The regression users reported against 0.7/0.8 — slow cold CLI, python, and optimization runs — is eliminated and inverted: v0.8.1 is faster cold than every previous VMEX release, including v0.3.0, on every deck measured, on both x86 and arm64.

Measured (cold CLI, full iteration budgets, all caches cleared): 36-core x86 — input.LandremanPaul2021_QA_lowres 43.1 s (v0.3.0) → 32.4 s, input.solovev 12.1 → 7.0 s; Apple M4 — QA 16.4–18.0 s (v0.3.0) / 22.3–23.3 s (v0.8.0) → 12.3–12.6 s, li383 → 4.0 s, solovev → 3.1 s. Warm reruns at parity or better, and the persistent compilation cache makes every rerun far cheaper than v0.3.0's flat cost. At QA resolution the cold start compiles 343 XLA programs (v0.8.0: 773; v0.3.0: 523).

What changed

  • The eager setup, WOUT-export, stage-interpolation, and printout passes each dispatched hundreds of single-op XLA programs per run; they are now module-level jitted lanes (#227, #230).
  • The iteration body traced the funct3d chain twice; it now traces once, halving every solver lane's HLO and compile time, and the ns4 preconditioner refresh runs only on its VMEC2000 cadence instead of every iteration (#229).
  • Print cadence, initial DELT, and ftol no longer key lane recompilation; host print loops and device transfers are gated (#231).
  • Non-finite iterations are detected from the residual scalars with full classification on trip (#232).
  • 3-D force-balance polishing no longer stalls in XLA constant folding: the Ruiz/probe jits stopped baking linearization residuals (~4.5 GB at QA scale) as constants; a small 3-D polish runs 2× faster and production-scale polish now executes end to end (#234).
  • CI pins compile budgets (lane HLO size, cold-solve program count) so this regression class fails loudly in review, and opt-in measured-memory auto chunk sizes land for the implicit multi-RHS paths (#233).
  • Guard repairs and CI hardening (#226, #228, #236, #238).

Numerical statement: per-iteration physics is unchanged; graph restructuring shifts XLA fusion, so trajectories can differ from v0.8.0 at 1 ULP per iteration with identical iteration counts and converged geometry agreeing at 1e-12. All golden, restart, step-control, parity, and end-to-end suites pass.