Enable compiler-managed CUDA graphs for UMA inference

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]))