cuML 25.06 spreads UMAP's kNN build across GPUs, lifting training off the single-GPU ceiling

NVIDIA has shipped, in cuML and cuVS 25.06, an all-neighbors kNN graph construction step that distributes across multiple GPUs, which the company says improves scale and cuts overall training runtime ^[5]. That matters because kNN construction is the expensive part of UMAP training, and in the previous out-of-core release the training stage stayed pinned to one GPU while only `transform()` could use several ^[3][4].

The mechanism is not new work, it is the same decomposition pushed one level further. The out-of-core design partitions the dataset into balanced clusters, overlaps vectors across nearby clusters so nearest-neighbor relationships survive the boundaries, computes a local kNN graph per cluster, and merges those into one global all-neighbors graph ^[7]. Because each cluster needs neither the full dataset nor coordination with its peers, the clusters can simply be handed out to GPUs, each of which gathers its own assigned data from CPU memory ^[8]. Each GPU builds local graphs and merges them into the global one, which is what lets the implementation skip the all-to-all communication that normally caps distributed all-neighbors kNN work ^[9]. The practical consequence: both stages of UMAP can now use multiple GPUs, where before only the transform side could ^[14].

For operators the interesting surface is two hyperparameters. `knn_n_clusters` sets how many partitions the data is split into, and `knn_overlap_factor` sets how many of the closest clusters each point is assigned to ^[10]. Clusters are approximately balanced, so vectors spread roughly evenly across available GPUs, and raising the cluster count shrinks each partition and therefore the amount of data that has to fit in any one GPU's memory ^[11]. Raising the overlap factor recovers more of the true cross-boundary neighbors, which generally improves graph quality and the resulting embedding ^[12]. In other words, per-GPU memory pressure is something you dial with the partition count rather than something dictated by the dataset size, and the quality knob is separate from it ^[15]. The balanced k-means underneath is documented in NVIDIA's paper, Massive-Scale Out-Of-Core UMAP on the GPU ^[13].

The framing NVIDIA uses is workflow, not throughput, and it is the right framing: UMAP runs are repeated while users tune parameters, and per-run cost has been what makes interactive exploration on large corpora impractical ^[2][1]. According to NVIDIA, the new path takes several-hundred-gigabyte workloads from hours or days to minutes ^[6]. Note the shape of that claim. It is a range, not a measurement, and the material supplied here contains no dataset, no GPU count, and no recall or trustworthiness figure to sit next to it ^[6].

Three things to watch. First, whether the quality claim holds under the settings you would actually pick for cost, since both knobs trade against each other and the guidance in the source is directional rather than quantified ^[11][12]. Second, what the gather from host memory costs at scale, given that each GPU pulls its clusters from CPU memory rather than from a shared device-side copy ^[8]. Third, whether the absence of all-to-all traffic means node-count scaling stays close to linear, or whether the merge into the global graph becomes the new serialisation point ^[9].