From 3206b850d7549ea50f387f049eb9949a0d1a8cac Mon Sep 17 00:00:00 2001 From: anon Date: Mon, 12 Oct 2026 15:27:38 +0000 Subject: [PATCH 1/9] fix-spectral-clustering-flaky --- .../distance/detail/distance_ops/l2_exp.cuh | 36 +++---------------- 1 file changed, 5 insertions(+), 31 deletions(-) diff --git a/cpp/src/distance/detail/distance_ops/l2_exp.cuh b/cpp/src/distance/detail/distance_ops/l2_exp.cuh index c0433c676a..505c94c5f0 100644 --- a/cpp/src/distance/detail/distance_ops/l2_exp.cuh +++ b/cpp/src/distance/detail/distance_ops/l2_exp.cuh @@ -12,22 +12,6 @@ namespace cuvs::distance::detail::ops { -/** - * Reserve 1 digit of precision from each floating-point type - * for round-off error tolerance. - * @tparam DataT - */ -template -__device__ constexpr AccT get_clamp_precision() -{ - switch (sizeof(DataT)) { - case 2: return AccT{1e-3}; - case 4: return AccT{1e-6}; - case 8: return AccT{1e-15}; - default: return AccT{0}; - } -} - // Epilogue operator for CUTLASS based kernel template struct l2_exp_cutlass_op { @@ -39,13 +23,9 @@ struct l2_exp_cutlass_op { { AccT outVal = aNorm + bNorm - AccT(2.0) * accVal; - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product (accVal) - * can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. - */ - outVal = - outVal * AccT(!((outVal * outVal < get_clamp_precision()) * (aNorm == bNorm))); - return sqrt ? raft::sqrt(outVal * static_cast(outVal > AccT(0))) : outVal; + // Expanded L2 can be slightly negative due to round-off; clamp to zero. + outVal = outVal > AccT{0} ? outVal : AccT{0}; + return sqrt ? raft::sqrt(outVal) : outVal; } __device__ AccT operator()(DataT aData) const noexcept @@ -113,14 +93,8 @@ struct l2_exp_distance_op { AccT accVal = acc[i][j]; AccT val = regxn[i] + regyn[j] - (AccT)2.0 * accVal; - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product - * (accVal) can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal - * instead. - */ - acc[i][j] = val * static_cast((val > AccT(0))) * - static_cast( - !((val * val < get_clamp_precision()) * (regxn[i] == regyn[j]))); + // Expanded L2 can be slightly negative due to round-off; clamp to zero. + acc[i][j] = val > AccT{0} ? val : AccT{0}; } } if (sqrt) { From b49edb5fe48931094186e189b95a2b4ab03ef1e8 Mon Sep 17 00:00:00 2001 From: anon Date: Mon, 12 Oct 2026 15:47:34 +0000 Subject: [PATCH 2/9] fix style --- cpp/src/distance/detail/distance_ops/l2_exp.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/distance/detail/distance_ops/l2_exp.cuh b/cpp/src/distance/detail/distance_ops/l2_exp.cuh index 505c94c5f0..e33b6bd976 100644 --- a/cpp/src/distance/detail/distance_ops/l2_exp.cuh +++ b/cpp/src/distance/detail/distance_ops/l2_exp.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From b06f50f571cc27596d379988599f04cf66356207 Mon Sep 17 00:00:00 2001 From: anon Date: Mon, 12 Oct 2026 16:31:54 +0000 Subject: [PATCH 3/9] address coderabbit NaN comment --- cpp/src/distance/detail/distance_ops/l2_exp.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/distance/detail/distance_ops/l2_exp.cuh b/cpp/src/distance/detail/distance_ops/l2_exp.cuh index e33b6bd976..a79bf2ba47 100644 --- a/cpp/src/distance/detail/distance_ops/l2_exp.cuh +++ b/cpp/src/distance/detail/distance_ops/l2_exp.cuh @@ -24,7 +24,7 @@ struct l2_exp_cutlass_op { AccT outVal = aNorm + bNorm - AccT(2.0) * accVal; // Expanded L2 can be slightly negative due to round-off; clamp to zero. - outVal = outVal > AccT{0} ? outVal : AccT{0}; + outVal = outVal < AccT{0} ? AccT{0} : outVal; return sqrt ? raft::sqrt(outVal) : outVal; } @@ -94,7 +94,7 @@ struct l2_exp_distance_op { AccT val = regxn[i] + regyn[j] - (AccT)2.0 * accVal; // Expanded L2 can be slightly negative due to round-off; clamp to zero. - acc[i][j] = val > AccT{0} ? val : AccT{0}; + acc[i][j] = val < AccT{0} ? AccT{0} : val; } } if (sqrt) { From b75a23d7345d86f755b0ecb5ba56b72819d841ab Mon Sep 17 00:00:00 2001 From: aamijar Date: Tue, 13 Oct 2026 16:49:29 +0000 Subject: [PATCH 4/9] fix self neighbor zeros for pairwise_distances --- cpp/src/distance/detail/distance.cuh | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/cpp/src/distance/detail/distance.cuh b/cpp/src/distance/detail/distance.cuh index 99a226a7d2..ced5db10b6 100644 --- a/cpp/src/distance/detail/distance.cuh +++ b/cpp/src/distance/detail/distance.cuh @@ -7,7 +7,9 @@ #include "distance_ops/all_ops.cuh" #include "pairwise_matrix/dispatch.cuh" + #include + #include #include #include @@ -15,13 +17,22 @@ #include #include #include // to_float +#include +#include #include namespace cuvs { namespace distance { namespace detail { +template +RAFT_KERNEL set_self_distances(OutT* out, IdxT diagonal_size, IdxT leading_dimension, FinOpT fin_op) +{ + IdxT idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < diagonal_size) { out[std::size_t(idx) * leading_dimension + idx] = fin_op(AccT{0}, 0); } +} + /** * @brief: A tag type for overload resolution based on DistanceType * @@ -591,6 +602,23 @@ void distance_impl_l2_expanded( // NOTE: different name ops::l2_exp_distance_op distance_op{perform_sqrt}; pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); + + // When both inputs are the same matrix, diagonal entries compare each row with itself. + // Set those known self-distances to zero. + if (x == y && m > 0 && n > 0) { + constexpr uint32_t block_size = 256; + IdxT diagonal_size = std::min(m, n); + IdxT leading_dimension = is_row_major ? n : m; + auto grid_size = raft::ceildiv(diagonal_size, IdxT{block_size}); + raft::launch_kernel(stream, + grid_size, + block_size, + set_self_distances, + out, + diagonal_size, + leading_dimension, + fin_op); + } } template From a5ced2d9d9c52591e813bb14ff9962fd1ed11ee6 Mon Sep 17 00:00:00 2001 From: aamijar Date: Tue, 13 Oct 2026 16:54:24 +0000 Subject: [PATCH 5/9] rename suffix --- cpp/src/distance/detail/distance.cuh | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cpp/src/distance/detail/distance.cuh b/cpp/src/distance/detail/distance.cuh index ced5db10b6..df9a21105e 100644 --- a/cpp/src/distance/detail/distance.cuh +++ b/cpp/src/distance/detail/distance.cuh @@ -27,7 +27,10 @@ namespace distance { namespace detail { template -RAFT_KERNEL set_self_distances(OutT* out, IdxT diagonal_size, IdxT leading_dimension, FinOpT fin_op) +RAFT_KERNEL set_self_distances_kernel(OutT* out, + IdxT diagonal_size, + IdxT leading_dimension, + FinOpT fin_op) { IdxT idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx < diagonal_size) { out[std::size_t(idx) * leading_dimension + idx] = fin_op(AccT{0}, 0); } @@ -613,7 +616,7 @@ void distance_impl_l2_expanded( // NOTE: different name raft::launch_kernel(stream, grid_size, block_size, - set_self_distances, + set_self_distances_kernel, out, diagonal_size, leading_dimension, From d61603a74ddec1f3531e0358e54ef78ce93226cf Mon Sep 17 00:00:00 2001 From: aamijar Date: Tue, 13 Oct 2026 21:10:00 +0000 Subject: [PATCH 6/9] fix col-major different size m and n layouts --- cpp/src/distance/detail/distance.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/distance/detail/distance.cuh b/cpp/src/distance/detail/distance.cuh index df9a21105e..e60674844b 100644 --- a/cpp/src/distance/detail/distance.cuh +++ b/cpp/src/distance/detail/distance.cuh @@ -608,7 +608,7 @@ void distance_impl_l2_expanded( // NOTE: different name // When both inputs are the same matrix, diagonal entries compare each row with itself. // Set those known self-distances to zero. - if (x == y && m > 0 && n > 0) { + if (x == y && m > 0 && n > 0 && (is_row_major || m == n)) { constexpr uint32_t block_size = 256; IdxT diagonal_size = std::min(m, n); IdxT leading_dimension = is_row_major ? n : m; From 7b4be11e22cdf37a00c0f6d8e6b5f8742a65bec6 Mon Sep 17 00:00:00 2001 From: aamijar Date: Thu, 15 Oct 2026 15:23:07 +0000 Subject: [PATCH 7/9] add regression tests --- cpp/tests/distance/dist_l2_exp.cu | 49 +++++++++++++++++++- cpp/tests/neighbors/distance_nn.cu | 73 ++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) diff --git a/cpp/tests/distance/dist_l2_exp.cu b/cpp/tests/distance/dist_l2_exp.cu index a91daff1b3..1c9dbe4e28 100644 --- a/cpp/tests/distance/dist_l2_exp.cu +++ b/cpp/tests/distance/dist_l2_exp.cu @@ -1,11 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../test_utils.cuh" #include "distance_base.cuh" +#include +#include + +#include +#include + namespace cuvs { namespace distance { @@ -162,5 +168,46 @@ INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucExpTestD, ::testing::ValuesIn( class BigMatrixEucExp : public BigMatrixDistanceTest {}; TEST_F(BigMatrixEucExp, Result) {} + +// Distinct points with equal norms must keep their (small) distances. The clamp of self-distance +// round-off to zero used to apply an absolute tolerance, which collapsed such points when the data +// has a small magnitude (e.g. a spectral embedding). +template +void test_small_scale_equal_norms(T scale) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n = 8; + std::vector h_x(n * n, T{0}); + for (int64_t i = 0; i < n; i++) { + h_x[i * n + i] = scale; + } + auto x = raft::make_device_matrix(handle, n, n); + auto d = raft::make_device_matrix(handle, n, n); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream); + + for (auto metric : {DistanceType::L2Expanded, DistanceType::L2SqrtExpanded}) { + cuvs::distance::pairwise_distance(handle, + raft::make_const_mdspan(x.view()), + raft::make_const_mdspan(x.view()), + d.view(), + metric); + std::vector h_d(n * n); + raft::update_host(h_d.data(), d.data_handle(), h_d.size(), stream); + raft::resource::sync_stream(handle, stream); + + const T expected = + metric == DistanceType::L2Expanded ? T(2) * scale * scale : std::sqrt(T(2)) * scale; + for (int64_t i = 0; i < n; i++) { + for (int64_t j = 0; j < n; j++) { + EXPECT_NEAR(h_d[i * n + j], i == j ? T(0) : expected, T(1e-3) * expected) + << "metric=" << static_cast(metric) << " i=" << i << " j=" << j; + } + } + } +} + +TEST(DistanceEucExpSmallScale, EqualNormsF) { test_small_scale_equal_norms(1e-2f); } +TEST(DistanceEucExpSmallScale, EqualNormsD) { test_small_scale_equal_norms(1e-5); } } // end namespace distance } // namespace cuvs diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 13b9f1d231..585ba379ba 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -16,8 +16,10 @@ #include #include #include +#include #include +#include namespace cuvs::neighbors { @@ -464,6 +466,77 @@ TEST(Top1nnPlan, RejectsMismatchedLaunch) plan))); } +// Distinct points with equal norms and a small magnitude (as in a spectral embedding) must not be +// treated as self-neighbors when finding the nearest neighbor. +TEST(Top1nnEqualNorms, SmallScale) +{ + using Backend = cuvs::distance::detail::Top1nnBackend; + using Kvp = raft::KeyValuePair; + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int m = 1000; + constexpr int n = 5; + constexpr int k = 5; + constexpr float s = 1e-2f; + std::vector h_x(m * k, 0.0f); + std::vector h_y(n * k, 0.0f); + for (int i = 0; i < m; i++) { + h_x[i * k + i % n] = s; + } + for (int j = 0; j < n; j++) { + h_y[j * k + j] = s; + } + auto x = raft::make_device_matrix(handle, m, k); + auto y = raft::make_device_matrix(handle, n, k); + auto x_norm = raft::make_device_vector(handle, m); + auto y_norm = raft::make_device_vector(handle, n); + auto output = raft::make_device_vector(handle, m); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream.get()); + raft::update_device(y.data_handle(), h_y.data(), h_y.size(), stream.get()); + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, m, stream.get()); + raft::linalg::rowNorm( + y_norm.data_handle(), y.data_handle(), k, n, stream.get()); + + cuvs::distance::detail::Top1nnTuning tuning{}; + for (auto backend : {Backend::Cutlass, Backend::Unfused}) { + if (!cuvs::distance::detail::is_top_1_nn_backend_available( + backend, x.data_handle(), y.data_handle(), m, n, k, DistanceType::L2Expanded)) { + continue; + } + const auto plan = cuvs::distance::probe_top_1_nn( + handle, x.data_handle(), y.data_handle(), m, n, k, tuning, DistanceType::L2Expanded, backend); + ASSERT_TRUE(plan.available); + auto workspace = raft::make_device_vector(handle, plan.workspace_bytes); + cuvs::distance::top_1_nn(handle, + output.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + tuning, + workspace.data_handle(), + workspace.size(), + false, + true, + true, + DistanceType::L2Expanded, + 0.0f, + plan); + std::vector h_out(m); + raft::update_host(h_out.data(), output.data_handle(), m, stream.get()); + raft::resource::sync_stream(handle, stream); + int n_wrong = 0; + for (int i = 0; i < m; i++) { + n_wrong += h_out[i].key != i % n; + } + EXPECT_EQ(n_wrong, 0) << "backend=" << static_cast(backend); + } +} + TEST(Top1nnPlan, ExplicitUnavailableBackendIsStrict) { raft::resources handle; From 2e76070e94bc88565a241ac3b5811b3da3664706 Mon Sep 17 00:00:00 2001 From: aamijar Date: Thu, 15 Oct 2026 19:37:14 +0000 Subject: [PATCH 8/9] loosen go test self-neighbor epsilon --- go/brute_force/brute_force_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/go/brute_force/brute_force_test.go b/go/brute_force/brute_force_test.go index 9525cc3691..842a8dd83f 100644 --- a/go/brute_force/brute_force_test.go +++ b/go/brute_force/brute_force_test.go @@ -13,7 +13,9 @@ func TestBruteForce(t *testing.T) { nFeatures = 16 nQueries = 4 k = 4 - epsilon = 0.001 + // Self-neighbors with L2SqrtExpanded may have small round-off error + // which can cause distances to be non-zero. + epsilon = 0.01 ) resource, _ := cuvs.NewResource(nil) From f3aefcbc824c989b8f443e5c748e77a7fc738703 Mon Sep 17 00:00:00 2001 From: aamijar Date: Fri, 16 Oct 2026 18:11:02 +0000 Subject: [PATCH 9/9] add brute_force and diagonal tests --- cpp/tests/distance/dist_l2_exp.cu | 46 ++++++++++++++++++ cpp/tests/neighbors/brute_force.cu | 78 ++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/cpp/tests/distance/dist_l2_exp.cu b/cpp/tests/distance/dist_l2_exp.cu index 1c9dbe4e28..c3082d6c23 100644 --- a/cpp/tests/distance/dist_l2_exp.cu +++ b/cpp/tests/distance/dist_l2_exp.cu @@ -209,5 +209,51 @@ void test_small_scale_equal_norms(T scale) TEST(DistanceEucExpSmallScale, EqualNormsF) { test_small_scale_equal_norms(1e-2f); } TEST(DistanceEucExpSmallScale, EqualNormsD) { test_small_scale_equal_norms(1e-5); } + +template +void test_exact_self_distance_diagonal() +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 50; + constexpr int64_t n_cols = 10; + constexpr bool row_major = std::is_same_v; + std::vector h_x(n_rows * n_cols); + for (int64_t i = 0; i < n_rows; ++i) { + for (int64_t j = 0; j < n_cols; ++j) { + auto index = row_major ? i * n_cols + j : j * n_rows + i; + h_x[index] = static_cast(((i + 3) * (j + 5)) % 37) / 37.0f; + } + } + + auto x = raft::make_device_matrix(handle, n_rows, n_cols); + auto d = raft::make_device_matrix(handle, n_rows, n_rows); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream); + + for (auto metric : {DistanceType::L2Expanded, DistanceType::L2SqrtExpanded}) { + cuvs::distance::pairwise_distance(handle, + raft::make_const_mdspan(x.view()), + raft::make_const_mdspan(x.view()), + d.view(), + metric); + std::vector h_d(n_rows * n_rows); + raft::update_host(h_d.data(), d.data_handle(), h_d.size(), stream); + raft::resource::sync_stream(handle, stream); + + for (int64_t i = 0; i < n_rows; ++i) { + EXPECT_EQ(h_d[i * n_rows + i], 0.0f) << "metric=" << static_cast(metric) << " row=" << i; + } + } +} + +TEST(DistanceEucExpSelfDistance, ExactDiagonalRowMajor) +{ + test_exact_self_distance_diagonal(); +} + +TEST(DistanceEucExpSelfDistance, ExactDiagonalColMajor) +{ + test_exact_self_distance_diagonal(); +} } // end namespace distance } // namespace cuvs diff --git a/cpp/tests/neighbors/brute_force.cu b/cpp/tests/neighbors/brute_force.cu index f8601859d2..1c09b9a739 100644 --- a/cpp/tests/neighbors/brute_force.cu +++ b/cpp/tests/neighbors/brute_force.cu @@ -18,6 +18,7 @@ #include #include +#include #include namespace cuvs::neighbors::brute_force { @@ -195,6 +196,83 @@ TEST_P(KNNTest_half_int64_t, BruteForce) { this->testBruteForce(); } INSTANTIATE_TEST_CASE_P(KNNTest, KNNTest_float_int64_t, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(KNNTest, KNNTest_half_int64_t, ::testing::ValuesIn(inputs)); +TEST(BruteForceEqualNorms, PreservesNeighborOrdering) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_database = 128; + constexpr int64_t n_queries = 1; + constexpr int64_t dim = 2; + constexpr double angle_step = 2e-3; + + std::vector h_database(n_database * dim); + for (int64_t i = 0; i < n_database; ++i) { + auto rank = (i * 73) % n_database; + double angle = static_cast(rank) * angle_step; + h_database[i * dim] = static_cast(std::cos(angle)); + h_database[i * dim + 1] = static_cast(std::sin(angle)); + } + std::vector h_queries{1.0f, 0.0f}; + + struct reference_neighbor { + double squared_distance; + int64_t index; + }; + std::vector reference(n_database); + for (int64_t i = 0; i < n_database; ++i) { + double dx = static_cast(h_queries[0]) - static_cast(h_database[i * dim]); + double dy = static_cast(h_queries[1]) - static_cast(h_database[i * dim + 1]); + reference[i] = {.squared_distance = dx * dx + dy * dy, .index = i}; + } + std::ranges::sort( + reference, [](const reference_neighbor& lhs, const reference_neighbor& rhs) -> bool { + return lhs.squared_distance < rhs.squared_distance || + (lhs.squared_distance == rhs.squared_distance && lhs.index < rhs.index); + }); + + auto database = raft::make_device_matrix(handle, n_database, dim); + auto queries = raft::make_device_matrix(handle, n_queries, dim); + raft::update_device(database.data_handle(), h_database.data(), h_database.size(), stream); + raft::update_device(queries.data_handle(), h_queries.data(), h_queries.size(), stream); + + for (auto metric : + {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::L2SqrtExpanded}) { + index_params index_params; + index_params.metric = metric; + auto index = build(handle, index_params, raft::make_const_mdspan(database.view())); + + for (int64_t k : {int64_t{15}, int64_t{65}}) { + auto neighbors = raft::make_device_matrix(handle, n_queries, k); + auto distances = raft::make_device_matrix(handle, n_queries, k); + search(handle, + search_params{}, + index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view(), + filtering::none_sample_filter{}); + + std::vector h_neighbors(k); + std::vector h_distances(k); + raft::update_host(h_neighbors.data(), neighbors.data_handle(), k, stream); + raft::update_host(h_distances.data(), distances.data_handle(), k, stream); + raft::resource::sync_stream(handle, stream); + + for (int64_t i = 0; i < k; ++i) { + auto expected_distance = metric == cuvs::distance::DistanceType::L2Expanded + ? reference[i].squared_distance + : std::sqrt(reference[i].squared_distance); + EXPECT_EQ(h_neighbors[i], reference[i].index) + << "metric=" << static_cast(metric) << " k=" << k << " rank=" << i; + EXPECT_NEAR(h_distances[i], + expected_distance, + metric == cuvs::distance::DistanceType::L2Expanded ? 5e-7 : 5e-5) + << "metric=" << static_cast(metric) << " k=" << k << " rank=" << i; + } + } + } +} + // Also test with larger random inputs, including col-major inputs struct RandomKNNInputs { int num_queries;