[tarang-jain]

Adds optional IVF-based centroid assignment to KMeans. Each Lloyd iteration rebuilds IVF lists over the current centroids, routes queries to their closest coarse clusters, sorts and gathers queries by list, calls the existing minClusterAndDistanceCompute / top_1_nn primitives within each list, and reduces the list results to one original centroid ID per query. No custom CUDA kernels are added.

centroid_ivf_params reuses the existing IVF-Flat build and search parameter structs from ivf_flat.hpp. The feature is available through C++, C, Python, and multi-GPU KMeans. Omitting it preserves exhaustive assignment. Final public fit inertia and prediction remain exhaustive.

Includes the pageable-host double-buffering changes from #2851; the batch loader is unchanged from that PR. The shared kmeans_type enum lives in neighbors/common.hpp, retaining its original namespace and values, so common.hpp no longer depends on the KMeans parameter definitions.

Single Lloyd iteration: checkpoint 5 → 6

Falcon: 960,000,000 FP32 vectors × 1,024 dimensions, 300,000 centroids, sharded evenly across 3 NVIDIA RTX PRO 6000 Blackwell Server Edition GPUs (devices 5–7, 320M rows each). All methods start from the same saved iteration-5 centroid checkpoint. One measured run per method on a shared host.

Implementation: d341aee19c533309a7a8ff2d54c0bf1d26a1934a. CUDA 13.3; driver 580.65.06. The benchmark links the current PR assignment primitives, current cuTile 1-NN kernels and current CAGRA search dispatch.

Assignment method	Iteration wall time (s)	Search time on GPUs 5 / 6 / 7 (s each)	Sum of search GPU-seconds	End-to-end M vectors/s	Final inertia, first 10M	Recall@1, uniform 1M
IVF, 512 lists / 32 probes	660.21	311.33 / 311.90 / 311.77	935.00	1.454	5,671,916.20	99.6710%
Exhaustive cuTile 1-NN	1419.20	1391.98 / 1392.67 / 1396.32	4180.97	0.676	5,671,827.34	99.9899%
CAGRA, NN-Descent	1441.03	1417.05 / 1412.58 / 1415.53	4245.17	0.666	5,672,070.52	99.0810%

IVF is 2.15× faster end to end than exhaustive assignment and 2.18× faster than CAGRA, with first-10M final inertia 0.00157% above exhaustive assignment.

Timing scope: one Lloyd update, including index construction, dataset reading/pageable-to-pinned staging, H2D transfers, norm computation, assignment, standard cuVS FP32 centroid accumulation, NCCL reduction, and centroid finalization. Initial context/buffer setup, JIT warmup, checkpoint writes, and separate final-inertia scoring are excluded; allocations performed within the measured iteration remain included. This measures the iteration primitives with the native batch-loader schedule, rather than the full public fit() call, which also performs final exhaustive scoring.

Search times are CUDA-event elapsed times around every assignment call, including internal host-submission gaps, accumulated separately on each GPU. Their sum is GPU-seconds, not parallel elapsed time. Iteration wall time is measured from the common start until all three GPUs finish the update. The IVF search measurement includes coarse routing, sorting/gathering, list scans, and cross-list top-1 reduction.

Settings:

• Common: device_buffer_samples=131072, batch_samples=131072, batch_centroids=0; the full staged batch is grouped by probed IVF list. PR Enable OOC KMeans transfer/compute overlap for pageable input #2851 uses two device buffers and two pinned-host staging buffers with a separate transfer stream.
• IVF: 512 coarse clusters over all 300k centroids, 20 coarse KMeans iterations, training fraction 1.0, seed 42, 32 probes. Each list is searched using the existing automatic top-1 backend, including cuTile where eligible.
• CAGRA: rebuild the centroid graph using NN-Descent; intermediate degree 128, graph degree 64; SINGLE_CTA, search_width=2, itopk_size=32, max_queries=16384, top-1 output, ordinary random initialization with no saved-neighbor seeds.
• Exhaustive assignment uses the same automatic cuTile 1-NN primitive. cuTile uses TF32 arithmetic for FP32 inputs where selected. Recall is measured against exhaustive FP32 reference assignments at checkpoint 5 on the same saved uniform 1M-row sample.
• Final inertia is evaluated against each method’s newly updated centroids on the first 10M rows, using the cuVS stable exhaustive FP32 top-1 primitive and FP64 accumulation of the returned squared distances. It is not full-960M inertia or the cost against the old centroids.

Method	Reset + index construction, slowest GPU (s)	Separate 10M inertia evaluation wall time (s)
IVF, 512 lists / 32 probes	0.095	39.76
Exhaustive cuTile 1-NN	0.001	37.20
CAGRA, NN-Descent	3.648	37.22

Checkpoint-5 SHA256: ac3b792623987bf04632ee3e5c3806d0422e371ba9d8c3294d283adf0bc1fef9. All runs audit exact shard coverage (960M rows total), save iteration-6 centroids, and retain per-batch CUDA-event timings.

Validation

19 focused C++/multi-GPU/C API regression tests passed, including all-probe equivalence, bounded-probe and tail batches, empty lists, original-ID tie handling, host double buffering, weighted fitting, coarse-training subsampling, parameter defaults, and multi-GPU agreement. Public-header include-order compilation passed; the 12 single-GPU cases were also rerun against the final benchmark runtime and passed. Python Cython translation and formatting checks passed earlier in the branch; the full repository CI and Python runtime suite have not been run.

[copy-pr-bot]

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.

[tarang-jain]

Closing in favor of cagra (discussed offline)