Fix IVF extend list resize bug by qwertyforce · Pull Request #2297 · NVIDIA/cuvs

qwertyforce · 2026-07-04T13:01:19Z

Description

Fixes a bug in IVF-Flat and IVF-SQ extend() where a list could receive new vectors but keep its old logical list::size.

The issue happens when extending a list that already has enough allocated capacity. Interleaved IVF list storage may be padded, so the old copy extent can be larger than the old logical size.

Example:

old logical size: 100
old padded copy size: 128
new logical size: 120

Previously, resize_list() received only the padded value, 128, and used it both for copying old storage and for deciding whether list::size needed to be updated. It checked 120 <= 128, skipped the update, and left list::size == 100 even though new vectors were inserted up to 120.

This stale size can break the shared-list / copy-on-write logic in later extends and corrupt search results.

Fix

Pass separate old sizes to ivf::resize_list():

old_logical_size: used for the list::size update/CAS decision
old_used_size:    used for copying old storage, may be padded

The existing overload is kept for call sites where both sizes are the same.

Tests

Added regression tests for:

IVF-Flat repeated extend from a shared base index with enough existing capacity
IVF-SQ in-place extend with enough existing capacity

copy-pr-bot · 2026-07-04T13:01:23Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

achirkin · 2026-07-05T03:48:20Z

old_used_size is the old logical size. old real size is checked via extends.

qwertyforce · 2026-07-05T05:30:36Z

old_used_size is the old logical size. old real size is checked via extends.

in ivf-flat extend, resize_list() receives raft::Pow2<kIndexGroupSize>::roundUp(old_list_sizes[label]) as an argument, right?

qwertyforce · 2026-07-05T06:06:45Z

here is poc

// ivf_flat_extend_size_poc.cpp

#include <cuvs/neighbors/ivf_flat.hpp>

#include <raft/core/copy.cuh>
#include <raft/core/device_mdarray.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/core/resource/cuda_stream.hpp>
#include <raft/core/resources.hpp>

#include <cstdint>
#include <iostream>
#include <numeric>
#include <optional>
#include <vector>

namespace ivf_flat = cuvs::neighbors::ivf_flat;

template <typename IndexT>
uint32_t visible_list_size(raft::resources const& handle, IndexT& index)
{
  auto stream = raft::resource::get_cuda_stream(handle);
  uint32_t h = 0;
  raft::copy(&h, index.list_sizes().data_handle(), 1, stream);
  raft::resource::sync_stream(handle);
  return h;
}

template <typename IndexT>
void print_sizes(char const* tag, raft::resources const& handle, IndexT& index)
{
  std::cout << tag << "\n"
            << "  index.list_sizes()[0] = " << visible_list_size(handle, index) << "\n"
            << "  list::size            = " << index.lists()[0]->get_size() << "\n"
            << "  capacity              = " << index.lists()[0]->indices_capacity() << "\n";
}

int main()
{
  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;

  std::vector<float> h_data(rows * dim);
  for (int64_t r = 0; r < rows; ++r) {
    for (int64_t c = 0; c < dim; ++c) {
      h_data[r * dim + c] = static_cast<float>(r + c);
    }
  }

  auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);
  raft::copy(data.data_handle(), h_data.data(), h_data.size(), stream);

  ivf_flat::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_view = raft::make_device_matrix_view<const float, int64_t, raft::row_major>(
    data.data_handle(), rows, dim);

  auto empty_index = ivf_flat::build(handle, params, all_view);

  auto base_view = raft::make_device_matrix_view<const float, int64_t, raft::row_major>(
    data.data_handle(), base_rows, dim);

  auto base_index = ivf_flat::extend(handle, base_view, std::nullopt, empty_index);
  raft::resource::sync_stream(handle);

  print_sizes("after base extend", handle, base_index);

  std::vector<int64_t> h_ids(grow_rows);
  std::iota(h_ids.begin(), h_ids.end(), base_rows);

  auto ids = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows);
  raft::copy(ids.data_handle(), h_ids.data(), h_ids.size(), stream);

  auto grow_view = raft::make_device_matrix_view<const float, int64_t, raft::row_major>(
    data.data_handle() + base_rows * dim, grow_rows, dim);

  auto ids_view =
    raft::make_device_vector_view<const int64_t, int64_t>(ids.data_handle(), grow_rows);

  auto grown_index = ivf_flat::extend(
    handle,
    grow_view,
    std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(ids_view),
    base_index);

  raft::resource::sync_stream(handle);

  print_sizes("after grow extend", handle, grown_index);

  auto visible  = visible_list_size(handle, grown_index);
  auto internal = grown_index.lists()[0]->get_size();

  if (visible == rows && internal == base_rows) {
    std::cout << "\nBUG: index-visible size is " << visible
              << " but internal list::size stayed " << internal << "\n";
    return 1;
  }

  std::cout << "\nNo mismatch observed\n";
  return 0;
}

after base extend
index.list_sizes()[0] = 100
list::size = 100
capacity = 128

after grow extend
index.list_sizes()[0] = 120
list::size = 100
capacity = 128

BUG: index-visible size is 120 but internal list::size stayed 100

achirkin · 2026-07-05T06:39:12Z

Ok, sorry for closing this prematurely, this does seem legit. I'm not sure though if this internal get_size() is used where it could cause any problems. I'll look into this in the following days.

qwertyforce · 2026-07-05T07:26:02Z

Thank you very much! I forgot to mention how i even stumbled upon this - i got very strange Illegal memory access errors (i am using python wrapper for multiple extends + search + torch for calculating the vectors) and believe that this bug is related (IMA was in calc_chunk_indices_kernel)
But I couldnt distill my production code to a small reproducible poc (the IMA is also non-deterministic), so all i can do is either the one poc i sent, or https://gist.github.com/qwertyforce/65d9bb2d108701240fc9ef1bc952e685 (tldr: if you create 2 new indexes by extend from the base one (not inplace), child_b data can leak into child_a. child_a is not modified after child_b is created, but child_a search results change)

qwertyforce · 2026-07-06T12:10:38Z

off-topic, figured out why i got IMA, NVIDIA/raft#3041, different issue.
should I create an issue/PR in cuVS to speed up the process, or is this comment enough?

achirkin · 2026-07-06T14:58:34Z

Everything that gets merged in the main branch before next burndown, gets automatically in the next release - the new release branch is forked from the main when burndown starts. Here's our schedule https://docs.rapids.ai/maintainers/

huuanhhuyn

Thank you for your contribution! I'd have 1 doubt.

huuanhhuyn · 2026-07-06T14:58:37Z

-      auto shared_list_size = old_used_size;
-      if (new_used_size <= old_used_size ||
+      auto shared_list_size = old_logical_size;
+      if (new_used_size <= old_logical_size ||


I guess there are 3 cases:

new_used_size <= old_logical_size -> no size update, no reallocation

old_logical_size < new_used_size <= old_used_size -> size update, no reallocation

new_used_size > old_used_size -> size update, reallocation

If I understand it correctly, your current version will do both size update and reallocation for the 2nd case.

If I understand it correctly, your current version will do both size update and reallocation for the 2nd case.

Only if orig_list->size.compare_exchange_strong(shared_list_size, new_used_size) returns false

it returns false if orig_list->size was already mutated before this resize_list call ( if the current shared list size is not old_logical_size anymore)

That's right. I misunderstood the logics.

huuanhhuyn · 2026-07-07T07:45:19Z

+           std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(grow_indices_view),
+           base_index);
+  raft::resource::sync_stream(handle);
+


Suggested change

ASSERT_EQ(first_grown_index.lists()[0].get(), base_index.lists()[0].get());

I'd suggest to test that the first grown list reuses the base list without reallocation.

huuanhhuyn · 2026-07-07T07:59:51Z

+  raft::resource::sync_stream(handle);
+
+  ASSERT_EQ(base_index.lists()[0]->get_size(), base_rows);
+  ASSERT_GE(base_index.lists()[0]->indices_capacity(), rows);


Suggested change

ASSERT_GE(base_index.lists()[0]->indices_capacity(), rows);

ASSERT_EQ(base_index.lists()[0]->indices_capacity(), raft::Pow2<kIndexGroupSize>::roundUp(rows));

I think this makes the test more expressive

huuanhhuyn · 2026-07-07T08:06:17Z

+  std::vector<int64_t> host_indices(grow_rows);
+  std::iota(host_indices.begin(), host_indices.end(), base_rows);
+  auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows);
+  raft::copy(indices.data_handle(), host_indices.data(), host_indices.size(), stream);


Suggested change

std::vector<int64_t> host_indices(grow_rows);

std::iota(host_indices.begin(), host_indices.end(), base_rows);

auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows);

raft::copy(indices.data_handle(), host_indices.data(), host_indices.size(), stream);

auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows);

raft::linalg::range(indices.data_handle(), base_rows, rows, stream);

raft has a device built-in method for it.
We should #include <raft/linalg/init.cuh> and remove the include <numeric>

huuanhhuyn · 2026-07-07T08:06:38Z

-      auto shared_list_size = old_used_size;
-      if (new_used_size <= old_used_size ||
+      auto shared_list_size = old_logical_size;
+      if (new_used_size <= old_logical_size ||


That's right. I misunderstood the logics.

huuanhhuyn · 2026-07-07T08:08:11Z


 INSTANTIATE_TEST_CASE_P(AnnIVFSQTest, AnnIVFSQTestF_float, ::testing::ValuesIn(inputs));

+TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity)


I'd have the same reviews for this test.

huuanhhuyn · 2026-07-07T08:23:02Z

+  std::vector<float> host_data(rows * dim);
+  for (int64_t row = 0; row < rows; row++) {
+    for (int64_t col = 0; col < dim; col++) {
+      host_data[row * dim + col] = static_cast<float>(row + col);
+    }
+  }
+
+  auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);
+  raft::copy(data.data_handle(), host_data.data(), host_data.size(), stream);


Suggested change

std::vector<float> host_data(rows * dim);

for (int64_t row = 0; row < rows; row++) {

for (int64_t col = 0; col < dim; col++) {

host_data[row * dim + col] = static_cast<float>(row + col);

}

}

auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);

raft::copy(data.data_handle(), host_data.data(), host_data.size(), stream);

auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim);

raft::matrix::fill(handle, data.view(), 0.0f);

I think we could initialize the data directly on device with zeros. We test the sizes instead of filled values. I assume the build() works fine with zeros.

After the changes, we should #include <raft/matrix/init.cuh> and remove #include <vector>

huuanhhuyn · 2026-07-07T08:35:13Z

Some suggestions to improve the tests. Otherwise LGTM! Thank you for the contribution!

huuanhhuyn · 2026-07-07T12:52:37Z

/ok to test f4cee4f

qwertyforce · 2026-07-07T15:49:22Z

fixed the tests

huuanhhuyn · 2026-07-07T18:41:09Z

@achirkin I can't trigger the full 85 CI checks. Can you trigger it or it is because of missing commit signing?

huuanhhuyn · 2026-07-07T18:37:59Z

  ASSERT_EQ(index.lists()[0]->get_size(), base_rows);
-  ASSERT_GE(index.lists()[0]->indices_capacity(), rows);
+  ASSERT_EQ(index.lists()[0]->indices_capacity(), raft::Pow2<kIndexGroupSize>::roundUp(rows));
+  auto* base_list = index.lists()[0].get();


Why do we need to get the base_list pointer here? Can't we later get it at line 69?

Fix IVF extend list resize

7d5f827

qwertyforce requested a review from a team as a code owner July 4, 2026 13:01

github-project-automation Bot added this to Unstructured Data Processing Jul 4, 2026

better comments + hpp

f4cee4f

achirkin closed this Jul 5, 2026

github-project-automation Bot moved this to Done in Unstructured Data Processing Jul 5, 2026

achirkin reopened this Jul 5, 2026

achirkin added bug Something isn't working non-breaking Introduces a non-breaking change labels Jul 5, 2026

achirkin self-assigned this Jul 5, 2026

huuanhhuyn reviewed Jul 6, 2026

View reviewed changes

huuanhhuyn suggested changes Jul 7, 2026

View reviewed changes

improve tests

7cd89ac

huuanhhuyn reviewed Jul 7, 2026

View reviewed changes


	ASSERT_EQ(first_grown_index.lists()[0].get(), base_index.lists()[0].get());

	ASSERT_GE(base_index.lists()[0]->indices_capacity(), rows);
	ASSERT_EQ(base_index.lists()[0]->indices_capacity(), raft::Pow2<kIndexGroupSize>::roundUp(rows));


		INSTANTIATE_TEST_CASE_P(AnnIVFSQTest, AnnIVFSQTestF_float, ::testing::ValuesIn(inputs));

		TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity)

Uh oh!

Conversation

qwertyforce commented Jul 4, 2026

Description

Fix

Tests

Uh oh!

copy-pr-bot Bot commented Jul 4, 2026

Uh oh!

achirkin commented Jul 5, 2026

Uh oh!

qwertyforce commented Jul 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

qwertyforce commented Jul 5, 2026

Uh oh!

achirkin commented Jul 5, 2026

Uh oh!

qwertyforce commented Jul 5, 2026

Uh oh!

qwertyforce commented Jul 6, 2026

Uh oh!

achirkin commented Jul 6, 2026

Uh oh!

huuanhhuyn left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

huuanhhuyn commented Jul 7, 2026

Uh oh!

huuanhhuyn commented Jul 7, 2026

Uh oh!

qwertyforce commented Jul 7, 2026

Uh oh!

huuanhhuyn commented Jul 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

qwertyforce commented Jul 5, 2026 •

edited

Loading