Optimize transfer/compute overlap in out-of-core KMeans

[viclafargue]

Summary

This PR improves data-transfer and compute overlap for host-resident out-of-core KMeans using a cyclic two-buffer pipeline.

• Batch 1 follows batch 0 on the copy stream while batch 0 compute starts as soon as its transfer completes.
• Subsequent transfers overlap computation of the current batch.
• Buffer recycling continues across iteration boundaries, allowing batch 0 of the next pass to be prefetched during the previous pass.
• Device-resident inputs remain zero-copy.

Implementation

• Adds a private KMeans batch loader with explicit staged, acquired, and reusable buffer states.
• Uses CUDA events so compute waits for H2D completion and a buffer cannot be overwritten until all of its consumers have been submitted.
• Applies the same dependency-driven scheduling to every batch, without special APIs or states for the first two batches.
• Uses persistent device scratch to avoid per-batch deallocation and its potential device-wide synchronization.
• Computes final inertia through the regular batched assignment and reduction pipeline, accumulates it on device, and copies only the final result to host.
• Leaves the shared ANN batch iterator unchanged.

Benchmark under similar configuration

10 GiB pinned-host FP32 dataset (10,485,760 × 256), 2,560 clusters, three iterations, ten 1 GiB out-of-core batches, and 131,072-sample assignment tiles.

Metric	main	PR
Median runtime	1.7795 s	0.7962 s
Speedup	1.00×	2.23×
Effective bulk throughput	22.48 GiB/s	50.24 GiB/s
Profiled GPU span	1773.12 ms	785.65 ms
H2D time	752.29 ms	752.38 ms
Kernel time	1003.08 ms	760.97 ms
H2D/kernel overlap	0.00 ms	727.91 ms
H2D overlapped by kernels	0.0%	96.75%
Kernels overlapped by H2D	0.0%	95.66%

Profile

Main branch :

This PR :

This PR (multi-GPU) :

[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.

[viclafargue]

/ok to test 4b8a5c5

[viclafargue]

/ok to test 9dd96bf

[tarang-jain]

if the data is already device accessible, why are we even using the workspace?

[tarang-jain]

Furthermore, if the large workspace is managed, the transfer speed is slightly slower (I found 55 GB/s versus 48 GB/s). I agree that the batchsize is not bounded here (it can be as large as the dataset) but I would argue for falling back to the large workspace only at the breaking point where allocating from the regular workspace is not possible. That calculation can get complicated to account for whether or not weights are present, so I'll tag @achirkin for some ideas.

[tarang-jain]

Actually since this is unbounded (only bounded by dataset size) lets just stick to large_mr @viclafargue. Since compute will overlap, its not a big deal.

[irina-resh-nvda]

A question about the python path: is there a way to set a stream pool for the single gpu ooc kmeans when using a python wrapper? If not, then there will be no h2d/compute overlap in this case. Is it intentional?

[viclafargue]

A question about the python path: is there a way to set a stream pool for the single gpu ooc kmeans when using a python wrapper? If not, then there will be no h2d/compute overlap in this case. Is it intentional?

Thanks for spotting this! I added some code to allow users to set up memory and stream pools from the C and Python cuVS APIs.