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

jamxia155 · 2026-07-22T22:10:09Z

/ok to test 350ed89

copy-pr-bot · 2026-07-22T22:10:13Z

/ok to test 350ed89

@jamxia155, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

jamxia155 · 2026-07-22T22:11:11Z

/ok to test 74caba9

jamxia155 · 2026-07-22T22:25:14Z

/ok to test b8238f3

jinsolp · 2026-07-23T00:41:15Z


-      num_neighbors = degree;
-      __syncthreads();
+      num_neighbors[warpIdx] = degree;


maybe make only lane 0 write to this?

Having just lane 0 write it would require a sync after. Seems simpler/cleaner to just have all threads write, since they have the same value.

jinsolp · 2026-07-23T02:49:45Z

+    __float2half(0.0f), __float2half(0.0f), __float2half(0.0f), __float2half(0.0f)};
+  for (int i = threadIdx.x; i < src_vec->Dim; i += 4 * blockDim.x) {
+    temp_dst[0] = dst_vec->coords[i];
+    if (i + 32 < src_vec->Dim) temp_dst[1] = dst_vec->coords[i + 32];


Should we use blockDim.x instead of hardcoding to 32 here?

I agree with this. We should not be hardcoding things because the kernel invocation can change (or the same kernel may be called somewhere else in the future).

This is run by a single warp within the block, and uses a warp width stride. I can change it to a define macro if that helps?

Replaced with raft::WarpSize

tarang-jain

There are a lot of changes in this PR where I want to basically tell the exact same things -- using raft macros and raft primtives wherever possible. We should certainly avoid writing small kernels which can easily be a lambda for raft::linalg. It helps keeps things readable and binary size increases can become slightly more predictable. For example, earlier we have found that raft::linalg::map kernels generally add less to the binary size than thrust::for_each (or similar thrust primitives).

tarang-jain · 2026-07-23T23:29:41Z

-                                                   Node<SUMTYPE> input_data,
-                                                   SUMTYPE* cur_max_val,
-                                                   int* max_idx)
+__inline__ __device__ void parallel_pq_max_enqueue_warp(Node<SUMTYPE>* pq,


Nit: we have macros to mark kernels: such as _RAFT_INLINE

I wouldnt say these macros should block the PR from merging.

I can change this throughout, but I can't find the docs in raft on the different INLINE macros.

tarang-jain · 2026-07-23T23:33:54Z

+__global__ void scatter_prefix_offsets(QueryCandidates<IdxT, accT>* query_list,
+                                       const int* edge_offsets,
+                                       int count)
+{


same here: we should be able to achieve these with raft primitives.

This scatter kernel is not so simple. I could do it but would need to re-work the code a bit more to use a strided mdspan or something. Can work on this if it's a blocker for the PR.

We can do this as a follow-up. @bkarsin do you mind creating an issue and referencing it in a comment here just to get this merged?

Created issue #2493 detailing that this kernel should be replaced with RAFT primitives.

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-24T08:52:20Z

    The following distance metrics are supported:
-        - L2Expanded
+        - L2Expanded (sqeuclidean)
+        - L2SqrtExpanded (l2 / euclidean distance)


Also update the C doc for cuvsVamanaBuild for the L2SqrtExpanded distance

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?

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

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 22, 2026 13:44

Fix launch bounds for older arch builds

e008a09

Add fp16 tests

350ed89

Merge branch 'main' into vamana-build-opt

74caba9

Add missing test file

b8238f3

jinsolp reviewed Jul 23, 2026

View reviewed changes

tarang-jain reviewed Jul 23, 2026

View reviewed changes

tarang-jain requested changes Jul 23, 2026

View reviewed changes

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

bkarsin added 7 commits July 24, 2026 13:59

Change distribution of fp16 vamana unit tests

70164ef

Replace hard-coded warp sizes with raft::warp and remove unused code

840860e

Switch to raft::copy and pre-commit check

71bf588

Merge branch 'main' into vamana-build-opt

9cbd829

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

Address review comments

c4d4d6c

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!

jamxia155 commented Jul 22, 2026

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

jamxia155 commented Jul 22, 2026

Uh oh!

jamxia155 commented Jul 22, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

tarang-jain left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tarang-jain Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

Uh oh!

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

bkarsin commented Jun 25, 2026 •

edited

Loading

tarang-jain Jul 23, 2026 •

edited

Loading