perf(dpa4c): accelerate compressed DPA4C CUDA inference in LAMMPS

deepmd/dpmodel/atomic_model/inner_potential.py

@@ -412,7 +412,8 @@ def call(
         r = xp.clip(r, min=1e-10)
         # Padded atoms carry type -1; their edges are masked below, so the
         # table lookup only needs a valid index for them.
-        atype_row = xp.clip(xp.astype(atype_flat, xp.int64), min=0)
+        atype_row = xp.astype(atype_flat, xp.int64)
+        atype_row = xp.maximum(atype_row, xp.zeros_like(atype_row))
         # Row ``center * (T + 1) + neighbor``, the ordered-pair index the
         # fused kernels use, with the destination as the center.
         row = xp.take(

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -933,7 +933,7 @@ def _compute_weight(
 
         # use the larger rr based on nlist
         nlist_larger = zbl_nlist if zbl_nnei >= dp_nnei else dp_nlist
-        masked_nlist = xp.clip(nlist_larger, 0, None)
+        masked_nlist = xp.maximum(nlist_larger, xp.zeros_like(nlist_larger))
         pairwise_rr = PairTabAtomicModel._get_pairwise_dist(
             extended_coord, masked_nlist
         )

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -434,7 +434,12 @@ def _extract_spline_coefficient(
         expanded_idx = xp.broadcast_to(
             idx[..., xp.newaxis, xp.newaxis], (*idx.shape, 1, 4)
         )
-        clipped_indices = xp.astype(xp.clip(expanded_idx, 0, nspline - 1), xp.int64)
+        clipped_indices = xp.astype(
+            xp.minimum(
+                xp.maximum(expanded_idx, xp.zeros_like(expanded_idx)), nspline - 1
+            ),
+            xp.int64,
+        )
 
         # (nframes, nloc, nnei, 4)
         final_coef = xp.squeeze(

deepmd/dpmodel/descriptor/dpa4c.py

@@ -245,6 +245,8 @@ class DescrptDPA4C(NativeOP, BaseDescriptor):
     _STAT_FRAMES_PER_SAMPLE = 64
     _COMPRESSION_BUFFER_NAMES = (
         "data",
+        "basis_table",
+        "basis_matrix",
         "info",
         "pair_film",
         "pair_mixing",
@@ -258,6 +260,10 @@ class DescrptDPA4C(NativeOP, BaseDescriptor):
         "output_mean",
         "output_inv_std",
     )
+    # The low-rank radial table and its expansion, which a payload compressed
+    # without that table does not carry; empty artifacts select the radial
+    # table itself, so such a payload loads as the snapshot it was.
+    _OPTIONAL_COMPRESSION_BUFFER_NAMES = ("basis_table", "basis_matrix")
 
     def __init__(
         self,
@@ -1931,11 +1937,36 @@ def deserialize(cls, data: dict) -> DescrptDPA4C:
             obj._load_compression(compression)
         return obj
 
+    @classmethod
+    def _compression_variables(cls, compression: dict) -> dict:
+        """Return every compressed-inference artifact of a payload.
+
+        Parameters
+        ----------
+        compression
+            The ``compress`` entry of a serialized descriptor.
+
+        Returns
+        -------
+        dict
+            The artifact of every name in ``_COMPRESSION_BUFFER_NAMES``, an
+            empty float32 array for an optional one the payload omits.
+        """
+        variables = compression["@variables"]
+        return {
+            name: (
+                np.empty(0, dtype=np.float32)
+                if name in cls._OPTIONAL_COMPRESSION_BUFFER_NAMES
+                and name not in variables
+                else variables[name]
+            )
+            for name in cls._COMPRESSION_BUFFER_NAMES
+        }
+
     def _load_compression(self, compression: dict) -> None:
         """Restore the compressed-inference artifact payload."""
-        variables = compression["@variables"]
-        for name in self._COMPRESSION_BUFFER_NAMES:
-            setattr(self, f"compress_{name}", variables[name])
+        for name, value in self._compression_variables(compression).items():
+            setattr(self, f"compress_{name}", value)
         self.compress = True
 
     @classmethod

deepmd/dpmodel/fitting/general_fitting.py

@@ -944,6 +944,49 @@ def readout_reference(self) -> Array:
             device=array_api_compat.device(self.bias_atom_e),
         )
 
+    def fold_case_embedding(self) -> None:
+        """Fold the case embedding into the first layer so the call needs no case columns.
+
+        The case embedding is the same trailing block of input columns for
+        every atom, so its product with the rows of the first-layer weight
+        that read it is a constant of the pre-activation. Adding that product
+        to the bias and removing the rows yields a fitting without the case
+        embedding that gives the same outputs. Export folds the fixed case of
+        a deployed model, so that the frozen branch of a multi-task model with
+        a shared fitting network is a plain fitting network. The fold rebuilds
+        the first layers in place, so it applies to a fitting that no other
+        branch shares. A first layer whose residual adds its input to the
+        output also carries the case embedding past the activation; such a
+        fitting keeps it.
+        """
+        if self.dim_case_embd == 0:
+            return
+        networks = [
+            self.nets[ii] for ii in range(self.ntypes if not self.mixed_types else 1)
+        ]
+        for net in networks:
+            width_in, width_out = net[0].w.shape
+            if net[0].resnet and width_out in (width_in, 2 * width_in):
+                return
+        case_embd = to_numpy_array(self.case_embd).astype(np.float64)
+        for net in networks:
+            data = net[0].serialize()
+            variables = data["@variables"]
+            weight = variables["w"]
+            kept = weight.shape[0] - self.dim_case_embd
+            variables["b"] = (
+                variables["b"].astype(np.float64)
+                + case_embd @ weight[kept:].astype(np.float64)
+            ).astype(weight.dtype)
+            variables["w"] = weight[:kept]
+            # The residual is inactive at the original input width, checked
+            # above; clearing it keeps the narrower input from activating it.
+            data["resnet"] = False
+            net[0] = type(net[0]).deserialize(data)
+            net.in_dim = kept
+        self.dim_case_embd = 0
+        self.case_embd = None
+
     def _call_common(
         self,
         descriptor: Array,

deepmd/pt/model/model/sezm_model.py

@@ -281,6 +281,11 @@
       Workaround for PyTorch <=2.11 bugs pytorch/pytorch#174379,
       #178080, #179494.  All three manifest only under data-dependent
       symbolic shapes -- exactly our edge count.
+* ``inplace_buffers=False``
+      Triton 3.7.1 can compute a reduction wrongly when the kernel stores
+      its result over an input that it loads in the same loop and the
+      tuned configuration replicates the block across warps; separate
+      output buffers leave no kernel that stores over its own inputs.
 
 NOTE 7 -- Multi-slot compile cache key
 --------------------------------------

deepmd/pt/utils/compile_compat.py

@@ -544,9 +544,11 @@ def build_inductor_compile_options(*, inference: bool = False) -> dict[str, Any]
     misbehaved on the combination of data-dependent edge counts and a
     second-order autograd graph -- most visibly the oversized fused Triton
     reduction kernels that fail ``make_ttgir`` (``PassManager::run failed``) on
-    some GPU/Triton combinations. Options absent from the running PyTorch's
-    configuration registry are dropped so the returned dictionary stays valid
-    across releases.
+    some GPU/Triton combinations. It also gives every generated kernel output
+    buffers separate from its inputs, without which Triton 3.7.1 can compute a
+    reduction wrongly. Options absent from the running PyTorch's configuration
+    registry are dropped so the returned dictionary stays valid across
+    releases.
 
     Parameters
     ----------
@@ -595,6 +597,18 @@ def build_inductor_compile_options(*, inference: bool = False) -> dict[str, Any]
         # access. A 1D grid keeps that axis on the x dimension (limit 2**31-1).
         # The option is shared by the training and evaluation graphs.
         "triton.max_tiles": 1,
+        # Inductor can generate a reduction kernel that loads an input inside
+        # its reduction loop and stores a result over that input in the same
+        # iteration. Where the tuned configuration launches more threads than
+        # the block has elements, Triton 3.7.1 replicates the block across
+        # warps: every copy loads, only the first copy stores, and no barrier
+        # separates the loads of an iteration from that store, so a lagging
+        # copy can load the stored value and the reduction comes out wrong
+        # intermittently. The Inductor cache keeps the tuned configuration, so
+        # every later process that loads the kernel from the cache can compute
+        # the reduction wrongly as well. Separate output buffers leave no
+        # kernel that stores over its own inputs.
+        "inplace_buffers": False,
     }
     # ``DP_TUNE_TRAIN`` grades the compile-time investment of the training
     # graphs (cumulative levels; inference graphs ignore it, the AOTI export

deepmd/pt_expt/descriptor/dpa4c.py

@@ -439,8 +439,8 @@ def deserialize(cls, data: dict) -> "DescrptDPA4C":
         if compression is not None:
             obj._set_compression(
                 {
-                    name: torch.as_tensor(compression["@variables"][name])
-                    for name in obj._COMPRESSION_BUFFER_NAMES
+                    name: torch.as_tensor(value)
+                    for name, value in obj._compression_variables(compression).items()
                 }
             )
         return obj

deepmd/pt_expt/infer/deep_eval.py

@@ -259,6 +259,17 @@ def __init__(
         # above (_load_pt2 / _load_pte / _load_pt), so this is the one
         # place that sees every model load regardless of archive kind.
         _warn_legacy_edge_vec(self.metadata)
+        if self.metadata.get("lower_input_kind") in (
+            "dpa1_canonical",
+            "dpa4c_canonical",
+        ) and (
+            self.metadata.get("canonical_index_dtype") != "auto"
+            or self.metadata.get("canonical_index_storage") != "byte_words"
+        ):
+            raise ValueError(
+                "Compact canonical graph artifacts require automatic-width byte "
+                "indices; re-freeze the model with the current DeePMD-kit version."
+            )
 
         # ``neighbor_graph_method`` selects searches for graph and edge-vector
         # inputs. Extended-coordinate inputs use ``nlist_backend`` instead.
@@ -2169,6 +2180,7 @@ def _prepare_graph_inputs(
         from deepmd.pt_expt.utils.canonical_graph import (
             CanonicalGraph,
             canonical_graph_from_neighbor_graph,
+            canonical_index_bytes,
         )
         from deepmd.pt_expt.utils.env import (
             DEVICE,
@@ -2246,7 +2258,7 @@ def _prepare_graph_inputs(
                 atype_t,
                 graph.n_node,
                 graph.n_local,
-                graph.source,
+                canonical_index_bytes(graph.source),
                 graph.edge_vec,
                 graph.destination_row_ptr,
                 graph.source_row_ptr,

deepmd/pt_expt/kernels/autotune.py

@@ -1,17 +1,20 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Freeze-time Triton launch-configuration autotuning.
+"""Freeze-time autotuning of custom-kernel launch choices.
 
 Custom-kernel families register a tuner through :func:`register_autotuner`; a
 model-freezing entry point calls :func:`run_autotune` once, before tracing, so
-the exported artifact bakes launch configurations tuned for the exact
-deployment shapes and GPU. A tuner sweeps only the shape keys its kernels need
-that are not already covered by the built-in tables, keeping repeat freezes
-cheap.
-
-Autotuning is a no-op below ``DP_TRITON_INFER`` level 2 (level-1 kernels run a
-single shape-independent configuration) and off CUDA -- Triton launch tables
-are GPU-specific and AOTInductor artifacts are not portable across GPU models,
-so the tuning target is always the local deployment GPU.
+the exported artifact bakes launch choices tuned for the exact deployment
+shapes and GPU: the Triton launch configurations of the DPA1 kernels, and the
+fitting routes of the fused energy-force operator of compressed DPA4C. A tuner
+measures only the keys its kernels need that are not already covered, keeping
+repeat freezes cheap.
+
+Each tuner runs from the ``DP_TRITON_INFER`` level it registers with (level-1
+Triton kernels run a single shape-independent configuration, so the Triton
+tuners start at level 2), for the lower kinds whose programs call its kernels,
+and only on CUDA: the tables are GPU-specific and AOTInductor artifacts are not
+portable across GPU models, so the tuning target is always the local
+deployment GPU.
 
 Adding a new custom-kernel family is one registration call at import time; the
 freeze paths need no change.
@@ -34,54 +37,77 @@
 
 log = logging.getLogger(__name__)
 
-# Family name -> tuner(model, level, device). Keying by name makes repeated
-# registration (e.g. module re-import) idempotent.
+# Family name -> (tuner(model, level, device), lowest level, lower kinds or
+# None for all). Keying by name makes repeated registration (e.g. module
+# re-import) idempotent.
 Autotuner = Callable[[torch.nn.Module, int, torch.device], None]
-_AUTOTUNERS: dict[str, Autotuner] = {}
+_AUTOTUNERS: dict[str, tuple[Autotuner, int, frozenset[str] | None]] = {}
 
 
-def register_autotuner(name: str, tuner: Autotuner) -> None:
-    """Register a kernel family's freeze-time launch-configuration tuner.
+def register_autotuner(
+    name: str,
+    tuner: Autotuner,
+    min_level: int = 2,
+    lower_kinds: tuple[str, ...] | None = None,
+) -> None:
+    """Register a kernel family's freeze-time tuner.
 
     Parameters
     ----------
     name : str
         Unique family identifier; a repeated name replaces the prior tuner.
     tuner : Callable[[torch.nn.Module, int, torch.device], None]
-        Sweeps the shape keys ``model`` needs for this family that are not
-        already covered and registers the winners into the family's launch
-        table. Receives the resolved ``DP_TRITON_INFER`` level and the target
-        CUDA device.
+        Measures the keys ``model`` needs for this family that are not already
+        covered and registers the winners into the family's table. Receives
+        the resolved ``DP_TRITON_INFER`` level and the target CUDA device.
+    min_level : int, optional
+        Lowest ``DP_TRITON_INFER`` level at which the tuner runs, by default 2,
+        that of the Triton launch tables.
+    lower_kinds : tuple[str, ...] or None, optional
+        Lower kinds whose programs call the kernels of the family; a freeze
+        to another lower kind skips the tuner. By default it runs for every
+        lower kind.
     """
-    _AUTOTUNERS[name] = tuner
+    _AUTOTUNERS[name] = (
+        tuner,
+        min_level,
+        None if lower_kinds is None else frozenset(lower_kinds),
+    )
 
 
-def run_autotune(model: torch.nn.Module, device: torch.device) -> None:
+def run_autotune(model: torch.nn.Module, device: torch.device, lower_kind: str) -> None:
     """Tune every registered kernel family for ``model`` on ``device``.
 
     Called by a freeze path after the model is built and before it is traced,
-    so ``resolve_*_config`` lookups made while tracing return freshly tuned
-    launches. A no-op below ``DP_TRITON_INFER`` level 2 or off CUDA.
+    so the lookups made while tracing return freshly tuned choices. A no-op
+    off CUDA; each tuner runs from the level it registered with, for the lower
+    kinds it registered with.
 
     Parameters
     ----------
     model : torch.nn.Module
         The model about to be frozen; tuners walk it for their shape keys.
     device : torch.device
         The target CUDA device the artifact will run on.
+    lower_kind : str
+        The lower kind of the program the freeze exports.
     """
     level = triton_infer_level()
-    if level < 2 or device.type != "cuda" or not torch.cuda.is_available():
+    if device.type != "cuda" or not torch.cuda.is_available():
         return
-    for name, tuner in _AUTOTUNERS.items():
-        # A tuning failure must degrade to the built-in / default launch tables
-        # rather than abort the freeze; the frozen artifact stays correct, only
-        # its launch configuration is untuned for uncovered shapes.
+    for name, (tuner, min_level, lower_kinds) in _AUTOTUNERS.items():
+        if level < min_level or (
+            lower_kinds is not None and lower_kind not in lower_kinds
+        ):
+            continue
+        # A tuning failure must degrade to the built-in / default tables rather
+        # than abort the freeze; the frozen artifact stays correct, only its
+        # launch choices are untuned for uncovered keys.
         try:
             tuner(model, level, device)
         except Exception:
             log.warning(
-                "Triton autotuner %r failed; keeping default launch configurations.",
+                "Autotuner %r failed; keeping the default launch choices.",
                 name,
                 exc_info=True,
             )

deepmd/pt_expt/kernels/cuda/dpa1/canonical.py

@@ -182,6 +182,11 @@ def _generic_topology(
     source: torch.Tensor,
     destination_row_ptr: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    from deepmd.pt_expt.utils.canonical_graph import (
+        canonical_index_words,
+    )
+
+    source = canonical_index_words(source)
     physical_edge_count = int(destination_row_ptr[-1].item())
     node_count = destination_row_ptr.shape[0] - 1
     destination = torch.repeat_interleave(