Optimize CAGRA_HNSW index build; add example - #195
Open
jamxia155 wants to merge 19 commits into
Open
Conversation
without loading the full set of data on the Java Heap, but instead allows us to stream the set of data to the Java Heap.
Address review feedback on duplicated code paths introduced by the host-streaming refactor: - Merge the two createMultiLayerHnswGraph overloads (List-based and CuVSMatrix-based) into a single CuVSMatrix-based method. The flush and quantized paths now build a matrix and sample rows from it, matching the merge path. - Merge the two writeFieldInternal overloads into one that takes a CuVSMatrix, with a thin List adapter for the flush/sorting path. - Drop the redundant `size` parameter throughout; it is derived from dataset.size(). - Remove the now-unused CuVSResources parameter from the Utils matrix builders (createFloatMatrix/createByteMatrix/createByteMatrixFromArray) and correct their docs, which still referred to device memory after the switch to hostBuilder. Updated all call sites (CuVS2510GPUVectorsWriter, binary/scalar quantized writers). While migrating the quantized writers to the unified graph method, fixed a latent byte-width bug: the quantized branch hardcoded (dimensions + 7) / 8 (binary packing), which is wrong for scalar quantization (one byte per dimension). The width is now taken from the matrix's column count, correct for both binary and scalar. Signed-off-by: Zack Meeks <zmeeks@nvidia.com>
Add ChunkedFbinIngestExample, a reference for streaming a large vector
file into an accelerated HNSW index without the two common pitfalls:
- fd churn: reopening/seeking/closing the source file per vector, which
is far slower than sequential I/O.
- doubled memory: pre-loading the whole file onto the JVM heap on top of
Lucene's own per-segment buffer.
The nested ChunkedFbinReader opens the file once and serves vectors from a
reusable buffer refilled by large sequential reads, holding at most one
chunk. Vectors are streamed one at a time into addDocument to build a
single-segment index. Because cuvs-lucene is a codec (below addDocument),
how the source is read is application code — the example is a template to
adapt to any vector source (DB, object store, stream).
Runnable with no arguments (generates a small demo .fbin) or against a real
file with a configurable chunk size. README updated with a run entry.
Level 0 (all N nodes) is now encoded in parallel: within memory-bounded waves, threads delta/VInt-encode contiguous node sub-ranges into per-thread buffers, which are concatenated to the IndexOutput in node order. The per-node encoding is independent and shared with the serial path, so the on-disk bytes are identical. Higher levels stay serial. Level-0 neighbors are read from the already-materialized layer0Neighbors array, so parallel reads are safe. Thread count comes from AcceleratedHNSWParams.getWriterThreads() (default 1 -> unchanged behavior; opt-in), and applies to CAGRA_HNSW and the quantized variants (shared writeGraph). Measured on deep1b-10M (degree 88): write-graph 36 s -> 12.8 s at 4 threads, 8.1 s at 16 (single-threaded concat is the ~5-6 s floor). Recall unchanged (within GPU-build run-to-run noise).
Stream vectors directly into a native host matrix during indexing instead of buffering them as a heap List<float[]>, so the full dataset is no longer held twice (Lucene's flat-writer list + the native CAGRA build input). This takes peak host memory from ~2x to ~1x, fixing OOM on out-of-core builds, and removes the per-vector createFloatMatrix copy (the matrix-assembly stage). Opt-in via AcceleratedHNSWParams.numInputVectors (default 0 = disabled → existing heap-buffered path unchanged). When set: - FieldWriter streams addValue() into a CuVSMatrix.hostBuilder host matrix (reusing the same builder the merge path already uses) and tracks its own DocsWithFieldSet. - NativeFlatVectorsWriter writes .vec/.vemf directly from that host matrix, byte-compatible with Lucene99FlatVectorsWriter (10.2.0) so the stock Lucene99FlatVectorsReader reads them; format constants are re-declared with a version-pinned attribution and an upgrade-check note. - The writer/format wire the native path into flush and skip the Lucene flat writer so it never opens the same .vec/.vemf. The host matrix is sized for exactly numInputVectors rows, so the value must equal the number of vectors added; the writer fails fast otherwise. Supported only for the unsorted single-segment flush path — merges and index-sorted segments throw. Validated on deep1b-10M: recall parity with the heap path (99.588 vs 99.594) confirms format correctness, and the matrix-assembly stage is eliminated.
fillNeighborArray materialized the CAGRA adjacency into on-heap NeighborArrays in a serial per-node loop (~8.5s at 10M x degree-88). Parallelize it under the existing writerThreads knob. CagraIndex.getGraph() returns a CuVSDeviceMatrix whose getRow uses a shared, stateful buffered reader that is not safe for concurrent access, so the parallel path pulls layer 0 to host once (bulk device->host copy) and materializes disjoint node ranges concurrently over the stateless host matrix, closing the temp copy after. The serial path (writerThreads=1) is unchanged and reads the device matrix directly. Applies to all CAGRA_HNSW variants. deep1b-10M @ 16 threads: hnsw-convert 8.5s -> 3.0s, recall unchanged.
Upgrade the large-.fbin ingestion reference from a single-buffered, fresh-array-per-vector reader to the optimized ingestion pattern used by the benchmark harness, and rename ChunkedFbinIngestExample -> OptimizedFbinIngestExample to reflect it. PrefetchingFbinReader now demonstrates the four properties that keep ingestion from bottlenecking the GPU build: - open once, read sequentially in large chunks (no fd churn) - bounded memory (at most two chunks resident) - overlap: a background reader thread fills the next chunk while the ingest thread feeds the current one into addDocument, hiding the disk read behind per-document indexing - reuse the vector array: get(int, float[]) unpacks into a caller-owned array, safe because Lucene copies each vector eagerly at addDocument Scope is ingestion only; the codec stays at defaults (no native flat buffering or other config-gated behavior). README updated (description, run command, jar version 26.08.0 -> 26.10.0).
HEURISTIC previously always chose the CAGRA build algo by row count (NN_DESCENT below 5M, else IVF_PQ) and ignored cagraGraphBuildAlgo. Now an explicit IVF_PQ or NN_DESCENT overrides that choice, while AUTO_SELECT keeps the row-count determination; params for the chosen algo stay auto-tuned (getIVFPQParams / getNNDescentParams). Flip DEFAULT_CAGRA_GRAPH_BUILD_ALGO from NN_DESCENT to AUTO_SELECT so callers that use the default keep the row-count behavior (the switch above would otherwise force NN_DESCENT for them). CUSTOM strategy is unchanged.
Rename OptimizedFbinIngestExample -> OptimizedCagraHnswBuildExample and
turn on every build-side optimization, so the example is a single golden
reference for building a large accelerated HNSW index rather than just an
ingestion pattern:
- Native flat buffering: size the codec's numInputVectors to each
segment, streaming vectors into the native host matrix instead of a
heap List<float[]> (single-segment; ~halves peak host memory).
- Automatic graph-build algorithm: HEURISTIC + AUTO_SELECT lets cuVS pick
NN_DESCENT vs IVF_PQ by dataset size and auto-tune its parameters.
- Partitioned multi-segment build with a user-specified segment count:
- sequential: K single-segment passes appended to one directory,
peak host = one slice (N/K);
- overlapped: a bounded pool (PIPELINE_DEPTH) builds segments into
their own dirs with the GPU commit serialized on a permit, then
hardlinks them into the final directory via addIndexes (no bulk
copy). Peak host = depth * (N/K).
The class Javadoc states the assumptions behind the segment count: native
flat buffering forces one segment per slice, the GPU build is serialized
(more segments buy host-memory headroom + ingest overlap, not device
parallelism), and search fans out across all K segment graphs.
Adds a lucene-misc dependency (HardlinkCopyDirectoryWrapper) and updates
the examples README with the new name and 4-arg usage.
- Add enableRMMAsyncMemory() call with a note that it must not be used with CPU-only codecs - Expose the primary tuning knobs in codecFor(): withMaxConn/withBeamWidth (recall/graph-size), withCuvsDistanceType, and withWriterThreads (seeded from availableProcessors()) - Drop explicit withCagraGraphBuildAlgo(AUTO_SELECT) — it is the default under HEURISTIC; update class Javadoc accordingly - Expand withNumInputVectors comment to cover the exact-count constraint, the filtered-ingest fallback, and index-sorted segment incompatibility - In Lucene99AcceleratedHNSWVectorsWriter, detect index-sorted segments at construction time rather than failing later at flush
jamxia155
added
improvement
…rsWriter NativeFlatVectorsWriter hand-transcribes Lucene 10.2.0's .vec/.vemf flat vector format so it can write directly from native memory, bypassing the per-vector FloatVectorValues indirection Lucene's own writer requires. That port had no automated safety net: a Lucene upgrade could silently change the format constants or write sequence with nothing catching it before files became unreadable by the stock Lucene99FlatVectorsReader. Add two tests, referenced from NativeFlatVectorsWriter's class javadoc: - TestNativeFlatVectorsWriterFormatConstants asserts the resolved lucene-core classpath version still matches the pinned 10.2.0 (via the jar's Specification-Version manifest attribute, not a pom.xml text parse). Fast and environment-independent, so it runs even without cuVS/GPU available. Fires on any version bump, by design — an upgrade always needs a manual re-verification pass regardless of whether it happens to touch the mirrored constants. - TestNativeFlatVectorsWriterRoundTrip positively verifies the format still works by building a single-segment index with numInputVectors set (native flat buffering) and asserting every vector round-trips byte-exact through the real, unmodified Lucene99FlatVectorsReader. Requires cuVS/GPU. Expand NativeFlatVectorsWriter's "on a Lucene upgrade" class javadoc into a concrete checklist that points at both tests.
…n/write Neither of the two writerThreads-tuned parallelizations had test coverage: materializing the CAGRA adjacency into NeighborArrays (GPUBuiltHnswGraph's constructor) and encoding level 0 to disk (AcceleratedHNSWUtils.writeGraph). Both are only supposed to reformat the same data across threads, with no change in output versus the serial path. Add TestWriterThreadsGraphEquivalence, which feeds a synthetic adjacency (CuVSMatrix.ofArray, matching the construction the higher-layer subset builder already uses) directly into both code paths at writerThreads=1 and writerThreads=4, and asserts identical results: matching per-node arc lists for the graph construction, and matching per-node byte-length offsets plus byte-identical encoded output for writeGraph. Using a synthetic adjacency rather than a real CAGRA build isolates these two parallelizations from CAGRA's own build-to-build variance, since writerThreads is also forwarded into the actual GPU build parameters. The synthetic dataset is sized just above PARALLEL_MIN_NODES (1 << 16, the threshold duplicated in both classes) so the parallel branches are actually exercised rather than silently falling through to serial.
TestNativeFlatVectorsWriterRoundTrip only verifies the flat .vec file round-trips through Lucene's reader -- it doesn't cover whether a natively-buffered index actually searches correctly, tolerates deletions, composes with the odd-graph-degree fix, or correctly rejects the usage patterns it doesn't support. None of that had test coverage before. Add two classes: - TestNativeFlatBufferingIndexAndSearch: end-to-end build + KNN search through the native-flat-buffered path; deletion applied after a natively buffered flush (with the javadoc spelling out that Lucene-level deletion is orthogonal to, and not caught by, the count-mismatch guard -- deletion only updates liveDocs and never touches FieldWriter's native buffer); and the M = ceil(cagraGraphDegree / 2) odd-degree fix combined with numInputVectors, a distinct call path (writeFieldNative) from the one TestAcceleratedHNSWOddGraphDegree exercises. - TestNativeFlatBufferingGuardRails: the three guard rails in Lucene99AcceleratedHNSWVectorsWriter around numInputVectors -- a count mismatch (simulating an ingest-time filter that skips addValue for some docs, as distinct from a delete), an index-sorted segment, and a merge attempt. The index-sort and count checks fire at different points (writer construction on the first addDocument, vs. flush), so each is asserted around the call that actually triggers it.
FloatVectorValues#size() on KnnVectorsWriter.MergedVectorValues's merged view is the raw sum of every source segment's on-disk vector count -- not the number of live (non-deleted) vectors the iterator actually yields (Lucene's own Lucene99FlatVectorsWriter never calls .size() for exactly this reason, only iterates). vectorBasedMerge used that count to pre-size CuVSMatrix.hostBuilder's fixed native buffer, so whenever a merge dropped deleted docs, the buffer ended up larger than the number of vectors actually written into it, and the resulting HNSW graph was built with more nodes than the flat vector file actually contains. Symptom: an intermittent EOFException in OffHeapFloatVectorValues during concurrent search, whenever traversal happened to reach one of the phantom out-of-bounds nodes. Fix: count live vectors via a throwaway iteration before allocating the host buffer, rather than trusting size() -- but only when at least one segment being merged actually has deletions (mergeState.liveDocs has a non-null entry). size() is trustworthy whenever no segment has deletions, since per-segment vector counts already exclude docs without a value for the field independent of deletions, so that common case (most merges have no prior deletes) skips the extra pass entirely and costs exactly what it did before this fix. Add TestMergedGraphOrdinalBounds: rather than depending on a concurrent search happening to traverse a bad node, this deterministically merges two segments with interspersed deletions and walks the entire resulting HNSW graph, asserting every neighbor ordinal is in-bounds relative to the merged segment's actual live vector count.
…+IVF_PQ Under Strategy.HEURISTIC with an explicit cagraGraphBuildAlgo override, getIVFPQParams honored every other caller-supplied field (graphDegree, intermediateGraphDegree, writerThreads, cuvsDistanceType) but always discarded a caller-configured CuVSIvfPqParams in favor of the auto-tuned one -- an inconsistency within the same branch, since nothing signaled that this one field behaved differently from the others. Add AcceleratedHNSWParams.isCuVSIvfPqParamsExplicit(), tracked via a new Builder flag set in withCuVSIvfPqParams(...), so getIVFPQParams can honor an explicitly-set value and only auto-tune when the caller left it unset. Also, while auditing this file: - Remove ALGO_SWITCH_THRESHOLD, dead since both create() overloads now fully delegate algorithm selection to cuVS's own fromDataset/ fromHnswParams heuristics rather than a local row-count threshold. - Restore getCuVSIvfPqParams' and roundUpSafe's original source-attribution comments (lost when getCuVSIvfPqParams was reintroduced in a prior commit after having been deleted upstream), rephrasing the ivf_pq_params reference to name the struct instead of a line-numbered URL that will drift from the real source. Drop roundUpSafe's comment; the arithmetic doesn't need it. - Add a TODO noting cuvs-java has no standalone binding for the C++ ivf_pq_params(dataset_extents, metric) constructor this method translates -- only fromHnswParams/fromDataset, which bundle algorithm selection with parameter tuning and can't be used to auto-tune IVF-PQ while pinning a different algorithm choice.
FieldWriter's native-buffering path always preallocates its host matrix as DataType.FLOAT (4 bytes/dimension), regardless of quantizationType. This is currently safe -- the only call site that ever passes numInputVectors > 0 (Lucene99AcceleratedHNSWVectorsWriter) always passes QuantizationType.NONE, and the quantized writers never enable native buffering -- but nothing enforced it: it held only because no caller combines them today. Add a check in FieldWriter's constructor so that combination fails loudly (matching the existing numInputVectors guard message convention) instead of silently mis-sizing the native buffer if a future caller ever does combine them. Document the constraint (currently unsupported, not fixable via a caller workaround) everywhere numInputVectors' requirements are already described: AcceleratedHNSWParams.withNumInputVectors's javadoc, FieldWriter's own field javadoc, and OptimizedCagraHnswBuildExample's inline comment alongside the existing index-sort/merge constraints.
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
2 participants
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.
This PR builds on the out-of-core host-streaming refactor from #173 by @nvzm123. This PR subsumes #173 to avoid stacking the PRs.
On top of that base, this PR adds the following.
Improvements
writerThreadsknob.IndexOutputin node order according to on-disk format in the serial path.New example: OptimizedCagraHnswBuildExample
A reference pattern for building a large accelerated HNSW index with all ingest- and build-side optimizations, showcasing:
addDocument).withNumInputVectors, avoiding the heap-buffered assembly copy.PIPELINE_DEPTH) builds segments concurrently into their own directories with the GPU commit serialized on a semaphore (ingest overlaps a prior segment's GPU commit), then combines the finished per-segment indexes by hardlinking their files into the final directory viaHardlinkCopyDirectoryWrapper+addIndexes(no bulk copy of vector data).enableRMMAsyncMemory()call (with a note that it must not be used with CPU-only codecs) to opt-in to RMM-managed memory resources, plus exposing the primary tuning knobs (withMaxConn, withBeamWidth, withCuvsDistanceType, withWriterThreads).