fix(pt): honor is_sorted in SE-Attention tabulation

deepmd/pt/model/descriptor/se_atten.py

@@ -813,7 +813,11 @@ def has_message_passing(self) -> bool:
 
     def need_sorted_nlist_for_lower(self) -> bool:
         """Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
-        return False
+        # Geometric compression uses the tabulate op's sorted-neighbor fold,
+        # which assumes padding and out-of-cutoff neighbors are trailing.
+        # `forward_lower` may receive an unsorted rcut+skin list from LAMMPS,
+        # so request the filtering/sorting pass whenever that op is active.
+        return self.geo_compress
 
 
 class NeighborGatedAttention(nn.Module):

source/op/pt/tabulate_multi_device.cc

@@ -25,6 +25,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
                               const torch::Tensor& em_tensor,
                               const torch::Tensor& two_embed_tensor,
                               int64_t last_layer_size,
+                              bool is_sorted,
                               torch::Tensor& descriptor_tensor) {
   // check input shape
   if (table_tensor.dim() != 2) {
@@ -60,15 +61,17 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor,
   if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em,
-                                     two_embed, nloc, nnei, last_layer_size);
+                                     two_embed, nloc, nnei, last_layer_size,
+                                     is_sorted);
 #else
     throw std::runtime_error(
         "The input tensor is on the GPU, but the GPU support for the "
         "customized OP library is not enabled.");
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   } else if (device == "CPU") {
     deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em,
-                                     two_embed, nloc, nnei, last_layer_size);
+                                     two_embed, nloc, nnei, last_layer_size,
+                                     is_sorted);
   }
 }
 
@@ -80,6 +83,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor,
                                   const torch::Tensor& two_embed_tensor,
                                   const torch::Tensor& dy_tensor,
                                   const torch::Tensor& descriptor_tensor,
+                                  bool is_sorted,
                                   torch::Tensor& dy_dem_x_tensor,
                                   torch::Tensor& dy_dem_tensor,
                                   torch::Tensor& dy_dtwo_tensor) {
@@ -111,18 +115,18 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor,
   // compute
   if (device == "GPU") {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    deepmd::tabulate_fusion_se_a_grad_gpu(dy_dem_x, dy_dem, dy_dtwo, table,
-                                          table_info, em_x, em, two_embed, dy,
-                                          nloc, nnei, last_layer_size);
+    deepmd::tabulate_fusion_se_a_grad_gpu(
+        dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy,
+        nloc, nnei, last_layer_size, is_sorted);
 #else
     throw std::runtime_error(
         "The input tensor is on the GPU, but the GPU support for the "
         "customized OP library is not enabled.");
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   } else if (device == "CPU") {
-    deepmd::tabulate_fusion_se_a_grad_cpu(dy_dem_x, dy_dem, dy_dtwo, table,
-                                          table_info, em_x, em, two_embed, dy,
-                                          nloc, nnei, last_layer_size);
+    deepmd::tabulate_fusion_se_a_grad_cpu(
+        dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy,
+        nloc, nnei, last_layer_size, is_sorted);
   }
 }
 
@@ -641,9 +645,13 @@ class TabulateFusionSeAGradOp
     torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
     torch::Tensor dy_dtwo_tensor = at::Tensor();
     // compute
+    // The non-attention se_a path invokes this op per type-pair block, so
+    // exclusions cannot interleave zero rows. Compressed forward_lower also
+    // requests a sorted nlist through the Python-side
+    // DescrptBlockSeA.need_sorted_nlist_for_lower() contract.
     TabulateFusionSeAGradForward<FPTYPE>(
         table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(),
-        dy_tensor, descriptor_tensor, dy_dem_x_tensor, dy_dem_tensor,
+        dy_tensor, descriptor_tensor, true, dy_dem_x_tensor, dy_dem_tensor,
         dy_dtwo_tensor);
     // save data
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
@@ -783,9 +791,11 @@ class TabulateFusionSeAOp
     torch::Tensor descriptor_tensor =
         torch::empty({em_tensor.size(0), 4, last_layer_size}, options);
     // compute
+    // Keep the sorted fold enabled: exclusions are uniform within each se_a
+    // type-pair invocation, and compressed forward_lower sorts its nlist first.
     TabulateFusionSeAForward<FPTYPE>(table_tensor, table_info_tensor,
                                      em_x_tensor, em_tensor, at::Tensor(),
-                                     last_layer_size, descriptor_tensor);
+                                     last_layer_size, true, descriptor_tensor);
     // save data
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
                             em_tensor, descriptor_tensor});
@@ -870,8 +880,8 @@ class TabulateFusionSeAttenGradOp
     torch::Tensor dy_dtwo_tensor = torch::zeros_like(two_embed_tensor);
     TabulateFusionSeAGradForward<FPTYPE>(
         table_tensor, table_info_tensor, em_x_tensor, em_tensor,
-        two_embed_tensor, dy_tensor, descriptor_tensor, dy_dem_x_tensor,
-        dy_dem_tensor, dy_dtwo_tensor);
+        two_embed_tensor, dy_tensor, descriptor_tensor, is_sorted,
+        dy_dem_x_tensor, dy_dem_tensor, dy_dtwo_tensor);
 
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
                             em_tensor, two_embed_tensor, descriptor_tensor});
@@ -969,9 +979,9 @@ class TabulateFusionSeAttenOp
     torch::Tensor descriptor_tensor =
         torch::empty({em_tensor.size(0), 4, last_layer_size}, options);
     // compute
-    TabulateFusionSeAForward<FPTYPE>(table_tensor, table_info_tensor,
-                                     em_x_tensor, em_tensor, two_embed_tensor,
-                                     last_layer_size, descriptor_tensor);
+    TabulateFusionSeAForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor,
+        two_embed_tensor, last_layer_size, is_sorted, descriptor_tensor);
     // save data
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
                             em_tensor, two_embed_tensor, descriptor_tensor});

source/tests/pt/model/test_compressed_se_atten_forward_lower.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Regression coverage for compressed DPA1 on an unsorted lower nlist.
+
+The geometric tabulation kernel folds trailing padding rows when its input is
+sorted. LAMMPS supplies ``forward_lower`` with an rcut+skin neighbor list that
+may put zero-switch, out-of-cutoff rows before real neighbors. A compressed
+DPA1 model must therefore request the extra filtering/sorting pass before the
+kernel applies that fold.
+"""
+
+import copy
+import unittest
+
+import torch
+
+from deepmd.pt.cxx_op import (
+    ENABLE_CUSTOMIZED_OP,
+)
+from deepmd.pt.model.model import (
+    get_model,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.nlist import (
+    extend_input_and_build_neighbor_list,
+)
+
+from ...seed import (
+    GLOBAL_SEED,
+)
+from .test_forward_lower import (
+    reduce_tensor,
+)
+from .test_permutation import (
+    model_dpa1,
+)
+
+dtype = torch.float64
+
+
+@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built")
+class TestCompressedSeAttenForwardLower(unittest.TestCase):
+    def setUp(self) -> None:
+        model_params = copy.deepcopy(model_dpa1)
+        # Geometric compression is available only for the strip representation
+        # without attention layers.
+        model_params["descriptor"]["tebd_input_mode"] = "strip"
+        model_params["descriptor"]["attn_layer"] = 0
+        self.model = get_model(model_params).to(env.DEVICE)
+
+    def _make_system(self):
+        """Create a sparse periodic system with neighbors in the skin region."""
+        natoms = 6
+        cell = 6.0 * torch.eye(3, dtype=dtype, device=env.DEVICE)
+        generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED)
+        coord = 5.5 * torch.rand(
+            [natoms, 3], dtype=dtype, device=env.DEVICE, generator=generator
+        )
+        atype = torch.tensor([0, 0, 1, 1, 2, 2], dtype=torch.int64, device=env.DEVICE)
+        return coord, atype, cell
+
+    def _min_nbor_dist(self, coord, cell) -> float:
+        """Return the periodic minimum pair distance for table construction."""
+        box = torch.diagonal(cell)
+        diff = coord[:, None, :] - coord[None, :, :]
+        diff = diff - torch.round(diff / box) * box
+        dist = torch.linalg.norm(diff, dim=-1)
+        dist = dist + torch.eye(coord.shape[0], device=coord.device) * 1e10
+        return float(dist.min())
+
+    def test_unsorted_overcut_nlist(self) -> None:
+        coord, atype, cell = self._make_system()
+        rcut = self.model.get_rcut()
+        sel = self.model.get_sel()
+
+        # Use a clean cutoff-bounded list as the uncompressed reference.
+        ec, ea, mapping, nlist = extend_input_and_build_neighbor_list(
+            coord.unsqueeze(0),
+            atype.unsqueeze(0),
+            rcut,
+            sel,
+            mixed_types=self.model.mixed_types(),
+            box=cell.unsqueeze(0),
+        )
+        ref = self.model.forward_lower(ec, ea, nlist, mapping, do_atomic_virial=False)
+
+        self.model.min_nbor_dist = torch.tensor(
+            0.9 * self._min_nbor_dist(coord, cell),
+            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
+        )
+        self.model.enable_compression()
+        self.assertTrue(self.model.need_sorted_nlist_for_lower())
+
+        # Mimic the unsorted rcut+skin list from LAMMPS and deliberately move
+        # its out-of-cutoff/padding rows ahead of the in-cutoff neighbors.
+        ec2, ea2, mapping2, nlist2 = extend_input_and_build_neighbor_list(
+            coord.unsqueeze(0),
+            atype.unsqueeze(0),
+            rcut + 2.0,
+            sum(sel),
+            mixed_types=True,
+            box=cell.unsqueeze(0),
+        )
+        safe_nlist = torch.clamp_min(nlist2, 0)
+        gather_index = safe_nlist.reshape(1, -1, 1).expand(-1, -1, 3)
+        neighbor_coord = torch.gather(ec2, 1, gather_index).view(
+            1, coord.shape[0], -1, 3
+        )
+        center_coord = ec2[:, : coord.shape[0], :].unsqueeze(2)
+        distance = torch.linalg.norm(neighbor_coord - center_coord, dim=-1)
+        real_neighbor = nlist2 >= 0
+        self.assertTrue(torch.any(real_neighbor & (distance <= rcut)).item())
+        self.assertTrue(torch.any(real_neighbor & (distance > rcut)).item())
+
+        nlist2 = torch.flip(nlist2, dims=[-1])
+        out = self.model.forward_lower(
+            ec2, ea2, nlist2, mapping2, do_atomic_virial=False
+        )
+
+        torch.testing.assert_close(out["energy"], ref["energy"], rtol=1e-10, atol=1e-10)
+        natoms = coord.shape[0]
+        ref_force = reduce_tensor(ref["extended_force"], mapping, natoms)
+        out_force = reduce_tensor(out["extended_force"], mapping2, natoms)
+        torch.testing.assert_close(out_force, ref_force, rtol=1e-10, atol=1e-10)
+
+
+if __name__ == "__main__":
+    unittest.main()

source/tests/pt/test_model_compression_se_atten.py

@@ -84,10 +84,13 @@ def _init_models_exclude_types():
     INPUT = str(tests_path / "input.json")
     jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json")))
 
-    # Configure se_atten descriptor with exclude_types
+    # Plain se_atten defaults to a zero average, so excluded rows share the
+    # padding sentinel and exercise the unsorted path in the compressed op.
     jdata["model"]["descriptor"] = {
-        "type": "se_atten_v2",
+        "type": "se_atten",
         "exclude_types": [[0, 1]],
+        "set_davg_zero": True,
+        "tebd_input_mode": "strip",
         "sel": 120,
         "rcut_smth": 0.50,
         "rcut": 6.00,