Repository navigation
GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) - #783
Open
foosekis4j wants to merge 15 commits into
Open
foosekis4j wants to merge 15 commits into
foosekis4j wants to merge 15 commits into
Conversation
foosekis4j
requested review from
MarkWolters,
ashkrisk,
jshook and
tlwillke
as code owners
November 3, 2026 17:25
foosekis4j
added a commit
to foosekis4j/jvector
that referenced
this pull request
Nov 3, 2026
Contributor
|
Before you submit for review:
If you did not complete any of these, then please explain below. |
Contributor
|
The compaction part is in-core only, right? As I read it, GpuCompaction loads every live vector into memory, rebuilds the graph from scratch on the GPU, and writes the result out. Our compaction is designed for the opposite case: large datasets or small machines, where the merge doesn’t fit in memory and has to run out of core. In that case, do we fall back to the CPU compactor? |
foosekis4j
requested review from
dian-lun-lin,
r-devulap and
vasimuddinmd
as code owners
November 6, 2026 12:47
Contributor
Author
Yes, GpuCompaction is in-core: it materialises every live vector on the heap, rebuilds from scratch on the GPU, and writes the result. It's meant for the case where the live set fits in host memory and in the GPU's FP16 budget (n × d × 2 bytes; 10M × 1024 fits on a 24 GB card). It does not replace your compactor's out-of-core design. To answer the fallback question precisely: it falls back to OnDiskGraphIndexCompactor only when the GPU build declines (no device, or not enough device memory). It doesn't check host memory first, so a merge larger than the heap would fail with an OOM instead of falling back. I'll fix that. The size checks will run before any vector is loaded, against a host-memory budget as well as the device budget, and anything that doesn't fit goes to OnDiskGraphIndexCompactor untouched. The rebuild will also read the live vectors from one off-heap region (or a temporary memory-mapped file) instead of heap objects, which removes one full copy. For the large-dataset case, a GPU should help your compactor rather than replace it: keep ts streaming cell join and its in-place merge, and offload the batched pieces. If it would help, I can split GpuCompaction out of this PR so the PR covers only graph construction. |
Contributor
Author
|
Following up on this: I moved The same question applied to the build, so 7f8f0ce adds a host-memory check to |
…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.
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.
…compiled PackedVectors.of first packs the dataset (often a memory-mapped file) and then the upper layers' rows (native memory). The hot loop read the rows through the MemorySegment, so C2 specialised it for the first segment class; the second call hit a class-check trap and, after repeated recompile/deoptimise rounds, the loop was left in C1 profiled code with every pool thread updating the same profile counters. Packing ~31k upper-layer rows of 1024 dimensions took 7-9 s, about half of a 1M-vector build. Each row is now copied to a per-thread float[] with MemorySegment.copy, which behaves the same for every segment class, and packed from the array. The packed output is identical. Cohere Wikipedia 1M x 1024: build 17.1 s -> 8.0 s; ada002-1M x 1536: 17.5 s -> 9.8 s (RTX 4090, TornadoVM 8.0.0), recall unchanged.
… PQ and NVQ out GpuCompaction rebuilds a merge in core: it needs the live vectors in host memory and the merged set in device memory. That suits small and medium merges but not OnDiskGraphIndexCompactor's out-of-core design, so it moves to a follow-up together with GPU PQ and NVQ encoding, where it can fall back to the out-of-core compactor before loading anything. This change keeps the GraphBuildAccelerator SPI and the GPU graph build. Removed: GpuCompaction, GpuProductQuantization, PqKernels, GpuNvq, NvqKernels, their benches (CompactionBench, PqBench, NvqBench, WritePhases) and TestGpuCompactionAndPqHost; GpuBuildCheck checks the graph build only; run.sh, README and the release note follow.
…e-edge merge GpuPrune.prune copied each batch into the TornadoVM arrays and read the keep flags back element by element on one thread, while the GPU waited. Batches are now laid out in host arrays by other threads while the GPU prunes the previous batch, moved with bulk MemorySegment copies, and their results applied in parallel behind it. A node is in one batch only, so the stages touch disjoint rows. The reverse-edge merge's two sequential passes are parallel, with each node's reverse list sorted by source so the result is the old one. Output is unchanged: built from exact candidates, the adjacency is identical to the previous version's (ada002-100k and sift-128, widths 48 and 96). RTX 4090, prune phase / whole build: Cohere Wikipedia 1M x 1024 2.56 -> 2.12 s / 8.0 -> 7.6 s; ada002-1M 3.54 -> 2.96 s / 10.3 -> 9.7 s; sift-128 2.99 -> 1.83 s / 6.7 -> 5.6 s.
…machine's architecture TornadoVM 8.1.0 (released 2026-11-06) has performance fixes; tornado-api and tornado-cuvs 8.1.0 are on Maven Central and the SDK on SDKMAN (8.1.0-jdk22plus-cuda). preflight.sh requires 8.1.0. cuvs-env.sh downloaded x86_64 wheels only; it now asks pip for the manylinux wheels of `uname -m`, so the module also sets up on aarch64 servers (checked on an A100 with a Neoverse-N1 host).
…(TornadoVM 8.1.0) Medians of 5 GPU and 3 CPU cold builds on an RTX 4090 against an i9-13900K, seven datasets from LastFM (2.3x) to ada002-1M (29.7x) and Cohere Wikipedia 10M x 1024 (186 s against 2320 s), with recall@10 of both graphs at search beams 40 and 160. Replaces the October tables measured on TornadoVM 7.x/8.0.0.
supports() checked only the device budget. A build that fits the GPU but not the host (a large dataset on a small machine) failed partway with an OutOfMemoryError or a failed native allocation instead of falling back to the CPU build. HostMemory estimates what the build holds at its peak from n, dim and the max degree alone: on the heap, the candidate and pruned lists and the builder's lists (~1.9 KB per node at M=32); in native memory, the FP16 rows, cuVS's output, the upper layers' vectors and, unless the vectors are a MappedVectors file, their float32 copy. supports() declines, and logs why, when the heap part exceeds the free heap or the native part plus the heap still to grow exceeds MemAvailable. -Djvector.gpu.hostMemory replaces both with one budget. The estimate errs high (it counts the FP16 rows and cuVS's output together, which a batched build never holds at once). Cohere Wikipedia 1M and 10M x 1024 still build on the GPU on a 62 GB host with run.sh's 48 GB heap.
The host-memory check compared the build's native memory with MemAvailable, which is the machine's: in a container or a memory-limited service, a build that fits the machine but not the cgroup passed the check and then hit the limit. availableNative() now also takes the least headroom of the process's cgroup and each ancestor (cgroup v2 memory.max / memory.current / memory.stat, or v1 memory.limit_in_bytes / memory.usage_in_bytes / memory.stat): each limit less the usage that is not page cache, which the kernel reclaims under the limit as MemAvailable counts it reclaimable. Found while running GPU compactions under systemd-run MemoryMax=8G: HostMemory saw the 62 GB host's free memory, not the 8 GB limit.
The install phase called addGraphNodeWithNeighbors from a parallel loop. Every call writes words all threads share (DenseIntMap's read lock and size, the highest node id, the entry point, the completion clock), so more threads made it slower: on a 128-core Neoverse N1 host it was the slowest phase of the build (3.7 s of 12.0 s on Cohere 1M, 37.9 s of 164 s on Cohere 10M). The NodeArrays are now built on all cores and handed over by one thread. The install step alone, 1M nodes with 32 neighbours: i9-13900K, 32 threads: 0.24 s -> 0.11 s Neoverse N1, 128 threads: 3.2-4.4 s -> 0.13 s
foosekis4j
force-pushed
the
gpu-graph-build
branch
from
November 10, 2026 03:14
96114fc to
b60dbef
Compare
The upper layers were built from exact candidates, m x m scores for a layer of m nodes. That is cheap for the 31k level-1 nodes of a 1M dataset and not for the 313k of a 10M one: 9.7 s of the Cohere 10M build on an RTX 4090 and 12.5 s on an A100. Layers of more than 65,536 nodes (-Djvector.gpu.upperExactMax) now take their candidates from NN-Descent, as the base layer does. Cohere Wikipedia 10M x 1024, upper layers and build: A100, 128-core host: 12.5 s -> 1.2 s, 138.95 s -> 129.95 s RTX 4090, i9-13900K: 9.7 s -> 5.4-7.6 s (host memory is tight at this size) recall@10 at beams 40 / 160: 0.868 / 0.963 -> 0.869 / 0.963 (A100). -Djvector.gpu.trace=true prints each upper layer's candidate, packing and pruning time.
…g of cuVS's buffers - The base layer was pruned from a narrowed copy of the candidate lists (Candidates.TopK.first, 0.8 GB at 1M x 96 on a COSINE dataset). The graph builder now reads the first width candidates of each list in place: the prune phase on Cohere 1M takes 2.05-2.10 s instead of 2.20 s, SIFT 1.69 s instead of 1.74 s (RTX 4090), recall unchanged. - cuVS's output buffers (1.2 GB at 1M rows) come from malloc instead of Arena.allocate, which zero-filled them before cuVS overwrote every element. - -Djvector.gpu.trace=true also prints the steps of each graph build (fill, prune, reverse merge, second prune).
Each batch's lists were copied into the TornadoVM arrays, and its kept flags out into a new array, by one thread: 2.1 s of each 4.3 s prune pass at 10M nodes. The copies now run in parallel slices, and the kept flags go into a ring of four reused buffers. -Djvector.gpu.trace=true prints each large prune's parts. RTX 4090, i9-13900K, cuVS candidates: deep-10M prune passes 4.3 s -> 3.2 s, prune phase 15.0 s -> 12.7 s, build 47.8 s -> 45.4 s; cohere-1M prune 2.05-2.10 s -> 1.96-2.03 s; sift-128 1.69 s -> 1.47-1.56 s; recall unchanged. Larger batches (16k, 32k nodes) gave nothing more.
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.Benchbuilds with the selectedaccelerator when
supports()says it can (AcceleratedBuild, used byCompressedOnDiskBuildandInMemoryBuild),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: cuVS NN-Descent candidates, then diversity pruning with the builder's
alpha on the GPU, using a tensor-core Gram matrix, and upper layers from the same pruning (exact candidates, or
NN-Descent for a layer of more than 65,536 nodes). The lists are installed from one thread, because
addGraphNodeWithNeighborswrites counters every thread shares. Without a GPU orTornadoVM, or when the build would not fit in device or host memory (the host's cgroup limit included),
supports()returns false and JVector buildson the CPU as usual.
GPU compaction, PQ and NVQ encoding, which an earlier version of this PR included, have moved out to a follow-up, so
this PR is only about graph construction.
Results
Same settings on both machines (M=32, efConstruction=100, overflow 1.2, alpha 1.2, hierarchy). Each build runs cold
in a fresh JVM, so GPU times include TornadoVM's kernel compilation and cuVS's start-up. Medians of 5 GPU and 3 CPU
builds (10M: 3 and 2), measured on this branch with TornadoVM 8.1.0. Recall@10 is from JVector's
GraphSearcheron each graph, at search beams 40 / 160.
RTX 4090 against JVector on an i9-13900K (32 threads):
A100 80 GB (PCIe) against JVector on the same server's host, 128 Arm Neoverse-N1 cores (aarch64, 128 threads):
GPU times on the A100 are medians of 3 cold builds on the current head (the 10M build: one), measured with the two
host-side changes above; the CPU times are the earlier medians of 3 (10M: 2).
searcher visits 4-16% fewer nodes on them than on its own graphs (on LastFM a search of the same beam reaches
half a point less recall).
Cohere 1M in 86 s against 180 s on the i9. JVector's build stops getting faster past about 64 threads.
addGraphNodeWithNeighborscall writescounters every thread shares (the map's lock and size, the highest node id, the entry point, the completion
clock), so 128 threads took 3.7 s on Cohere 1M where one thread takes 0.27 s. That change, and NN-Descent
candidates for the 313k-node first upper layer at 10M (12.5 s to 1.2 s), took the A100 builds from 12.0 s to
8.7 s on Cohere 1M, 8.9 s to 5.4 s on SIFT and 164 s to 130 s on Cohere 10M.
MappedVectors. cuVS builds its candidates in 3 batches on the80 GB A100 and in 9 on the 24 GB 4090. The A100 build peaks at 32.6 GB of device memory, the 4090 build at 19.7 GB.
Trying it
tornado-apiandtornado-cuvs8.1.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), for x86_64 oraarch64.
Testing
TestExternalGraphBuild, plus the existingGraphIndexBuilderTest,TestGraphIndexBuilderConfigandTestOnDiskGraphIndexCompactor: 45 tests, all pass.jvector-examplestests on the refactored build classes (TestGridSpec,TestCompressedBuildPreflight,TestCompressorCaching,TestReportedParameters): 15 tests, all pass.run.sh bench ada002-100kbuilds with theaccelerator through them.
jvector-gpuhost tests (no GPU needed): 13 tests, all pass, covering:jvector-gpu/run.sh checkon an RTX 4090 prints OK. It builds COSINE, DOT_PRODUCT and EUCLIDEAN graphs on theGPU and on the CPU and compares their recall.
mvn apache-rat:checkpasses with-Pgpu.Notes and questions
docs/release notes/4.1.0/783.feature.md.separately and keep only the interface here?
getGraph(), once a JVector release has the interface.mainafter the build refactor (Refactor grid #776).build_native_lib.sh, which on x86_64 needs GCC 12 for Highway'sAVX3_SPR target and should check
$CXX.