Repository navigation
OnDiskGraphIndexCompactor: CompactionAccelerator hook, with the wide code, probes and cell walk on the GPU (follow-up to #794) - #795
Draft
foosekis4j wants to merge 26 commits into
Draft
foosekis4j wants to merge 26 commits into
foosekis4j wants to merge 26 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.
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.
Contributor
|
Before you submit for review:
If you did not complete any of these, then please explain below. |
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.
…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.
…the cell walk elsewhere The full-precision merge's cell walk scores each cell's vectors against the nodes that probe it and keeps each node's top k per larger source. It is the largest phase of a merge (70.7 s of 230.0 s on ada002-1M, 4 segments) and already batched, so it can run on other hardware without changing anything else in the compactor. CompactionAccelerator (@experimental) is found with ServiceLoader through CompactionAccelerators.find() and selected with -Djvector.compaction.accelerator=<name>. For each target source of a range, the compactor hands it a CellWalk: the range's queries, the sorted (cell, query) pairs, each cell's ordinal range, the raw rows through PreEncodedCodeCache.copyRows (a bulk copy of contiguous ordinals), JVector's score function, and setTopK for the results. If walk() returns false, or no accelerator is selected, the compactor walks as before. ReferenceCompactionAccelerator does the walk in plain Java through the public contract only; testCellJoinMergeRecallWithAccelerator runs the cell-join merge with it and checks that it walked and that recall meets the existing bound.
…oVM) TornadoCompactionAccelerator implements CompactionAccelerator with two TornadoVM kernels in Java (CellWalkKernels): score, one work-group per (cell, query) pair with the query in local memory and one warp per member row, writing FP32 scores in JVector's units at offsets the host assigns; and topk, which merges each query's candidates into its running top k. A range's queries are uploaded once and kept on the device (persistOnDevice, with their norms computed there); the target's rows stream from the merge's scratch region in fixed-size chunks of whole cells. Every buffer has a fixed size, so a larger merge runs as more launches, never as a larger allocation. ada002-1M, 4 segments, RTX 4090 / i9-13900K: cell walk 70.7 s -> 23.9 s, the compaction 230.0 s -> 181.0 s, recall@10 at beams 20-160 within 0.002. A CPU and a GPU merge agree on 0.92% of base-layer neighbour lists, two CPU merges on 0.94% (mean overlap 0.865 in both): the merge itself is not deterministic. WalkBench (run.sh walk) runs the compactor with the CPU and the GPU walk on CompactionBench's segments and compares time, lists and recall.
…p through the hook The probes choose, for each node of a range, the cells it scans: a beam search over the merged hierarchy's level-1 nodes (ef 64), from whose best-first list the compactor takes cells until each target source's scan budget is spent. At 57-59 s of 230 s on ada002-1M they are, with the walk on the GPU, the largest phase. CompactionAccelerator.probe(CellProbe) (a default method returning false) ranks every query's cells at once: CellProbe gives the range's queries, the level-1 vectors (HubMap.copyLevel1Rows, a bulk copy), the similarity and the most probes to return, and takes each query's ranked level-1 positions. With an accelerator, the compactor decodes the range's queries first, offers them to probe(), and then applies its scan budgets to the returned lists exactly as to the beam's. ReferenceCompactionAccelerator ranks every center exactly, and the accelerator test checks that both steps ran.
TornadoCompactionAccelerator.probe scores every query of a range against every level-1 center with a tiled FP32 matrix product (ProbeKernels.centerScores, 64 x 64 tiles staged through local memory, in JVector's similarity order) and keeps each query's best 1,024 centers (selectTop: the 1,024th score bracketed by bisection, the centers above it gathered in position order, a bitonic sort in local memory). The centers and their norms are uploaded once per range and stay on the device; the queries go in blocks whose score matrix fits 1 GB. ada002-1M, 4 segments, RTX 4090 / i9-13900K: probes 58.8 s -> 13.4 s, with the GPU walk the compaction 230.0 s -> 134.8 s (1.7x). Recall@10 at beams 20-160 is 0.917 / 0.959 / 0.983 / 0.993 against 0.914 / 0.958 / 0.982 / 0.993: exact ranking orders the cells at least as well as the beam. -Djvector.gpu.probe=false keeps the probes on the CPU.
…ide code through the hook The full-precision merge trains a wide code (a PQ with four dimensions per subspace, 256 centroids) on a balanced 128,000-vector sample, for the finalize's diversity comparisons. On ada002-1M that is 37.3 s of the merge, of which 0.5 s reads the sample. trainWideCode(sample, subspaces, clusters, globallyCenter), a default method returning null, lets an accelerator train it as ProductQuantization.compute would; PQRetrainer.train passes the sample to it and falls back to ProductQuantization.compute. The log line now says how long reading the sample took and who trained the code. The reference accelerator trains it with ProductQuantization.compute, and the accelerator test checks all three steps ran.
TornadoCompactionAccelerator.trainWideCode trains the wide code with GpuProductQuantization.compute, which follows ProductQuantization.compute (same subspaces, global centering, k-means++ seeding and Lloyd iterations). -Djvector.gpu.wideCode=false keeps it on the CPU. ada002-1M, 4 segments, RTX 4090 / i9-13900K: the wide code 37.3 s -> 2.8 s; with the probes and the walk on the GPU too, the compaction 230.0 s -> 99.9 s (2.3x), recall@10 at beams 20-160 0.917 / 0.959 / 0.983 / 0.993 against 0.914 / 0.958 / 0.982 / 0.993 on the CPU.
…, and an out-of-core benchmark The walk's and the probes' buffers had fixed sizes (up to 2 GB of queries and 1 GB of scores). They now take shares of TornadoVM's device memory budget (-Dtornado.device.memory): 60% for the walk, of which at most 40% queries, 10% rows, 15% candidates and 5% pairs, and 25% for the probes, with the score matrix capped at 256 MB. The query block is no larger than the compactor's range: TornadoVM keeps a resident host copy of every buffer, so a larger block only takes host memory from the merge's page cache. That mattered: under an 8 GB host limit, the first version's 3.5 GB of host copies pushed the merge's mapped scratch region and code cache out of memory, and the finalize slowed from about 4 s to up to 53 s a range. OocBench runs the merge in a process that opens only the segments (compact), so it can run under a cgroup memory limit, and evaluates afterwards (eval, evalraw). It also builds segments from a raw float32 file, on the GPU or the CPU (segments). ada002-1M, 4 segments, systemd-run MemoryMax=8G (no swap), -Xmx4g, 4 GB device budget: CPU 219.7 s, GPU 125.0 s (1.76x), recall@10 at beams 20-160 0.915 / 0.959 / 0.982 / 0.993 against 0.915 / 0.958 / 0.981 / 0.993, no OOM kill. Heap plus native memory peaked at 4.0 GB (CPU) and 7.3 GB (GPU): the GPU path's host copies are a fixed cost of the device budget, not of the merge's size.
…o 100k centers Setting C (Cohere 10M x 1024, 8 segments, 32 GB cgroup limit, 21 GB device budget) made the GPU path slower than the CPU's (4,272 s against 3,180 s), for two reasons that grow with the merge: - A launch's row chunk spanned every row from its first to its last probed cell. At 10M a range's queries probe a sparse set of cells, so the chunk read most of each target's rows from disk. The chunk now holds only the launch's distinct cells, packed (copied in parallel); the score kernel takes each pair's chunk rows and the first ordinal of its cell. - Exact probes rank every query against every level-1 center: 313k centers at 10M make it 42 TFLOP a range, slower than the compactor's beam. Above -Djvector.gpu.probe.maxCenters (default 100,000) the probes stay on the CPU. Setting C with both: GPU 2,986 s against the CPU's 3,180 s, recall@10 at beams 20-160 0.767 / 0.863 / 0.923 / 0.960 against 0.766 / 0.854 / 0.918 / 0.957 (the first GPU run's outputs). The merge is I/O-bound at this size: the walk's row copies (476 s) and launches (456 s) still run one after the other. ada002-1M (unconstrained): 109.7 s, recall unchanged.
Each launch was planned (pairs, offsets, the distinct cells) and had its rows read from the merge's scratch region before it ran, and only then did the GPU start; at 10M x 1024 under a 32 GB limit the row reads (476 s) and the launches (456 s) ran one after the other. A background thread now plans the next launch and reads its rows into one host staging set (pairs, offsets and a direct buffer of the chunk's rows) while the GPU runs the current one; the staging set is copied into TornadoVM's buffers, then reused. One set is enough: TornadoVM's buffers are the second. ada002-1M, 4 segments: the GPU walk 3.1-3.4 s -> 2.6 s a target (the GPU now waits about 60 ms a target for the next launch), the merge 109.7 s -> 98.9 s, recall unchanged.
foosekis4j
force-pushed
the
ooc-compaction
branch
from
November 9, 2026 17:10
8d9037d to
1f41a99
Compare
Exact probes rank every query against every level-1 center; at 10M x 1024 (313k centers) that costs 42 TFLOP a range and lost to the compactor's beam, so the probes had stayed on the CPU there (872 s of setting C). Past -Djvector.gpu.probe.maxCenters (100,000) they now go through an inverted-list search over the centers: every 32nd center is a coarse center and every center joins the list of its nearest coarse center (once per merge, on the GPU: centerScores + argmaxRows); a query scores the coarse centers, keeps its 128 best lists (selectTop), scores their members exactly (gatherScores, one warp a center, at most 4,096) and keeps the best 1,024 (selectListed). -Djvector.gpu.probe.coarseRatio, coarseProbes, listCap and twoLevel=false tune or disable it. ada002-1M, forced with maxCenters=1000 (31,381 centers in 981 lists): 2.7 s a source (exact: 3.5 s); recall@10 at beams 20-160 0.913 / 0.957 / 0.982 / 0.992 against the CPU beam's 0.914 / 0.958 / 0.982 / 0.993 (64 lists: 0.912 / 0.955 / 0.980 / 0.992).
At 10M x 1024 (setting C: 8 segments, 32 GB cgroup limit, 21 GB device budget), 128 lists of 9,773 cost recall: 0.753 / 0.851 / 0.914 / 0.954 at beams 20-160 against the CPU merge's 0.766 / 0.854 / 0.918 / 0.957. With 256 lists and up to 8,192 centers a query: 0.763 / 0.855 / 0.918 / 0.956, the merge in 2,183 s against the CPU's 3,180 s (1.46x), probes 314 s against 868 s. A two-level launch now takes at most 16,384 queries (-Djvector.gpu.probe.block): its score matrices and candidate buffers have resident host copies, and 65k-query blocks made the merge's heap plus native memory peak at 24.5 GB; with 16k blocks, 20.3 GB.
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 follows the discussion in #783 about large datasets and small machines. It accelerates
OnDiskGraphIndexCompactoritself instead of replacing it: the compactor keeps its out-of-core design (source order, region ordinals, scratch region, reverse offers, in-place writer), and three of its phases can run on a GPU: the wide-code training, the probes and the cell walk.Stacked on #794 (which is stacked on #783): the branch contains their commits. The changes to review are the last eleven commits.
OnDiskGraphIndexCompactor:CompactionAccelerator, a hook for running the cell walk elsewhere (core, no new dependencies).@Experimentalservice interface, found withServiceLoaderthroughCompactionAccelerators.find()and selected with-Djvector.compaction.accelerator=<name>.CellWalk:PreEncodedCodeCache.copyRows, a bulk copy of contiguous ordinals;setTopKfor the results.walk()returns false, the compactor walks exactly as before.ReferenceCompactionAccelerator, in the tests, does the walk in plain Java through the public contract only.jvector-gpu: the cell walk on the GPU (TornadoVM).CompactionAccelerator.probe, the probe step through the same hook (core).probe(CellProbe)is a default method returning false. It ranks every query's cells at once from the level-1 vectors (HubMap.copyLevel1Rows).jvector-gpu: the probe step on the GPU.-Djvector.gpu.probe=falsekeeps the probes on the CPU.CompactionAccelerator.trainWideCode, the wide code through the same hook (core).trainWideCodeis a default method returning null. An accelerator may train the code asProductQuantization.computewould.jvector-gpu: the wide code on the GPU, withGpuProductQuantization.computefrom jvector-gpu: GPU compaction, PQ and NVQ encoding (follow-up to #783) #794.-Djvector.gpu.wideCode=falsekeeps it on the CPU.jvector-gpu: buffers sized from the device budget, and an out-of-core benchmark (OocBench).-Dtornado.device.memoryand the probes 25%. The query block is no larger than the compactor's range.Walk only the probed cells' rows; GPU probes only up to 100k centers.
-Djvector.gpu.probe.maxCenters(100,000 by default) the compactor's beam is cheaper, and the probes stay on the CPU.The walk's launches are pipelined. A background thread plans the next launch and reads its rows while the GPU runs the current one.
Two-level GPU probes for merges with many centers. Past 100k centers, the probes become an inverted-list search over the centers.
Results
ada002-1M (982,790 x 1536, COSINE), 4 segments built by JVector on the CPU (M=32, hierarchy). RTX 4090 against an i9-13900K (32 threads), TornadoVM 8.1.0,
run.sh walk ada002-1M 4 <workDir> both:Out of core
Each merge runs in its own process under a cgroup memory limit (
systemd-run --user --scope -p MemoryMax=... -p MemorySwapMax=0). Recall is evaluated afterwards in a separate process.No run was OOM-killed.
Testing
TestOnDiskGraphIndexCompactor, 23 tests, all pass, on top of jvector-gpu: GPU compaction, PQ and NVQ encoding (follow-up to #783) #794's (63 in all with thejvector-gpuhost tests).testCellJoinMergeRecallWithAcceleratoris new: it runs the cell-join merge through the reference accelerator, and checks that the accelerator trained the wide code, ranked the probes and walked, and that recall meets the existing bound.TestExternalGraphBuild,GraphIndexBuilderTest,TestGraphIndexBuilderConfigand thejvector-gpuhost tests, 62 tests, all pass.run.sh walkon an RTX 4090: the GPU wide code, probes and walk on ada002-1M, with the results above.mvn apache-rat:checkpasses with-Pgpu.Questions
jvector-gpu, which would be much harder to keep in sync with yours.