fix(vacuum): bound isolated graph references and preserve export kernels

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -413,11 +413,11 @@ def vacuum_descriptor(self) -> Array:
         conditioned as the neutral ground-state atom: zero charge with the
         ground-state multiplicity when the descriptor takes the charge/spin
         condition, and a spin vector of one Bohr magneton per unpaired
-        electron when it takes the native spin. A descriptor with a graph
-        lower is evaluated on the neighbor graph of those single-atom frames,
-        which holds no real edge; a dense descriptor on a neighbor list of its
-        own capacity filled with padding, since its equations are sized by
-        that capacity.
+        electron when it takes the native spin.
+
+        Descriptors that opt into graph vacuum evaluation use isolated nodes,
+        independently of their advertised neighbor capacity; other descriptors
+        retain their configured all-masked dense neighbor list.
 
         Returns
         -------
@@ -429,39 +429,36 @@ def vacuum_descriptor(self) -> Array:
         xp = array_api_compat.array_namespace(bias)
         device = array_api_compat.device(bias)
         ntypes = self.get_ntypes()
-        coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device)
-        atype = xp.reshape(
-            xp.arange(ntypes, dtype=xp.int64, device=device), (ntypes, 1)
-        )
+        atype = xp.arange(ntypes, dtype=xp.int64, device=device)
         conditions = {
             name: xp.asarray(table, dtype=bias.dtype, device=device)
             for name, table in self.vacuum_conditions().items()
         }
-        if self.descriptor.uses_graph_lower():
+        if self.uses_graph_lower() and self.descriptor.supports_graph_vacuum():
             from deepmd.dpmodel.utils.neighbor_graph import (
-                build_neighbor_graph,
+                NeighborGraph,
             )
 
-            graph = build_neighbor_graph(
-                coord, atype, None, self.descriptor.get_rcut(), with_csr=True
-            )
-            # Only descriptors whose ``call_graph`` declares a condition take
-            # the keyword; a declared condition the model does not use is None.
-            spin_kwargs = (
-                {"spin": conditions.get("spin")} if self.supports_native_spin() else {}
-            )
-            charge_spin_kwargs = (
-                {"charge_spin": conditions.get("charge_spin")}
-                if self.supports_charge_spin
-                else {}
+            # One isolated node per frame, with the same two masked guard
+            # edges as the graph builders. Never size this graph from sel:
+            # graph-only descriptors may report an unbounded sentinel there.
+            # The graph interface also carries native spin, which the dense
+            # compatibility adapters need not support.
+            graph = NeighborGraph(
+                n_node=xp.ones((ntypes,), dtype=xp.int64, device=device),
+                edge_index=xp.zeros((2, 2), dtype=xp.int64, device=device),
+                edge_vec=xp.zeros((2, 3), dtype=bias.dtype, device=device),
+                edge_mask=xp.zeros((2,), dtype=xp.bool, device=device),
             )
             return self.descriptor.call_graph(
                 graph,
-                xp.reshape(atype, (ntypes,)),
+                atype,
                 type_embedding=self.descriptor.graph_type_embedding_table(),
-                **spin_kwargs,
-                **charge_spin_kwargs,
+                **conditions,
             )[0]
+
+        coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device)
+        atype = xp.reshape(atype, (ntypes, 1))
         nlist = xp.full(
             (ntypes, 1, self.descriptor.get_nnei()),
             -1,
@@ -639,10 +636,11 @@ def forward_atomic_graph(
                 atype.shape[0],
             )
         # === Step 2. Vacuum reference nodes ===
-        # The reference atoms are appended as single-node frames, so the same
-        # descriptor call yields the vacuum descriptor of every type.
+        # Opted-in descriptors evaluate the reference in the same graph call.
+        # Legacy descriptors must retain the dense reference used by folding.
         vacuum_ref = self.fitting_net.needs_vacuum_descriptor()
-        if vacuum_ref:
+        graph_vacuum = vacuum_ref and self.descriptor.supports_graph_vacuum()
+        if graph_vacuum:
             graph, atype, charge_spin, spin = self.append_vacuum_frames(
                 graph, atype, charge_spin, spin
             )
@@ -671,12 +669,14 @@ def forward_atomic_graph(
         gg, rot_mat = outputs[0], outputs[1]
         node_gate = outputs[2] if len(outputs) > 2 else None
         vacuum_kwargs = {}
-        if vacuum_ref:
+        if graph_vacuum:
             vacuum_kwargs["vacuum_descriptor"] = gg[n_real:]
             gg = gg[:n_real]
             atype = atype[:n_real]
             rot_mat = None if rot_mat is None else rot_mat[:n_real]
             node_gate = None if node_gate is None else node_gate[:n_real]
+        elif vacuum_ref:
+            vacuum_kwargs["vacuum_descriptor"] = self.vacuum_descriptor()
         # === Step 4. Fitting ===
         # The source gate of a bridged descriptor fades the learned part of
         # the fitting output: an atom of a frozen pair keeps only the bias of

deepmd/dpmodel/descriptor/dpa4.py

@@ -2512,6 +2512,10 @@ def dense_lower_supports_comm(self) -> bool:
         """
         return False
 
+    def supports_graph_vacuum(self) -> bool:
+        """The dense adapter uses the same graph core; native spin needs it."""
+        return True
+
     def uses_graph_lower(self) -> bool:
         """Whether this descriptor supports the sel-free NeighborGraph lower.
 

deepmd/dpmodel/descriptor/dpa4c.py

@@ -2118,6 +2118,10 @@ def get_env_protection(self) -> float:
         """Return the direction regularization scale in Å."""
         return self._EPS
 
+    def supports_graph_vacuum(self) -> bool:
+        """Isolated references use the graph core without a dense sel capacity."""
+        return True
+
     def uses_graph_lower(self) -> bool:
         """Return whether graph-native lowering is supported."""
         return True

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -287,6 +287,18 @@ def graph_type_embedding_table(self) -> Any | None:
             """
             return None
 
+        def supports_graph_vacuum(self) -> bool:
+            """Whether isolated-atom references can use the graph interface.
+
+            Opting in requires ``call_graph`` to produce the same isolated
+            reference as the dense compatibility path, when one exists;
+            graph-only descriptors may also opt in. This is distinct from
+            ``uses_graph_lower``: some legacy descriptors retain different
+            padding semantics on their dense and graph paths, so switching
+            their reference would change the model's energy zero.
+            """
+            return False
+
         @abstractmethod
         def need_sorted_nlist_for_lower(self) -> bool:
             """Returns whether the descriptor needs sorted nlist when using `forward_lower`."""

deepmd/pt_expt/common.py

@@ -273,6 +273,28 @@ def _try_convert_list(name: str, value: list) -> torch.nn.Module | None:
     return None
 
 
+def _module_buffer_device(obj: torch.nn.Module, name: str) -> torch.device:
+    """Device for a numpy array assigned as a buffer on ``obj``.
+
+    Prefer the existing buffer of the same name when it already holds a
+    tensor, then any parameter or non-None buffer on the module, then
+    ``env.DEVICE``. Assignments after ``Module.to`` therefore keep residency
+    with the module rather than jumping back to the process default device.
+    """
+    from deepmd.pt_expt.utils import env  # deferred - avoids circular import
+
+    if "_buffers" in obj.__dict__:
+        existing = obj._buffers.get(name)
+        if existing is not None:
+            return existing.device
+    for tensor in obj.parameters():
+        return tensor.device
+    for tensor in obj.buffers():
+        if tensor is not None:
+            return tensor.device
+    return env.DEVICE
+
+
 def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, Any]:
     """Common __setattr__ logic for pt_expt wrappers around dpmodel classes.
 
@@ -319,6 +341,10 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer
     - In dpmodel, all persistent arrays use `self.xxx = np.array(...)`. Scalars
       use `.item()`, lists use `.tolist()`. So `isinstance(value, np.ndarray)`
       reliably identifies buffer-worthy attributes.
+    - A numpy array becomes a buffer on the module's current device (see
+      :func:`_module_buffer_device`), not unconditionally on ``env.DEVICE``, so
+      a later assignment after ``Module.to`` does not pull tensors back to the
+      process default device.
     - torch.Tensor values assigned to existing buffers fall through to
       torch.nn.Module.__setattr__, which correctly updates them.
     - dpmodel objects are identified by registry lookup (exact type match), so
@@ -328,10 +354,10 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer
 
     **Circular import resolution:**
 
-    The function uses a deferred import `from deepmd.pt_expt.utils import env`
-    inside the function body. This breaks the circular dependency chain:
-    common.py → utils/__init__.py → exclude_mask.py → common.py. The import is
-    cached by Python after the first call, so there's no performance penalty.
+    Device resolution and related helpers defer
+    ``from deepmd.pt_expt.utils import env`` so this module does not import
+    ``utils`` at load time (common.py → utils/__init__.py → exclude_mask.py →
+    common.py). The import is cached after the first call.
 
     **Usage pattern:**
 
@@ -354,11 +380,9 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer
     >>> # Case 3: dpmodel object → pt_expt module
     >>> obj.emask = AtomExcludeMaskDP(...)  # becomes AtomExcludeMask module
     """
-    from deepmd.pt_expt.utils import env  # deferred - avoids circular import
-
     # numpy array → torch buffer
     if isinstance(value, np.ndarray) and "_buffers" in obj.__dict__:
-        tensor = torch.as_tensor(value, device=env.DEVICE)
+        tensor = torch.as_tensor(value, device=_module_buffer_device(obj, name))
         if name in obj._buffers:
             obj._buffers[name] = tensor
             return True, tensor

deepmd/pt_expt/utils/serialization.py

@@ -1750,6 +1750,70 @@ def _trace_and_export(
         )
 
 
+@contextlib.contextmanager
+def _cpu_reference_descriptors(model: torch.nn.Module) -> Iterator[None]:
+    """Temporarily use CPU reference descriptors for the eager vacuum fold.
+
+    Deployment descriptors bind accelerator operators at construction, before
+    they move to CPU for fake tracing. Changing their device or the kernel
+    environment cannot unbind those operators. Rebuild only descriptors whose
+    fittings still need a reference, leaving the folded bias and non-persistent
+    vacuum table on the original fitting. Restore the deployment descriptors
+    before tracing so their optimized bindings remain intact.
+    """
+    from deepmd.dpmodel.atomic_model.dp_atomic_model import (
+        DPAtomicModel,
+    )
+    from deepmd.pt_expt.utils import (
+        env,
+    )
+
+    atomic_models = [
+        module
+        for module in model.modules()
+        if isinstance(module, DPAtomicModel)
+        and module.fitting_net.needs_vacuum_descriptor()
+    ]
+    if not atomic_models:
+        yield
+        return
+
+    levels = (
+        "DP_CUDA_INFER",
+        "DP_TRITON_INFER",
+        "DP_CUTILE_INFER",
+        "DP_CUTE_INFER",
+        "DP_CUDA_TRAIN",
+        "DP_TRITON_TRAIN",
+        "DP_AMP_INFER",
+    )
+    saved_levels = {name: os.environ.get(name) for name in levels}
+    saved_device = env.DEVICE
+    descriptors = []
+    try:
+        env.DEVICE = torch.device("cpu")
+        os.environ.update(dict.fromkeys(levels, "0"))
+        for atomic_model in atomic_models:
+            descriptor = atomic_model.descriptor
+            reference = type(descriptor).deserialize(descriptor.serialize()).eval()
+            # Virtual-spin backbones can disable this route at runtime; that
+            # choice is deliberately absent from the portable descriptor.
+            if not descriptor.uses_graph_lower() and reference.uses_graph_lower():
+                reference.disable_graph_lower()
+            descriptors.append((atomic_model, descriptor))
+            atomic_model.descriptor = reference
+        yield
+    finally:
+        for atomic_model, descriptor in descriptors:
+            atomic_model.descriptor = descriptor
+        env.DEVICE = saved_device
+        for name, value in saved_levels.items():
+            if value is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = value
+
+
 def _trace_and_export_impl(
     data: dict,
     model_json_override: dict | None = None,
@@ -1833,12 +1897,15 @@ def _trace_and_export_impl(
         # pt_expt BaseModel registry returns this backend's torch class.
         model = BaseModel.deserialize(data["model"])
     model.eval()
-    # The vacuum reference is resolved on the export object, folded into the
-    # fitting bias or stored as a per-type table, so the exported graph
-    # carries no reference atoms; the move to the tracing device follows, so
-    # the table lands there with the rest of the model.
-    model.fold_vacuum_reference()
+    # Fold the vacuum reference on the CPU tracing device. vacuum_descriptor
+    # evaluates an isolated atom of every type through the descriptor; on CUDA
+    # that path can allocate several GiB even when export already holds most of
+    # the GPU. Move first so the fold runs where tracing will, and so the
+    # folded bias or vacuum table stays on that device with the rest of the
+    # model.
     model.to("cpu")
+    with torch.no_grad(), _cpu_reference_descriptors(model):
+        model.fold_vacuum_reference()
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)