perf(mta): GPU-resident metatomic force provider

[HaoZeke]

Summary

• Pure-ML GPU path (MetatomicGpuForceProvider) bypassing IForceProvider for zero-copy coordinate wrapping, GPU NL construction from nbnxm GpuPairlist, and GPU-resident model evaluation
• PME GPU integration pattern: DeviceBuffer<RVec> + GpuEventSynchronizer registered with GpuForceReduction
• CUDA kernel converts nbnxm GpuPairlist to metatensor NL format on device; no D2H copies in single-rank hot path
• DD supported: GPU model eval + CPU sparse MPI force exchange (same pattern as CPU MetatomicForceProvider)
• New nbnxm accessors gpuGetPairlist() and gpuGetAtomIndex() for GPU pairlist access
• Activation via metatomic-gpu MDP section or GMX_METATOMIC_GPU env var; CUDA build required (stub for CPU-only)

Single-rank GPU-resident only; no multi-GPU

The zero-copy force output via GpuForceReduction is gated on !hasDomainDecomposition(). With DD, forces go to CPU for sparse MPI exchange through distributeNonHomeForces(). No multi-GPU support exists: LAMMPS Kokkos handles this via MPI_Comm_split_type round-robin GPU assignment, but the metatomic GPU path just uses whatever device GROMACS assigns through DeviceStreamManager. Multi-rank multi-GPU is untested.

Not upstream-ready

The GPU force provider bypasses IForceProvider and wires directly into do_force() in sim_util.cpp (coordinate sync, atom mapping, NL build, model evaluation, output application) and runner.cpp (provider construction, model path resolution from KVT, topology iteration). It also adds a metatomicGpu member to t_forcerec, a useGpuMetatomic flag to SimulationWork, and nbnxm accessors that expose internal GPU data structures.

Upstreaming would require a proper GPU force provider interface (analogous to PME GPU integration) rather than the current direct wiring. Either IForceProvider needs GPU buffer support, or a separate IGpuForceProvider interface with DeviceBuffer input/output.

Phase 1 audit (latest commit)

• Device validation: GMX_RELEASE_ASSERT that DeviceContext matches cudaGetDevice()
• Coordinate lifetime: debug-mode checksum of d_x before model eval, verified unchanged after backward() (fires with GMX_METATOMIC_DEBUG=1)
• Interaction range: GPU module registers interaction_range + max_nl_cutoff (prerequisite for Newton-mode DD)
• Stub fix: CPU builds failed to link because the stub was missing the xReadyOnDevice parameter
• Parity test: TestGPUCPUParity asserts GPU/CPU step-0 energy within 1 kJ/mol

Test plan

• CPU build (metatomic-cpu env) compiles with stub
• Existing CPU metatomic tests pass
• CI benchmark workflow fixed (stub signature mismatch was the CPU link failure)
• CUDA build (metatomic-cuda env) on remote workstation
• Single-rank GPU energy matches CPU reference (TestGPUCPUParity)
• DD GPU energy matches serial GPU reference
• nsys profile confirms no GPU-CPU copies in single-rank hot path

[HaoZeke]

Benchmark: 4-Way Metatomic Scaling (CPU + GPU vs LAMMPS + LAMMPS/kk)

Hardware: cosmopc5 -- RTX 4070 Ti SUPER, Threadripper PRO 3955WX
Model: PET-MAD v1.0.2 | Steps: 500 | Threading: -ntmpi 1 -ntomp 4

Results (ms/step, lower is better)

Atoms	GMX CPU	GMX GPU	LAMMPS	LAMMPS/kk
648	68.3	69.4	86.0	80.0
1536	149.1	144.7	196.0	180.0
3000	272.1	260.3	354.0	328.0

Speedup

Atoms	GMX CPU vs LAMMPS	GMX GPU vs LAMMPS/kk	GPU vs CPU
648	21% faster	13% faster	2% slower
1536	24% faster	20% faster	3% faster
3000	23% faster	21% faster	4% faster

Energy parity (step 0)

Engine	648 atoms (kJ/mol)	3000 atoms (kJ/mol)
GMX CPU	-309,456	-1,407,420
GMX GPU	-309,109	--
LAMMPS	-309,022	-1,407,437
LAMMPS/kk	-309,022	-1,407,437

All within 0.14%. The difference is from coordinate precision: GROMACS uses float32 (-DGMX_DOUBLE=OFF), LAMMPS metal units use float64. The GMX GPU path has an additional source of difference from the CUDA NL kernel using float32 positions (nbnxm d_xq stores xyz as float).

Analysis

GROMACS vs LAMMPS (20-24% faster): Both call the same metatensor C++ API and PyTorch model. The gap comes from atom representation -- GROMACS uses shift-vector NL (648 atoms) vs LAMMPS explicit ghost atoms (~6700 atoms for 648 local). PET-MAD's O(N^2) transformer attention makes this 10x atom count difference significant.

GPU-resident vs CPU IForceProvider (1-4% faster): The CPU pairlist mode already reuses the GROMACS NL via MDModulesPairlistConstructedSignal. The GPU path only saves the D2H coordinate copy and H2D force copy around model evaluation. At these sizes that is ~5-10 ms vs ~260 ms model evaluation. The benefit grows with system size (the GPU advantage increases from -2% to +4% going from 648 to 3000 atoms).

Reproduction

cd mta_bench/water_scaling && python prepare.py

# All 4 engines:
python3 run_benchmarks.py --engines gmx_cpu,gmx_gpu,lmp,lmp_kk --sizes 216,512,1000

# Or individually:
# GMX CPU
$GMX_BIN grompp -f gmx_cpu.mdp -c conf.gro -p topol.top -o topol.tpr -maxwarn 10
$GMX_BIN mdrun -s topol.tpr -ntomp 4 -ntmpi 1 -g md.log -pin off

# GMX GPU-resident
$GMX_BIN grompp -f gmx_gpu.mdp -c conf.gro -p topol.top -o topol.tpr -maxwarn 10
$GMX_BIN mdrun -s topol.tpr -ntomp 4 -ntmpi 1 -g md.log -pin off -nb gpu -pme cpu

[HaoZeke]

Some changes since I ran these but.. single GPU numbers with PET seem good