Enable compiler-managed CUDA graphs for UMA inference

Torch compile can capture UMA compute regions, but internal neighbor generation changes the exact edge shape from one frame to the next. Those shape changes prevent reliable CUDA graph reuse even when the atom count is fixed.

Expose reduce-overhead as an inference compile mode and mark prediction boundaries for CUDA graph trees. For internal NVIDIA v3 graphs, generate the neighbor list before model execution and pad its edges to a configurable bucket. Padded edges are moved beyond the cutoff so they contribute zero energy, force, and stress while keeping the compiled input shape stable. Neighbor construction remains outside capture because the released NVIDIA implementation allocates temporary buffers.

The internal path is limited to CUDA, graph version 3, and non-distributed inference. External graphs can use reduce-overhead without padding.

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
```

Author: mlazos

docs/core/common_tasks/ase_calculator.md

@@ -92,8 +92,10 @@ The advanced user might quickly see that **default**, **batch**, and **turbo** m
 | activation_checkpointing | this uses a custom chunked activation checkpointing algorithm and allows significant savings in memory for a small inference speed penalty. If you are predicting on systems >1000 atoms, we recommend keeping this on. However, if you want the absolute fastest inference possible for small systems, you can turn this off |
 | merge_mole | This is useful in long rollout applications where the system composition stays constant. By pre-merge the MoLE weights, we can save both memory and compute. |
 | compile | This uses torch.compile to significantly speed up computation. Due to the way pytorch traces the internal graph, it requires a long compile time during the first iteration and can even recompile anytime it detected a significant change in input dimensions. It is not recommended if you are computing frequently on very different atomic systems. |
+| compile_mode | Set to `"reduce-overhead"` to use compiler-managed CUDA graphs. With internal graph generation this requires version 3, a CUDA device, and no graph parallelism. It is intended for repeated evaluations of similarly sized systems. |
 | external_graph_gen | Only use this if you want to use an external graph generator. This should be rarely used except for development |
 | internal_graph_gen_version | currently we support v2[default], an internal implementation that is better suited for parallelism and v3 the neighborlist from Nvidia Alchemi library which is faster for single gpu operations. |
+| internal_graph_edge_bucket_size | Number of edges per padded bucket when `compile_mode="reduce-overhead"` uses internal graph version 3. Larger buckets reduce shape changes but perform more work on masked edges. Default 1024. |
 | edge_chunk_size | Experimental. Used for padding edge sizes. This helps reduce re-compilations from torch compile, default to None |
 | use_quaternion_wigner | enable quaternion-based Wigner D matrix computation. If false we fall back to euler-angle based rotations. default True. |
 | base_precision_dtype | governs the main precision type of the computation, default to FP32, FP64 is also supported |
@@ -109,8 +111,9 @@ settings = InferenceSettings(
     activation_checkpointing=False,
     merge_mole=True,
     compile=True,
+    compile_mode="reduce-overhead",
     external_graph_gen=False,
-    internal_graph_gen_version=2,
+    internal_graph_gen_version=3,
 )
 
 predictor = pretrained_mlip.get_predict_unit(

src/fairchem/core/graph/padded_nvidia_graph.py

@@ -0,0 +1,54 @@
+"""Bucketed NVIDIA neighbor graphs for compiled inference."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from fairchem.core.common import gp_utils
+from fairchem.core.graph.radius_graph_pbc_nvidia import radius_graph_pbc_nvidia
+
+
+class PaddedNvidiaGraphGenerator:
+    """Generate v3 neighbor graphs and replace data's edges with a padded bucket."""
+
+    def __init__(self, settings, backbone):
+        if gp_utils.initialized():
+            raise ValueError(
+                "padded NVIDIA graph generation does not support graph parallelism"
+            )
+        if not getattr(backbone, "supports_padded_edges", False):
+            raise ValueError(
+                f"{type(backbone).__name__} does not support padded neighbor edges"
+            )
+        if torch.device(next(backbone.parameters()).device).type != "cuda":
+            raise ValueError("padded NVIDIA graph generation requires CUDA")
+        self.cutoff = float(backbone.cutoff)
+        self.max_neighbors = int(backbone.max_neighbors)
+        self.enforce_max_neighbors_strictly = bool(
+            backbone.enforce_max_neighbors_strictly
+        )
+        self.edge_bucket_size = settings.internal_graph_edge_bucket_size
+
+    def generate(self, data):
+        edge_index, cell_offsets, neighbors = radius_graph_pbc_nvidia(
+            data,
+            self.cutoff,
+            self.max_neighbors,
+            self.enforce_max_neighbors_strictly,
+            pbc=data.pbc,
+        )
+        num_edges = edge_index.shape[1]
+        edge_capacity = (
+            (num_edges + self.edge_bucket_size - 1) // self.edge_bucket_size
+        ) * self.edge_bucket_size
+        pad = edge_capacity - num_edges
+
+        data.edge_index = F.pad(edge_index, (0, pad))
+        data.cell_offsets = F.pad(cell_offsets.to(data.pos.dtype), (0, 0, 0, pad))
+        data.nedges = neighbors.clone()
+        data.nedges[-1] += pad
+        data.edge_valid_mask = (
+            torch.arange(edge_capacity, device=edge_index.device) < num_edges
+        )
+        return data

src/fairchem/core/models/uma/escn_md.py

@@ -278,6 +278,8 @@ def resolve_dataset_mapping(
 
 @registry.register_model("escnmd_backbone")
 class eSCNMDBackbone(nn.Module, MOLEInterface):
+    supports_padded_edges = True
+
     def __init__(
         self,
         max_num_elements: int = 100,
@@ -663,10 +665,12 @@ def _generate_graph(self, data_dict):
 
             # Compute shifts from cell offsets
             if len(data_dict["natoms"]) == 1:
-                # Single system: use matmul (compile-friendly, no data-dependent ops)
-                shifts = data_dict["cell_offsets"].to(
-                    data_dict["cell"].dtype
-                ) @ data_dict["cell"].squeeze(0)
+                offsets = data_dict["cell_offsets"].to(data_dict["cell"].dtype)
+                if "edge_valid_mask" in data_dict:
+                    cell_per_edge = data_dict["cell"].expand(offsets.shape[0], -1, -1)
+                    shifts = torch.bmm(offsets.unsqueeze(1), cell_per_edge).squeeze(1)
+                else:
+                    shifts = offsets @ data_dict["cell"].squeeze(0)
             else:
                 # Batched: need repeat_interleave for variable edges per system
                 cell_per_edge = data_dict["cell"].repeat_interleave(
@@ -684,6 +688,13 @@ def _generate_graph(self, data_dict):
                 - data_dict["pos"][data_dict["edge_index"][1]]
                 + shifts
             )  # [n_edges, 3]
+            edge_valid_mask = data_dict.get("edge_valid_mask", None)
+            if edge_valid_mask is not None:
+                padding_vec = edge_distance_vec.new_zeros(3)
+                padding_vec[0].fill_(self.cutoff + 1.0)
+                edge_distance_vec = torch.where(
+                    edge_valid_mask.unsqueeze(1), edge_distance_vec, padding_vec
+                )
             # pylint: disable=E1102
             edge_distance = torch.linalg.norm(
                 edge_distance_vec, dim=-1, keepdim=False
@@ -952,7 +963,14 @@ def build_inference_settings(cls, settings: InferenceSettings) -> dict:
         if settings.edge_chunk_size is not None:
             overrides["edge_chunk_size"] = settings.edge_chunk_size
         if settings.external_graph_gen is not None:
-            overrides["otf_graph"] = not settings.external_graph_gen
+            padded_internal_graph = (
+                settings.compile_mode == "reduce-overhead"
+                and not settings.external_graph_gen
+                and settings.internal_graph_gen_version == 3
+            )
+            overrides["otf_graph"] = not (
+                settings.external_graph_gen or padded_internal_graph
+            )
         if settings.internal_graph_gen_version is not None:
             overrides["radius_pbc_version"] = settings.internal_graph_gen_version
         if settings.use_quaternion_wigner is not None:

src/fairchem/core/units/mlip_unit/api/inference.py

@@ -201,7 +201,30 @@ class InferenceSettings:
     # VRAM but allow bigger systems; reduce if you run into OOM errors.
     max_atoms: int | None = None
 
+    # Optional torch.compile mode. "reduce-overhead" enables Inductor CUDA graphs.
+    compile_mode: str | None = None
+
+    # Edge bucket used by padded NVIDIA internal graph generation.
+    internal_graph_edge_bucket_size: int = 1024
+
     def __post_init__(self):
+        if self.compile_mode not in (None, "reduce-overhead"):
+            raise ValueError("compile_mode must be None or 'reduce-overhead'")
+        if self.compile_mode is not None and not self.compile:
+            raise ValueError("compile_mode requires compile=True")
+        if (
+            type(self.internal_graph_edge_bucket_size) is not int
+            or self.internal_graph_edge_bucket_size < 1
+        ):
+            raise ValueError("internal_graph_edge_bucket_size must be positive")
+        if (
+            self.compile_mode == "reduce-overhead"
+            and not self.external_graph_gen
+            and self.internal_graph_gen_version != 3
+        ):
+            raise ValueError(
+                "internal reduce-overhead requires internal_graph_gen_version=3"
+            )
         if isinstance(self.base_precision_dtype, str):
             self.base_precision_dtype = getattr(torch, self.base_precision_dtype)
             assert (

src/fairchem/core/units/mlip_unit/predict.py

@@ -38,6 +38,7 @@
 )
 from fairchem.core.components.batch_server import get_app_handle_with_retry
 from fairchem.core.datasets.atomic_data import AtomicData, warn_if_upcasting
+from fairchem.core.graph.padded_nvidia_graph import PaddedNvidiaGraphGenerator
 from fairchem.core.models.uma.nn.execution_backends import (
     ExecutionMode,
     maybe_update_settings_backend,
@@ -150,6 +151,7 @@ def __init__(
         self.lazy_model_intialized = False
         self.assert_on_nans = assert_on_nans
         self._warned_upcast = False
+        self._padded_graph_generator = None
 
         if self.model.module.backbone.regress_config.direct_forces:
             logging.warning(
@@ -448,6 +450,8 @@ def _fall_back_from_fast_path(
         # 1. change flags
         self.inference_settings.merge_mole = False
         self.inference_settings.compile = False
+        self.inference_settings.compile_mode = None
+        self._padded_graph_generator = None
         if self.inference_settings.execution_mode == ExecutionMode.UMAS_FAST_GPU:
             self.inference_settings.execution_mode = ExecutionMode.GENERAL
         self.lazy_model_intialized = False
@@ -493,9 +497,17 @@ def predict(
             if single_atom_result is not None:
                 return single_atom_result
 
+        if self.inference_settings.compile_mode == "reduce-overhead":
+            torch.compiler.cudagraph_mark_step_begin()
+
+        use_padded_internal_graph = (
+            self.inference_settings.compile_mode == "reduce-overhead"
+            and not self.inference_settings.external_graph_gen
+        )
+
         # Regular model prediction path
         # this needs to be .clone() to avoid issues with graph parallel modifying this data with MOLE
-        data_device = data.to(self.device).clone()
+        data_device = data.clone().to(self.device)
 
         dtype = self.inference_settings.base_precision_dtype
         if not self._warned_upcast:
@@ -505,6 +517,12 @@ def predict(
                 data_device[key] = val.to(dtype)
 
         backbone = self.model.module.backbone
+        if use_padded_internal_graph:
+            if self._padded_graph_generator is None:
+                self._padded_graph_generator = PaddedNvidiaGraphGenerator(
+                    self.inference_settings, backbone
+                )
+            data_device = self._padded_graph_generator.generate(data_device)
         _prepare_inference_gradients(backbone, data_device)
 
         # Model handles any per-prediction checks (e.g., MOLE consistency)
@@ -514,13 +532,28 @@ def predict(
             self._fall_back_from_fast_path(error)
             self._lazy_init(data)
             self.model.module.on_predict_check(data_device)
-
         return self._run_inference(data_device, undo_element_references)
 
     def _lazy_init(self, data: AtomicData) -> None:
         """
         Lazy initialization on first predict call.
         """
+        if (
+            self.inference_settings.compile_mode == "reduce-overhead"
+            and not self.inference_settings.external_graph_gen
+        ):
+            if gp_utils.initialized():
+                raise ValueError(
+                    "internal reduce-overhead does not support graph parallelism"
+                )
+            if not getattr(self.model.module.backbone, "supports_padded_edges", False):
+                raise ValueError(
+                    "internal reduce-overhead is not supported by "
+                    f"{type(self.model.module.backbone).__name__}"
+                )
+            if torch.device(self._requested_device).type != "cuda":
+                raise ValueError("internal reduce-overhead requires CUDA")
+
         # Model handles its own preparation (MOLE merge, eval mode, etc.)
         self.model.module.prepare_for_inference(data, self.inference_settings)
         # Inference differentiates outputs with respect to inputs, not weights.
@@ -539,7 +572,11 @@ def _lazy_init(self, data: AtomicData) -> None:
             # The model's scalars are fixed at inference, so this skips dynamo's
             # TensorifyScalarRestartAnalysis retrace during compile.
             torch._dynamo.config.specialize_float = True
-            self.model = torch.compile(self.model, dynamic=True)
+            self.model = torch.compile(
+                self.model,
+                dynamic=True,
+                mode=self.inference_settings.compile_mode,
+            )
 
         self.lazy_model_intialized = True
 

tests/core/graph/test_padded_nvidia_graph.py

@@ -0,0 +1,76 @@
+"""Tests for bucketed NVIDIA neighbor graphs."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+
+import fairchem.core.graph.padded_nvidia_graph as padded_graph
+from fairchem.core.graph.padded_nvidia_graph import PaddedNvidiaGraphGenerator
+
+
+class _GraphData(SimpleNamespace):
+    def clone(self):
+        return _GraphData(
+            **{
+                key: value.clone() if torch.is_tensor(value) else value
+                for key, value in vars(self).items()
+            }
+        )
+
+
+def _generator(bucket_size=4):
+    generator = object.__new__(PaddedNvidiaGraphGenerator)
+    generator.cutoff = 5.0
+    generator.max_neighbors = 32
+    generator.enforce_max_neighbors_strictly = False
+    generator.edge_bucket_size = bucket_size
+    return generator
+
+
+def test_generate_pads_edges_to_bucket(monkeypatch):
+    edge_index = torch.tensor([[1, 0, 2], [0, 1, 1]])
+    cell_offsets = torch.arange(9, dtype=torch.float64).reshape(3, 3)
+    neighbors = torch.tensor([2, 1])
+
+    def graph(*args, **kwargs):
+        return edge_index, cell_offsets, neighbors
+
+    monkeypatch.setattr(padded_graph, "radius_graph_pbc_nvidia", graph)
+    data = _GraphData(
+        pos=torch.empty(3, 3, dtype=torch.float32),
+        pbc=torch.ones(2, 3, dtype=torch.bool),
+    )
+
+    result = _generator().generate(data)
+
+    assert result is data
+    torch.testing.assert_close(result.edge_index[:, :3], edge_index)
+    torch.testing.assert_close(
+        result.edge_index[:, 3], torch.zeros(2, dtype=torch.long)
+    )
+    torch.testing.assert_close(result.cell_offsets[:3], cell_offsets.float())
+    torch.testing.assert_close(result.cell_offsets[3], torch.zeros(3))
+    torch.testing.assert_close(result.nedges, torch.tensor([2, 2]))
+    torch.testing.assert_close(
+        result.edge_valid_mask, torch.tensor([True, True, True, False])
+    )
+
+
+def test_generate_keeps_exact_bucket_shape(monkeypatch):
+    edge_index = torch.tensor([[1, 0, 2, 1], [0, 1, 1, 2]])
+    cell_offsets = torch.zeros(4, 3)
+    neighbors = torch.tensor([4])
+    monkeypatch.setattr(
+        padded_graph,
+        "radius_graph_pbc_nvidia",
+        lambda *args, **kwargs: (edge_index, cell_offsets, neighbors),
+    )
+    data = _GraphData(pos=torch.empty(3, 3), pbc=torch.ones(1, 3, dtype=torch.bool))
+
+    result = _generator().generate(data)
+
+    torch.testing.assert_close(result.edge_index, edge_index)
+    torch.testing.assert_close(result.nedges, neighbors)
+    assert result.edge_valid_mask.all()

tests/core/models/uma/test_padded_edges.py

@@ -0,0 +1,49 @@
+"""Tests for padded neighbor edges in UMA."""
+
+from __future__ import annotations
+
+import torch
+
+from fairchem.core.models.uma.escn_md import eSCNMDBackbone
+
+
+def test_precomputed_padding_is_moved_beyond_cutoff():
+    backbone = object.__new__(eSCNMDBackbone)
+    torch.nn.Module.__init__(backbone)
+    backbone.cutoff = 5.0
+    backbone.otf_graph = False
+    data = {
+        "pos": torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]),
+        "cell": torch.eye(3).unsqueeze(0),
+        "natoms": torch.tensor([2]),
+        "edge_index": torch.tensor([[0, 0], [1, 0]]),
+        "cell_offsets": torch.zeros(2, 3),
+        "edge_valid_mask": torch.tensor([True, False]),
+    }
+
+    graph = backbone._generate_graph(data)
+
+    torch.testing.assert_close(graph["edge_distance"][0], torch.tensor(1.0))
+    torch.testing.assert_close(graph["edge_distance"][1], torch.tensor(6.0))
+    torch.testing.assert_close(data["scatter_target"], torch.tensor([1, 0]))
+
+
+def test_precomputed_periodic_shifts_match_internal_graph_math():
+    backbone = object.__new__(eSCNMDBackbone)
+    torch.nn.Module.__init__(backbone)
+    backbone.cutoff = 5.0
+    backbone.otf_graph = False
+    data = {
+        "pos": torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]),
+        "cell": torch.diag(torch.tensor([3.0, 4.0, 5.0])).unsqueeze(0),
+        "natoms": torch.tensor([2]),
+        "edge_index": torch.tensor([[0], [1]]),
+        "cell_offsets": torch.tensor([[1.0, 0.0, 0.0]]),
+    }
+
+    graph = backbone._generate_graph(data)
+
+    torch.testing.assert_close(
+        graph["edge_distance_vec"], torch.tensor([[2.0, 0.0, 0.0]])
+    )
+    torch.testing.assert_close(graph["edge_distance"], torch.tensor([2.0]))