perf(dpa1): optimize graph CUDA inference and deployment (#5758)

DPA1 graph inference previously decomposed descriptor, fitting, and
force/virial work into many generic tensor operations, retained large
edge-scale autograd state, and lacked a shared sparse topology contract
across
Python export, the C++ runtime, and LAMMPS. Build an inference-oriented
graph
pipeline that keeps eligible attention-free DPA1 models and their
neighbor data
on the GPU while retaining explicit fallbacks for unsupported
configurations.

- Add inference-only CUDA operators for attention-free concat and strip
DPA1.
  The uncompressed path fuses environment construction, the three-layer
embedding network, type handling, moment reduction, Gram contraction,
and
analytical edge-gradient backward. The geometrically compressed strip
path
fuses quintic table interpolation, type-pair gating, moments,
descriptors,
and analytical gradients. It supports masked edges, canonical or
permutation
  CSR, int32 or int64 addressing, width buckets from 8 through 256,
  non-bucket padding and slicing, and axis widths up to 16.

- Add descriptor-agnostic fused energy fitting and force/virial
operators.
Run eligible fitting networks through pedantic FP32 cuBLAS GEMMs with
TF32
disabled and fused bias, activation, timestep, residual, and derivative
epilogues, while retaining atomic and per-frame energies in FP64. Reduce
both
edge incidences with one warp per node and no global floating-point
atomics,
  then accumulate frame virials through FP64 partial reductions without
  materializing an (E, 9) outer-product tensor.

- Define cumulative DP_CUDA_INFER levels. Level 1 uses separately
registered
descriptor, fitting, and CSR force operators with first-order autograd.
Level 2 returns force and virial as values, using one opaque end-to-end
operator for eligible uncompressed DPA1 and an explicit custom-operator
chain
for compressed DPA1. This removes the inference autograd tape,
suppresses the
  unused rotation output, avoids retaining the descriptor only for shape
  metadata, and shortens saved-state lifetimes for memory reuse.

- Add make_fx- and export-composable Triton kernels for the PyTorch
  se-family environment matrix and the dense and graph DPA1 environment
convolutions. Their closed-form first backward preserves the force path
while
  avoiding decomposed gather, switch, normalization, gating, and segment
  reduction chains. DP_TRITON_INFER=2 resolves per-GPU launch tables and
freeze-time tuning fills uncovered model shapes; level 3 optionally
enables
the compensated FP16x3 final embedding GEMM. The default CUDA graph path
  remains FP32.

- Select compressed-kernel resource policy at first uncaptured use from
balanced or occupancy-oriented 128/256-thread launches, cache the result
by
device, direction, descriptor shape, topology, index width, and workload
class, and retain architecture-specific fallbacks for small inputs,
capture,
or tuning failures. Build the native operators only with a CUDA-enabled
PyTorch and include a lowest-supported portable PTX target alongside
native
  code.

- Extend NeighborGraph with keyword-only destination/source CSR orders,
int64 row pointers, and an explicit destination_sorted property without
changing its original positional constructor. Generic builders preserve
payload order and omit CSR by default; consumers opt into CSR or stable
destination-major canonicalization, which moves masked guards to the
suffix
and makes destination order the identity. Edge masks remain
authoritative
inside every row, and export validates permutations, row membership,
bounds,
  and canonical identity before tracing.

- Make graph-form AOTInductor deployment select the optimized path
  automatically for eligible .pt2 conversion and export compressed graph
models directly. Keep frame, node, and edge axes dynamic; preserve
custom
operators through export; record the graph edge-vector dtype in
metadata;
and make per-atom virial part of the graph artifact contract required by
the
Kokkos consumer. Eligible compressed artifacts accept FP32 edge geometry
directly, while generic and uncompressed graph artifacts retain the FP64
geometry ABI. Use int64 Inductor indexing and bounded Triton launch
tiling for
  large dynamic graphs.

- Extend DeepEval and the C++ API to consume the same canonical graph
ABI.
Add device-edge capability and dtype queries, FP32 and FP64 edge-vector
overloads, runtime frame and atomic parameters, total-versus-owned node
counts, and optional communication metadata. Host ingestion
canonicalizes
arbitrary payloads, while the device path constructs destination
identity
  CSR and source order with histogram, prefix-sum, and counting scatter.

- Add pair_style deepmd/kk for device-resident edge and graph .pt2
inference. Build compact model-cutoff edges from the Kokkos full
neighbor
list with count, scan, and fill passes; compact NULL-mapped atom types;
emit
FP32 or FP64 vectors according to artifact metadata; and scatter
model-node
outputs back on device. Single-rank execution folds periodic images onto
local owners, whereas domain decomposition retains local-plus-halo nodes
and
explicitly reverse-communicates force and centroid per-atom virial
through
  either device-aware or host-staged communication. Message-passing edge
artifacts use their with-comm forward, including empty-rank phantom
inputs;
message-passing graph artifacts fail fast because that multi-rank ABI is
not
  supported.

- Preserve graph correctness at ownership and masking boundaries.
Exclude halo
fitting outputs from energy, zero halo atomic parameters, retain
virtual-atom
and pair-exclusion masks in fused level 2, handle zero-node graphs, and
reduce
frame virial from node virials instead of a contended edge scatter.
Validate
and broadcast multi-frame fparam/aparam layouts, fix wide
atomic-parameter
allocation when filtering virtual atoms, and avoid copying incompletely
constructed LAMMPS model members. Reduce graph-builder overhead with
bounded
  GPU neighbor-capacity estimates and direct length-three periodic-shift
  reductions.

- Require regeneration of graph-form .pt2 artifacts because their
positional
ABI now includes n_local and both CSR views. Compressed graph export is
  limited to FP32, geometrically compressed, attention-free strip models
  without excluded type pairs, with output width at most 256 and
axis_neuron <= 16. The uncompressed CUDA descriptor requires three FP32
  embedding layers whose widths stay equal or double within its compiled
bounds. Other configurations continue through Triton, reference graph,
or
nlist lowering. deepmd/kk additionally requires a GPU edge/graph
artifact,
  one model, and a valid atom map.

- Add regression coverage for CUDA and Triton forward/gradient parity,
compressed and uncompressed end-to-end energy, force, global virial,
atomic
energy, and atom virial; smooth and non-smooth one- and two-sided type
gates;
  residual and activation variants; non-power-of-two widths; int32/int64
addressing; masked cached edges; canonical and permuted CSR; ownership;
zero-node and many-frame reductions; make_fx, torch.compile, and dynamic
  graph export. Extend DeepEval and C++ tests for dynamic edge counts,
multi-rank graph ingestion, parameter validation and broadcasting,
atomic
  outputs, CSR construction, and virtual-atom parameter preservation.

Author: OutisLi

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,