diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..0d62ecb4a3 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -394,34 +394,17 @@ void initScalableKMeansPlusPlus(raft::resources const& handle, int niter = std::min(8, (int)ceil(log(psi))); RAFT_LOG_DEBUG("KMeans||: psi = %g, log(psi) = %g, niter = %d ", psi, log(psi), niter); + // Buffer for d(x, C') when incrementally updating min distances after each round. + auto newMinClusterDistanceVec = raft::make_device_vector(handle, n_samples); + // <<<< Step-3 >>> : for O( log(psi) ) times do + // minClusterDistanceVec / psi already hold d(x, C) / phi_X(C) from Step-2 (and are + // refreshed at the end of each round against only the newly sampled C'). for (int iter = 0; iter < niter; ++iter) { RAFT_LOG_DEBUG("KMeans|| - Iteration %d: # potential centroids sampled - %d", iter, potentialCentroids.extent(0)); - cuvs::cluster::kmeans::detail::minClusterDistanceCompute( - handle, - X, - potentialCentroids, - minClusterDistanceVec.view(), - L2NormX.view(), - L2NormBuf_OR_DistBuf, - params.metric, - params.batch_samples, - params.batch_centroids, - workspace); - - cuvs::cluster::kmeans::detail::computeClusterCost( - handle, - minClusterDistanceVec.view(), - workspace, - raft::make_device_scalar_view(clusterCost.data()), - raft::identity_op{}, - raft::add_op{}); - - psi = clusterCost.value(stream); - // <<<< Step-4 >>> : Sample each point x in X independently and identify new // potentialCentroids raft::random::uniform( @@ -458,6 +441,38 @@ void initScalableKMeansPlusPlus(raft::resources const& handle, potentialCentroids = raft::make_device_matrix_view(centroidsBuf.data(), tot_centroids, n_features); /// <<<< End of Step-5 >>> + + // Refresh d(x, C) = min(d(x, C), d(x, C')) for the next sampling round. + // Skip when C' is empty or this was the last oversampling iteration. + if (Cp.extent(0) > 0 && iter + 1 < niter) { + cuvs::cluster::kmeans::detail::minClusterDistanceCompute( + handle, + X, + Cp, + newMinClusterDistanceVec.view(), + L2NormX.view(), + L2NormBuf_OR_DistBuf, + params.metric, + params.batch_samples, + params.batch_centroids, + workspace); + + raft::linalg::map(handle, + minClusterDistanceVec.view(), + raft::min_op{}, + raft::make_const_mdspan(minClusterDistanceVec.view()), + raft::make_const_mdspan(newMinClusterDistanceVec.view())); + + cuvs::cluster::kmeans::detail::computeClusterCost( + handle, + minClusterDistanceVec.view(), + workspace, + raft::make_device_scalar_view(clusterCost.data()), + raft::identity_op{}, + raft::add_op{}); + + psi = clusterCost.value(stream); + } } /// <<<< Step-6 >>> RAFT_LOG_DEBUG("KMeans||: total # potential centroids sampled - %d", diff --git a/cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh b/cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh index 034bfa0f62..dff2196361 100644 --- a/cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh +++ b/cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -294,21 +295,14 @@ void initKMeansPlusPlus_distributed( RAFT_LOG_DEBUG( "Distributed KMeans||: rank=%d, psi=%g, niter=%d", rank, static_cast(psi), niter); + // Buffer for d(x, C') when incrementally updating min distances after each round. + auto newMinClusterDistance = + raft::make_device_vector(handle, std::max(n_local, IndexT{1})); + // Steps 3-6: sample candidates `O(log psi)` times, gathering across ranks. + // minClusterDistance / psi already hold d(x, C) / phi_X(C) from Step-2 (and are + // refreshed at the end of each round against only the newly sampled C'). for (int iter = 0; iter < niter; ++iter) { - if (iter > 0) - psi = compute_global_cluster_cost(handle, - params, - X_parts, - part_offsets, - n_local, - L2NormX.view(), - potentialCentroids, - minClusterDistance.view(), - L2NormBuf_OR_DistBuf, - workspace, - comms); - if (n_local > 0) { auto rands_view = raft::make_device_vector_view(uniformRands.data_handle(), n_local); @@ -390,6 +384,58 @@ void initKMeansPlusPlus_distributed( IndexT tot_centroids = static_cast(potentialCentroids.extent(0)) + total_new; potentialCentroids = raft::make_device_matrix_view(centroidsBuf.data(), tot_centroids, n_features); + + // Refresh d(x, C) = min(d(x, C), d(x, C')) for the next sampling round. + // Skip when C' is empty or this was the last oversampling iteration. + if (total_new > 0 && iter + 1 < niter) { + auto Cp = raft::make_device_matrix_view( + centroidsBuf.data() + (static_cast(potentialCentroids.extent(0)) - total_new) * + n_features, + total_new, + n_features); + + auto d_partial = raft::make_device_scalar(handle, DataT{0}); + for (std::size_t p = 0; p < X_parts.size(); ++p) { + auto part_rows = static_cast(X_parts[p].extent(0)); + if (part_rows == 0) { continue; } + auto x_slice = raft::make_device_matrix_view( + X_parts[p].data_handle(), part_rows, n_features); + auto mcd_slice = raft::make_device_vector_view( + minClusterDistance.data_handle() + part_offsets[p], part_rows); + auto new_mcd_slice = raft::make_device_vector_view( + newMinClusterDistance.data_handle() + part_offsets[p], part_rows); + auto norm_slice = raft::make_device_vector_view( + L2NormX.data_handle() + part_offsets[p], part_rows); + + cuvs::cluster::kmeans::min_cluster_distance(handle, + x_slice, + Cp, + new_mcd_slice, + norm_slice, + L2NormBuf_OR_DistBuf, + params.metric, + params.batch_samples, + params.batch_centroids, + workspace); + + raft::linalg::map(handle, + mcd_slice, + raft::min_op{}, + raft::make_const_mdspan(mcd_slice), + raft::make_const_mdspan(new_mcd_slice)); + } + + if (n_local > 0) { + auto mcd_view = + raft::make_device_vector_view(minClusterDistance.data_handle(), n_local); + cuvs::cluster::kmeans::cluster_cost( + handle, mcd_view, workspace, d_partial.view(), raft::add_op{}); + } + + comms.allreduce(d_partial.data_handle(), d_partial.data_handle(), 1); + raft::copy(&psi, d_partial.data_handle(), 1, stream); + raft::resource::sync_stream(handle); + } } RAFT_LOG_DEBUG("Distributed KMeans||: rank=%d, total candidates = %d", diff --git a/fern/pages/rust_api/rust-api-cuvs-test-utils.md b/fern/pages/rust_api/rust-api-cuvs-test-utils.md new file mode 100644 index 0000000000..d3c680e734 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-test-utils.md @@ -0,0 +1,15 @@ +--- +slug: api-reference/rust-api-cuvs-test-utils +--- + +# Test Utils Module + +_Rust module: `cuvs::test_utils`_ + +_Source: `rust/cuvs/src/test_utils.rs`_ + +Test-only tensor adapters. + +[`DeviceTensor`] is an RMM-backed device matrix, and the `ndarray` host +adapters below implement [`AsDlTensor`]/[`AsDlTensorMut`] for plain host +arrays. We use `ndarray` only as a dev-dependency to assist with unit tests. diff --git a/python/cuvs/benchmarks/bench_pq_build.py b/python/cuvs/benchmarks/bench_pq_build.py new file mode 100644 index 0000000000..ce55dd73ae --- /dev/null +++ b/python/cuvs/benchmarks/bench_pq_build.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +""" +End-to-end benchmark for cuVS Product Quantizer build. + +Mirrors the C++ setup: + - 256k train vectors (sampled from SIFT1M .fbin) + - pq_bits=8 (256 centers), pq_dim=128 subspaces, use_vq=False + - classic kmeans with max_iter=12 (default init = k-means|| / scalable++) + +Example: + python bench_pq_build.py \\ + --dataset /path/to/sift-128-euclidean/base.fbin \\ + --num-train 262144 --repeats 3 +""" + +from __future__ import annotations + +import argparse +import statistics +import time +from pathlib import Path + +import numpy as np + +from cuvs.common import Resources +from cuvs.preprocessing.quantize import pq + + +def read_fbin(path: Path) -> np.ndarray: + """Load a .fbin file (int32 n, int32 dim, then n*dim float32).""" + with open(path, "rb") as f: + n, dim = np.fromfile(f, dtype=np.int32, count=2) + data = np.fromfile(f, dtype=np.float32, count=int(n) * int(dim)) + return np.ascontiguousarray(data.reshape(int(n), int(dim))) + + +def load_train(dataset: str, num_train: int, dim: int, seed: int) -> np.ndarray: + path = Path(dataset) + if not path.exists(): + raise FileNotFoundError(path) + + data = read_fbin(path) + if data.ndim != 2: + raise ValueError(f"expected 2D dataset, got shape {data.shape}") + if data.shape[1] != dim: + raise ValueError(f"dataset dim={data.shape[1]} does not match --dim={dim}") + if data.shape[0] < num_train: + raise ValueError( + f"dataset has {data.shape[0]} rows, need at least --num-train={num_train}" + ) + + if data.shape[0] == num_train: + train = data + else: + rng = np.random.default_rng(seed) + idx = rng.choice(data.shape[0], size=num_train, replace=False) + train = np.ascontiguousarray(data[idx]) + + print(f"Loaded {path}: using {train.shape[0]} / {data.shape[0]} rows, dim={train.shape[1]}") + return train + + +def time_pq_build( + train: np.ndarray, + *, + pq_bits: int, + pq_dim: int, + kmeans_n_iters: int, + pq_kmeans_type: str, + use_subspaces: bool, + warmup: int, + repeats: int, +) -> list[float]: + """Return wall times (seconds) for end-to-end pq.build calls.""" + resources = Resources() + params = pq.QuantizerParams( + pq_bits=pq_bits, + pq_dim=pq_dim, + use_subspaces=use_subspaces, + use_vq=False, + vq_n_centers=0, + kmeans_n_iters=kmeans_n_iters, + pq_kmeans_type=pq_kmeans_type, + # Use the full trainset (matches C++ max_train_points_per_pq_code=num_train). + max_train_points_per_pq_code=train.shape[0], + max_train_points_per_vq_cluster=train.shape[0], + ) + + print( + "QuantizerParams(" + f"pq_bits={params.pq_bits}, pq_dim={params.pq_dim}, " + f"use_subspaces={params.use_subspaces}, use_vq={params.use_vq}, " + f"kmeans_n_iters={params.kmeans_n_iters}, pq_kmeans_type={pq_kmeans_type}, " + f"max_train_points_per_pq_code={params.max_train_points_per_pq_code})" + ) + + def once() -> float: + resources.sync() + t0 = time.perf_counter() + quantizer = pq.build(params, train, resources=resources) + resources.sync() + elapsed = time.perf_counter() - t0 + _ = quantizer.pq_bits + return elapsed + + for i in range(warmup): + t = once() + print(f" warmup[{i}]: {t:.3f}s") + + times: list[float] = [] + for i in range(repeats): + t = once() + times.append(t) + print(f" repeat[{i}]: {t:.3f}s") + return times + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--dataset", + type=str, + required=True, + help="Path to SIFT (or other) base.fbin", + ) + p.add_argument("--num-train", type=int, default=262144, help="Train vectors (default 256k)") + p.add_argument("--dim", type=int, default=128) + p.add_argument("--pq-bits", type=int, default=8, help="Implies n_centers = 2**pq_bits") + p.add_argument("--pq-dim", type=int, default=128, help="Number of PQ subspaces / chunks") + p.add_argument("--kmeans-n-iters", type=int, default=12) + p.add_argument( + "--pq-kmeans-type", + choices=("kmeans", "kmeans_balanced"), + default="kmeans", + help="'kmeans' uses classic kmeans (default init = scalable k-means++)", + ) + p.add_argument("--no-subspaces", action="store_true", help="Disable per-subspace codebooks") + p.add_argument("--warmup", type=int, default=1) + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--seed", type=int, default=42) + return p.parse_args() + + +def main() -> None: + args = parse_args() + + train = load_train(args.dataset, args.num_train, args.dim, args.seed) + n_centers = 1 << args.pq_bits + print( + f"Benchmark: n_train={train.shape[0]}, dim={train.shape[1]}, " + f"pq_dim={args.pq_dim}, n_centers={n_centers}, max_iter={args.kmeans_n_iters}" + ) + + times = time_pq_build( + train, + pq_bits=args.pq_bits, + pq_dim=args.pq_dim, + kmeans_n_iters=args.kmeans_n_iters, + pq_kmeans_type=args.pq_kmeans_type, + use_subspaces=not args.no_subspaces, + warmup=args.warmup, + repeats=args.repeats, + ) + + mean = statistics.mean(times) + stdev = statistics.stdev(times) if len(times) > 1 else 0.0 + print("-" * 60) + print( + f"pq.build: mean={mean:.3f}s stdev={stdev:.3f}s " + f"min={min(times):.3f}s max={max(times):.3f}s (n={len(times)})" + ) + + +if __name__ == "__main__": + main()