fix(gpu): preserve exact-capacity neighbor rows

source/lib/src/gpu/prod_env_mat.cu

@@ -110,6 +110,12 @@ __device__ inline void decoding_nbor_info(int& type,
   index = key & 0xFFFFFF;
 }
 
+__device__ inline bool is_padding_nbor_info(const uint_64 key) {
+  // Empty sort slots are initialized bytewise to 0xff. Compare with that
+  // sentinel directly because a full row has a valid key in its final slot.
+  return key == static_cast<uint_64>(-1);
+}
+
 __global__ void get_i_idx(int* i_idx,
                           const int nloc,
                           const int nframes,
@@ -175,7 +181,7 @@ __global__ void fill_nei_iter(int* nei_iter_dev,
   const FPTYPE* key_out = key + nloc * max_nbor_size + row * max_nbor_size;
   int nei_type_cur = -1, nbor_idx_cur = 0;
   int nei_type_pre = -1, nbor_idx_pre = 0;
-  if (col < max_nbor_size && key_out[col] != key_out[max_nbor_size - 1]) {
+  if (col < max_nbor_size && !is_padding_nbor_info(key_out[col])) {
     if (col >= 1) {
       decoding_nbor_info(nei_type_pre, nbor_idx_pre, key_out[col - 1]);
     }
@@ -201,7 +207,7 @@ __global__ void format_nlist_fill_b(int* nlist,
   FPTYPE* key_out = key + nloc * max_nbor_size + row * max_nbor_size;
   int* row_nlist = nlist + row * nlist_size;
   if (col < max_nbor_size) {
-    if (key_out[col] != key_out[max_nbor_size - 1]) {
+    if (!is_padding_nbor_info(key_out[col])) {
       int nei_type = 0, nbor_idx = 0;
       decoding_nbor_info(nei_type, nbor_idx, key_out[col]);
       int out_indx = col - nei_iter[nei_type] + sec[nei_type];

source/lib/tests/test_fmt_nlist.cc

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #include <gtest/gtest.h>
 
+#include <string>
+
 #include "fmt_nlist.h"
 #include "neighbor_list.h"
 
@@ -434,6 +436,69 @@ TEST_F(TestFormatNlistShortSel, gpu) {
   }
 }
 
+TEST(FormatNlistGpu, preserves_exact_capacity_rows) {
+  // Exercise every supported radix-sort size without leaving a padding slot.
+  // These exact boundaries are reachable after the caller rounds row capacity.
+  for (const int max_nbor_size : {256, 512, 1024, 2048, 4096}) {
+    SCOPED_TRACE("max_nbor_size=" + std::to_string(max_nbor_size));
+    const int nloc = 1;
+    const int nall = max_nbor_size + 1;
+    const float rcut = 2.0f;
+    const std::vector<int> sec = {0, max_nbor_size - 1, max_nbor_size};
+    std::vector<double> coord(static_cast<size_t>(nall) * 3, 0.0);
+    std::vector<int> type(nall, 0);
+    std::vector<int> neighbors(max_nbor_size);
+    for (int ii = 0; ii < max_nbor_size; ++ii) {
+      neighbors[ii] = ii + 1;
+      coord[static_cast<size_t>(ii + 1) * 3] = 1.0;
+    }
+    // Put the second type boundary in the final occupied slot so the test also
+    // requires fill_nei_iter to process the last valid key.
+    type[max_nbor_size] = 1;
+
+    std::vector<int> ilist = {0};
+    std::vector<int> numneigh = {max_nbor_size};
+    std::vector<int*> firstneigh = {neighbors.data()};
+    deepmd::InputNlist in_nlist(nloc, ilist.data(), numneigh.data(),
+                                firstneigh.data()),
+        gpu_inlist;
+    std::vector<int> formatted(max_nbor_size, -1);
+
+    double* coord_dev = NULL;
+    int *type_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL,
+        *memory_dev = NULL;
+    uint_64* array_longlong_dev = NULL;
+    deepmd::malloc_device_memory_sync(coord_dev, coord);
+    deepmd::malloc_device_memory_sync(type_dev, type);
+    deepmd::malloc_device_memory_sync(nlist_dev, formatted);
+    deepmd::malloc_device_memory(array_int_dev,
+                                 sec.size() + nloc * sec.size() + nloc);
+    deepmd::malloc_device_memory(array_longlong_dev,
+                                 static_cast<size_t>(nloc) * max_nbor_size * 2);
+    deepmd::malloc_device_memory(memory_dev,
+                                 static_cast<size_t>(nloc) * max_nbor_size);
+    deepmd::convert_nlist_gpu_device(gpu_inlist, in_nlist, memory_dev,
+                                     max_nbor_size);
+
+    format_nbor_list_gpu(nlist_dev, coord_dev, type_dev, gpu_inlist,
+                         array_int_dev, array_longlong_dev, max_nbor_size, nloc,
+                         nall, 1, rcut, sec);
+    deepmd::memcpy_device_to_host(nlist_dev, formatted);
+
+    deepmd::delete_device_memory(nlist_dev);
+    deepmd::delete_device_memory(coord_dev);
+    deepmd::delete_device_memory(type_dev);
+    deepmd::delete_device_memory(array_int_dev);
+    deepmd::delete_device_memory(array_longlong_dev);
+    deepmd::delete_device_memory(memory_dev);
+    deepmd::free_nlist_gpu_device(gpu_inlist);
+
+    for (int ii = 0; ii < max_nbor_size; ++ii) {
+      EXPECT_EQ(formatted[ii], ii + 1) << "neighbor slot " << ii;
+    }
+  }
+}
+
 TEST_F(TestEncodingDecodingNborInfo, valid_nbor_info_gpu) {
   int *valid_type_dev = NULL, *valid_index_dev = NULL, *out_type_dev = NULL,
       *out_index_dev = NULL;