Repository navigation
Compute the expanded L2 distance in fp32 instead of fp64 - #2531
rapids-bot[bot] merged 4 commits into
Conversation
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 NVIDIA#991) is untouched. This is not architecture specific; any part with a low fp64:fp32 ratio pays it.
|
@lowener thank you :) |
lowener
added
improvement
|
/ok to test 36d5e5d |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
@lowener merged Locally (RTX 5090, sm_120, CUDA 13.3 env from Could you |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Summary
Merge Risk: ⚪ Minimal · up to No actionable merge-blocking risk is evident in the described single-precision multiplier change. The PR is ready to merge after normal checks. Pre-merge checks |
 |
| Check name | Status | Explanation |
|---|---|---|
| Title check |  |
The title clearly and concisely describes the main change: computing the expanded L2 distance in fp32 instead of fp64. |
| Description check | |
The description is directly related to the code change and provides relevant implementation details, performance measurements, accuracy results, testing information, and scope. |
| Docstring Coverage | |
No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0… |
| Linked Issues check | |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
@lowener I merged I rebuilt and ran the complete This is approved but CI has never run on a current head. Could you |
|
/ok to test 2b4015e |
|
/merge |
3 participants
Problem
In
calculate_metric(cpp/src/neighbors/detail/nn_descent.cuh), theL2Expanded/L2SqrtExpandedbranch computess_distances[i] = l2_norms[..] + l2_norms[..] - 2.0 * s_distances[i];2.0is adoubleliteral, so the whole expression is promoted to fp64: bothl2_normsloads ands_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_wmmaon those parts. Counting fp64 opcodes (DADD/DMUL/DFMA/DSETP/F2D/D2F) in the SASS for this kernel:Measurements
RTX 5090 (sm_120a), CUDA 13.2, SIFT1M (1M × 128),
graph_degree64,intermediate_graph_degree128, 20 iterations, clocks locked.local_jointime fromnsys, build wall clock is the best of three repeats.local_join_kernel_wmma(20 iters)nn_descent::buildwall clockBaseline 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_joinis ~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 fp16wmma, andl2_normsis computed in fp32 from fp16 data — so evaluating the final subtraction in fp64 cannot recover precision that was lost upstream.mainacross 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.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