perf(dpa1): optimize graph CUDA inference and deployment

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -320,6 +320,51 @@ def forward_common_atomic(
         atom_mask = xp_take_first_n(ext_atom_mask, 1, nloc)
         return self._finalize_atomic_ret(ret_dict, atom_mask, atype)
 
+    def _prepare_graph_nodes(
+        self,
+        n_node: Array,
+        n_local: Array | None,
+        atype: Array,
+        reference: Array,
+    ) -> tuple[Array, Array]:
+        """Apply node masks shared by generic and compact graph forwards."""
+        xp = array_api_compat.array_namespace(reference)
+        atype = xp.asarray(atype, device=array_api_compat.device(reference))
+        atom_mask = self.make_atom_mask(atype)
+        atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype))
+        output_mask = atom_mask
+        if n_local is not None:
+            from deepmd.dpmodel.utils.neighbor_graph import (
+                node_ownership_mask,
+            )
+
+            output_mask = output_mask & node_ownership_mask(
+                n_node,
+                n_local,
+                atype.shape[0],
+            )
+        if self.atom_excl is not None:
+            output_mask = xp.logical_and(
+                output_mask,
+                self.atom_excl.build_type_exclude_mask(atype_clamped),
+            )
+        return atype_clamped, output_mask
+
+    def _prepare_graph_inputs(
+        self,
+        graph: "NeighborGraph",
+        atype: Array,
+    ) -> tuple["NeighborGraph", Array, Array]:
+        """Apply graph masks shared by standard and fused atomic forwards."""
+        atype_clamped, output_mask = self._prepare_graph_nodes(
+            graph.n_node,
+            graph.n_local,
+            atype,
+            graph.edge_vec,
+        )
+        self._assert_graph_pair_excluded(graph, atype_clamped)
+        return graph, atype_clamped, output_mask
+
     def forward_common_atomic_graph(
         self,
         graph: "NeighborGraph",
@@ -359,25 +404,15 @@ def forward_common_atomic_graph(
             the result dict on the flat node axis, defined by the `FittingOutputDef`.
 
         """
-        xp = array_api_compat.array_namespace(graph.edge_vec)
-        atype = xp.asarray(atype, device=array_api_compat.device(graph.edge_vec))
-        atom_mask = self.make_atom_mask(atype)  # (N,) bool
-        atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype))
-        # NOTE: model-level ``pair_exclude_types`` is NOT applied here. It is a
-        # graph-BUILD transform (decision #18) already folded into
-        # ``graph.edge_mask`` by the NeighborGraph builder (Python) or
-        # ``applyPairExclusion`` (C++); this method consumes a pre-excluded graph.
-        # Fail-safe (eager only): guard against a caller that skipped the build
-        # seam, which would silently INCLUDE excluded pairs (fail-open).
-        self._assert_graph_pair_excluded(graph, atype_clamped)
+        graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype)
         ret_dict = self.forward_atomic_graph(
             graph,
             atype_clamped,
             fparam=fparam,
             aparam=aparam,
             charge_spin=charge_spin,
         )
-        return self._finalize_atomic_ret(ret_dict, atom_mask, atype)
+        return self._finalize_atomic_ret(ret_dict, output_mask, atype)
 
     def _assert_nlist_pair_excluded(self, nlist: Array, extended_atype: Array) -> None:
         """Fail-safe: assert the nlist reaching the dense seam is pre-excluded.
@@ -491,10 +526,16 @@ def _finalize_atomic_ret(
 
         """
         xp = array_api_compat.array_namespace(atype)
-        ret_dict = self.apply_out_stat(ret_dict, atype)
+        safe_atype = xp.where(
+            self.make_atom_mask(atype),
+            atype,
+            xp.zeros_like(atype),
+        )
+        ret_dict = self.apply_out_stat(ret_dict, safe_atype)
         if self.atom_excl is not None:
             atom_mask = xp.logical_and(
-                atom_mask, self.atom_excl.build_type_exclude_mask(atype)
+                atom_mask,
+                self.atom_excl.build_type_exclude_mask(safe_atype),
             )
         lead = atom_mask.shape  # (nf, nloc) dense | (N,) graph
         for kk in ret_dict.keys():

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -37,6 +37,46 @@
 )
 
 
+def _extend_graph_aparam(
+    aparam: Array,
+    n_node: Array,
+    n_local: Array,
+    n_total: int,
+) -> Array:
+    """Expand frame-local atomic parameters onto a local-plus-halo node axis."""
+    import array_api_compat
+
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        frame_id_from_n_node,
+        node_ownership_mask,
+    )
+
+    xp = array_api_compat.array_namespace(aparam, n_node, n_local)
+    frame_id = frame_id_from_n_node(n_node, n_total=n_total)
+    frame_end = xp.cumulative_sum(n_node)
+    frame_start = frame_end - n_node
+    node_index = xp.arange(
+        n_total,
+        dtype=n_node.dtype,
+        device=array_api_compat.device(n_node),
+    )
+    index_in_frame = node_index - xp.take(frame_start, frame_id, axis=0)
+    local_capacity = aparam.shape[1]
+    sentinel = xp.zeros(
+        (aparam.shape[0], 1, aparam.shape[2]),
+        dtype=aparam.dtype,
+        device=array_api_compat.device(aparam),
+    )
+    padded_aparam = xp.concat([aparam, sentinel], axis=1)
+    padded_capacity = local_capacity + 1
+    local_index = index_in_frame % padded_capacity
+    flat_index = frame_id * padded_capacity + local_index
+    flat_aparam = xp.reshape(padded_aparam, (-1, aparam.shape[-1]))
+    gathered = xp.take(flat_aparam, flat_index, axis=0)
+    ownership = node_ownership_mask(n_node, n_local, n_total)
+    return xp.where(ownership[:, None], gathered, xp.zeros_like(gathered))
+
+
 @BaseAtomicModel.register("standard")
 class DPAtomicModel(BaseAtomicModel):
     r"""Model give atomic prediction of some physical property.
@@ -302,10 +342,27 @@ def forward_atomic_graph(
         )
         fparam_node = None
         if fparam is not None:
-            frame_id = frame_id_from_n_node(graph.n_node)
+            frame_id = frame_id_from_n_node(
+                graph.n_node,
+                n_total=atype.shape[0],
+            )
             fparam_node = xp.take(fparam, frame_id, axis=0)  # (N, ndf)
+        aparam_node = aparam
+        if aparam is not None and graph.n_local is not None and aparam.ndim == 3:
+            aparam_node = _extend_graph_aparam(
+                aparam,
+                graph.n_node,
+                graph.n_local,
+                atype.shape[0],
+            )
         return self.fitting_net.call_graph(
-            gg, atype, gr=rot_mat, g2=None, h2=None, fparam=fparam_node, aparam=aparam
+            gg,
+            atype,
+            gr=rot_mat,
+            g2=None,
+            h2=None,
+            fparam=fparam_node,
+            aparam=aparam_node,
         )
 
     def compute_or_load_stat(

deepmd/dpmodel/descriptor/dpa1.py

@@ -434,12 +434,13 @@ def uses_graph_lower(self) -> bool:
         """Returns whether this descriptor supports the graph-native lower.
 
         The graph-native lower (``call_graph``) covers the factorizable path
-        AND transformer attention (``attn_layer >= 0``, NeighborGraph PR-D)
-        with concat OR strip type-embedding.  ``exclude_types`` is fully
-        supported via
+        and transformer attention (``attn_layer >= 0``) with concat or strip
+        type embedding. ``exclude_types`` is applied through
         :func:`~deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`.
-        Compressed descriptors are the remaining ineligible config and fall
-        back to the legacy dense path, so those models keep working unchanged.
+        Geo-compressed strip models
+        (``geo_compress=True``, ``attn_layer == 0``) are also graph-eligible;
+        tebd-only compression and compressed descriptors with descriptor-level
+        exclusions stay on the legacy dense path.
 
         Eligibility does NOT imply numerical interchangeability with the
         dense route for every config: with ``smooth_type_embedding=True``
@@ -449,13 +450,13 @@ def uses_graph_lower(self) -> bool:
         """
         if self._graph_lower_disabled:
             return False
-        # compressed descriptors have no graph kernel (geo/tebd tabulation is
-        # dense-only); keep them on the legacy dense path.
         if self.compress:
-            return False
-        # strip is graph-eligible (per-edge factorized embedding, no neighbor
-        # coupling); exclude_types is graph-native via ``apply_pair_exclusion``
-        # (owned at this seam). Only compression / the disable flag force dense.
+            return (
+                self.geo_compress
+                and self.se_atten.tebd_input_mode == "strip"
+                and self.se_atten.attn_layer == 0
+                and not self.se_atten.exclude_types
+            )
         return self.se_atten.tebd_input_mode in ("concat", "strip")
 
     def disable_graph_lower(self) -> None:
@@ -1775,7 +1776,7 @@ def call_graph(
         Notes
         -----
         Known limitations:
-        - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense;
+        - ``tebd_input_mode`` in {"concat", "strip"};
         - ``exclude_types`` is applied graph-natively via ``apply_pair_exclusion``.
         """
         from deepmd.dpmodel.utils.neighbor_graph import (

deepmd/dpmodel/model/make_model.py

@@ -666,28 +666,25 @@ def forward_common_atomic_graph(
             Parameters
             ----------
             atype
-                (N,) flat LOCAL atom types, ``N == sum(n_node)``.
+                (N,) flat local-plus-halo atom types, ``N == sum(n_node)``.
             n_node
-                (nf,) per-frame local atom counts.
+                (nf,) per-frame total node counts.
             edge_index
                 (2, E) ``[src, dst]`` edge endpoints (flat local indices).
             edge_vec
                 (E, 3) neighbor-minus-center edge vectors.
             edge_mask
                 (E,) boolean/0-1 valid-edge mask.
             n_local
-                Per-rank local atom counts for multi-rank inference. Ignored in
-                PR-A (single-rank); accepted for ABI stability.
+                Per-frame owned node counts. Halo fitting outputs are masked.
             fparam
                 Frame parameter, ``(nf, ndf)``.
             aparam
                 Atomic parameter, ``(N, nda)``.
             comm_dict
-                MPI communication metadata. Ignored in PR-A; accepted for ABI
-                stability.
+                Optional MPI communication metadata.
             charge_spin
-                charge/spin conditioning. Ignored in PR-A; accepted for ABI
-                stability with charge/spin-conditioned descriptors.
+                Charge/spin conditioning.
 
             Returns
             -------
@@ -701,6 +698,7 @@ def forward_common_atomic_graph(
                 edge_index=edge_index,
                 edge_vec=edge_vec,
                 edge_mask=edge_mask,
+                n_local=n_local,
             )
             atomic_ret = self.atomic_model.forward_common_atomic_graph(
                 graph, atype, fparam=fparam, aparam=aparam, charge_spin=charge_spin
@@ -739,28 +737,25 @@ def call_common_lower_graph(
             Parameters
             ----------
             atype
-                (N,) flat LOCAL atom types, ``N == sum(n_node)``.
+                (N,) flat local-plus-halo atom types, ``N == sum(n_node)``.
             n_node
-                (nf,) per-frame local atom counts.
+                (nf,) per-frame total node counts.
             edge_index
                 (2, E) ``[src, dst]`` edge endpoints (flat local indices).
             edge_vec
                 (E, 3) neighbor-minus-center edge vectors.
             edge_mask
                 (E,) boolean/0-1 valid-edge mask.
             n_local
-                Per-rank local atom counts for multi-rank inference. Ignored in
-                PR-A (single-rank); accepted for ABI stability.
+                Per-frame owned node counts. Halo fitting outputs are masked.
             fparam
                 Frame parameter, ``(nf, ndf)``.
             aparam
                 Atomic parameter, ``(N, nda)``.
             comm_dict
-                MPI communication metadata. Ignored in PR-A; accepted for ABI
-                stability.
+                Optional MPI communication metadata.
             charge_spin
-                charge/spin conditioning. Ignored in PR-A; accepted for ABI
-                stability with charge/spin-conditioned descriptors.
+                Charge/spin conditioning.
 
             Returns
             -------

deepmd/dpmodel/utils/neighbor_graph/__init__.py

@@ -3,8 +3,9 @@
 
 The unified edge/graph neighbor-list contract and its supporting machinery:
 ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity
-+ edge padding), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher +
-the ``from_dense_quartet`` legacy converter), ``segment`` (mask-aware
++ edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization),
+``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the
+``from_dense_quartet`` legacy converter), ``segment`` (mask-aware
 segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly).
 See the design discussion wanghan-iapcm/deepmd-kit#4.
 """
@@ -25,6 +26,11 @@
     from_dense_quartet,
     graph_from_dense_quartet,
 )
+from .csr import (
+    attach_edge_csr,
+    build_edge_csr,
+    canonicalize_neighbor_graph,
+)
 from .derivatives import (
     edge_force_virial,
 )
@@ -39,6 +45,7 @@
     NeighborGraph,
     apply_pair_exclusion,
     frame_id_from_n_node,
+    node_ownership_mask,
     node_validity_mask,
     pad_and_guard_angles,
     pad_and_guard_edges,
@@ -61,9 +68,12 @@
     "angle_to_node_sum",
     "apply_pair_exclusion",
     "attach_angles",
+    "attach_edge_csr",
     "build_angle_index",
+    "build_edge_csr",
     "build_neighbor_graph",
     "build_neighbor_graph_ase",
+    "canonicalize_neighbor_graph",
     "center_edge_pairs",
     "edge_env_mat",
     "edge_force_virial",
@@ -72,6 +82,7 @@
     "graph_angle_cos",
     "graph_from_dense_quartet",
     "neighbor_graph_from_ijs",
+    "node_ownership_mask",
     "node_validity_mask",
     "pad_and_guard_angles",
     "pad_and_guard_edges",

deepmd/dpmodel/utils/neighbor_graph/angles.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """3-body angle graph: pairs of edges sharing a center within a_rcut.
 
-Angles reference EDGES (angle_index into [0,E)); edge_vec stays the only
-geometry leaf. a_sel is normalization-only (not a truncation). Reuses PR-D's
-center_edge_pairs; a_rcut filters the participating edges.
+Angles reference edges (angle_index into [0,E)); edge_vec stays the only
+geometry leaf. a_sel is normalization-only (not a truncation), and a_rcut
+filters the edges passed to center_edge_pairs.
 """
 
 from __future__ import (
@@ -89,8 +89,8 @@ def build_angle_index(
     # the normalization below.
     dist = safe_for_vector_norm(edge_vec, axis=-1)  # (E,)
     a_edge_mask = xp.astype(edge_mask, xp.bool) & (dist < a_rcut)
-    # compact eager form only (static_nnei not exposed until angle export is
-    # needed, PR-G). dst = edge_index[1, :] per the [src, dst] SoA convention.
+    # The compact eager form groups by destination under the [src, dst] SoA
+    # convention.
     q_e, k_e, pair_mask = center_edge_pairs(
         edge_index[1, :],
         a_edge_mask,