perf(pt/dpa1): optimize graph CUDA inference and deployment

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

@@ -313,6 +313,42 @@ 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_inputs(
+        self,
+        graph: "NeighborGraph",
+        atype: Array,
+    ) -> tuple["NeighborGraph", Array, Array]:
+        """Apply graph masks shared by standard and fused atomic forwards."""
+        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))
+        output_mask = atom_mask
+        if graph.n_local is not None:
+            from deepmd.dpmodel.utils.neighbor_graph import (
+                node_ownership_mask,
+            )
+
+            output_mask = output_mask & node_ownership_mask(
+                graph.n_node,
+                graph.n_local,
+                atype.shape[0],
+            )
+        if self.pair_excl is not None:
+            keep = self.pair_excl.build_edge_exclude_mask(
+                graph.edge_index, atype_clamped
+            )
+            graph = dataclasses.replace(
+                graph,
+                edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype),
+            )
+        if self.atom_excl is not None:
+            output_mask = xp.logical_and(
+                output_mask,
+                self.atom_excl.build_type_exclude_mask(atype_clamped),
+            )
+        return graph, atype_clamped, output_mask
+
     def forward_common_atomic_graph(
         self,
         graph: "NeighborGraph",
@@ -352,26 +388,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))
-        if self.pair_excl is not None:
-            keep = self.pair_excl.build_edge_exclude_mask(
-                graph.edge_index, atype_clamped
-            )
-            graph = dataclasses.replace(
-                graph,
-                edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype),
-            )
+        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 _finalize_atomic_ret(
         self, ret_dict: dict, atom_mask: Array, atype: Array
@@ -399,10 +424,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

@@ -431,21 +431,25 @@ 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. Remaining ineligible configs
-        (``exclude_types``, and compressed descriptors) fall back to the legacy
-        dense path, so those models keep working unchanged.
+        and transformer attention (``attn_layer >= 0``) with concat or strip
+        type-embedding. Geo-compressed strip models
+        (``geo_compress=True``, ``attn_layer == 0``) are also graph-eligible;
+        tebd-only compression (``geo_compress=False``) and ``exclude_types``
+        stay on the legacy dense path.
 
         Eligibility does NOT imply numerical interchangeability with the
         dense route for every config: with ``smooth_type_embedding=True``
         the carry-all graph attention is sel-independent by design and
         differs from the dense lower by up to ~1e-4 (see the Notes of
         :meth:`call_graph`).
         """
-        # compressed descriptors have no graph kernel (geo/tebd tabulation is
-        # dense-only); keep them on the legacy dense path.
         if self.compress:
-            return False
+            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
+            )
         # exclude_types stays dense (graph exclusion is owned elsewhere); strip is
         # now graph-eligible (per-edge factorized embedding, no neighbor coupling).
         return (
@@ -1756,7 +1760,7 @@ def call_graph(
         -----
         Known limitations:
         - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense;
-        - ``exclude_types`` is not yet supported and raises (lands in a later PR).
+        - ``exclude_types`` is not supported and raises.
         """
         from deepmd.dpmodel.utils.neighbor_graph import (
             edge_env_mat,
@@ -1768,10 +1772,7 @@ def call_graph(
                 f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}"
             )
         if self.exclude_types:
-            raise NotImplementedError(
-                "graph path does not yet apply exclude_types (NeighborGraph PR-A); "
-                "type exclusion lands in a later PR"
-            )
+            raise NotImplementedError("graph path does not support exclude_types")
         if type_embedding is None:
             raise ValueError("type_embedding is required for the graph path")
         xp = array_api_compat.array_namespace(graph.edge_vec)

deepmd/dpmodel/model/make_model.py

@@ -641,28 +641,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
             -------
@@ -676,6 +673,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
@@ -714,28 +712,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.
 """
@@ -24,6 +25,10 @@
     build_neighbor_graph,
     from_dense_quartet,
 )
+from .csr import (
+    build_edge_csr,
+    canonicalize_neighbor_graph,
+)
 from .derivatives import (
     edge_force_virial,
 )
@@ -37,6 +42,7 @@
     GraphLayout,
     NeighborGraph,
     frame_id_from_n_node,
+    node_ownership_mask,
     node_validity_mask,
     pad_and_guard_angles,
     pad_and_guard_edges,
@@ -59,15 +65,18 @@
     "angle_to_node_sum",
     "attach_angles",
     "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",
     "frame_id_from_n_node",
     "from_dense_quartet",
     "graph_angle_cos",
     "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,