Vamana build optimization set and fp16 support by bkarsin · Pull Request #2264 · NVIDIA/cuvs

bkarsin · 2026-06-25T00:11:03Z

Series of GPU Vamana build performance optimizations that addresses #2178 and #1757. Initial estimates from #1757 were not accurate, so many other optimizations were tried (some abandoned, some successful). This PR includes:

GreedySearch optimizations:

Multi-warp blocks - increases occupancy of GreedySearch
Reduce shared memory used per warp
fp16 query approximation (along with fp16 support)

RobustPrune optimizations:

Multi-warp block support (block size depends on graph degree)
Caching accepted candidate vectors during occlusion loop
Remove redundant merge of candidate lists and avoid syncs

General optimizations:

Replace prefix_sums kernel with cub variant
Re-use distances from GreedySearch in RobustPrune kernel
fp16 support added

Together these optimizations give significant speedups across all configs with minimal recall variance compared to the current baseline. I benchmarked performance across a range of synthetic datasets and two real-world datasets. (NOTE: finishing benchmarks and will update tables below once they are all collected).

Synthetic dataset build tests (all 1M vector datasets)

		RTX pro 6000	RTX pro 6000	H100	H100	L4	L4
TYPE	Dim/Degree	BASELINE	OPT	BASELINE	OPT	BASELINE	OPT
fp32	64/32	1.11	1.07	1.35	1.19	9.83	3.57
fp32	768 / 32	5.98	4.03	6.01	4.92	36.6	25.8
fp32	960 / 32	7.82	5.61	8.23	6.52	50.1	38.7
fp32	64/64	4.4	4.05	6.14	5.4	16.4	14.8
fp32	768/64	28.8	18.2	33.7	25.3	194	125
fp32	960/64	39.4	24.9	50.7	36.5	251	185
fp16	64/32		1.05				3.8
fp16	768 / 32		2.5				11.1
fp16	960 / 32		3.13				14.1
fp16	64/64		4.01				13.6
fp16	768/64		10.4				53.8
fp16	960/64		14.2				65.6
int8	64/32	1.06	0.966	1.69	1.74	3.5	3.35
int8	768 / 32	2.98	2.38	4.82	3.69	11.2	9.21
int8	960 / 32	3.85	3	6.39	4.32	15.1	12
int8	64/64	4.15	3.61	6.36	6.05	14.4	13.3
int8	768/64	12.9	9.74	22.4	14.9	51.7	39.7
int8	960/64	17.1	12.5	28.8	19.9	72.1	52.1

Also tested real-world BIGANN 10M (uint8 128D) and GIST (fp32 960D) datasets:

		RTX pro 6000	RTX pro 6000	H100	H100	L4	L4
Dataset	deg / iters	BASELINE	OPT	BASELINE	OPT	BASELINE	OPT
BIGANN	32 / 1.0	12.017	10.736	13.1152	12.58	40.3484	39.887
BIGANN	32 / 2.0	25.942	23.229	32.0109	30.4	90.6138	89.994
BIGANN	64 / 1.0	36.642	30.688	48.1399	43.65	126.248	120.394
BIGANN	64 / 2.0	86.221	72.026	118.444	107.91	299.2	285.543
GIST (fp32)	32 / 1.0	6.39807	4.9804	6.773	5.51	41.48	32.3789
GIST (fp32)	32 / 2.0	16.2624	12.1026	19.0785	13.7675	103.197	76.84
GIST (fp32)	64 / 1.0	22.9393	16.74	27.2307	20.1767	150.2	109.98
GIST (fp32)	64 / 2.0	60.1051	41.452	71.1611	53.6482	380.185	262.1
GIST (fp16)	32 / 1.0		2.597		3.80264
GIST (fp16)	32 / 2.0		6.176		8.692
GIST (fp16)	64 / 1.0		9.965		13.793
GIST (fp16)	64 / 2.0		23.95		34.1922

(cherry picked from commit 14e36b3)

(cherry picked from commit d8b547b)

(cherry picked from commit 4d970e0)

…L2 comparators. (cherry picked from commit 0746d15)

…e vector in shared memory in the RobustPrune occlusion loop (cherry picked from commit f45fd1b49283434eb4a3017da069ead501e938c3)

…efix sum with cub scan and hoist per-batch reverse-edge allocations (cherry picked from commit 3b8650f5ebea52421f187332a2f6f3bdd599c42e)

…lusion across multiple warps per query (raise occupancy) (cherry picked from commit 2e02f938f97e65ca073daf07397d538432b52867)

… query->existing-edge distances in the RobustPrune merge (avoid recompute) (cherry picked from commit 041d355c585b7f98302d054b80ac287d637a07e4)

…oords in FP16 smem for dim>=512 to raise GreedySearch occupancy (salvage of N8)

…nce (one warp) instead of redundantly on all 128 threads, then broadcast

…ck (4 vs 8) to raise occupancy/MLP on the degree-64 occlusion sweep

copy-pr-bot · 2026-06-25T00:11:06Z

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.

cjnolet · 2026-07-08T14:36:56Z

 #define KERNEL_TIMING (RAFT_LOG_ACTIVE_LEVEL <= RAPIDS_LOGGER_LOG_LEVEL_DEBUG)

+template <typename accT, typename IdxT>
+__global__ void gather_query_sizes(QueryCandidates<IdxT, accT>* query_list,


Do you have a sense for how much this is adding to the binary size? @divyegala should be able to explain how to see the deployed metrics for PRs.

jamxia155 · 2026-07-17T20:01:15Z

/ok to test c055c60

launch_bounds

tarang-jain

I have not looked into the implementation itself (what the code inside the kernels does) and my review continues to be along the same direction as my previous comments -- to improve overall code quality and readability. Pointing out a couple of very small things.

tarang-jain · 2026-08-08T00:00:48Z

+
+  const int laneId    = threadIdx.x & 31;
+  const int warpId    = threadIdx.x >> 5;
+  const int num_warps = blockDim.x >> 5;


Lets keep this simple. I have been advised by @tfeher in the past to not use bit arithmetic for simple division / mod operations, since the compiler is typically smart enough to automatically simplify these. So we can do again things like laneId = threadIdx.x % raft::WarpSize and warpId = threadIdx.x / 32.

I think this can also be ignored @bkarsin . Its too much nit-picking.

tarang-jain · 2026-08-08T00:05:53Z

+  auto scan_temp_storage = raft::make_device_mdarray<uint8_t>(
+    res, large_ws, raft::make_extents<int64_t>(std::max(scan_temp_bytes, size_t{1})));
+
+  thrust::device_vector<IdxT> edge_dest_vec(max_total_edges);


Use raft::device_vector here

tarang-jain · 2026-08-19T23:45:31Z

/ok to test e1e0b7e

jamxia155 · 2026-08-20T22:24:36Z

/ok to test f0579f9

lowener · 2026-08-24T09:10:59Z

+__device__ SUMTYPE l2_ILP2_half(Point<__half, SUMTYPE>* src_vec, Point<__half, SUMTYPE>* dst_vec)
+{
+  __half temp_dst[2]    = {__float2half(0.0f), __float2half(0.0f)};
+  __half partial_sum[2] = {__float2half(0.0f), __float2half(0.0f)};
+  for (int i = threadIdx.x; i < src_vec->Dim; i += 2 * blockDim.x) {
+    temp_dst[0] = dst_vec->coords[i];
+    if (i + raft::WarpSize < src_vec->Dim) temp_dst[1] = dst_vec->coords[i + raft::WarpSize];
+
+    l2_half_fma_sq(partial_sum[0], src_vec[0].coords[i], temp_dst[0]);
+    if (i + raft::WarpSize < src_vec->Dim)
+      l2_half_fma_sq(partial_sum[1], src_vec[0].coords[i + raft::WarpSize], temp_dst[1]);
+  }
+  partial_sum[0] = __hadd(partial_sum[0], partial_sum[1]);
+
+  return l2_half_warp_reduce<SUMTYPE>(partial_sum[0]);
+}


Would it make sense to also provide an overload of l2_half_fma_sq with half2 type? I see that __hfma2 and __hsub2 exist.

It might benefit l2_ILP4_half below too

Good point, though this is not a super simple change. Would have to choose what versions to replace with this (only ILP variants?) and deal with edge cases like odd dimensions. Would also require some re-benchmarking. To avoid making this PR even bigger, would adding at issue for this be better?

Okay then can you create that issue, and add a comment in the code linking to that issue for better tracking please?

Created issue #2510 and added comment in code that references it

lowener · 2026-08-24T09:48:07Z

+__device__ __forceinline__ void l2_fma_sq2(SUMTYPE& acc, float sx, float sy, float2 dst2)
+{
+  float dx = sx - dst2.x;
+  float dy = sy - dst2.y;
+  acc      = fmaf(dx, dx, acc);
+  acc      = fmaf(dy, dy, acc);
+}


This might also benefit from fsub2/fmaf2. The data loaded is already float2

Created issue #2510 and added comment

lowener

LGTM

jamxia155 · 2026-08-28T18:46:56Z

/ok to test 882b781

divyegala · 2026-08-31T22:27:35Z

/ok to test f70b984

lowener · 2026-09-02T22:20:43Z

/ok to test 0c1a48e

lowener · 2026-09-03T15:59:04Z

/merge

bkarsin added 15 commits June 15, 2026 14:11

Initial multi-warp greedy search optimization

a4d82e1

(cherry picked from commit 14e36b3)

Add fp16 support for Vamana build/serialize

05cc8cc

(cherry picked from commit d8b547b)

Fix recall issue with multi-warp optimization

5bb5e6c

(cherry picked from commit 4d970e0)

Speed up fp16 vamana build with promoted query vectors and optimized …

3dca1b0

…L2 comparators. (cherry picked from commit 0746d15)

perf(vamana): P1_cache_candidate_vector — Cache the accepted candidat…

02b5638

…e vector in shared memory in the RobustPrune occlusion loop (cherry picked from commit f45fd1b49283434eb4a3017da069ead501e938c3)

perf(vamana): O1_parallel_scan_buffer_pool — Replace single-thread pr…

9ec9b99

…efix sum with cub scan and hoist per-batch reverse-edge allocations (cherry picked from commit 3b8650f5ebea52421f187332a2f6f3bdd599c42e)

perf(vamana): N3_multiwarp_robust_prune — Parallelize RobustPrune occ…

696d3d2

…lusion across multiple warps per query (raise occupancy) (cherry picked from commit 2e02f938f97e65ca073daf07397d538432b52867)

perf(vamana): N7_reuse_search_distances_in_prune — Reuse GreedySearch…

24f2ed7

… query->existing-edge distances in the RobustPrune merge (avoid recompute) (cherry picked from commit 041d355c585b7f98302d054b80ac287d637a07e4)

perf(vamana): M3_fp16_query_smem_occupancy — Store the cached query c…

bbe3d50

…oords in FP16 smem for dim>=512 to raise GreedySearch occupancy (salvage of N8)

perf(vamana): M5_prune_single_warp_merge — Do the RobustPrune merge o…

eeca286

…nce (one warp) instead of redundantly on all 128 threads, then broadcast

perf(vamana): M6_prune_warps_tuning — Sweep RobustPrune warps-per-blo…

4348e3e

…ck (4 vs 8) to raise occupancy/MLP on the degree-64 occlusion sweep

Clean up and remove dead code

4a0c179

pre-commit fixes

cb91c99

Fix bug with odd dimension data for some cases

2dc5fa9

Merge branch 'NVIDIA:main' into vamana-build-opt

a9b8a8d

github-project-automation Bot added this to Unstructured Data Processing Jun 25, 2026

josephine-wolf-oberholtzer added this to Unstructured Data Processing Jul 1, 2026

cjnolet added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels Jul 1, 2026

cjnolet moved this to In Progress in Unstructured Data Processing Jul 1, 2026

cjnolet reviewed Jul 8, 2026

View reviewed changes

Merge branch 'main' into vamana-build-opt

5da19e0

cjnolet requested review from a team as code owners July 15, 2026 15:04

cjnolet assigned bkarsin Jul 15, 2026

Merge branch 'main' into vamana-build-opt

c055c60

bkarsin added 3 commits July 24, 2026 18:04

Fix bug with raft::copy

f269c3b

replace kernel with map_offset and add raft directives

03a29e4

Merge branch 'main' into vamana-build-opt

3cc19c4

tarang-jain requested changes Jul 28, 2026

View reviewed changes

Comment thread cpp/src/neighbors/detail/vamana/greedy_search.cuh

bkarsin added 3 commits July 29, 2026 16:44

Replace CUDA_ARCH with raft::resource::get_device_properties to choose

38d7f1a

launch_bounds

Fix bug with RAFT directives

6cac7b1

Merge branch 'main' into vamana-build-opt

1db9190

tarang-jain approved these changes Aug 8, 2026

View reviewed changes

bkarsin added 2 commits August 19, 2026 13:01

Merge branch 'main' into vamana-build-opt

ecc61d1

Replace thrust vector with raft

e1e0b7e

Fix build error for some architectures

f0579f9

lowener reviewed Aug 24, 2026

View reviewed changes

bkarsin added 2 commits August 25, 2026 12:40

Address review comments

c4d4d6c

Add comments and update C doc for vamana

db5c3fe

lowener approved these changes Aug 27, 2026

View reviewed changes

lowener and others added 2 commits August 27, 2026 11:19

Merge branch 'main' into vamana-build-opt

9f5baca

Merge branch 'main' into vamana-build-opt

882b781

Increase binary size limitation slightly to accomodate PR

46dd9bc

bkarsin requested a review from a team as a code owner August 31, 2026 21:01

bkarsin requested a review from KyleFromNVIDIA August 31, 2026 21:01

Merge branch 'main' into vamana-build-opt

f70b984

Merge branch 'main' into vamana-build-opt

0c1a48e

bdice approved these changes Sep 2, 2026

View reviewed changes

rapids-bot Bot merged commit 6c31e37 into NVIDIA:main Sep 3, 2026
169 of 171 checks passed

Uh oh!

Conversation

bkarsin commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jun 25, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jamxia155 commented Jul 17, 2026

Uh oh!

Uh oh!

tarang-jain left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tarang-jain commented Aug 19, 2026

Uh oh!

jamxia155 commented Aug 20, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lowener left a comment

Choose a reason for hiding this comment

Uh oh!

jamxia155 commented Aug 28, 2026

Uh oh!

divyegala commented Aug 31, 2026

Uh oh!

lowener commented Sep 2, 2026

Uh oh!

lowener commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

bkarsin commented Jun 25, 2026 •

edited

Loading