From 8321c8d1daf37cf7ccdccb2c31a92cf8e107330f Mon Sep 17 00:00:00 2001 From: nethum529 Date: Fri, 17 Jul 2026 22:32:40 -0500 Subject: [PATCH 1/4] Fix KNN neighbor buffer layout issues Signed-off-by: nethum529 --- .../cuml/cuml/neighbors/nearest_neighbors.pyx | 3 +- python/cuml/cuml/neighbors/weights.py | 6 ++- .../cuml/tests/test_kneighbors_regressor.py | 36 ++++++++++++- python/cuml/tests/test_nearest_neighbors.py | 54 ++++++++++++++++++- 4 files changed, 94 insertions(+), 5 deletions(-) diff --git a/python/cuml/cuml/neighbors/nearest_neighbors.pyx b/python/cuml/cuml/neighbors/nearest_neighbors.pyx index 32e4d67603..f0f9861949 100644 --- a/python/cuml/cuml/neighbors/nearest_neighbors.pyx +++ b/python/cuml/cuml/neighbors/nearest_neighbors.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # import warnings @@ -192,6 +192,7 @@ void swap_kernel(long long int* I, float* D, int n_rows, int n_cols) { def _drop_self_edges(distances, indices): """Drop edges between a point and itself in the knn graph""" + indices = cp.ascontiguousarray(indices, dtype=cp.int64) rows, cols = indices.shape # Launch config diff --git a/python/cuml/cuml/neighbors/weights.py b/python/cuml/cuml/neighbors/weights.py index ca369c2ed3..6774710aeb 100644 --- a/python/cuml/cuml/neighbors/weights.py +++ b/python/cuml/cuml/neighbors/weights.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import cupy as cp @@ -56,7 +56,9 @@ def compute_weights(distances, weights): return raw_weights elif callable(weights): # Custom callable weights (raw, not normalized) - raw_weights = cp.asarray(weights(distances), dtype=cp.float32) + raw_weights = cp.ascontiguousarray( + cp.asarray(weights(distances), dtype=cp.float32) + ) # Return raw weights - normalization will be done in C++ kernel return raw_weights else: diff --git a/python/cuml/tests/test_kneighbors_regressor.py b/python/cuml/tests/test_kneighbors_regressor.py index a0402c8b3b..59b91ef000 100644 --- a/python/cuml/tests/test_kneighbors_regressor.py +++ b/python/cuml/tests/test_kneighbors_regressor.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -173,6 +173,40 @@ def test_weights_predict(weights, n_neighbors): np.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) +def test_callable_weights_non_c_contiguous_cupy_view(): + X, y = make_regression( + n_samples=64, + n_features=6, + n_informative=4, + random_state=42, + ) + X = X.astype(np.float32) + y = y.astype(np.float32) + + X_train, X_test = X[:48], X[48:] + y_train = y[:48] + + def non_c_contiguous_weights(distances): + xp = cp if isinstance(distances, cp.ndarray) else np + weights = xp.asarray(1.0 / (1.0 + distances), dtype=xp.float32) + base = xp.empty((weights.shape[1], weights.shape[0]), dtype=xp.float32) + base[...] = weights.T + result = base.T + assert result.dtype == xp.float32 + assert not result.flags.c_contiguous + return result + + knn_cu = cuKNN(n_neighbors=5, weights=non_c_contiguous_weights) + knn_cu.fit(X_train, y_train) + pred_cu = knn_cu.predict(cp.asarray(X_test)) + + knn_sk = skKNN(n_neighbors=5, weights=non_c_contiguous_weights) + knn_sk.fit(X_train, y_train) + pred_sk = knn_sk.predict(X_test) + + cp.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) + + @pytest.mark.parametrize("weights", ["uniform", "distance"]) def test_weights_multioutput(weights): """Test weights parameter with multioutput regression.""" diff --git a/python/cuml/tests/test_nearest_neighbors.py b/python/cuml/tests/test_nearest_neighbors.py index 5f72a7dfa0..b8a7037a51 100644 --- a/python/cuml/tests/test_nearest_neighbors.py +++ b/python/cuml/tests/test_nearest_neighbors.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -712,6 +712,58 @@ def test_nearest_neighbors_sparse( assert (len(diffs[diffs > 0]) / len(np.ravel(skI))) <= 0.005 +def test_nearest_neighbors_sparse_x_none_self_edge_swap_indices_int64(): + X_dense = np.array( + [ + [1.0, 0.0, 2.0, 0.0], + [1.0, 0.0, 2.0, 0.0], + [0.0, 3.0, 0.0, 4.0], + [2.0, 0.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + X = cupyx.scipy.sparse.csr_matrix(cp.asarray(X_dense)) + + nn = cuKNN( + metric="euclidean", + n_neighbors=2, + algorithm="brute", + output_type="cupy", + ) + nn.fit(X) + + explicit_distances, explicit_indices = nn.kneighbors(X, n_neighbors=2) + assert explicit_indices.dtype == cp.int32 + assert explicit_distances.dtype == cp.float32 + + distances, indices = nn.kneighbors(X=None, n_neighbors=2) + assert indices.dtype == cp.int64 + assert indices.flags.c_contiguous + assert distances.dtype == cp.float32 + assert distances.flags.c_contiguous + + indices_np = cp.asnumpy(indices) + distances_np = cp.asnumpy(distances) + for row, row_indices in enumerate(indices_np): + assert row not in row_indices + + sk_distances, sk_indices = ( + skKNN( + metric="euclidean", + n_neighbors=2, + algorithm="brute", + ) + .fit(X_dense) + .kneighbors(X=None, n_neighbors=2) + ) + + np.testing.assert_allclose( + distances_np, sk_distances, atol=1e-5, rtol=1e-5 + ) + for row_indices, expected_indices in zip(indices_np, sk_indices): + assert set(row_indices.tolist()) == set(expected_indices.tolist()) + + @pytest.mark.parametrize("n_neighbors", [1, 5, 6]) def test_haversine(n_neighbors): hoboken_nj = [40.745255, -74.034775] From 27827bbd52b8d4717f1173daab0385404eb98298 Mon Sep 17 00:00:00 2001 From: nethum529 Date: Fri, 17 Jul 2026 22:46:30 -0500 Subject: [PATCH 2/4] Stabilize sparse KNN self-edge regression Signed-off-by: nethum529 --- python/cuml/tests/test_nearest_neighbors.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/cuml/tests/test_nearest_neighbors.py b/python/cuml/tests/test_nearest_neighbors.py index b8a7037a51..1cba505752 100644 --- a/python/cuml/tests/test_nearest_neighbors.py +++ b/python/cuml/tests/test_nearest_neighbors.py @@ -715,10 +715,10 @@ def test_nearest_neighbors_sparse( def test_nearest_neighbors_sparse_x_none_self_edge_swap_indices_int64(): X_dense = np.array( [ - [1.0, 0.0, 2.0, 0.0], - [1.0, 0.0, 2.0, 0.0], - [0.0, 3.0, 0.0, 4.0], - [2.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [10.0, 0.0, 0.0, 0.0], + [10.25, 0.0, 0.0, 0.0], ], dtype=np.float32, ) @@ -736,7 +736,7 @@ def test_nearest_neighbors_sparse_x_none_self_edge_swap_indices_int64(): assert explicit_indices.dtype == cp.int32 assert explicit_distances.dtype == cp.float32 - distances, indices = nn.kneighbors(X=None, n_neighbors=2) + distances, indices = nn.kneighbors(X=None, n_neighbors=1) assert indices.dtype == cp.int64 assert indices.flags.c_contiguous assert distances.dtype == cp.float32 @@ -750,11 +750,11 @@ def test_nearest_neighbors_sparse_x_none_self_edge_swap_indices_int64(): sk_distances, sk_indices = ( skKNN( metric="euclidean", - n_neighbors=2, + n_neighbors=1, algorithm="brute", ) .fit(X_dense) - .kneighbors(X=None, n_neighbors=2) + .kneighbors(X=None, n_neighbors=1) ) np.testing.assert_allclose( From ee12f2c1adc281dcb12c7f835196fd04dbe06fbe Mon Sep 17 00:00:00 2001 From: nethum529 Date: Fri, 17 Jul 2026 23:09:34 -0500 Subject: [PATCH 3/4] Fix callable KNN weight alignment Signed-off-by: nethum529 --- python/cuml/cuml/neighbors/weights.py | 21 ++++++-- .../cuml/tests/test_kneighbors_regressor.py | 53 +++++++++++++++++++ 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/python/cuml/cuml/neighbors/weights.py b/python/cuml/cuml/neighbors/weights.py index 6774710aeb..0d38cd2269 100644 --- a/python/cuml/cuml/neighbors/weights.py +++ b/python/cuml/cuml/neighbors/weights.py @@ -56,11 +56,26 @@ def compute_weights(distances, weights): return raw_weights elif callable(weights): # Custom callable weights (raw, not normalized) - raw_weights = cp.ascontiguousarray( - cp.asarray(weights(distances), dtype=cp.float32) + raw_weights = cp.asarray(weights(distances), dtype=cp.float32) + + # Copy as bytes so an unaligned float32 view is never dereferenced. + # The destination is a CuPy-owned, aligned C-contiguous allocation. + aligned_weights = cp.empty( + raw_weights.shape, dtype=cp.float32, order="C" + ) + byte_shape = raw_weights.shape + (raw_weights.itemsize,) + raw_weight_bytes = cp.ndarray( + byte_shape, + dtype=cp.uint8, + memptr=raw_weights.data, + strides=raw_weights.strides + (1,), ) + aligned_weight_bytes = aligned_weights.view(cp.uint8).reshape( + byte_shape + ) + cp.copyto(aligned_weight_bytes, raw_weight_bytes) # Return raw weights - normalization will be done in C++ kernel - return raw_weights + return aligned_weights else: raise ValueError( f"weights must be 'uniform', 'distance', or a callable, got {weights}" diff --git a/python/cuml/tests/test_kneighbors_regressor.py b/python/cuml/tests/test_kneighbors_regressor.py index 59b91ef000..d74433251c 100644 --- a/python/cuml/tests/test_kneighbors_regressor.py +++ b/python/cuml/tests/test_kneighbors_regressor.py @@ -12,6 +12,7 @@ from sklearn.utils.validation import check_random_state from cuml.neighbors import KNeighborsRegressor as cuKNN +from cuml.neighbors.weights import compute_weights from cuml.testing.utils import array_equal @@ -207,6 +208,58 @@ def non_c_contiguous_weights(distances): cp.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) +def test_callable_weights_unaligned_c_contiguous_cupy_view(): + float32_alignment = np.dtype(np.float32).alignment + + def unaligned_weights(distances): + values = cp.asarray(1.0 / (1.0 + distances), dtype=cp.float32) + storage = cp.empty(values.nbytes + 1, dtype=cp.uint8) + result = storage[1:].view(cp.float32).reshape(values.shape) + cp.cuda.runtime.memcpyAsync( + result.data.ptr, + values.data.ptr, + values.nbytes, + cp.cuda.runtime.memcpyDeviceToDevice, + cp.cuda.get_current_stream().ptr, + ) + cp.cuda.get_current_stream().synchronize() + assert result.flags.c_contiguous + assert result.data.ptr % float32_alignment != 0 + return result + + distances = cp.arange(15, dtype=cp.float32).reshape(3, 5) + computed_weights = compute_weights(distances, unaligned_weights) + + assert computed_weights.dtype == cp.float32 + assert computed_weights.flags.c_contiguous + assert computed_weights.data.ptr % float32_alignment == 0 + assert computed_weights.flags.owndata + cp.testing.assert_allclose( + computed_weights, 1.0 / (1.0 + distances), rtol=1e-6 + ) + + X, y = make_regression( + n_samples=64, + n_features=6, + n_informative=4, + random_state=42, + ) + X = X.astype(np.float32) + y = y.astype(np.float32) + X_train, X_test = X[:48], X[48:] + y_train = y[:48] + + knn_cu = cuKNN(n_neighbors=5, weights=unaligned_weights) + knn_cu.fit(X_train, y_train) + pred_cu = knn_cu.predict(cp.asarray(X_test)) + + knn_sk = skKNN(n_neighbors=5, weights=lambda d: 1.0 / (1.0 + d)) + knn_sk.fit(X_train, y_train) + pred_sk = knn_sk.predict(X_test) + + cp.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) + + @pytest.mark.parametrize("weights", ["uniform", "distance"]) def test_weights_multioutput(weights): """Test weights parameter with multioutput regression.""" From 0d98baac95b70d3096a73da8f1c5ca97fdd0de89 Mon Sep 17 00:00:00 2001 From: nethum529 Date: Mon, 20 Jul 2026 20:18:00 -0500 Subject: [PATCH 4/4] Simplify callable neighbor weights Signed-off-by: nethum529 --- python/cuml/cuml/neighbors/weights.py | 21 ++------ .../cuml/tests/test_kneighbors_regressor.py | 53 ------------------- 2 files changed, 3 insertions(+), 71 deletions(-) diff --git a/python/cuml/cuml/neighbors/weights.py b/python/cuml/cuml/neighbors/weights.py index 0d38cd2269..1a1f058abd 100644 --- a/python/cuml/cuml/neighbors/weights.py +++ b/python/cuml/cuml/neighbors/weights.py @@ -56,26 +56,11 @@ def compute_weights(distances, weights): return raw_weights elif callable(weights): # Custom callable weights (raw, not normalized) - raw_weights = cp.asarray(weights(distances), dtype=cp.float32) - - # Copy as bytes so an unaligned float32 view is never dereferenced. - # The destination is a CuPy-owned, aligned C-contiguous allocation. - aligned_weights = cp.empty( - raw_weights.shape, dtype=cp.float32, order="C" - ) - byte_shape = raw_weights.shape + (raw_weights.itemsize,) - raw_weight_bytes = cp.ndarray( - byte_shape, - dtype=cp.uint8, - memptr=raw_weights.data, - strides=raw_weights.strides + (1,), + raw_weights = cp.asarray( + weights(distances), dtype="float32", order="C" ) - aligned_weight_bytes = aligned_weights.view(cp.uint8).reshape( - byte_shape - ) - cp.copyto(aligned_weight_bytes, raw_weight_bytes) # Return raw weights - normalization will be done in C++ kernel - return aligned_weights + return raw_weights else: raise ValueError( f"weights must be 'uniform', 'distance', or a callable, got {weights}" diff --git a/python/cuml/tests/test_kneighbors_regressor.py b/python/cuml/tests/test_kneighbors_regressor.py index d74433251c..59b91ef000 100644 --- a/python/cuml/tests/test_kneighbors_regressor.py +++ b/python/cuml/tests/test_kneighbors_regressor.py @@ -12,7 +12,6 @@ from sklearn.utils.validation import check_random_state from cuml.neighbors import KNeighborsRegressor as cuKNN -from cuml.neighbors.weights import compute_weights from cuml.testing.utils import array_equal @@ -208,58 +207,6 @@ def non_c_contiguous_weights(distances): cp.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) -def test_callable_weights_unaligned_c_contiguous_cupy_view(): - float32_alignment = np.dtype(np.float32).alignment - - def unaligned_weights(distances): - values = cp.asarray(1.0 / (1.0 + distances), dtype=cp.float32) - storage = cp.empty(values.nbytes + 1, dtype=cp.uint8) - result = storage[1:].view(cp.float32).reshape(values.shape) - cp.cuda.runtime.memcpyAsync( - result.data.ptr, - values.data.ptr, - values.nbytes, - cp.cuda.runtime.memcpyDeviceToDevice, - cp.cuda.get_current_stream().ptr, - ) - cp.cuda.get_current_stream().synchronize() - assert result.flags.c_contiguous - assert result.data.ptr % float32_alignment != 0 - return result - - distances = cp.arange(15, dtype=cp.float32).reshape(3, 5) - computed_weights = compute_weights(distances, unaligned_weights) - - assert computed_weights.dtype == cp.float32 - assert computed_weights.flags.c_contiguous - assert computed_weights.data.ptr % float32_alignment == 0 - assert computed_weights.flags.owndata - cp.testing.assert_allclose( - computed_weights, 1.0 / (1.0 + distances), rtol=1e-6 - ) - - X, y = make_regression( - n_samples=64, - n_features=6, - n_informative=4, - random_state=42, - ) - X = X.astype(np.float32) - y = y.astype(np.float32) - X_train, X_test = X[:48], X[48:] - y_train = y[:48] - - knn_cu = cuKNN(n_neighbors=5, weights=unaligned_weights) - knn_cu.fit(X_train, y_train) - pred_cu = knn_cu.predict(cp.asarray(X_test)) - - knn_sk = skKNN(n_neighbors=5, weights=lambda d: 1.0 / (1.0 + d)) - knn_sk.fit(X_train, y_train) - pred_sk = knn_sk.predict(X_test) - - cp.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) - - @pytest.mark.parametrize("weights", ["uniform", "distance"]) def test_weights_multioutput(weights): """Test weights parameter with multioutput regression."""