jvector-gpu: GPU compaction, PQ and NVQ encoding (follow-up to #783) by foosekis4j · Pull Request #794 · datastax/jvector

foosekis4j · 2026-11-08T15:11:05Z

Summary

Follow-up to #783, which now covers only graph construction. This PR adds the other GPU paths that were split out of
it: compaction by rebuilding, PQ training and encoding, and NVQ encoding. All of them are in the optional
jvector-gpu module (mvn -Pgpu), and every one falls back to JVector's CPU code without a GPU or TornadoVM.

This PR is stacked on #783: its branch contains #783's commits. The changes to review are the last four
commits:

GPU compaction, PQ and NVQ encoding:
- GpuCompaction: same inputs as OnDiskGraphIndexCompactor, with ordinalMappers() as there;
- GpuProductQuantization;
- GpuNvq;
- their benches (run.sh compact, pq, nvq, phases) and host tests.
GpuCompaction decides from sizes, and reads the live vectors from a mapped file. This answers the question
about large datasets and small machines in GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) #783. The rebuild is in-core, so compact() first decides from the
sizes alone, before it reads any vector:
- the build must fit in host memory, by the GPU build's own estimate (HostMemory, added in GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) #783);
- the merged vectors must fit the device in FP16.
Anything else goes to OnDiskGraphIndexCompactor untouched, so out-of-core merges are unchanged. On the GPU path,
the live vectors are streamed to a temporary raw float32 file next to the output and read through a memory
mapping, never through the Java heap.
README numbers measured on this branch.
Release note (docs/release notes/4.1.0/794.feature.md).

Results

RTX 4090 against JVector on an i9-13900K (32 threads), ada002-1M (982,790 x 1536, COSINE), TornadoVM 8.1.0:

	JVector (CPU)	GPU	speed-up	quality
compaction, 4 segments of ~246k (M=32, hierarchy)	226.4 s	16.0 s	14.2x	recall@10 at beam 20 / 160: 0.925 / 0.994 vs 0.915 / 0.993
PQ train + encode (M=192, K=256)	15.79 s	2.59 s	6.1x	recall@10 of a PQ scan 0.8023 vs 0.7999
PQ encode with JVector's codebooks	5.32 s	1.01 s	5.3x	all 188,695,680 codes identical
NVQ encode (8 bits, 2 sub-vectors)	4.02 s	1.29 s	3.1x	recall@10 0.9974 vs 0.9974; 99.994% of code bytes identical

Compaction: the merged graph is a fresh build rather than a stitched one, so at the same beam it has slightly
higher recall, at the same search latency.
PQ: identical codes come from the same codebooks; where two centroids are equally close up to float rounding,
the GPU may pick the other one.
NVQ: the code bytes that differ are near-ties, because the GPU sums each loss in a different order.

Testing

TestGpuCompactionAndPqHost (no GPU needed), 4 tests, all pass:
- live vectors at their new ordinals, and the fallback to the CPU compactor;
- PQ falls back to JVector without the runtime;
- PQ codebooks round-trip through JVector's format;
- compaction declines from the sizes alone when the host cannot hold the build.
With GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) #783's tests (TestGpuBuildHost, TestExternalGraphBuild, GraphIndexBuilderTest,
TestGraphIndexBuilderConfig, TestOnDiskGraphIndexCompactor), 62 tests, all pass.
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 (20000/20000), and a four-segment compaction with 10% deletes.
mvn apache-rat:check passes with -Pgpu.

Notes

Release note: docs/release notes/4.1.0/794.feature.md.
What compaction is for: GpuCompaction doesn't replace the out-of-core compactor. It is a faster path for
merges a GPU can hold; larger ones go to OnDiskGraphIndexCompactor.
Merge order: once GPU graph construction: GraphBuildAccelerator SPI and an optional jvector-gpu module (TornadoVM + cuVS) #783 is merged, I'll rebase this onto main so only these commits remain.

github-actions · 2026-11-08T15:11:16Z

@experimental

…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

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).

…U graph build) Restores GpuCompaction, GpuProductQuantization, PqKernels, GpuNvq, NvqKernels, their benches and host tests, and the run.sh modes and README sections, on top of the graph-build change they were split from.

…from a mapped file GpuCompaction's rebuild is in-core, and it used to load every live vector onto the heap before asking whether the GPU could take the merge, so a merge larger than memory failed instead of falling back. compact() now decides first, from the sizes alone: - the build must fit in host memory, by the GPU build's own estimate (HostMemory: ~1.9 KB of heap per live node at M=32, plus native memory for the FP16 rows and cuVS's output; -Djvector.gpu.hostMemory sets a fixed budget); - the merged vectors must fit the device in FP16 and the degree must be one the GPU build takes (TornadoGraphBuildAccelerator.supports(n, dim, maxDegree, mapped)). Anything else goes to OnDiskGraphIndexCompactor before a vector is read. The GPU path streams the live vectors, source by source in parallel, to a temporary raw float32 file next to the output and builds from a memory mapping of it, so the vectors are never held on the heap or copied a second time; the file is deleted afterwards. README says which merges the rebuild is for.

… 8.1.0) Remeasured on an RTX 4090 / i9-13900K with run.sh compact, pq and nvq on ada002-1M: compaction 16.0 s against 226.4 s (14.2x), PQ 2.59 s against 15.79 s (6.1x, codes with JVector's codebooks all identical), NVQ 1.29 s against 4.02 s (3.1x). Replaces the October tables.

foosekis4j requested review from MarkWolters, ashkrisk, dian-lun-lin, jshook, r-devulap, tlwillke and vasimuddinmd as code owners November 8, 2026 15:11

foosekis4j added a commit to foosekis4j/jvector that referenced this pull request Nov 8, 2026

Release note for datastax#794: GPU compaction, PQ and NVQ encoding

cffe0fb

foosekis4j mentioned this pull request Nov 9, 2026

OnDiskGraphIndexCompactor: CompactionAccelerator hook, with the wide code, probes and cell walk on the GPU (follow-up to #794) #795

Open

foosekis4j added a commit to foosekis4j/jvector that referenced this pull request Nov 9, 2026

Release note for datastax#794: GPU compaction, PQ and NVQ encoding

eecdbf3

foosekis4j force-pushed the gpu-compaction branch from cffe0fb to eecdbf3 Compare November 9, 2026 16:35

foosekis4j added 12 commits November 10, 2026 00:01

Release note for datastax#783: GPU graph construction

4146227

jvector-gpu: license headers on pom.xml and README.md

f5556c9

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.

foosekis4j added a commit to foosekis4j/jvector that referenced this pull request Nov 10, 2026

Release note for datastax#794: GPU compaction, PQ and NVQ encoding

0496433

foosekis4j force-pushed the gpu-compaction branch from eecdbf3 to 0496433 Compare November 10, 2026 06:40

foosekis4j added a commit to foosekis4j/jvector that referenced this pull request Nov 10, 2026

Release note for datastax#794: GPU compaction, PQ and NVQ encoding

9442ca3

foosekis4j force-pushed the gpu-compaction branch from 0496433 to 9442ca3 Compare November 10, 2026 07:28

foosekis4j added 4 commits November 10, 2026 11:45

Release note for datastax#794: GPU compaction, PQ and NVQ encoding

c7f0db3

foosekis4j force-pushed the gpu-compaction branch from 9442ca3 to c7f0db3 Compare November 10, 2026 11:47

Navigation Menu

Repository navigation

jvector-gpu: GPU compaction, PQ and NVQ encoding (follow-up to #783) - #794

jvector-gpu: GPU compaction, PQ and NVQ encoding (follow-up to #783)#794
foosekis4j wants to merge 18 commits into
datastax:mainfrom
foosekis4j:gpu-compaction

foosekis4j commented Nov 8, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Nov 8, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

foosekis4j commented Nov 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Results

Testing

Notes

Uh oh!

github-actions Bot commented Nov 8, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

foosekis4j commented Nov 8, 2026 •

edited

Loading