feat(pt_expt): O(N) on-device NeighborGraph builders — vesin & nvalchemiops (PR-C)

deepmd/dpmodel/model/make_model.py

@@ -463,7 +463,8 @@ def _call_common_graph(
                 ng = build_neighbor_graph_ase(cc, atype, bb, self.get_rcut())
             else:
                 raise ValueError(
-                    f"unknown neighbor_graph_method {method!r}; use 'dense' or 'ase'"
+                    f"unknown neighbor_graph_method {method!r}; the dpmodel/jax backend "
+                    "supports 'dense'/'ase' only ('vesin'/'nv' require the pt_expt backend)."
                 )
             xp = array_api_compat.array_namespace(atype)
             nf, nloc = atype.shape[:2]

deepmd/dpmodel/utils/neighbor_graph/ase_builder.py

@@ -58,7 +58,8 @@ def build_neighbor_graph_ase(
     coord
         (nf, nloc, 3) local coordinates.
     atype
-        (nf, nloc) local atom types (unused for the search; carried for API parity).
+        (nf, nloc) local atom types; ``type < 0`` marks a virtual atom, excluded
+        as center and neighbor (the search itself is type-blind).
     box
         (nf, 3, 3) simulation cell, or ``None`` for non-periodic.
     rcut
@@ -127,6 +128,14 @@ def _to_cpu_numpy(x: Any) -> np.ndarray:
         np.concatenate(nframe_parts) if nframe_parts else np.zeros((0,), dtype=np.int64)
     )
 
+    # virtual atoms (atype < 0) are excluded as centers AND neighbors -- the
+    # World-2 builder contract shared with the dense reference builder; the
+    # geometric search above cannot know about them.
+    atype_np = _to_cpu_numpy(atype).reshape(nf, nloc)
+    keep = (atype_np[nframe_all, i_all] >= 0) & (atype_np[nframe_all, j_all] >= 0)
+    i_all, j_all = i_all[keep], j_all[keep]
+    S_all, nframe_all = S_all[keep], nframe_all[keep]
+
     return neighbor_graph_from_ijs(
         i_all, j_all, S_all, coord, box, nframe_all, nloc, layout=layout
     )

deepmd/pt/utils/nv_nlist.py

@@ -33,9 +33,6 @@
     EdgeNeighborList,
     NeighborList,
 )
-from deepmd.pt.utils.region import (
-    normalize_coord,
-)
 from deepmd.pt_expt.utils.edge_schema import (
     edge_schema_from_neighbor_matrix,
 )
@@ -164,74 +161,24 @@ def build(
         See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The
         returned ``nlist`` is distance-sorted and truncated to ``sum(sel)``.
         """
-        from nvalchemiops.torch.neighbors import (
-            neighbor_list,
+        device = coord.device
+        nf, nloc = atype.shape[:2]
+        target_neighbors = int(sum(sel))
+        coord = coord.reshape(nf, nloc, 3)
+
+        # Delegate the raw search to the shared helper in nv_graph_builder.
+        # Function-level import avoids a module-level pt -> pt_expt cycle while
+        # keeping the search logic in exactly one place (graph-builder primary,
+        # legacy strategy is the deprecation-bound caller).
+        from deepmd.pt_expt.utils.nv_graph_builder import (
+            nv_search_matrix,
         )
 
-        device = coord.device
-        with _input_device_context(device):
-            nf, nloc = atype.shape[:2]
-            target_neighbors = int(sum(sel))
-            search_capacity = target_neighbors
-            total_atoms = nf * nloc
-            coord = coord.reshape(nf, nloc, 3)
-            periodic = box is not None
-            if not periodic:
-                cell = None
-                pbc = None
-            else:
-                cell = box.reshape(nf, 3, 3).to(device=device, dtype=coord.dtype)
-                coord = normalize_coord(coord, cell)
-                pbc = torch.ones((nf, 3), dtype=torch.bool, device=device)
-            positions_for_nlist = coord.reshape(total_atoms, 3).detach()
-            batch_idx = torch.arange(
-                nf, dtype=torch.int32, device=device
-            ).repeat_interleave(nloc)
-            batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc
-            method = choose_nv_nlist_method(nloc, periodic=periodic, device=device)
-
-            # ``batch_naive`` otherwise derives ``max_atoms_per_system`` from
-            # ``batch_ptr`` with a ``.max().item()`` device->host sync on every
-            # call. Our batches are homogeneous (``nloc`` atoms per frame), so the
-            # value is known on the host; passing it explicitly removes that
-            # per-call sync. ``batch_cell_list`` neither accepts the argument nor
-            # has a ``**kwargs`` catch-all, so the override is guarded on method.
-            extra_nl_kwargs: dict[str, Any] = {}
-            if method == "batch_naive":
-                extra_nl_kwargs["max_atoms_per_system"] = int(nloc)
-
-            # Grow the search capacity until all neighbors fit so the distance-sort
-            # below selects the true nearest ``sum(sel)``.
-            while True:
-                nlist_result = neighbor_list(
-                    positions_for_nlist,
-                    float(rcut),
-                    cell=cell,
-                    pbc=pbc,
-                    batch_idx=batch_idx,
-                    batch_ptr=batch_ptr,
-                    method=method,
-                    max_neighbors=int(search_capacity),
-                    return_neighbor_list=False,
-                    wrap_positions=False,
-                    **extra_nl_kwargs,
-                )
-                if len(nlist_result) == 2:
-                    neighbor_matrix, num_neighbors = nlist_result
-                    shifts = torch.zeros(
-                        (*neighbor_matrix.shape, 3),
-                        dtype=torch.int32,
-                        device=device,
-                    )
-                else:
-                    neighbor_matrix, num_neighbors, shifts = nlist_result
-                max_found = (
-                    int(num_neighbors.max().item()) if num_neighbors.numel() > 0 else 0
-                )
-                if max_found <= search_capacity:
-                    break
-                search_capacity = max(max_found, _grow_search_capacity(search_capacity))
+        coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix(
+            coord, box, rcut, start_capacity=target_neighbors
+        )
 
+        with _input_device_context(device):
             if return_mode == "edges":
                 return edge_schema_from_neighbor_matrix(
                     coord=coord,

deepmd/pt_expt/infer/deep_eval.py

@@ -65,6 +65,10 @@
 if TYPE_CHECKING:
     import ase.neighborlist
 
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        NeighborGraph,
+    )
+
 
 # Public output keys emitted by the graph-form AOTI forward
 # (``forward_lower_graph_exportable``) keyed by the output-variable category that
@@ -125,6 +129,19 @@ class DeepEval(DeepEvalBackend):
     neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
         The ASE neighbor list class to produce the neighbor list. If None, the
         neighbor list will be built natively in the model.
+    nlist_backend : str, default: "auto"
+        Neighbor-list builder for the NLIST/extended lower path (``.pte`` and
+        nlist-form ``.pt2``): ``"auto"`` / ``"vesin"`` / ``"native"``. Not
+        used by graph-form ``.pt2`` artifacts.
+    neighbor_graph_method : str, default: "dense"
+        Carry-all graph builder for GRAPH-FORM ``.pt2`` artifacts ONLY
+        (``metadata["lower_input_kind"] == "graph"``): ``"dense"`` / ``"ase"``
+        (backend-agnostic) or ``"vesin"`` / ``"nv"`` (on-device O(N)). A
+        non-default value on any other artifact raises at construction — the
+        knob would silently do nothing there; use ``nlist_backend`` for the
+        nlist path instead. All builders emit the same neighbor set, so the
+        choice is performance-only. Consolidating the two knobs into a single
+        backend-selection API is deferred to the dense-nlist deprecation.
     **kwargs : dict
         Keyword arguments.
     """
@@ -137,11 +154,15 @@ def __init__(
         auto_batch_size: bool | int | AutoBatchSize = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
         nlist_backend: str = "auto",
+        neighbor_graph_method: str = "dense",
         **kwargs: Any,
     ) -> None:
         self.output_def = output_def
         self.model_path = model_file
         self.neighbor_list = neighbor_list
+        # World-2 graph-form ``.pt2`` (lower_input_kind == "graph") builder select:
+        # "dense"/"ase" (backend-agnostic) or "vesin"/"nv" (on-device O(N)).
+        self._neighbor_graph_method = neighbor_graph_method
         self._is_pt2 = model_file.endswith(".pt2")
 
         if self._is_pt2:
@@ -157,6 +178,20 @@ def __init__(
                 "`.pt` (training checkpoint)."
             )
 
+        # neighbor_graph_method is consumed ONLY by graph-form .pt2 eval
+        # (_eval_model_graph); fail fast instead of silently ignoring it on
+        # nlist-form artifacts (there, the builder knob is nlist_backend).
+        if (
+            neighbor_graph_method != "dense"
+            and getattr(self, "metadata", {}).get("lower_input_kind") != "graph"
+        ):
+            raise ValueError(
+                f"neighbor_graph_method={neighbor_graph_method!r} only applies to "
+                "graph-form .pt2 artifacts (lower_input_kind == 'graph'); this "
+                f"model is not graph-form. Use nlist_backend to select the "
+                "neighbor-list builder for the nlist path."
+            )
+
         self._setup_nlist_backend(nlist_backend)
 
         if isinstance(auto_batch_size, bool):
@@ -1660,9 +1695,6 @@ def _eval_model_graph(
         forward returns the LOCAL public keys directly, so results are reshaped
         without ``communicate_extended_output``.
         """
-        from deepmd.dpmodel.utils.neighbor_graph import (
-            build_neighbor_graph,
-        )
         from deepmd.pt_expt.utils.env import (
             DEVICE,
         )
@@ -1678,28 +1710,19 @@ def _eval_model_graph(
         box_input = cells.reshape(nframes, 9) if cells is not None else None
         # Dynamic edge axis (B2.0): build the carry-all graph at its exact edge
         # count (no static padding); the AOTI artifact accepts any E.
-        graph = build_neighbor_graph(
-            coord_input,
-            atom_types,
-            box_input,
-            self._rcut,
-        )
+        graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE)
 
         atype_t = torch.tensor(
             np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE
         )
-        n_node_t = torch.tensor(
-            np.asarray(graph.n_node), dtype=torch.int64, device=DEVICE
-        )
-        edge_index_t = torch.tensor(
-            np.asarray(graph.edge_index), dtype=torch.int64, device=DEVICE
-        )
-        edge_vec_t = torch.tensor(
-            np.asarray(graph.edge_vec), dtype=torch.float64, device=DEVICE
-        )
-        edge_mask_t = torch.tensor(
-            np.asarray(graph.edge_mask), dtype=torch.bool, device=DEVICE
+        # graph fields may be numpy (dense/ase) or torch, possibly on CUDA
+        # (vesin/nv) -- torch.as_tensor handles both and moves to DEVICE.
+        n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE)
+        edge_index_t = torch.as_tensor(
+            graph.edge_index, dtype=torch.int64, device=DEVICE
         )
+        edge_vec_t = torch.as_tensor(graph.edge_vec, dtype=torch.float64, device=DEVICE)
+        edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE)
 
         fparam_t, aparam_t = self._prepare_optional_lower_inputs(
             fparam, aparam, nframes, natoms, DEVICE
@@ -1734,6 +1757,61 @@ def _eval_model_graph(
                 )
         return tuple(results)
 
+    def _build_eval_graph(
+        self,
+        coord_input: np.ndarray,
+        atom_types: np.ndarray,
+        box_input: np.ndarray | None,
+        device: "torch.device",
+    ) -> "NeighborGraph":
+        """Build the carry-all NeighborGraph for graph-form ``.pt2`` inference.
+
+        Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run
+        backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)).
+        All backends emit the SAME neighbor set (carry-all, sel-free), so the
+        selection is a pure performance choice and results are unchanged.
+        """
+        method = self._neighbor_graph_method
+        if method == "dense":
+            from deepmd.dpmodel.utils.neighbor_graph import (
+                build_neighbor_graph,
+            )
+
+            return build_neighbor_graph(coord_input, atom_types, box_input, self._rcut)
+        if method == "ase":
+            from deepmd.dpmodel.utils.neighbor_graph import (
+                build_neighbor_graph_ase,
+            )
+
+            return build_neighbor_graph_ase(
+                coord_input, atom_types, box_input, self._rcut
+            )
+        if method in ("vesin", "nv"):
+            cc = torch.as_tensor(coord_input, dtype=torch.float64, device=device)
+            aa = torch.as_tensor(
+                np.asarray(atom_types), dtype=torch.int64, device=device
+            )
+            bb = (
+                torch.as_tensor(box_input, dtype=torch.float64, device=device)
+                if box_input is not None
+                else None
+            )
+            if method == "vesin":
+                from deepmd.pt_expt.utils.vesin_graph_builder import (
+                    build_neighbor_graph_vesin,
+                )
+
+                return build_neighbor_graph_vesin(cc, aa, bb, self._rcut)
+            from deepmd.pt_expt.utils.nv_graph_builder import (
+                build_neighbor_graph_nv,
+            )
+
+            return build_neighbor_graph_nv(cc, aa, bb, self._rcut)
+        raise ValueError(
+            f"unknown neighbor_graph_method {method!r}; "
+            "use 'dense', 'ase', 'vesin', or 'nv'"
+        )
+
     def _get_output_shape(
         self, odef: OutputVariableDef, nframes: int, natoms: int
     ) -> list[int]:

deepmd/pt_expt/model/make_model.py

@@ -456,14 +456,38 @@ def _call_common_graph(
                 build_neighbor_graph_ase,
             )
 
+            # mirror the dpmodel guard: _resolve_graph_method's eligibility
+            # check only protects the default (None) path; an EXPLICIT
+            # neighbor_graph_method would otherwise reach the builders for
+            # descriptors without a graph lower.
+            descriptor = getattr(self.atomic_model, "descriptor", None)
+            uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False)
+            if not (self.mixed_types() and uses_graph_lower()):
+                raise NotImplementedError(
+                    "neighbor_graph_method requires a mixed_types descriptor with a "
+                    "graph lower (e.g. dpa1 attn_layer=0)"
+                )
             rcut = self.get_rcut()
             if method == "dense":
                 ng = build_neighbor_graph(cc, atype, bb, rcut)
             elif method == "ase":
                 ng = build_neighbor_graph_ase(cc, atype, bb, rcut)
+            elif method == "vesin":
+                from deepmd.pt_expt.utils.vesin_graph_builder import (
+                    build_neighbor_graph_vesin,
+                )
+
+                ng = build_neighbor_graph_vesin(cc, atype, bb, rcut)
+            elif method == "nv":
+                from deepmd.pt_expt.utils.nv_graph_builder import (
+                    build_neighbor_graph_nv,
+                )
+
+                ng = build_neighbor_graph_nv(cc, atype, bb, rcut)
             else:
                 raise ValueError(
-                    f"unknown neighbor_graph_method {method!r}; use 'dense' or 'ase'"
+                    f"unknown neighbor_graph_method {method!r}; "
+                    "use 'dense', 'ase', 'vesin', or 'nv'"
                 )
             nf, nloc = atype.shape[:2]
             atype_flat = atype.reshape(nf * nloc)