[BUG] KMeans crashes for larger Datasets

[Intron7]

cuML 26.06.00 aborts in KMeans.fit for a 2,000,000 × 40 float32 array.
The identical script succeeds with cuML 26.04.
Actual: cudaErrorIllegalAddress; process exit code 134.
Expected: KMeans.fit completes without aborting.
GPU: RTX PRO 6000 Blackwell (97,887 MiB), driver 610.43.02, CUDA 13.2.

Reproducer

import cupy as cp
from cuml.cluster import KMeans

X = cp.random.RandomState(0).standard_normal((2_000_000, 40), dtype=cp.float32)
KMeans(
    n_clusters=100,
    init="scalable-k-means++",
    n_init=1,
    max_iter=1,
    random_state=0,
).fit(X)
cp.cuda.runtime.deviceSynchronize()

[viclafargue]

I could not reproduce on my end (on a RTX Quadro 8000), but this is probably related to the fact that the 1-NN assignment step is currently unfused on Blackwell which causes a spike in VRAM use. There is a lot of ongoing work to optimize this step (see NVIDIA/cuvs#2249).

In the meantime, you may want to make use of the recently introduced out-of-core KMeans API or adjust the batch_samples and batch_centroids parameters exposed in the C++ API to modify the tiling of the distance matrix calculations.

cc @tarang-jain @vinaydes @tfeher @divyegala

[Intron7]

Quick update i can also reproduce this on an DGX-Spark

[viclafargue]

Fixed by NVIDIA/cuvs#2346.

import cupy as cp
import rmm.statistics
from cuml.cluster import KMeans
from rmm.allocators.cupy import rmm_cupy_allocator

# Ensure CuPy allocations are also visible to RMM.
cp.cuda.set_allocator(rmm_cupy_allocator)

with rmm.statistics.statistics():
    X = cp.random.RandomState(0).standard_normal(
        (2_000_000, 40), dtype=cp.float32
    )

    KMeans(
        n_clusters=100,
        init="scalable-k-means++",
        n_init=1,
        max_iter=1,
        random_state=0,
    ).fit(X)

    cp.cuda.runtime.deviceSynchronize()
    stats = rmm.statistics.get_statistics()

print(f"Peak RMM memory: {stats.peak_bytes / 1024**3:.3f} GiB")

Before on B200 :

$ python repro.py
terminate called after throwing an instance of 'rmm::_RMM_26_8::cuda_error'
  what():  CUDA error at: /tmp/conda-bld-output/bld/rattler-build_librmm/work/cpp/src/cuda_stream_view.cpp:45: cudaErrorIllegalAddress an illegal memory access was encountered
Aborted (core dumped)

After still on B200 :

$ python repro.py
Peak RMM memory: 0.581 GiB