Accelerating high-fidelity gyrokinetic fusion simulations with iGENE

Abstract

We present the latest version of iGENE, a differentiable implementation of the local (flux-tube) gyrokinetic code GENE, now fully rewritten in JAX. iGENE features a multispecies, collisional, fully electromagnetic nonlinear gyrokinetic model while preserving end-to-end differentiability. The code is verified against GENE for linear ITG, TEM and KBM physics and for nonlinear turbulent fluxes with adiabatic and kinetic electrons. Forward-mode derivatives of the linear growth rate agree with finite differences for all five parameters tested, resolving the 30–50% discrepancy in $\hat{s}$ and $ε$ reported for the earlier TensorFlow version. All collectives are native JAX operations, and the code strong-scales from 2 to 32 H100 GPUs with 86% parallel efficiency. We then assess three routes to a shorter time to solution. Custom CUDA FFT kernels operating in bfloat16 are up to $3.97\times$ faster than fp32 cuFFT on the transforms, and on a single H100 iGENE is $2.63\times$ faster per time step than single-precision GENE. Multi-rate sub-cycling of the parallel dynamics yields a further $1.89\times$ speedup with kinetic electrons. Finally, a Hermite-Laguerre velocity-space representation is implemented, offering tunable fidelity.