Skip to content

GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) - #783

Open
foosekis4j wants to merge 4 commits into
datastax:mainfrom
foosekis4j:gpu-graph-build
Open

foosekis4j wants to merge 4 commits into
datastax:mainfrom
foosekis4j:gpu-graph-build

Conversation

@foosekis4j

Copy link
Copy Markdown

Summary

This PR adds a way to build a JVector graph index outside GraphIndexBuilder's incremental insertion, and an
optional module that uses it to build graphs on NVIDIA GPUs. Nothing changes unless the module is built and selected:
GraphIndexBuilder, search, writers and the on-disk format behave as before.

It is two commits, which can be reviewed (or split into two PRs) separately:

  1. GraphBuildAccelerator (core, ~460 lines, no new dependencies). An @Experimental service interface
    (name(), supports(builder, vectors, similarity), build(builder, vectors, similarity)), found with
    ServiceLoader through GraphBuildAccelerators.find() and selected with -Djvector.graph.accelerator=<name>. An
    accelerator computes every node's diversity-pruned neighbour list and installs it with new @Experimental methods
    on GraphIndexBuilder:

    • addGraphNodeWithNeighbors(level, node, neighbors), safe for distinct (node, level) pairs from many threads;
    • sampleGraphLevels(n), which draws levels as addGraphNode does, so a hierarchy keeps its distribution;
    • completeExternalBuild(entryNode);
    • getMaxDegree() and getAlpha().

    cleanup() then runs as usual; it skips refineFinalGraph for an external build. Bench (Grid) builds with the
    selected accelerator when supports() says it can, and exactly as before otherwise. TestExternalGraphBuild checks
    the contract with a brute-force reference accelerator.

  2. jvector-gpu (optional module, mvn -Pgpu). It implements the interface for NVIDIA GPUs with
    TornadoVM (GPU kernels written in Java) and
    cuVS:

    • How the build works: cuVS NN-Descent candidates; diversity pruning with the builder's alpha on the GPU, using
      a tensor-core Gram matrix; upper layers from the same pruning.
    • Also on the GPU:
      • PQ training and encoding (GpuProductQuantization; with JVector's codebooks the codes are identical);
      • NVQ encoding (GpuNvq);
      • compaction by rebuilding (GpuCompaction, with ordinalMappers() as in OnDiskGraphIndexCompactor).
    • Without a GPU or TornadoVM, every entry point falls back to JVector's CPU code.

Results

RTX 4090 against JVector on an i9-13900K (32 threads), same settings (M=32, efConstruction=100, overflow 1.2,
alpha 1.2, hierarchy). Cohere Wikipedia v3 (embed-multilingual-v3, English), 1,000,000 x 1024, COSINE, cold builds,
measured on this branch with TornadoVM 8.0.0:

graph build recall@10 at beam 20 / 40 / 80 / 160
JVector (CPU) 181.7 s 0.8117 / 0.8934 / 0.9459 / 0.9741
jvector-gpu 17.4 s (median of 4, 16.0-19.0 s) 0.8060-0.8073 / 0.8943-0.8960 / 0.9476-0.9485 / 0.9761-0.9766

That's about 10x. Recall is higher from beam 40 up and about 0.005 lower at beam 20.

jvector-gpu/README.md has more, measured in October on JVector 4.0.1 with TornadoVM 7.x:

  • other datasets: ada002-1M 19x, a 10M x 1024 build in 207 s against 2271 s;
  • compaction: 15x;
  • PQ: 3.6x;
  • NVQ: 3.0x;
  • what the speed-up depends on.

Trying it

export JAVA_HOME=<JDK 22+>
sdk install tornadovm 8.0.0-jdk22plus-cuda       # TornadoVM 8.0.0 with the CUDA backend
export TORNADOVM_HOME=$HOME/.sdkman/candidates/tornadovm/8.0.0-jdk22plus-cuda PATH=$TORNADOVM_HOME/bin:$PATH
mvn -Pgpu -pl jvector-gpu,jvector-examples -am package -DskipTests
jvector-gpu/run.sh check                         # every GPU path against JVector's CPU code; prints OK
jvector-gpu/run.sh bench ada002-100k             # JVector's Bench with the graph built on the GPU
  • Dependencies: tornado-api and tornado-cuvs 8.0.0 come from Maven Central. They're provided, because the
    TornadoVM SDK supplies them at run time.
  • cuVS: installed from NVIDIA's Python wheels by jvector-gpu/cuvs-env.sh (about 2.1 GB, once).
  • Hardware: CUDA only, compute capability 8.0 or newer.

Testing

  • TestExternalGraphBuild, plus the existing GraphIndexBuilderTest, TestGraphIndexBuilderConfig and
    TestOnDiskGraphIndexCompactor: 45 tests, all pass.
  • jvector-gpu host tests (no GPU needed): 12 tests, covering:
    • FP16 packing and normalisation;
    • scores against JVector's;
    • candidate selection;
    • batching and the device-memory budget;
    • a PQ codebook round-trip through JVector's format;
    • the fallbacks to the CPU code (build, PQ, compaction).
  • jvector-gpu/run.sh check on an RTX 4090 prints OK. It covers COSINE, DOT_PRODUCT and EUCLIDEAN builds, PQ codes
    identical with the CPU's codebooks, and a four-segment compaction with 10% deletes.
  • mvn apache-rat:check passes with -Pgpu.

Notes and questions

  • Release note: I'll add docs/release notes/4.1.0/<this PR>.feature.md once this PR has a number.
  • Where the module lives: would you rather keep it in this repository behind a profile, as here, or publish it
    separately and keep only the interface here?
  • OpenSearch: using the accelerator from OpenSearch's JVector plugin would be a small follow-up there, in its
    getGraph(), once a JVector release has the interface.
  • GCC 12 fix: a separate small PR fixes build_native_lib.sh, which on x86_64 needs GCC 12 for Highway's
    AVX3_SPR target and should check $CXX.

…builds

GraphIndexBuilder builds a graph by inserting nodes one at a time. GraphBuildAccelerator lets an
implementation build the same graph another way, for example on a GPU, and hand JVector the
result. Nothing changes unless one is installed and selected: the search path and the on-disk
format are untouched.

- GraphBuildAccelerator (@experimental): name(), supports(builder, vectors, similarity) and
  build(builder, vectors, similarity). GraphBuildAccelerators.find() loads the implementations with
  ServiceLoader and returns the one -Djvector.graph.accelerator=<name> selects, if any.
- GraphIndexBuilder, for builds from outside:
  - addGraphNodeWithNeighbors(level, node, neighbors) installs a node's diversity-pruned list in
    one layer; safe for distinct (node, level) pairs from many threads;
  - sampleGraphLevels(n) draws node levels as addGraphNode does, so a hierarchy keeps its
    distribution;
  - completeExternalBuild(entryNode) sets the entry point of a flat graph;
  - getMaxDegree() and getAlpha() expose the build parameters;
  - cleanup() skips refineFinalGraph for an external build, which already chose its lists.
- Bench (Grid) builds with the selected accelerator when supports() says it can, alongside the
  inline feature writes, and otherwise exactly as before.
- TestExternalGraphBuild: a reference accelerator (brute-force neighbors, JVector's own diversity
  pruning) builds flat and hierarchical graphs that search with recall@10 above 0.9, keep within
  the max degree and put the entry point on the top layer; find() returns nothing unless an
  accelerator is installed and selected.
…ornadoVM + cuVS)

An optional module, built only with -Pgpu, that implements GraphBuildAccelerator for NVIDIA GPUs
through TornadoVM (Java GPU kernels) and NVIDIA cuVS. Without -Pgpu nothing changes; with it but
without a GPU, every entry point falls back to JVector's CPU code.

- TornadoGraphBuildAccelerator ("tornadovm"): cuVS NN-Descent candidates, diversity pruning with
  the builder's alpha on the GPU (tensor-core Gram matrix), upper layers from the same pruning, then
  addGraphNodeWithNeighbors and completeExternalBuild. COSINE, DOT_PRODUCT and EUCLIDEAN, any
  dimension (padded to 16), M <= 64.
- GpuProductQuantization: k-means++ training and encoding on the GPU; with JVector's codebooks the
  codes are identical to the CPU's.
- GpuCompaction: merges on-disk segments by rebuilding the graph on the GPU; ordinalMappers() as
  OnDiskGraphIndexCompactor.
- GpuNvq: NVQ encoding on the GPU (8-bit).
- GpuBuildCheck (run.sh check): every feature against JVector on synthetic data; prints OK.
- Benchmarks (run.sh bench|cpu|pq|compact|nvq|wikipedia|groundtruth|phases), cuvs-env.sh (cuVS
  from NVIDIA's wheels), preflight.sh, and a README with results.
- Host tests (no GPU): candidate width, packing, fallbacks to the CPU code.

Wikipedia 1M x 1024 (Cohere embed-multilingual-v3), RTX 4090, cold build: GPU 17.4 s (median of
4, 16.0-19.0 s) vs JVector 181.7 s on the CPU; recall@10 higher at search beams 40-160, 0.005
lower at beam 20.
@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown
Contributor

Before you submit for review:

  • Does your PR follow guidelines from CONTRIBUTIONS.md?
  • Did you summarize what this PR does clearly and concisely?
  • Did you include performance data for changes which may be performance impacting?
  • Did you include useful docs for any user-facing changes or features?
  • Did you include useful javadocs for developer oriented changes, explaining new concepts or key changes?
  • Did you rebase your branch onto the latest main for regression testing and PR submission?
  • Did you trigger regression testing via Run Bench Main and review results?
  • Did you adhere to the code formatting guidelines (TBD)
  • Did you group your changes for easy review, providing meaningful descriptions for each commit?
  • Did you ensure that all files contain the correct copyright header?
  • Did you add documentation for this feature to the release notes directory?

If you did not complete any of these, then please explain below.

@foosekis4j

Copy link
Copy Markdown
Author

Added the release note as docs/release notes/4.1.0/783.feature.md (fb6267f).

The gpu module is only listed under the gpu profile, so the RAT check of
the default build does not skip it as a sub-project and reads every file
in it; the two files without a header failed the check.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant