[FEA] Faster, leaner NN-Descent / CAGRA builds for Elasticsearch-shaped workloads (RTX PRO 6000 / sm_120, 1024-d) #2837

cc @cjnolet

Summary

Elastic recently published GPU-accelerated vector indexing in Elasticsearch with cuVS: 138M 1024-d vectors on 8× RTX PRO 6000 Blackwell, 7× indexing throughput and ~48× faster force-merge versus CPU. Their GPU codec builds CAGRA graphs with NN-Descent (below 5M vectors per segment) from device-resident fp32 on flush and host fp32 on merge.

We profiled exactly that path on an RTX PRO 6000 (sm_120) and found a set of independent, mostly small changes. Together they make the build substantially faster with the same recall and lower memory use:

• Same recall. kNN-graph recall and CAGRA graph overlap with exact neighbours are unchanged within run-to-run noise on every config (e.g. NN-Descent recall@64 0.9675 → 0.9676).
• Bit-identical arithmetic. Every kernel change produces byte-identical join output on a deterministic kernel-level harness: 42 cases covering dims 64/128/130/768/1024, L2/IP/cosine, and the WMMA-half, WMMA-float and SIMT kernels. The end-to-end graph is not run-to-run deterministic even on main, because of the atomic reverse-edge slot claim, so equivalence there is shown statistically.
• Less memory. Peak transient VRAM on the device path drops from ~16 GB to ~1.1 GB at 4M × 1024 (Don't duplicate a device-resident dataset in CAGRA sort_knn_graph #2834). About 1.9 GB of per-build pinned host memory goes away at 5M rows (Keep NN-Descent per-iteration sample lists on device #2836).

These changes are not Elastic-specific. Every CAGRA build that uses NN-Descent, and every direct nn_descent::build caller, benefits on every architecture. Only the sm_120 occupancy observation is GPU-specific. The gains are largest for high-dimensional (768–1024-d) data and device-resident inputs.

Why the RTX PRO 6000 matters

GB202 parts (RTX PRO 6000, RTX 5090) are consumer/workstation-class: FP64 runs at 1/64 of the FP32 rate. Any fp64 left on a hot path is disproportionately expensive there compared with A100/H100/B200. #2531 already removes the last fp64 from the NN-Descent device build path (2.0 → 2.0f in the expanded-L2 epilogue). Everything below is measured on top of #2531 and none of it is a precision change.

PRs

The four new PRs are independent of each other and of #2531: each applies cleanly on main, so they can be reviewed and merged in any order. All measurements are on main + #2531.

The changes

Each one is a separate PR against main (+#2531), independently reviewable and revertible. Speedups are end-to-end CAGRA build time at the Elastic config (4M × 1024, device fp32), from interleaved runs with 95% CIs, unless noted.

1. Downcast device-resident fp32 to fp16 once before the local join (#2833)

GNND::build downcasts host fp32 into d_data_half_, but device fp32 was passed through, so the WMMA join converted every gathered element on the fly with 4-byte scalar loads. Now device input is converted once device-to-device, with a fallback to the old path if the buffer does not fit (has_enough_device_memory already budgets for it).

• +7.6% (Elastic flush), +12% with cuVS default degrees; join kernel −8%. Host path unchanged.
• Bit-identical: the WMMA-float kernel on fp32 equals the WMMA-half kernel on pre-converted data in all 15 dim/metric cases.
• Cost: a transient n × dim fp16 buffer.

2. sort_knn_graph: don't duplicate a device-resident dataset (#2834)

sort_knn_graph always allocated a full fp32 copy of the dataset, even for device mdspans. It now reads device-accessible datasets in place and only stages host datasets.

• Peak extra VRAM on the device path: 15.9 GB → 1.1 GB (4M × 1024), 4.0 GB → 0.27 GB (1M × 1024).
• Time is within noise: the removed device-to-device copy takes about 20 ms at 1.8 TB/s. This is a memory fix, and it also lets larger flushes stay on NN-Descent.

3. Fuse the two WMMA join phases into one pass over the features (#2835)

Phases 1 (new×new) and 2 (new×old) each walked the dimension in 128-column tiles, gathering the new vectors twice. Now the kernel does one pass with two accumulator fragments.

• +7% (Elastic flush, 4M × 1024), +2–3% on wiki_all 768-d, no change at dim ≤ 128 (single tile). Join kernel −11% at 4M × 1024.
• Registers drop 64 → 51; occupancy is unchanged.
• Bit-identical (per-fragment accumulation order is unchanged).
• It also moves the per-k-step __syncthreads out of the k loop. Improvements on NND kernels: u8 MMA and warp guard #2751 makes the same barrier change and adds per-warp MMA guards; the two are compatible, and Fuse the two WMMA local-join phases in NN-Descent #2835 will rebase onto whichever lands first.

4. Keep the per-iteration sample lists on device (#2836)

Each iteration the join kernels read four sampled lists from pinned host memory over PCIe. The reverse lists were built on device only to be copied back to host. The old lists were copied host-to-host so the worker thread could resample. Now the lists live on device (4 × n × 32 ids), and reverse lists are built in place. There is a pinned-host fallback if they don't fit; that path does a swap instead of a copy. The dead sentinel fill_n passes are dropped.

• +21% (Elastic flush), +15–31% (merge), +24–36% (standalone NN-Descent). This is the largest single item, about 200 ms per iteration.
• Removes ~1.9 GB of pinned host allocation per build at 5M rows; costs 2.56 GB of VRAM at 5M rows when it fits.
• 870/870 NN-Descent tests pass with both the device path and the forced fallback.

5. Investigated, not proposed as PRs

• Per-k-step __syncthreads removal: correct and bit-identical, but no measurable gain alone. It is part of Fuse the two WMMA local-join phases in NN-Descent #2835 (and of Improvements on NND kernels: u8 MMA and warp guard #2751).
• __launch_bounds__(512, 2) for sm_120/121: ptxas already allocates exactly 64 registers for local_join_kernel_wmma on sm_120 (CUDA 13.0), and the SASS is byte-identical with and without the bound. Occupancy is already at the shared-memory ceiling of 2 blocks/SM. Nothing to gain today.
• Pre-lock reject in insert_to_global_graph: as usually written, with a plain load, it is unsafe. The load caches the line in L1, and the locked read-modify-write can then see stale data, desynchronising ids and distances; our kernel harness found 111 corrupted (id, dist) pairs. With an L2-only load (__ldcg) it is correct but ~1–2% slower. We do not recommend it. It may also be worth auditing the lock protocol's acquire semantics.
• sort_lists per-thread scratch / staging norms in shared memory: correct, with no measurable effect.
• IVF-PQ graph build batch prefetch: no measurable gain at 2M–6M × 1024 host input, and +0.5 GB VRAM.

Related open PRs

• Compute the expanded L2 distance in fp32 instead of fp64 #2531: fp32 expanded-L2 epilogue (fp64 removal), see above.
• Improvements on NND kernels: u8 MMA and warp guard #2751: NND kernel improvements (u8 MMA, per-warp MMA guards, the same barrier move as Fuse the two WMMA local-join phases in NN-Descent #2835).
• Fix NN-descent L2 norm shared memory opt-in #2310: dynamic shared-memory opt-in for compute_l2_norms_kernel at very high dims. Downcast device-resident fp32 input to fp16 once in NN-Descent #2833 adds a second launch site that should get the same opt-in.

Methodology

• Hardware and software: RTX PRO 6000 Blackwell Server Edition (sm_120a), CUDA 13.0, driver 580; main @ 24b0b17 + Compute the expanded L2 distance in fp32 instead of fp64 #2531.
• Elastic parameters, taken from their GPU codec: graph_degree 12, intermediate 22, NN-Descent iterations 11, cosine.
• Data: synthetic clustered low-intrinsic-dim 1024-d at 100K–6M rows, plus wiki_all_1M (768-d) and SIFT1M.
• Timing: warm-up build, then 3–8 timed builds per variant. The headline 4M numbers come from interleaved runs (variant order rotated across rounds). NVML-sampled peak VRAM; nsys kernel breakdowns.
• Correctness: full NEIGHBORS_ANN_NN_DESCENT_TEST (870) on every variant, NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST (1456) on the CAGRA-touching ones, plus the kernel-equivalence harness.

Happy to share the harness and raw results, and to split or combine PRs however reviewers prefer.