Repository navigation
GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) - #783
foosekis4j wants to merge 10 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.
|
Before you submit for review:
If you did not complete any of these, then please explain below. |
|
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.
|
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? |
…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.
|
Pushed f1ba7ac, a performance fix in The FP16 packing loop is JIT-compiled first for the memory-mapped dataset. When the upper layers then pass rows in native memory, a class check in the compiled loop fails. C2 then went through rounds of recompiling and deoptimising, and the loop ended up in C1 profiled code. Packing the ~31k upper-layer rows took 7-9 s at 1024-1536 dimensions, about half of a 1M build. Each row is now copied to a Cold builds on the same RTX 4090 / i9-13900K, same settings as in the description (M=32, efConstruction 100, alpha 1.2, hierarchy):
Recall is unchanged. On Cohere 1M, recall@10 at beam 40 / 80 is 0.898 / 0.952 for the GPU graph and 0.893 / 0.949 for JVector's. |
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. |
… 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.
|
Following up on this: I moved The same question applied to the build, so 7f8f0ce adds a host-memory check to |
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: 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. Without a GPU or
TornadoVM, or when the build would not fit in device or host memory,
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):
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.
addGraphNodeWithNeighborstakes 3.7 s on the 128-core host against 0.3 s on the i9 (Cohere 1M). That step runson the CPU.
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-gpuhost tests (no GPU needed): 12 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.build_native_lib.sh, which on x86_64 needs GCC 12 for Highway'sAVX3_SPR target and should check
$CXX.