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

## Summary
- Replace CPU-side `buildTypeSortedNlist` with `createNlistTensor(data,
expected_nnei)` — avoids distance computation and type sorting every
step; model's compiled `format_nlist` handles this on-device
- Export with `do_atomic_virial=False` by default — avoids 3 extra
`torch.autograd.grad` backward passes; add `--atomic-virial` flag to `dp
convert-backend`
- Cache `mapping_tensor` as member variable, only rebuild when `ago ==
0`
- Store `nnei` and `do_atomic_virial` in .pt2 metadata for C++ to read
at init
- Make nnei dynamic in torch.export — compiled graph accepts
variable-size neighbor lists via internal padding + sort branch

## Benchmark (V100-SXM2-16GB, 192-atom water, LAMMPS MD)

### Before this PR

.pt2 was **9x slower** than .pth due to CPU-side nlist sorting, baked-in
atomic virial backward passes, and excessive clones:

| Atoms | .pth (ms/step) | .pt2 (ms/step) | .pt2/.pth |
|------:|---------------:|---------------:|:---------:|
| 192 | 11 | 97 | **8.8x** |

### After this PR

#### DPA1 L0 (se_atten nlayer=0)

| Atoms | .pth (ms) | .pt2 (ms) | .pt2/.pth |
|------:|----------:|----------:|:---------:|
| 192 | 5.60 | 4.93 | **0.88x** |
| 384 | 6.69 | 8.45 | **1.26x** |
| 768 | 10.9 | 16.3 | **1.49x** |
| 1536 | 19.3 | 31.2 | **1.62x** |
| 3072 | 36.7 | 58.8 | **1.60x** |
| 6144 | 72.0 | 116 | **1.62x** |
| 12288 | 140 | 229 | **1.63x** |

#### DPA1 L2 (se_atten nlayer=2)

| Atoms | .pth (ms) | .pt2 (ms) | .pt2/.pth |
|------:|----------:|----------:|:---------:|
| 192 | 13.0 | 9.17 | **0.71x** |
| 384 | 22.2 | 16.2 | **0.73x** |
| 768 | 41.0 | 30.4 | **0.74x** |
| 1536 | 77.8 | 58.8 | **0.76x** |

#### DPA2 (repinit + repformer)

| Atoms | .pth (ms) | .pt2 (ms) | .pt2/.pth |
|------:|----------:|----------:|:---------:|
| 192 | 28.5 | 15.6 | **0.55x** |
| 384 | 34.6 | 28.2 | **0.81x** |
| 768 | 60.5 | 53.4 | **0.88x** |
| 1536 | 112.9 | 104 | **0.92x** |

For models with more compute (DPA1 L2, DPA2), .pt2 is **24-45% faster**
than .pth. For the smallest model (DPA1 L0), .pt2 has higher per-call
overhead that dominates at large atom counts.

## Test plan
- [x] All Python export/make_fx tests pass (74 tests)
- [x] All Python model tests pass
- [x] All C++ ctest pass (0 failures)
- [x] All 37 LAMMPS .pt2 tests pass
- [x] V100 benchmark confirms speedup for DPA1 L2 and DPA2


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added `--atomic-virial` command-line flag to enable atomic virial
correction during model conversion and export operations
* Models exported with this feature now include per-atom virial
contributions for improved computational accuracy
  * Atomic virial support available for all exportable model formats

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

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