Repository navigation
GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) - #783
Open
foosekis4j wants to merge 4 commits into
Open
foosekis4j wants to merge 4 commits into
foosekis4j wants to merge 4 commits into
Conversation
…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.
foosekis4j
requested review from
MarkWolters,
ashkrisk,
jshook and
tlwillke
as code owners
November 3, 2026 17:25
Contributor
|
Before you submit for review:
If you did not complete any of these, then please explain below. |
Author
|
Added the release note as |
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds a way to build a JVector graph index outside
GraphIndexBuilder's incremental insertion, and anoptional 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:
GraphBuildAccelerator(core, ~460 lines, no new dependencies). An@Experimentalservice interface(
name(),supports(builder, vectors, similarity),build(builder, vectors, similarity)), found withServiceLoaderthroughGraphBuildAccelerators.find()and selected with-Djvector.graph.accelerator=<name>. Anaccelerator computes every node's diversity-pruned neighbour list and installs it with new
@Experimentalmethodson
GraphIndexBuilder:addGraphNodeWithNeighbors(level, node, neighbors), safe for distinct (node, level) pairs from many threads;sampleGraphLevels(n), which draws levels asaddGraphNodedoes, so a hierarchy keeps its distribution;completeExternalBuild(entryNode);getMaxDegree()andgetAlpha().cleanup()then runs as usual; it skipsrefineFinalGraphfor an external build.Bench(Grid) builds with theselected accelerator when
supports()says it can, and exactly as before otherwise.TestExternalGraphBuildchecksthe contract with a brute-force reference accelerator.
jvector-gpu(optional module,mvn -Pgpu). It implements the interface for NVIDIA GPUs withTornadoVM (GPU kernels written in Java) and
cuVS:
a tensor-core Gram matrix; upper layers from the same pruning.
GpuProductQuantization; with JVector's codebooks the codes are identical);GpuNvq);GpuCompaction, withordinalMappers()as inOnDiskGraphIndexCompactor).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:
That's about 10x. Recall is higher from beam 40 up and about 0.005 lower at beam 20.
jvector-gpu/README.mdhas more, measured in October on JVector 4.0.1 with TornadoVM 7.x:Trying it
tornado-apiandtornado-cuvs8.0.0 come from Maven Central. They'reprovided, because theTornadoVM SDK supplies them at run time.
jvector-gpu/cuvs-env.sh(about 2.1 GB, once).Testing
TestExternalGraphBuild, plus the existingGraphIndexBuilderTest,TestGraphIndexBuilderConfigandTestOnDiskGraphIndexCompactor: 45 tests, all pass.jvector-gpuhost tests (no GPU needed): 12 tests, covering:jvector-gpu/run.sh checkon an RTX 4090 prints OK. It covers COSINE, DOT_PRODUCT and EUCLIDEAN builds, PQ codesidentical with the CPU's codebooks, and a four-segment compaction with 10% deletes.
mvn apache-rat:checkpasses with-Pgpu.Notes and questions
docs/release notes/4.1.0/<this PR>.feature.mdonce this PR has a number.separately and keep only the interface here?
getGraph(), once a JVector release has the interface.build_native_lib.sh, which on x86_64 needs GCC 12 for Highway'sAVX3_SPR target and should check
$CXX.