Use batched unfused 1-NN for KMeans on Blackwell

[viclafargue]

An investigation of large-scale KMeans failures on Blackwell revealed that the unfused 1-NN path bypasses KMeans distance batching and materializes the complete n_samples × n_clusters distance matrix. This can require hundreds of gigabytes of temporary workspace and may overflow int32 indexing.

This change batches the unfused 1-NN implementation on Blackwell, restoring bounded workspace usage. It also renames the dispatch booleans to make the metric-capability and implementation-selection decisions clearer.

[tfeher]

I don't think this is not the right solution. Missing the tiling logic was an oversight in #1768. It was actually discussed, and was thought to be addressed, but apparently there is another call path where this is not working.

I think the correct solution would add a simple 1D tiling along n_samples. Wouldn't that solve the problem?

[tarang-jain]

I think the correct solution would add a simple 1D tiling along n_samples. Wouldn't that solve the problem?

A 1D tiling would also have to be done on the centroids right?

[tfeher]

I think the correct solution would add a simple 1D tiling along n_samples

Having said that, this code path will be reworked in #2249, and what we do here will be deprecated. Therefore we should not develop anything complicated to fix this.

[tfeher]

A 1D tiling would also have to be done on the centroids right?

Depends on what are typical parameter ranges we expect this to work on. If the target use case would not work with tiling only on n_samples, then I would just go for the simpler solution proposed by Victor.

[tarang-jain]

Depends on what are typical parameter ranges we expect this to work on. If the target use case would not work with tiling only on n_samples, then I would just go for the simpler solution proposed by Victor.

The behabior of the API is to acknowledge tile parameters on both the centroids and the samples. That said, kmeans_balanced might have different requirements. Changing the use_fused function could mean affecting the behavior there as well, which may not be intended for balanced kmeans.

[tarang-jain]

Here, we should think about the effect on balanced kmeans. Avoidance of tiling there may have be intentional. Then if balanced and regular kmeans are expected to behave differently, we must solve this problem in another way -- kmeans should be able to override this function and force the use of fused kernel always.

[viclafargue]

Balanced KMeans uses the same assignment function but does not pass the regular KMeans tiling parameters. Instead, it bounds memory through its own outer minibatching and calls the assignment with batch_samples = 0 and batch_centroids = 0. Its use of the unfused Blackwell path was therefore intentional (normally executes as a single inner tile).

My latest commit fix restores Blackwell to the shared unfused heuristic and makes regular KMeans’s unfused assignment honor batch_samples and batch_centroids, avoiding the full distance-matrix allocation.

[achirkin]

Thanks for the PR! The allocation/batch sizing seems right. I'd suggest to add some tests to cover specifically this codepath - it's not clear whether it is ever tested in the CI with the current test suite.

[viclafargue]

Added a test that exercises no batching, sample-only batching, centroid-only batching, and combined batching for both 32-bit and 64-bit indices. It validates labels and inertia on every GPU and verifies that batching reduces device-memory allocations.
We use predict because minClusterAndDistanceCompute is an internal, non-exported template and cannot be linked directly from the test executable.

[achirkin]

Thank for the new tests! LGTM

[tarang-jain]

Thanks. LGTM except one small comment.

[tfeher]

Thanks Victor, LGTM!

[viclafargue]

/merge