Enable compiler-managed CUDA graphs for UMA inference

[mlazos]

Summary

This adds an opt-in compile_mode="reduce-overhead" inference setting so
torch.compile can use its built-in CUDA graph support for UMA.

Internal NVIDIA v3 neighbor generation remains outside the captured model
regions. Before the model runs, its variable-length edge list is padded to a
configurable bucket. Padded edges are moved beyond the cutoff, so they
contribute zero energy, force, and stress while keeping the compiled tensor
shape stable across nearby molecular-dynamics frames.

This uses released nvalchemiops 0.4.0; it does not require a patched NVIDIA
library or a custom outer CUDA graph runner. It depends on #2154 for the UMA
fast-GPU kernel baseline.

Performance

Measured on one H100 with:

• UMA-S-1p2, 1,000 Cu atoms
• energy, forces, and stress
• FP32 model state with TF32 enabled
• activation_checkpointing=False
• merge_mole=True
• compile=True, dynamic=True
• external_graph_gen=False
• internal_graph_gen_version=3
• execution_mode="umas_fast_gpu"
• edge bucket size 1,024
• TunableOp disabled
• PR Optimize UMA SO2 and Wigner inference kernels #2154 as the baseline

Both predictors were compiled and warmed in the same process. Ten alternating
20-step blocks used identical position sequences.

Variant	Median QPS	Mean QPS	Median latency
Dynamic compile baseline	25.5091	25.8275	39.202 ms
reduce-overhead	27.3169	27.2558	36.607 ms

This is a 7.09% median QPS improvement and a 5.53% mean QPS improvement. There
were no timed recompiles. Inductor reported 29 recorded CUDA graph regions.

The measured frame had 78,424 real edges and was padded to 78,848, adding 424
masked edges (0.54%). Separate-process steady reserved memory was 9.285 GiB for
the baseline and 7.533-7.660 GiB with reduce-overhead.

Numerical checks

At the same displaced 1,000-atom frame, maximum absolute differences versus
the dynamic compile baseline were:

• energy: 0.0022302
• forces: 0.0004342
• stress: 0.000001186

NVIDIA v3 returns the same edge multiset in nondeterministic order. A
same-order padded-versus-unpadded check stayed within the variation seen when
repeating the unpadded path. Direct eager checks on single- and two-system
batches were also within 5.96e-8 for energy, 8.38e-9 for forces, and 1.64e-11
for stress.

The added CUDA regression uses the real UMA-S-1p2 checkpoint. It compares
eager and reduce-overhead energy, forces, and stress across repeated calls,
retains an earlier output to detect replay-buffer overwrites, and crosses to a
different edge bucket before returning to the original bucket.

Limitations

• Internal reduce-overhead currently requires CUDA and NVIDIA graph version
  3.
• Graph-parallel inference is rejected for this path.
• Neighbor construction itself is not captured because the released NVIDIA
  implementation allocates temporary buffers.
• Crossing an edge bucket records another compiler-managed CUDA graph. Returning
  to a previous bucket reuses its existing compiled graph.

Test plan

PYTHONPATH="$PWD/src" pytest -q -s tests/core/graph/test_padded_nvidia_graph.py tests/core/models/uma/test_padded_edges.py tests/core/units/mlip_unit/test_inference_settings.py
PYTHONPATH="$PWD/src" pytest -q -s tests/core/units/mlip_unit/test_predict.py::test_reduce_overhead_internal_graph_predict
pre-commit run --files docs/core/common_tasks/ase_calculator.md src/fairchem/core/graph/padded_nvidia_graph.py src/fairchem/core/models/uma/escn_md.py src/fairchem/core/units/mlip_unit/api/inference.py src/fairchem/core/units/mlip_unit/predict.py tests/core/graph/test_padded_nvidia_graph.py tests/core/models/uma/test_padded_edges.py tests/core/units/mlip_unit/test_inference_settings.py tests/core/units/mlip_unit/test_predict.py

[github-actions]

This PR has been marked as stale because it has been open for 30 days with no activity.