perf(pt2): optimize .pt2 C++ inference path

deepmd/dpmodel/model/make_model.py

@@ -614,7 +614,13 @@ def _format_nlist(
                     axis=-1,
                 )
 
-            if n_nnei > nnei or extra_nlist_sort:
+            # Order matters for torch.export: Python evaluates `or` left-to-right
+            # with short-circuit.  When `extra_nlist_sort=True` (Python bool) is
+            # on the left, the right-hand `n_nnei > nnei` is not evaluated, so no
+            # symbolic guard is registered on the dynamic `n_nnei` dimension.
+            # Swapping the operands would force the SymInt comparison to run and
+            # emit an `_assert_scalar` node in the exported graph.
+            if extra_nlist_sort or n_nnei > nnei:
                 n_nf, n_nloc, n_nnei = nlist.shape
                 # make a copy before revise
                 m_real_nei = nlist >= 0

deepmd/entrypoints/convert_backend.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import logging
 from typing import (
     Any,
 )
@@ -7,11 +8,14 @@
     Backend,
 )
 
+log = logging.getLogger(__name__)
+
 
 def convert_backend(
     *,  # Enforce keyword-only arguments
     INPUT: str,
     OUTPUT: str,
+    atomic_virial: bool = False,
     **kwargs: Any,
 ) -> None:
     """Convert a model file from one backend to another.
@@ -20,12 +24,31 @@ def convert_backend(
     ----------
     INPUT : str
         The input model file.
-    INPUT : str
+    OUTPUT : str
         The output model file.
+    atomic_virial : bool
+        If True, export .pt2/.pte models with per-atom virial correction.
+        This adds ~2.5x inference cost.  Default False.  Silently ignored
+        (with a warning) for backends that don't support the flag.
     """
     inp_backend: Backend = Backend.detect_backend_by_model(INPUT)()
     out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)()
     inp_hook = inp_backend.serialize_hook
     out_hook = out_backend.deserialize_hook
     data = inp_hook(INPUT)
-    out_hook(OUTPUT, data)
+    # Forward atomic_virial to pt_expt deserialize_to_file if applicable;
+    # warn and skip the flag for backends that don't accept it so that
+    # scripts passing --atomic-virial indiscriminately don't break.
+    import inspect
+
+    sig = inspect.signature(out_hook)
+    if "do_atomic_virial" in sig.parameters:
+        out_hook(OUTPUT, data, do_atomic_virial=atomic_virial)
+    else:
+        if atomic_virial:
+            log.warning(
+                "--atomic-virial is only meaningful for pt_expt .pt2/.pte "
+                "outputs; ignoring it for output backend %s",
+                out_backend.name,
+            )
+        out_hook(OUTPUT, data)

deepmd/main.py

@@ -915,6 +915,15 @@ def main_parser() -> argparse.ArgumentParser:
     )
     parser_convert_backend.add_argument("INPUT", help="The input model file.")
     parser_convert_backend.add_argument("OUTPUT", help="The output model file.")
+    parser_convert_backend.add_argument(
+        "--atomic-virial",
+        action="store_true",
+        default=False,
+        help="Export .pt2/.pte models with per-atom virial correction. "
+        "This adds ~2.5x inference cost but is required for "
+        "LAMMPS compute/atom virial output. "
+        "Ignored (with a warning) for other output backends.",
+    )
 
     # * show model ******************************************************************
     parser_show = subparsers.add_parser(

deepmd/pt_expt/model/dipole_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import types
 from typing import (
     Any,
 )
@@ -16,6 +17,7 @@
 )
 
 from .make_model import (
+    _pad_nlist_for_export,
     make_model,
 )
 from .model import (
@@ -137,6 +139,7 @@ def fn(
             aparam: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
             extended_coord = extended_coord.detach().requires_grad_(True)
+            nlist = _pad_nlist_for_export(nlist)
             return model.forward_lower(
                 extended_coord,
                 extended_atype,
@@ -147,6 +150,13 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn, **make_fx_kwargs)(
-            extended_coord, extended_atype, nlist, mapping, fparam, aparam
-        )
+        # See make_model.py for the rationale of the pad + monkeypatch.
+        _orig_need_sort = model.need_sorted_nlist_for_lower
+        model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model)
+        try:
+            traced = make_fx(fn, **make_fx_kwargs)(
+                extended_coord, extended_atype, nlist, mapping, fparam, aparam
+            )
+        finally:
+            model.need_sorted_nlist_for_lower = _orig_need_sort
+        return traced

deepmd/pt_expt/model/dos_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import types
 from typing import (
     Any,
 )
@@ -16,6 +17,7 @@
 )
 
 from .make_model import (
+    _pad_nlist_for_export,
     make_model,
 )
 from .model import (
@@ -117,6 +119,7 @@ def fn(
             aparam: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
             extended_coord = extended_coord.detach().requires_grad_(True)
+            nlist = _pad_nlist_for_export(nlist)
             return model.forward_lower(
                 extended_coord,
                 extended_atype,
@@ -127,6 +130,13 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn, **make_fx_kwargs)(
-            extended_coord, extended_atype, nlist, mapping, fparam, aparam
-        )
+        # See make_model.py for the rationale of the pad + monkeypatch.
+        _orig_need_sort = model.need_sorted_nlist_for_lower
+        model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model)
+        try:
+            traced = make_fx(fn, **make_fx_kwargs)(
+                extended_coord, extended_atype, nlist, mapping, fparam, aparam
+            )
+        finally:
+            model.need_sorted_nlist_for_lower = _orig_need_sort
+        return traced

deepmd/pt_expt/model/dp_linear_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import types
 from typing import (
     Any,
 )
@@ -19,6 +20,7 @@
 )
 
 from .make_model import (
+    _pad_nlist_for_export,
     make_model,
 )
 from .model import (
@@ -142,6 +144,7 @@ def fn(
             aparam: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
             extended_coord = extended_coord.detach().requires_grad_(True)
+            nlist = _pad_nlist_for_export(nlist)
             return model.forward_lower(
                 extended_coord,
                 extended_atype,
@@ -152,9 +155,16 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn, **make_fx_kwargs)(
-            extended_coord, extended_atype, nlist, mapping, fparam, aparam
-        )
+        # See make_model.py for the rationale of the pad + monkeypatch.
+        _orig_need_sort = model.need_sorted_nlist_for_lower
+        model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model)
+        try:
+            traced = make_fx(fn, **make_fx_kwargs)(
+                extended_coord, extended_atype, nlist, mapping, fparam, aparam
+            )
+        finally:
+            model.need_sorted_nlist_for_lower = _orig_need_sort
+        return traced
 
     @classmethod
     def update_sel(

deepmd/pt_expt/model/dp_zbl_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import types
 from typing import (
     Any,
 )
@@ -16,6 +17,7 @@
 )
 
 from .make_model import (
+    _pad_nlist_for_export,
     make_model,
 )
 from .model import (
@@ -139,6 +141,7 @@ def fn(
             aparam: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
             extended_coord = extended_coord.detach().requires_grad_(True)
+            nlist = _pad_nlist_for_export(nlist)
             return model.forward_lower(
                 extended_coord,
                 extended_atype,
@@ -149,6 +152,15 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn, **make_fx_kwargs)(
-            extended_coord, extended_atype, nlist, mapping, fparam, aparam
-        )
+        # Force `_format_nlist`'s sort branch into the compiled graph so the
+        # exported model tolerates oversized nlists at runtime — see
+        # make_model.py for the full rationale.
+        _orig_need_sort = model.need_sorted_nlist_for_lower
+        model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model)
+        try:
+            traced = make_fx(fn, **make_fx_kwargs)(
+                extended_coord, extended_atype, nlist, mapping, fparam, aparam
+            )
+        finally:
+            model.need_sorted_nlist_for_lower = _orig_need_sort
+        return traced

deepmd/pt_expt/model/make_model.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import math
+import types
 from typing import (
     Any,
 )
@@ -28,6 +29,28 @@
 )
 
 
+def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor:
+    """Append a single ``-1`` column to ``nlist`` for export-time tracing.
+
+    Used inside ``forward_common_lower_exportable`` (and its spin counterpart)
+    so that ``_format_nlist``'s terminal slice ``ret[..., :nnei]`` truncates
+    to a statically sized output.  Without the extra column, torch.export
+    cannot prove the ``ret.shape[-1] == nnei`` assertion at trace time and
+    would specialise the dynamic ``nnei`` dim to the sample value.
+
+    Combined with the short-circuit order in ``_format_nlist``
+    (``extra_nlist_sort`` on the left) and the ``need_sorted_nlist_for_lower``
+    override during tracing, this keeps the compiled graph's ``nnei`` axis
+    fully dynamic and free of symbolic shape guards.
+    """
+    pad = -torch.ones(
+        (*nlist.shape[:2], 1),
+        dtype=nlist.dtype,
+        device=nlist.device,
+    )
+    return torch.cat([nlist, pad], dim=-1)
+
+
 def _cal_hessian_ext(
     model: Any,
     kk: str,
@@ -346,6 +369,7 @@ def fn(
                 aparam: torch.Tensor | None,
             ) -> dict[str, torch.Tensor]:
                 extended_coord = extended_coord.detach().requires_grad_(True)
+                nlist = _pad_nlist_for_export(nlist)
                 return model.forward_common_lower(
                     extended_coord,
                     extended_atype,
@@ -356,13 +380,26 @@ def fn(
                     do_atomic_virial=do_atomic_virial,
                 )
 
-            return make_fx(fn, **make_fx_kwargs)(
-                extended_coord,
-                extended_atype,
-                nlist,
-                mapping,
-                fparam,
-                aparam,
+            # Force `_format_nlist`'s sort branch into the compiled graph so the
+            # exported model tolerates oversized nlists at runtime (LAMMPS builds
+            # nlists with rcut+skin).  Combined with the short-circuit order in
+            # `_format_nlist`, no symbolic guard on the dynamic `nnei` axis is
+            # emitted.
+            _orig_need_sort = model.need_sorted_nlist_for_lower
+            model.need_sorted_nlist_for_lower = types.MethodType(
+                lambda self: True, model
             )
+            try:
+                traced = make_fx(fn, **make_fx_kwargs)(
+                    extended_coord,
+                    extended_atype,
+                    nlist,
+                    mapping,
+                    fparam,
+                    aparam,
+                )
+            finally:
+                model.need_sorted_nlist_for_lower = _orig_need_sort
+            return traced
 
     return CM

deepmd/pt_expt/model/polar_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import types
 from typing import (
     Any,
 )
@@ -16,6 +17,7 @@
 )
 
 from .make_model import (
+    _pad_nlist_for_export,
     make_model,
 )
 from .model import (
@@ -117,6 +119,7 @@ def fn(
             aparam: torch.Tensor | None,
         ) -> dict[str, torch.Tensor]:
             extended_coord = extended_coord.detach().requires_grad_(True)
+            nlist = _pad_nlist_for_export(nlist)
             return model.forward_lower(
                 extended_coord,
                 extended_atype,
@@ -127,6 +130,13 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn, **make_fx_kwargs)(
-            extended_coord, extended_atype, nlist, mapping, fparam, aparam
-        )
+        # See make_model.py for the rationale of the pad + monkeypatch.
+        _orig_need_sort = model.need_sorted_nlist_for_lower
+        model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model)
+        try:
+            traced = make_fx(fn, **make_fx_kwargs)(
+                extended_coord, extended_atype, nlist, mapping, fparam, aparam
+            )
+        finally:
+            model.need_sorted_nlist_for_lower = _orig_need_sort
+        return traced