Compute the expanded L2 distance in fp32 instead of fp64

[maxwbuckley]

Problem

In calculate_metric (cpp/src/neighbors/detail/nn_descent.cuh), the L2Expanded / L2SqrtExpanded branch computes

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 executes 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 ends up dominating local_join_kernel_wmma on those parts. Counting fp64 opcodes (DADD/DMUL/DFMA/DSETP/F2D/D2F) in the SASS for this kernel:

	fp64 instructions
before	18
after	0

Measurements

RTX 5090 (sm_120a), CUDA 13.2, SIFT1M (1M × 128), graph_degree 64, intermediate_graph_degree 128, 20 iterations, clocks locked. local_join time from nsys, build wall clock is the best of three repeats.

	before	after	speedup
local_join_kernel_wmma (20 iters)	1578.0 ms	1412.0 ms	1.118x
nn_descent::build wall clock	3627.9 ms	3452.9 ms	1.051x

Baseline reproducibility across three separate builds of unmodified main: 1577.99 / 1578.21 / 1577.91 ms (0.02% spread).

The end-to-end figure is smaller than the kernel figure because local_join is ~84% of GPU kernel time but only ~44% of build wall clock; the rest is host-side work.

Accuracy

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 that was lost upstream.

• recall@10 on SIFT1M: 0.9992–1.0000, versus 0.9996–1.0000 for unmodified main across repeats. The spread is run-to-run nondeterminism in graph construction, not an effect of this change: an unrelated control patch that cannot alter any computed value shows the same variation.
• A standalone harness driving this kernel directly produces a bit-identical graph at d=96, d=128 and d=960.
• The negative-distance clamp immediately below ([BUG] NN Descent resulting in negative and NaN distances #991) is untouched and still behaves the same way.

Scope

This is not architecture specific — any part with a low fp64:fp32 ratio pays it, so sm_80/sm_90 should benefit too, by an amount that scales with their fp64 ratio. I only have Blackwell consumer hardware to measure on, so I have not quantified it elsewhere.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxwbuckley]

@lowener thank you :)

[lowener]

/ok to test 36d5e5d