Gyaradax: Local Gyrokinetics JAX Code

Motivation

• Known: Gyrokinetic simulation은 핵융합 플라즈마 난류 이해에 필수적이며, GKW, GENE 등 CPU 기반 legacy 코드가 수십 년간 표준으로 사용되어 왔다. 최근 CGYRO나 GENE의 GPU 지원 버전이 등장했지만 여전히 ML 워크플로우 통합에는 유연성이 부족하다.
• Gap: 기존 gyrokinetics 코드들은 복잡한 MPI 병렬화와 legacy Fortran 구조로 인해 유지보수가 어렵고 automatic differentiation이나 최적화·ML 워크플로우와 근본적으로 호환되지 않는다는 문제가 있다.
• Why: gyrokinetics 시뮬레이션에 automatic differentiation과 GPU 가속을 결합하면 inverse problem, sensitivity analysis, ML 기반 대체모델 개발 등 플라즈마 물리와 ML의 교차 연구가 가능해지므로 핵융합 연구 가속화에 중요한 인프라가 된다.
• Approach: GKW의 물리 모델을 기반으로 JAX로 재구현하면서 GPU native 가속, mixed precision, custom CUDA kernel(LTO callback을 통한 cuFFT 융합) 등을 도입하고, 이 전 과정을 coding agent를 활용한 vibecoding 방식으로 수행하였다.

Achievement

1. 경량화된 코드베이스: 약 30,000줄의 Fortran 코드(GKW)를 기능적으로 동등한 약 3,000줄의 JAX 코드로 축소하여 구현하였다.
2. 검증된 정확성: Rosenbluth & Hinton zonal flow test, Cyclone Base Case 등 analytical case와 GKW 참조 데이터에 대한 광범위한 empirical benchmark를 통해 formal agreement 및 statistical parity를 입증하였다.
3. 성능 향상: mixed precision과 custom CUDA kernel(Z2Z packing, LTO callback 기반 cuFFT 융합)을 통해 기본 JAX 대비 약 2배의 추가 속도 향상을 달성하였으며, GKW 대비 substantial speedup을 보였다.
4. AI-driven 과학 소프트웨어 개발 사례 제시: unit testing을 통한 측정 가능한 진행상황과 구조화된 프롬프팅을 결합해 coding agent가 복잡한 Fortran 코드를 빠르게 JAX로 번역하고 최적화할 수 있음을 보였다.
5. ML-plasma physics 교차 연구 예시: differentiable programming을 활용해 R/LT gradient 기반 복원과 growth rate에 대한 sensitivity analysis 예제를 시연하였다.

How

• Vlasov-Poisson 방정식을 δf 형태로 분해하여 parallel advection, magnetic drift, mirror term, equilibrium drive, field drive, nonlinear E×B term, dissipation term 등 물리적 구성요소를 명시적으로 구현
• 시간 적분은 explicit 4th-order Runge-Kutta(RK4)를 사용하며 jax.lax.scan으로 timestep을 fuse하여 Python interpreter overhead 감소
• 평행/속도 방향은 4th-order central/upwinded stencil, 수직 방향은 pseudospectral method(3/2-rule dealiasing)로 이산화
• species-dependent 계수(Bessel function, Maxwellian, drift velocity 등)를 time loop 이전에 precompute하여 RK4의 4회 RHS 평가에서 중복 연산 제거
• mixed precision(Float32/Float64 혼합)으로 FFT 및 nonlinear Poisson bracket 계산의 메모리 대역폭 절감
• Z2Z packing으로 inverse FFT 횟수를 4회에서 2회로 축소하고, Hermitian symmetrization 보정을 적용
• custom CUDA kernel과 Link-Time Optimization(LTO) callback을 사용해 cuFFT 호출 전후 연산을 butterfly pass에 직접 임베딩하여 HBM round-trip 제거
• Rosenbluth & Hinton zonal flow test, Cyclone Base Case(Dimits et al., 2000)로 formal verification 수행
• GKW 참조 데이터와 time-averaged transport flux, wavenumber spectrum, growth rate 등 통계량을 다양한 equilibrium configuration(adiabatic/kinetic electron)에서 비교
• differentiable programming을 활용한 R/LT gradient 기반 recovery 및 growth rate sensitivity analysis 실험 수행
• 개발 과정에서 coding agent(vibecoding)를 인간 전문가의 감독, 구조화된 프롬프팅, unit testing 기반 측정 가능한 진행 지표와 결합하여 활용

Originality

• legacy Fortran 기반 gyrokinetics 코드(GKW)를 JAX/CUDA로 완전히 재구현하여 automatic differentiation과 GPU native 가속을 동시에 지원하는 최초의 사례 중 하나로 제시됨
• custom CUDA kernel과 LTO callback을 활용해 XLA의 일반적 stencil/FFT lowering 한계를 넘어서는 fine-grained 성능 최적화(Z2Z packing, single-pass fused stencil)를 도입
• 과학 컴퓨팅 코드베이스 개발에 coding agent를 인간 전문가·구조화된 프롬프팅·unit testing과 결합해 대규모 Fortran→JAX 번역 및 최적화를 수행한 방법론적 사례를 제시
• inverse problem(R/LT recovery)과 sensitivity analysis 등 differentiable gyrokinetics의 실용적 활용 예시를 시연

Limitation & Further Study

• 현재 electrostatic collisionless plasma만 지원하며 electromagnetic perturbation, collisionality, Coriolis effect(Term VI), global gyrokinetics는 미지원
• 시간 적분이 explicit RK4로 제한되어 fast kinetic species에서 CFL 제약으로 인해 속도가 느려질 수 있으며, GKW가 제공하는 implicit/semi-implicit integrator가 없음
• 기하학적으로 analytic circular(Lapillonne)와 s-alpha geometry만 구현되어 있고 Miller parameterization이나 EFIT/CHEASE 기반 numerical MHD equilibria를 요구하는 shaped plasma는 미지원
• dense array 기반 5D domain 처리로 GKW의 sparse format 대비 메모리 효율이 떨어질 수 있으며, grid parallelism(MPI 기반 domain decomposition) 미지원으로 대규모 확장성에 제한
• 후속 연구로 electromagnetic perturbation, collisional effect, global gyrokinetics, 병렬화 지원 확장이 필요

Evaluation

총평: gyrokinetics 시뮬레이션을 JAX 생태계로 성공적으로 이식하여 GPU 가속과 automatic differentiation을 결합한 실용적이고 검증된 도구를 제시했으며, coding agent를 활용한 과학 소프트웨어 개발 방법론 자체도 흥미로운 기여로 평가된다. 다만 electromagnetic effect, collisionality, global geometry 등 물리적 커버리지의 한계는 향후 확장이 필요한 부분이다.