Repository navigation
Compute the expanded L2 distance in fp32 instead of fp64 - #2531
maxwbuckley wants to merge 2 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⚙️ Run configuration
📒 Files selected for processing (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe expanded L2 distance calculation changes its multiplier literal from ChangesL2 distance calculation
Priority: ➖ Normal Estimated code review effort: 1 (Trivial) | ~2 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The change keeps the expanded L2 calculation in single precision, with possible rounding differences. No additional changed files or concrete regressions are identified, so no material merge risk remains. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
2 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