Reuse minClusterAndDistance Helper for Balanced KMeans

[tarang-jain]

The norm computation + fused reduction is already present in the minClusterDistanceCompute function. We can reuse that for balanced kmeans.

Furthermore, this PR updates the minClusterDistanceCompute function to also use the fused kernel for the cosine metric.

Binary size savings (conda-cpp-build check with CUDA 12.9.1 + amd):
main: 660.72 MB
This PR: 648.83 MB
(conda-cpp-build check with CUDA 13.1.1 + amd):
main: 305.70 B
This PR: 300.63 MB

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[aamijar]

Thanks for the refactor Tarang! Could you add a description to the PR?

[jinsolp]

Thanks @tarang-jain ! Suggesting small changes and adding a question:

[jinsolp]

should we be computing norms like this here too?

if (metric == cuvs::distance::DistanceType::CosineExpanded) {
      raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(
        handle, centroids, centroidsNorm, raft::sqrt_op{});
    } else {
      raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(
        handle, centroids, centroidsNorm);
    }

[tarang-jain]

Good catch! I have fixed this.

[jinsolp]

Looking at the deleted code in kmeans_balanced.cuh, this used to be false for the CosineExpended metric. However, this condition passes true for the CosineExpended metric and sqrt-s the distances output.

[tarang-jain]

Yes that is because Cosine is now supported in our fused kernel.

[tarang-jain]

This wasn't being caught because we were lacking tests for the cosine metric. We always test L2 (default). I am trying to add cosine inputs to the tests next.

[achirkin]

Could you please refresh my understanding on couple math questions?

1. By the time we call minClusterDistanceCompute, the data is normalized, so we can use the same L2 norm everywhere, right?
2. Why is sqrt needed for the cosine case?

[tarang-jain]

By the time we call minClusterDistanceCompute, the data is normalized, so we can use the same L2 norm everywhere, right?

Cosine is only supported in balanced kmeans. When we call minClusterAndDistance, the data is not normalized, but the centroids are.

Why is sqrt needed for the cosine case?

sqrt is needed because the cosine distance op directly divides by the norms (it does not do the sqrt):

cuvs/cpp/src/distance/detail/distance_ops/cosine.cuh

Line 20 in 988ebf5

return static_cast<AccT>(1.0) - static_cast<AccT>(accVal / (aNorm * bNorm));