Fix IVF extend list resize bug

cpp/include/cuvs/neighbors/ivf_flat.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -11,6 +11,7 @@
 #include <cuvs/neighbors/common.hpp>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/host_mdspan.hpp>
+#include <raft/util/integer_utils.hpp>
 #include <sstream>
 #include <string>
 
@@ -110,7 +111,8 @@ struct list_spec {
   /** Determine the extents of an array enough to hold a given amount of data. */
   constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
   {
-    return raft::make_extents<SizeT>(n_rows, dim);
+    return raft::make_extents<SizeT>(raft::round_up_safe<SizeT>(n_rows, SizeT{kIndexGroupSize}),
+                                     dim);
   }
 };
 

cpp/include/cuvs/neighbors/ivf_sq.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -118,7 +118,8 @@ struct CUVS_EXPORT list_spec {
   constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
   {
     uint32_t padded = raft::round_up_safe<uint32_t>(dim, kVecLen);
-    return raft::make_extents<SizeT>(n_rows, padded);
+    return raft::make_extents<SizeT>(raft::round_up_safe<SizeT>(n_rows, SizeT{kIndexGroupSize}),
+                                     padded);
   }
 };
 

cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -282,11 +282,8 @@ void extend(raft::resources const& handle,
     raft::resource::sync_stream(handle);
     auto& lists = index->lists();
     for (uint32_t label = 0; label < n_lists; label++) {
-      ivf::resize_list(handle,
-                       lists[label],
-                       list_device_spec,
-                       new_list_sizes[label],
-                       raft::Pow2<kIndexGroupSize>::roundUp(old_list_sizes[label]));
+      ivf::resize_list(
+        handle, lists[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]);
     }
   }
   // Update the pointers and the sizes

cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -16,7 +16,6 @@
 #include <raft/core/numpy_serializer.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/serialize.hpp>
-#include <raft/util/pow2_utils.cuh>
 
 #include <fstream>
 
@@ -27,7 +26,7 @@ namespace cuvs::neighbors::ivf_flat::detail {
 // backward compatibility.
 // TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward
 //             compatible fashion.
-constexpr int serialization_version = 4;
+constexpr int serialization_version = 5;
 
 /**
  * Save the index to file.
@@ -72,11 +71,7 @@ void serialize(raft::resources const& handle, std::ostream& os, const index<T, I
 
   list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list(handle,
-                        os,
-                        index_.lists()[label],
-                        list_store_spec,
-                        raft::Pow2<kIndexGroupSize>::roundUp(sizes_host(label)));
+    ivf::serialize_list(handle, os, index_.lists()[label], list_store_spec, sizes_host(label));
   }
   raft::resource::sync_stream(handle);
 }

cpp/src/neighbors/ivf_list.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 

cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -354,11 +354,8 @@ void extend_inplace(raft::resources const& handle,
     raft::resource::sync_stream(handle);
     auto& lists = index->lists();
     for (uint32_t label = 0; label < n_lists; label++) {
-      ivf::resize_list(handle,
-                       lists[label],
-                       list_device_spec,
-                       new_list_sizes[label],
-                       raft::Pow2<kIndexGroupSize>::roundUp(old_list_sizes[label]));
+      ivf::resize_list(
+        handle, lists[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]);
     }
   }
   ivf::detail::recompute_internal_state(handle, *index);

cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -15,13 +15,12 @@
 #include <raft/core/mdarray.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/serialize.hpp>
-#include <raft/util/pow2_utils.cuh>
 
 #include <fstream>
 
 namespace cuvs::neighbors::ivf_sq::detail {
 
-constexpr int serialization_version = 1;
+constexpr int serialization_version = 2;
 
 template <typename CodeT>
 void serialize(raft::resources const& handle, std::ostream& os, const index<CodeT>& index_)
@@ -63,11 +62,7 @@ void serialize(raft::resources const& handle, std::ostream& os, const index<Code
 
   list_spec<uint32_t, CodeT, int64_t> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list(handle,
-                        os,
-                        index_.lists()[label],
-                        list_store_spec,
-                        raft::Pow2<kIndexGroupSize>::roundUp(sizes_host(label)));
+    ivf::serialize_list(handle, os, index_.lists()[label], list_store_spec, sizes_host(label));
   }
   raft::resource::sync_stream(handle);
 }

cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu

@@ -1,12 +1,18 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <gtest/gtest.h>
 
 #include "../ann_ivf_flat.cuh"
 
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/linalg/init.cuh>
+#include <raft/matrix/init.cuh>
+
+#include <optional>
+
 namespace cuvs::neighbors::ivf_flat {
 
 typedef AnnIVFFlatTest<float, float, int64_t> AnnIVFFlatTestF_float;
@@ -19,4 +25,65 @@ TEST_P(AnnIVFFlatTestF_float, AnnIVFFlat)
 
 INSTANTIATE_TEST_CASE_P(AnnIVFFlatTest, AnnIVFFlatTestF_float, ::testing::ValuesIn(inputs));
 
+TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity)
+{
+  raft::resources handle;
+  auto stream = raft::resource::get_cuda_stream(handle);
+
+  constexpr int64_t base_rows = 100;
+  constexpr int64_t grow_rows = 20;
+  constexpr int64_t rows      = base_rows + grow_rows;
+  constexpr int64_t dim       = 4;
+
+  auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);
+  raft::matrix::fill(handle, data.view(), 0.0f);
+
+  index_params params;
+  params.n_lists                        = 1;
+  params.metric                         = cuvs::distance::DistanceType::L2Expanded;
+  params.add_data_on_build              = false;
+  params.kmeans_trainset_fraction       = 1.0;
+  params.adaptive_centers               = false;
+  params.conservative_memory_allocation = false;
+
+  auto all_data_view =
+    raft::make_device_matrix_view<const float, int64_t>(data.data_handle(), rows, dim);
+  auto empty_index = build(handle, params, all_data_view);
+
+  auto base_data_view =
+    raft::make_device_matrix_view<const float, int64_t>(data.data_handle(), base_rows, dim);
+  auto base_index = extend(handle, base_data_view, std::nullopt, empty_index);
+  raft::resource::sync_stream(handle);
+
+  ASSERT_EQ(base_index.lists()[0]->get_size(), base_rows);
+  ASSERT_EQ(base_index.lists()[0]->indices_capacity(), raft::Pow2<kIndexGroupSize>::roundUp(rows));
+
+  auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows);
+  raft::linalg::range(indices.data_handle(), base_rows, rows, stream);
+
+  auto grow_data_view = raft::make_device_matrix_view<const float, int64_t>(
+    data.data_handle() + base_rows * dim, grow_rows, dim);
+  auto grow_indices_view =
+    raft::make_device_vector_view<const int64_t, int64_t>(indices.data_handle(), grow_rows);
+  auto first_grown_index =
+    extend(handle,
+           grow_data_view,
+           std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(grow_indices_view),
+           base_index);
+  raft::resource::sync_stream(handle);
+
+  ASSERT_EQ(first_grown_index.lists()[0].get(), base_index.lists()[0].get());
+  ASSERT_EQ(first_grown_index.lists()[0]->get_size(), rows);
+
+  auto second_grown_index =
+    extend(handle,
+           grow_data_view,
+           std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(grow_indices_view),
+           base_index);
+  raft::resource::sync_stream(handle);
+
+  EXPECT_NE(first_grown_index.lists()[0].get(), second_grown_index.lists()[0].get());
+  EXPECT_EQ(second_grown_index.lists()[0]->get_size(), rows);
+}
+
 }  // namespace cuvs::neighbors::ivf_flat

cpp/tests/neighbors/ann_ivf_sq/test_float_int64_t.cu

@@ -1,17 +1,73 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <gtest/gtest.h>
 
 #include "../ann_ivf_sq.cuh"
 
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/linalg/init.cuh>
+#include <raft/matrix/init.cuh>
+
+#include <optional>
+
 namespace cuvs::neighbors::ivf_sq {
 
 typedef AnnIVFSQTest<float, float, int64_t> AnnIVFSQTestF_float;
 TEST_P(AnnIVFSQTestF_float, AnnIVFSQ) { this->testAll(); }
 
 INSTANTIATE_TEST_CASE_P(AnnIVFSQTest, AnnIVFSQTestF_float, ::testing::ValuesIn(inputs));
 
+TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity)
+{
+  raft::resources handle;
+  auto stream = raft::resource::get_cuda_stream(handle);
+
+  constexpr int64_t base_rows = 100;
+  constexpr int64_t grow_rows = 20;
+  constexpr int64_t rows      = base_rows + grow_rows;
+  constexpr int64_t dim       = 4;
+
+  auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);
+  raft::matrix::fill(handle, data.view(), 0.0f);
+
+  index_params params;
+  params.n_lists                        = 1;
+  params.metric                         = cuvs::distance::DistanceType::L2Expanded;
+  params.add_data_on_build              = false;
+  params.max_train_points_per_cluster   = 256;
+  params.conservative_memory_allocation = false;
+
+  auto all_data_view =
+    raft::make_device_matrix_view<const float, int64_t>(data.data_handle(), rows, dim);
+  auto index = build(handle, params, all_data_view);
+
+  auto base_data_view =
+    raft::make_device_matrix_view<const float, int64_t>(data.data_handle(), base_rows, dim);
+  extend(handle, base_data_view, std::nullopt, &index);
+  raft::resource::sync_stream(handle);
+
+  ASSERT_EQ(index.lists()[0]->get_size(), base_rows);
+  ASSERT_EQ(index.lists()[0]->indices_capacity(), raft::Pow2<kIndexGroupSize>::roundUp(rows));
+  auto* base_list = index.lists()[0].get();
+
+  auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows);
+  raft::linalg::range(indices.data_handle(), base_rows, rows, stream);
+
+  auto grow_data_view = raft::make_device_matrix_view<const float, int64_t>(
+    data.data_handle() + base_rows * dim, grow_rows, dim);
+  auto grow_indices_view =
+    raft::make_device_vector_view<const int64_t, int64_t>(indices.data_handle(), grow_rows);
+  extend(handle,
+         grow_data_view,
+         std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(grow_indices_view),
+         &index);
+  raft::resource::sync_stream(handle);
+
+  EXPECT_EQ(index.lists()[0].get(), base_list);
+  EXPECT_EQ(index.lists()[0]->get_size(), rows);
+}
+
 }  // namespace cuvs::neighbors::ivf_sq