Exposing opt-in double buffering for MG batched KMeans

[viclafargue]

Partially answers #2292

• GPUs: 2xQuadro RTX 8000
• Input: float32 host memory; RMM pool: 90%
• Timing: median of 3 runs after one warmup per mode; max iterations: 5
• Speedup is single-buffer time / double-buffer time; values above 1 favor prefetch.
• Device input buffers are shown per rank as single → double; common KMeans workspace is excluded.

Scenario	Host	Rows/GPU	Features	Clusters	Batch rows	Batches/GPU	Device input buffers/rank	Iterations	Single (s)	Double (s)	Speedup
Copy-dominated	pageable	4,000,000	32	2	524,288	8	64 MiB → 128 MiB	5	0.767	0.698	1.099x
Overlap-friendly	pinned	4,000,000	32	2	524,288	8	64 MiB → 128 MiB	5	0.594	0.491	1.209x
Compute-dominated	pinned	4,000,000	32	16384	524,288	8	64 MiB → 128 MiB	5	5.028	4.594	1.094x
Single batch	pinned	4,000,000	32	32	4,000,000	1	488 MiB → 488 MiB	5	0.172	0.178	0.970x
Large buffer	pinned	8,000,000	128	2	4,000,000	2	1.91 GiB → 3.81 GiB	5	4.726	4.234	1.116x

Prefetch is beneficial but not universal: it gives up to 1.21× speedup when pinned H2D copies can overlap KMeans work, but only ~1.09× in copy- or compute-limited cases and no gain for a single batch.

[cjnolet]

Thanks for giving this a first pass. I'd also like for @tfeher and the devtech team to take a look at this as well. Also, please post benchmarks on here.

[cjnolet]

I really don't like the use of the term "streaming" here. It carries altogether too much baggage with it. Can we please just remove it here? I think we need to consider removing it from "streaming_batch_size" as well. Streaming to me means a k-means that runs through a streaming mechanism, meaning it can only go forward as an iterator and can restart from where it left off when more data is received. There are variants of k-means that satisfy this, but the batched variant we've implemented does not.

[tarang-jain]

The problem with batch_size is that we have two other parameters -- batch_centroids and batch_samples. So batch_size can be confusing. Perhaps we should reconsider the names of those two parameters then? Or maybe think of some other prefix for batch_size?

[cjnolet]

What are the meanings of batch samples vs batch size?

[tarang-jain]

the tile sizes for the distance computation (local reductions are done on this tile) -- i.e. batch_samples * batch_centroids distances are computed at once. So I was saying that we are already doing what the first part of the flash-kmeans paper mentions (at least that was my understanding)

[tarang-jain]

  /**
   * batch_samples and batch_centroids are used to tile 1NN computation which is
   * useful to optimize/control the memory footprint
   * Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0
   * then don't tile the centroids
   *
   * NB: These parameters are unrelated to streaming_batch_size, which controls how many
   * samples to transfer from host to device per batch when processing out-of-core
   * data.
   */
  int batch_samples = 1 << 15;

  /**
   * if 0 then batch_centroids = n_clusters
   */
  int batch_centroids = 0;

[cjnolet]

If one is for tiling then let's call it like it is, that's a start in the right direction. The longer we make the names, the more a user has to remember.

Let's trying and think of how we can make this easier to use. "Streaming" is not a good word because it has too many other meanings.

Out-of-core buffered (or batched) is more aligned with the actual purpose but that's also pretty verbose.

There's really 2 levels of batching going on here- one is to control intermediate computation of the 1nn itself, the other is to control the device footprint of each "batching" of data vectors from host.

Having 3 options with batch in the name is super confusing. But I think we can alleviate this (and avoid having to break existing apis) if we change batch to device_buffer just for the non-tiled version.

So that would be device_buffer_samples and device_buffer_prefetch. That seems to be more indicative of "I'm buffering up some number of samples on device".

But above and beyond this change- what would be the negative impact to just making prefetch always on? I can imagine the gpu utilization would fill up with work before the memory would fill up with batches, which would suggest to me that there's a buffer buffer size that can remain pretty constant for most gpus.

[viclafargue]

Prefetch is not universally beneficial: it only improves throughput when the GPU can overlap host-to-device copies with meaningful computation. If either copies or compute dominate completely, the second stream and buffer provide little or no gain.

Making it unconditional would also make the memory tradeoff unconditional. Every applicable rank would reserve space for two device batches instead of one, even for workloads where overlap cannot help. That extra allocation can be large enough to cause an out-of-memory failure for configurations that otherwise fit.

[cjnolet]

Let's get some benchmarks attached to this PR so that we can evaluate. Again , I suspect the actual device utilization is likely to cap long before the memory does, which would indicate that even if there's no gain in performance, so long as there is not a drop in performance, the simplicity for the user could be justified in the default. But let's back that up with numbers.

[cjnolet]

BTW I'm not necessarily say we should or shouldn't make it the default, but I am saying we should look into the impact and provide real evidence either way (concrete numbers / formula always helps but benchmarks should ultimately guide us).