[BUG] cuVS KMeans error during model training: `reduce_cols_by_key.cuh line=140`

Ran cuVS KMeans clustering on 4M vector dataset (fp32, 1024d) on an A100 80GB system. Run failed to complete with the following error:

Starting cuvs_kmeans sweeps.... 

Loading dataset: miracl-fp32-1024d-4M
Data shape: (3999994, 1024) 

{'load_data_avg_cpu_util': 1.1333333333333333, 'load_data_max_cpu_util': 2.9, 'load_data_max_ram_gb': 27.70366668701172, 'load_data_avg_gpu_util': 0.0, 'load_data_max_gpu_util': 0, 'load_data_max_vram_gb': 0.749267578125, 'load_data_duration_sec': 3.746360223274678}

Starting cuvs_kmeans parameter sweep....
n_clusters = 10
Training cuvs_kmeans model....
Traceback (most recent call last):
  File "/myworkspace/kmeans_sweep.py", line 292, in <module>
    telem, kmeans_model = train_kmeans(X, n_clusters, algorithm, time_delay)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/myworkspace/kmeans_sweep.py", line 142, in wrapper
    result = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/myworkspace/kmeans_sweep.py", line 216, in train_kmeans
    centroids, inertia, n_iter = cuvs_kmeans.fit(cuvs_kmeans_params, X)
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "cuvs/common/resources.pyx", line 110, in cuvs.common.resources.auto_sync_resources.wrapper
  File "/opt/conda/lib/python3.11/site-packages/pylibraft/common/outputs.py", line 83, in wrapper
    ret_value = f(*args, **kwargs)
                ^^^^^^^^^^^^^^^^^^
  File "cuvs/cluster/kmeans/kmeans.pyx", line 240, in cuvs.cluster.kmeans.kmeans.fit
  File "cuvs/cluster/kmeans/kmeans.pyx", line 241, in cuvs.cluster.kmeans.kmeans.fit
  File "cuvs/common/exceptions.pyx", line 37, in cuvs.common.exceptions.check_cuvs
cuvs.common.exceptions.CuvsException: CUDA error encountered at: file=/tmp/conda-bld-output/bld/rattler-build_libcuvs/host_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehol/include/raft/linalg/detail/reduce_cols_by_key.cuh line=140: call='cudaPeekAtLastError()', Reason=cudaErrorInvalidConfiguration:invalid configuration argument
Obtained 27 stack frames
#1 in /opt/conda/lib/python3.11/site-packages/cuml/cluster/../../../../libcuvs.so(+0x4976fd) [0x7f2a11f4d6fd]
#2 in /opt/conda/lib/python3.11/site-packages/cuml/cluster/../../../../libcuvs.so: void raft::linalg::detail::reduce_cols_by_key<float, thrust::transform_iterator<cuvs::cluster::kmeans::detail::KeyValueIndexOp<int, float>, raft::KeyValuePair<int, float>*, thrust::use_default, thrust::use_default>, int>(float const*, thrust::transform_iterator<cuvs::cluster::kmeans::detail::KeyValueIndexOp<int, float>, raft::KeyValuePair<int, float>*, thrust::use_default, thrust::use_default>, float*, int, int, int, CUstream_st*, bool) +0x8e3 [0x7f2a1204aec3]
#3 in /opt/conda/lib/python3.11/site-packages/cuml/cluster/../../../../libcuvs.so(+0x7ab623) [0x7f2a12261623]
#4 in /opt/conda/lib/python3.11/site-packages/cuml/cluster/../../../../libcuvs.so(+0x7c92a7) [0x7f2a1227f2a7]
#5 in /opt/conda/lib/python3.11/site-packages/cuvs/cluster/kmeans/../../../../../libcuvs_c.so: cuvsKMeansFit +0xc26 [0x7f29ad08a3a6]
#6 in /opt/conda/lib/python3.11/site-packages/cuvs/cluster/kmeans/kmeans.cpython-311-x86_64-linux-gnu.so(+0x11ac1) [0x7f29ad11dac1]
#7 in python3: PyObject_Call +0x134 [0x55a8349c4964]
#8 in python3: _PyEval_EvalFrameDefault +0x4a29 [0x55a83499a049]
#9 in python3: _PyFunction_Vectorcall +0x17f [0x55a8349baa4f]
#10 in python3: PyVectorcall_Call +0xda [0x55a834943b1f]
#11 in /opt/conda/lib/python3.11/site-packages/cuvs/common/resources.cpython-311-x86_64-linux-gnu.so(+0xb44f) [0x7f29ad00644f]
#12 in python3: _PyObject_MakeTpCall +0x29b [0x55a8349886ab]
#13 in python3: _PyEval_EvalFrameDefault +0x70a [0x55a834995d2a]
#14 in python3: _PyFunction_Vectorcall +0x17f [0x55a8349baa4f]
#15 in python3: _PyEval_EvalFrameDefault +0x4a29 [0x55a83499a049]
#16 in python3(+0x29b0ad) [0x55a834a4d0ad]
#17 in python3: PyEval_EvalCode +0x9f [0x55a834a4c7ef]
#18 in python3(+0x2b872a) [0x55a834a6a72a]
#19 in python3(+0x2b43b3) [0x55a834a663b3]
#20 in python3(+0x2c9780) [0x55a834a7b780]
#21 in python3: _PyRun_SimpleFileObject +0x1bc [0x55a834a7b11c]
#22 in python3: _PyRun_AnyFileObject +0x44 [0x55a834a7aef4]
#23 in python3: Py_RunMain +0x383 [0x55a834a75643]
#24 in python3: Py_BytesMain +0x37 [0x55a834a3ca17]
#25 in /lib/x86_64-linux-gnu/libc.so.6(+0x29d90) [0x7f2f1ef00d90]
#26 in /lib/x86_64-linux-gnu/libc.so.6: __libc_start_main +0x80 [0x7f2f1ef00e40]
#27 in python3(+0x28a8ca) [0x55a834a3c8ca]

Code repo: https://github.com/btong04/cuvs-bench-runner. Error can be reproduced by creating a dataset and running kmeans_sweep.py against it using algorithm=cuvs_kmeans. The notebook provides an example of usage with a custom dataset converted into the cuvs-bench data format.

Failure occurred during model training with the entire dataset loaded into system memory. The same code worked for 2M vectors.

Reproduce the failure with kmeans_sweep.py from the linked cuvs-bench-runner repository, using algorithm=cuvs_kmeans and the cuVS benchmark-format dataset described in the notebook. Start at the cuvs_kmeans.fit call and the reduce_cols_by_key.cuh line=140 error, comparing the failing 4M-vector run with the successful 2M-vector run. Done means training completes for the larger dataset without the invalid configuration error.