From 36d5e5dd9a4a65ef3d3492901281bec130b793b7 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Mon, 31 Aug 2026 13:43:41 +0200 Subject: [PATCH] Compute the expanded L2 distance in fp32 instead of fp64 In calculate_metric the L2Expanded / L2SqrtExpanded branch reads s_distances[i] = l2_norms[..] + l2_norms[..] - 2.0 * s_distances[i]; `2.0` is a double literal, so the whole expression is promoted to fp64: both l2_norms loads and s_distances[i] are widened, the arithmetic runs on the fp64 pipe, and the result is narrowed back on assignment. This happens once per Gram matrix element, twice per CTA, for every row of the dataset on every NN-Descent iteration. Consumer GPUs run fp64 at a small fraction of their fp32 rate (1/64 on GB202), so this dominates local_join_kernel_wmma on those parts. Counting fp64 opcodes (DADD/DMUL/DFMA/DSETP/F2D/D2F) in the SASS for this kernel gives 18 before the change and 0 after. Measured on an RTX 5090 (sm_120a), SIFT1M, 1M x 128, graph_degree 64, intermediate_graph_degree 128, 20 iterations, clocks locked: local_join_kernel_wmma 1578.0 ms -> 1412.0 ms (1.118x) nn_descent::build 3627.9 ms -> 3452.9 ms (1.051x) The fp64 intermediate was not buying accuracy. Both operands already carry fp16-level error: s_distances[i] comes out of an fp16 wmma, and l2_norms is computed in fp32 from fp16 data, so evaluating the final subtraction in fp64 cannot recover precision lost upstream. recall@10 is unchanged at 0.9992-1.0000, which is the run-to-run spread of the unmodified build, and a standalone harness produces a bit-identical graph at d=96, d=128 and d=960. The negative-distance clamp below (issue #991) is untouched. This is not architecture specific; any part with a low fp64:fp32 ratio pays it. --- cpp/src/neighbors/detail/nn_descent.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 432cde7ffc..ce8b089be3 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -549,7 +549,7 @@ __device__ __forceinline__ void calculate_metric(float* s_distances, } else if (metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded) { s_distances[i] = - l2_norms[row_neighbors[row_id]] + l2_norms[col_neighbors[col_id]] - 2.0 * s_distances[i]; + l2_norms[row_neighbors[row_id]] + l2_norms[col_neighbors[col_id]] - 2.0f * s_distances[i]; // for fp32 vs fp16 precision differences resulting in negative distances when distance // should be 0 related issue: https://github.com/nvidia/cuvs/issues/991 s_distances[i] = s_distances[i] < 0.0f ? 0.0f : s_distances[i];