Skip to content

Speed up batched IVF Flat routing - #229

Draft
jay23606 wants to merge 1 commit into
VectorDB-NTU:mainfrom
jay23606:fix/ivf-batched-flat-routing-topn
Draft

jay23606 wants to merge 1 commit into
VectorDB-NTU:mainfrom
jay23606:fix/ivf-batched-flat-routing-topn

Conversation

@jay23606

Copy link
Copy Markdown

Closes #218

This draft adds a blocked top-nprobe assignment kernel for batched Flat IVF searches. It rotates query blocks, screens centroid tiles with the transposed matrix-product kernel, and recomputes candidates near each row's top-n boundary with the existing distance functions. The scalar search path is unchanged; the batch path falls back to scalar routing for one query and non-Flat initializers.

The IVF batch regression test now checks nprobe values 1, 2, 3, and all clusters across L2/IP, quantization settings, HACC modes, and thread counts, comparing IDs and distances to repeated scalar searches.

Validation is incomplete in this environment: git diff --check passed, but C++ tests and formatting were not run because this Windows environment has no CMake, compiler, Bash, or clang-format 23 installed. No before/after GIST1M benchmark or Recall@10 measurement has been run. Please treat this as an early implementation draft; CI and the required benchmark/quality report remain outstanding.

@gouyt13

gouyt13 commented Nov 13, 2026

Copy link
Copy Markdown
Member

Thanks for opening this early. The approach looks right: the GEMM only screens candidates, and the ones kept are recomputed with the same float euclidean_sqr / dot_product_dis the scalar path uses. I built it locally with the missing include added. On random data (L2 and IP, 1024 and 4096 clusters, nprobe 1/16/64/all, 1 and 8 threads), batched IDs and distances matched repeated search calls exactly, and the existing IVF and k-means C++ tests passed.

Must fix

  1. Build. kmeans.hpp calls index_detail::parallel_for without including rabitqlib/index/detail/parallel_for.hpp, so every file that includes kmeans.hpp fails to compile.
  2. Formatting. ./scripts/check-format.sh reports 13 violations in kmeans.hpp and 6 in ivf.cpp.
  3. Nested threads. exact_assign_topn calls simd::matrix_product_transposed (an Eigen GEMM) inside parallel_for without calling omp_set_num_threads(1) first, so Eigen starts its own OpenMP team inside each worker. exact_assign makes that call for this reason (kmeans.hpp line 141). Measured with OMP_WAIT_POLICY=passive: num_threads=1 used about 1.6 cores (1.0 with OMP_NUM_THREADS=1), and num_threads=8 used about 10. Python's IvfIndex.search always goes through search_batch, so this also affects Python's default num_threads=1.
  4. Ties. Batched and scalar routing can choose different clusters when centroid distances tie. The scalar path uses partial_sort and compares distance only; the batch path orders by squared distance, then ID. With duplicated centroids (each odd centroid a copy of the one before it), nprobe=15 gave 6,022 of 20,000 differing result entries for both L2 and IP. Please make both paths order by the returned float distance, then by ID (this means changing FlatInitializer::centroids_distances too), and add a duplicate-centroid test. Ordering by squared distance is not enough, because distinct squared distances can round to the same float after sqrt.

Performance

  • When nprobe equals the number of clusters, every candidate passes the screen, is recomputed, and is inserted into a sorted vector, which costs O(k·nprobe). With 2,000 queries, dim 128, 4,096 clusters and 1 thread, the batched path took 3.0 s versus 2.0 s for scalar search. This was a diagnostic run on a loaded shared host, not a benchmark, but the direction is clear. Please skip the screen when nprobe is close to the cluster count, or use a heap.
  • Centroid norms are recomputed in every 32-row block. Please compute them once per call.
  • Every batch with more than one query takes the new path. Please measure small batches and add a threshold if needed. Also consider streaming smaller row blocks through rotate, route and scan, instead of rotating up to 128 MiB of queries before any routing.
  • Without reuse_workers, parallel_for creates new std::threads on each of its three calls per block. The existing batch path uses omp_for_dynamic, which avoids this.

Tests. The regression test has 4 clusters and 13 queries, so it never covers more than one centroid tile (more than 256 clusters) or more than one row block (more than 32 queries). Please add cases that cross both, plus duplicate centroids, k larger than the number of points, and removed points. A direct test of exact_assign_topn would also help.

Minor. The comment says distances are "recomputed in double precision", but they are float euclidean_sqr / dot_product_dis results widened to double, which is what keeps them identical to the scalar path. Please say that, and note that the error bound covers the float distance error.

Toolchain. Please set up a Linux build environment; WSL2 or Docker works. The benchmark in #218 needs one anyway, and CI on this repository needs manual approval for each push. Before the next push, please run the IVF C++ tests and check-format.sh. Then report GIST1M throughput and Recall@10 before and after the change, plus single-query latency and batch sizes 8 and 64, along with the CPU, compiler, ISA and thread count.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

IVF-RaBitQ: route batched queries with a blocked top-nprobe assignment

2 participants