Skip to content

[FEA] Optimize Vamana index build kernels #1757

Description

@bkarsin

There are a number of perf optimizations that can improve the GPU Vamana index build in cuVS. This issue tracks these optimizations that may be applied with several PRs to limit scope and maintain stability. Perf estimates for each optimization are rough and attempt to be conservative.
Note: speedups are speculative and should not be relied on until optimizations are tested

High Impact / Priority Optimizations

  • Increase blockDim of GreedySearch and RobustPrune kernels from 32 to 128-256. Block size limits theoretical occupancy to 50%, estimated 1.2x-1.5x kernel perf improvement - got 0% to 38% kernel speedup
  • Parallelize all distance computations in enqueue_all_neighbors for improved parallelism. Could significantly reduce sync overheads on method that GreedySearch spends 30-40% of runtime on. Expected 1.4-1.6x GreedySearch improvement - Optimization unsuccessful, no speedup achieved
  • Reduce syncthreads required during merge loop of RobustPrune. Can reduce syncs by 10x during merge loop. RobustPrune spends 20% of time here, so estimate 1.3-1.4x improvement of RobustPrune. - Achieved 0% to 28% kernel spedup
  • Replace serial prefix sum execution with CUB parallel prefix sums. Minimal time is spent here, so roughly 5% E2E build time improvement. - Achieved 2-5% E2E speedup
  • Optimize the check_visited data structure of RobustPrune (hash set or bitset). Currently a linear search, estimate a 1.2x speedup to GreedySearch.

Overall perf gain estimate

  • 2-2.5x speedup of GreedySearch
  • 1.5-2.2x speedup of RobustPrune
    Depending on the dataset, the time breakdown can vary between these kernels. On narrow vectors, GreedySearch dominates, while RobustPrune is more costly for wide vectors.

Expected E2E perf gains: 1.6x-2.4x

E2E perf achieved thus far: 1x-1.4x

Additional more challenging / speculative optimizations

  • Try to improve cache reuse of candidate-candidate distances computed during RobustPrune. Shared memory is a major occupancy limiter for these kernels, so this may not be feasible.
  • Reduce synchronizations used by PriorityQueue structure by replacing them with warp primitives.

Metadata

Metadata

Assignees

Labels

Fields

No fields configured for Enhancement.

Projects

Status
Todo
Status
Todo

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions