Multi segment cagra search by jamxia155 · Pull Request #133 · NVIDIA/cuvs-lucene

jamxia155 · 2026-04-22T23:26:04Z

Companion PR to cuvs!2035, addresses #124.

Existing CAGRA search code path for each search query:

Call CAGRA search API on one index segment
Copy results back to host
Add results into host-side global top-k priority queue
Repeat for all index segments

Proposed change:

Leverage new multi-segment CAGRA search API to launch all per-segment searches in one API call
Leave results on device and run GPU-accelerated select-k API to compute global top-k
Copy final top-k results to host

Overrides rewrite() to run all segment searches into a shared device buffer and merge with cuvsSelectK entirely on GPU, eliminating per-segment D2H copies and CPU-side TopDocs.merge(). Falls back to the standard Lucene per-segment path when any segment lacks a CAGRA index, an explicit filter is set, or k > 1024. Also adds ordToDoc() and getCagraIndexForField() helpers to CuVS2510GPUVectorsReader to support result decoding. Fixes for Lucene 10.2 API changes: CodecReader moved to org.apache.lucene.index; createRewrittenQuery() removed and replaced with an inline docAndScoreQuery() implementation using the public Weight/ScorerSupplier API.

- CuVS2510GPUVectorsFormat: call CuVSProvider.provider().enableRMMAsyncMemory() in the static initializer so that cuda_async_memory_resource is active for the lifetime of the codec. This makes CAGRA workspace deallocations stream-ordered and non-blocking, which is required for the CudaStreamPool to provide any parallelism benefit. - GPUKnnFloatVectorQuery: upload the query vector to device once before the per-segment loop and share the resulting CuVSMatrix across all CagraQuery instances, reducing host-to-device copies from O(numSegments) to 1 per query. Wrap the shared device matrix in try-with-resources to close the RMM allocation promptly after MultiSegmentCagraSearch.search() returns. - FilterCuVSProvider: delegate enableRMMAsyncMemory() to the wrapped provider.

GPUKnnFloatVectorQuery / GPUPerLeafCuVSKnnCollector: - Add persistent, persistentLifetime, and persistentDeviceUsage parameters, threaded through all constructor overloads and forwarded to CagraSearchParams.Builder in both the multi-segment rewrite() path and the per-segment approximateSearch() fallback path. - Add threadBlockSize parameter (0 = auto) to allow tuning of the persistent kernel's worker_queue_size, which determines how many concurrent query threads can run without latency increase. Fix persistent-runner hash instability across segments (rewrite() path): - When max_iterations is 0 (auto), CAGRA computes it from each segment's dataset size. Different-sized segments produce different values, causing a distinct runner hash per segment and a destroy/recreate cycle on every search call. - Add computeMaxIterations(), which mirrors adjust_search_params() from search_plan.cuh, and call it once using the largest segment's graph size and degree. All segments then share the same max_iterations, producing a stable runner hash across the full multi-segment query. CuVS2510GPUVectorsReader: - Forward threadBlockSize, persistent, persistentLifetime, and persistentDeviceUsage from GPUPerLeafCuVSKnnCollector to CagraSearchParams.Builder in the per-segment fallback path.

Remove persistent kernel mode: - Drop persistent, persistentLifetime, and persistentDeviceUsage fields and parameters from GPUKnnFloatVectorQuery, GPUPerLeafCuVSKnnCollector, and CuVS2510GPUVectorsReader. The persistent kernel is superseded by the native multi-segment kernel (cuvsCagraSearchMultiSegment) which achieves better concurrency without the per-runner lifecycle overhead. - Collapse the 11-argument GPUKnnFloatVectorQuery constructor (which only existed to accept persistent parameters) into the standard 8-argument form. - Remove stale comments that described max_iterations uniformity in terms of persistent-runner hash stability; replace with accurate explanation (consistent search quality across segments of different sizes). Add workspace pool configuration: - Add WORKSPACE_POOL_SIZE_PROPERTY constant to ThreadLocalCuVSResourcesProvider. - On resources creation, read com.nvidia.cuvs.workspacePoolSize system property and call setWorkspacePool() if set, so callers can pre-warm the per-thread RMM pool without modifying cuvs-lucene source.

copy-pr-bot · 2026-04-22T23:26:08Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

…earch

GPUKnnFloatVectorQuery previously fell back to the per-segment CPU path whenever a filter query was present. This change extends the GPU fast path to handle filters by packing the accepted-ordinal bitset for every segment into a single FilterBitsetHandle and passing it to MultiSegmentCagraSearch. FilterBitsetCache (new): 16-entry shared LRU cache keyed by (filter Query, per-segment reader-cache keys, field). Per-reader keys are used rather than core keys so that liveDocs changes caused by a reader reopen automatically invalidate the cached bitset. GPUKnnFloatVectorQuery changes: - Remove the early bail-out for filter != null - Add buildOrGetCachedFilterHandle: checks FilterBitsetCache, falls back to buildFilterHandle on miss - Add buildFilterHandle: evaluates the filter Weight per segment via FloatVectorValues.getAcceptOrds (intersecting with liveDocs), then packs the result into combinedLongs / segBitOffsets for FilterBitsetHandle - When a handle is present, CagraQuery is built without a per-query prefilter (the bitset already encodes the intersection)

cuVS's search_multi_partition was refactored to take a single queries matrix and return globally merged top-k outputs (partition_ids, neighbors, distances) instead of per-partition vectors of matrix views. This commit updates GPUKnnFloatVectorQuery to match. Changes in rewrite(): - Build one shared CagraQuery (queryVector + searchParams), not one per Lucene segment; drop the per-segment construction loop. - The remaining loop only collects the CagraIndex per segment for the cuVS partition list. - Single call to MultiPartitionCagraSearch.search(...), passing the optional FilterBitsetHandle as the last argument (no more filterHandle != null branching between overloads). Filter handling: - Since a single shared CagraQuery cannot carry per-segment liveDocs via withPrefilter anymore, fold liveDocs into the FilterBitsetHandle alongside any explicit Lucene filter. The handle is now built whenever filter != null OR any segment has deletes. - buildFilterHandle extended to accept filter == null, in which case acceptDocs is just liveDocs per segment. - Removed buildCagraQuery helper that constructed per-segment liveDocs bitsets; the unused java.util.BitSet import is removed with it.

The k>1024 short-circuit at the top of rewrite() forced large-k queries through Lucene's per-leaf approximateSearch fallback, where they fail either at the codec reader's k<=1024 gate or — when that gate doesn't fire — by trying to use a brute-force index that callers typically do not build. The multi-partition path is the right place to serve these queries: with N segments at SINGLE_CTA's itopk_size cap of 512, the per-partition × num_segments candidate pool is more than enough to assemble a global top-k well above 1024. Companion to the cuVS change that decouples the multi-partition kernel's per-partition output count from the global topk, allowing each partition to emit up to itopk_size candidates and merging across partitions to produce the final top-k. - Drop the k>1024 early-return in rewrite(). Feasibility is now enforced by cuVS's itopk_size * num_partitions >= topk RAFT_EXPECTS in search_multi_partition. - Clamp itopk_size to 512 when building CagraSearchParams, since SINGLE_CTA rejects larger values. The clamp is a no-op when the caller-requested itopk_size is already <= 512. - Update the class Javadoc to reflect that the multi-partition path is no longer k-capped. The per-leaf approximateSearch fallback's k<=1024 gate and brute-force- index assumption are unchanged; they fire only when some segment lacks a GPU reader, which is independent of this change.

This reverts commit 7f394de.

computeMaxIterations was added to stabilize the persistent-kernel runner hash across segments by pre-computing max_iterations on the Java side from the largest segment, so every segment received the same value. Persistent-kernel support was subsequently removed, eliminating that rationale. The cuVS multi-partition C++ entry (search_multi_partition in cagra_search.cuh) already iterates partitions to find max_dataset_size and max_graph_degree, then calls adjust_search_params(), which produces the same max_iterations value the Java side was computing. Passing max_iterations = 0 (the CagraSearchParams default) lets that path run naturally. - Delete the per-segment loop that computed max_dataset_size and graph_degree. - Drop .withMaxIterations() from the multi-partition search builder. - Delete the computeMaxIterations helper and its Javadoc. No behavior change; the C++ side computes the identical value.

The gate forced k>1024 queries through Lucene's per-leaf fallback, where they hit the codec reader's k<=1024 check and tried to invoke a brute-force index that GPU-indexed segments typically do not write — producing a confusing NPE. cuVS now handles k>1024 in the multi-partition path via MULTI_KERNEL, routed automatically by the AUTO heuristic on itopk_size or selected explicitly via cagraSearchAlgo. Feasibility is validated at the C++ side per algo: SINGLE_CTA still requires itopk_size * num_partitions >= topk, MULTI_KERNEL requires topk <= itopk_size. Failures surface as RAFT_EXPECTS messages instead of NPEs.

…earch

cuvs-java moved FilterBitsetHandle/MultiPartitionCagraSearch to the base source set behind CuVSProvider: - GPUKnnFloatVectorQuery: new FilterBitsetHandle(...) -> FilterBitsetHandle.create(...). - FilterCuVSProvider: forward the two new CuVSProvider methods (newFilterBitsetHandle, searchCagraMultiPartition) to the delegate.

…earch

The merge of origin/main (c840848) combined two independently-added <repositories> blocks: the GCS-mirror + Maven Central block from NVIDIA#166 (releases only, snapshots disabled) and the central-snapshots block from the snapshot-build change (snapshots enabled). Two <repositories> tags is invalid XML, so the POM failed to parse. Consolidate both into a single <repositories> block, preserving all three repositories: the GCS mirror and Maven Central for releases, and central-snapshots for snapshot resolution.

max_iterations was already plumbed through cuvs-java (CagraSearchParams.Builder.withMaxIterations) and the C-API, but cuvs-lucene always left it at the default (0 = auto). Thread it through so callers can bound CAGRA search iterations, following the existing threadBlockSize pattern: - GPUKnnFloatVectorQuery: new maxIterations field; the full constructor takes it (between threadBlockSize and searchAlgo). The 6-arg constructor is unchanged and delegates with 0. - GPUPerLeafCuVSKnnCollector: carries maxIterations with a getter. - Both search paths (multi-segment rewrite and per-leaf fallback) pass it into CagraSearchParams.Builder.withMaxIterations. 0 = auto preserves the previous behavior, so existing callers are unaffected.

jameslamb · 2026-07-01T16:51:42Z

I manually re-triggered CI here, it was failing with the issues fixed by rapidsai/ci-imgs#437

cuVS now recomputes per-partition bit offsets from the index sizes, so the combined bitset is passed as a plain BITSET filter with no offsets or total- bit count. Update the consumers of `FilterBitsetHandle` accordingly: - `GPUKnnFloatVectorQuery.buildFilterHandle` calls the 1-arg `FilterBitsetHandle.create(combinedLongs)`. The per-segment offsets are still computed locally to pack each segment's slice into the combined bitset (64-bit word-aligned, matching cuVS's recomputation) — they're just no longer transported. - `FilterCuVSProvider.newFilterBitsetHandle` delegate override updated to the 1-arg signature.

imotov

Sorry it took me a while. I'm still trying to wrap my head around this codebase. I left a few suggestions.

imotov · 2026-07-15T21:08:43Z

+   * @throws IOException if the vector values cannot be read
+   */
+  public int ordToDoc(String field, int ordinal) throws IOException {
+    return flatVectorsReader.getFloatVectorValues(field).ordToDoc(ordinal);


This is a minor nit, but ordToDoc() is called for every result. Although getFloatVectorValues() is fairly inexpensive, it still performs some allocations. Since the number of results may be significantly larger than the number of underlying segments, it could be worth stashing the FloatVectorValues per segment instead of repeatedly creating and discarding them.

Updated as suggested. Thanks.

imotov · 2026-07-15T21:17:14Z

+
+                      @Override
+                      public int advance(int target) {
+                        while (pos < n && sortedDocs[pos] < target) pos++;


This might cause NPE if advance(int target) is called before nextDoc(), which sometimes happens when you iterate over multiple DocIdSetIterators at the same time.

Fixed. Thanks.

imotov · 2026-07-15T22:03:41Z

+    FilterBitsetHandle cached = FilterBitsetCache.get(filter, segReaderKeys, field);
+    if (cached != null) return cached;
+
+    FilterBitsetHandle handle = buildFilterHandle(indexSearcher, leaves, gpuReaders);
+    FilterBitsetCache.put(filter, segReaderKeys, field, handle);


I think there are a few potential race conditions in this cache that we should address. Because get() and put() are synchronized independently, it's possible for two threads operating on the same key to reach buildFilterHandle() at roughly the same time. In that case, both threads will perform the work and create a FilterBitsetHandle, which is somewhat wasteful. A bigger problem though is that when the second thread calls put(), it will overwrite the FilterBitsetHandle in the CACHE created by the first thread. Unless the replaced handle is explicitly closed, this appears to create a potential resource leak.

Implemented reference counting as resolution.

imotov · 2026-07-15T22:05:10Z

+        @Override
+        protected boolean removeEldestEntry(Map.Entry<FilterCacheKey, FilterBitsetHandle> eldest) {
+          if (size() > MAX_HOST_ENTRIES) {
+            eldest.getValue().close();


What will happen if this handle is still in use?

Resolved with ref counting solution.

imotov · 2026-07-15T22:15:30Z

+          public Explanation explain(LeafReaderContext ctx, int doc) {
+            for (ScoreDoc sd : scoreDocs) {
+              if (sd.doc == ctx.docBase + doc) {
+                return Explanation.match(sd.score, "GPU multi-segment CAGRA search");


Since we're applying the boost to the score earlier, we should probably include that boost in the explanation as well. One option is to make it part of the explanation details; alternatively, we could simply use: return Explanation.match(sd.score * boost, "GPU multi-segment CAGRA search");

Added "boost" to explanation details.

imotov · 2026-07-15T22:30:39Z

+
+                  @Override
+                  public float getMaxScore(int upTo) {
+                    return Float.MAX_VALUE;


We can probably do a bit better here without much overhead by returning a real max value and ignoring upTo.

Indeed. Updated as suggested.

imotov · 2026-07-15T22:52:59Z

+      if (t instanceof IOException) throw (IOException) t;
+      if (t instanceof RuntimeException) throw (RuntimeException) t;
+      throw new RuntimeException("Multi-segment GPU search failed", t);


I think, for now, this can be replaced with Utils.handleThrowable. More generally, though, I don't think this is a great pattern. We shouldn't indiscriminately wrap everything in a RuntimeException, as that prevents applications using this library from handling serious conditions such as OutOfMemoryError or AssertionError appropriately. In general, it's better to propagate Error subclasses unchanged and only wrap exceptions that are actually intended to be converted into unchecked exceptions.

Replaced with handleThrowable for now as a proper fix as suggested would be too pervasive for the scope of this PR.

Can one of you create a GitHub issue to follow up on this, just so it doesn't get lost in the noise?

Created #178, @imotov please share more context as needed. Thanks.

imotov · 2026-07-15T22:57:33Z

+    for (int i = 0; i < numSegments; i++) {
+      segBitOffsets[i] = totalBits;
+      int numOrds = gpuReaders.get(i).getFloatVectorValues(field).size();
+      totalBits += ((long) (numOrds + 63) / 64) * 64;


I think, + here might cause an int overflow.

Thanks for the spot. Updated to upcast numOrds before any subsequent arithmetic.

imotov · 2026-07-15T23:03:16Z

+      return new Bits() {
+        @Override
+        public boolean get(int i) {
+          return false;
+        }
+
+        @Override
+        public int length() {
+          return maxDoc;
+        }
+      };


return new Bits.MatchNoBits(maxDoc);

Updated as suggested. Thanks.

The result-mapping loop in GPUKnnFloatVectorQuery called ordToDoc() once per result, each call allocating a fresh FloatVectorValues via the reader. Since the number of results can be far larger than the number of segments, cache one FloatVectorValues per segment and reuse it across all results landing in that segment, reducing allocations from O(results) to O(distinct segments touched). Removes the now-unused CuVS2510GPUVectorsReader.ordToDoc() helper.

FilterBitsetCache had two concurrency problems. get() and put() were synchronized independently, so concurrent misses on the same key both built a FilterBitsetHandle and the second put() overwrote the first without closing it, leaking device memory. And eviction closed a handle that another thread's in-flight search could still be using, freeing its device allocation underneath the search. Rework the cache around FilterBitsetHandle's new reference counting. Values are now per-key CompletableFutures, so a key is built exactly once while builds for different keys still run in parallel. A cached handle carries one cache-owned reference, released via close() on eviction; acquire() takes an additional reference for the caller, which GPUKnnFloatVectorQuery releases in a finally after the search. Because the allocation is freed only when the last reference drops, eviction can no longer free a handle still in use, and a replaced entry cannot leak. Add TestFilterBitsetCache (compute-once, eviction refcount, failed-build retry, concurrent balance) and TestMultiSegmentGPUFilterConcurrency (end-to-end filtered multi-segment GPU search under cache eviction).

The per-segment search decoder skipped empty result slots by checking ord < 0, which only catches the 0xFFFFFFFF sentinel index. Single-CTA emits 0x7FFFFFFF for empty slots, which is positive and slipped through as a bogus ordinal under high filter selectivity. Gate on the FLT_MAX sentinel distance instead, which every search algorithm honors and which is robust regardless of how the index is mapped downstream.

The test built its index with TestUtil.alwaysKnnVectorsFormat, whose wrapped reader is not a CuVS2510GPUVectorsReader, so unwrapGpuReader returned null and every query fell back to per-segment search — the multi-partition path was never tested. Use CuVS2510GPUSearchCodec, which exposes the CuVS reader directly, so GPUKnnFloatVectorQuery takes the multi-partition path the test is meant to cover.

…filters The per-segment search passed BitSet.length() as the prefilter's numDocs (the "total dataset vectors" cuVS uses to size the filter). BitSet.length() is the highest set bit plus one, so under a selective filter — where the trailing ordinals are rejected — it is smaller than the vector count. cuVS then treats ordinals beyond that length as outside the filter and default-accepts them, returning filtered-out vectors. Under a permissive filter the last ordinal is usually accepted, so length() matched the vector count and the bug was masked. Pass the vector count (acceptedOrds.length()) instead, so the prefilter covers every ordinal. Add TestPerSegmentGPUFilterSearch, which forces the per-segment path (alwaysKnnVectorsFormat wraps the reader) with selective filters and asserts every hit satisfies its filter.

…ph degree unwrapGpuReader only recognized a directly-wrapped CuVS reader, so a real index -- where PerFieldKnnVectorsFormat.FieldsReader wraps the per-field reader -- never matched and every query silently fell back to per-segment search, leaving the multi-partition path (and its filter cache) unused. Unwrap the PerField reader so those segments take the multi-partition path. That path requires every partition to share one built CAGRA graph degree; cuvsCagraSearchMultiPartition rejects a mismatch, and a small segment can have its degree truncated at build time. Gate the multi-partition path on all segments reporting the same CagraIndex.getGraphDegree(), falling back to the per-segment path (which handles heterogeneous segments, liveDocs, and gaps natively) otherwise. Rework TestPerSegmentGPUFilterSearch to drive leaf.searchNearestVectors directly, since the unwrap change now routes its former setup through the multi-partition path; the reworked test exercises the per-segment prefilter regardless of higher-level routing.

Weight.explain returned the raw CAGRA score while the scorer multiplies it by the query boost, so the explained value disagreed with the actual score whenever boost != 1.0 (and would fail Lucene's explain-consistency checks). Report sd.score * boost, with the raw score and boost as detail factors.

getMaxScore returned Float.MAX_VALUE, which disables TOP_SCORES/MaxScore pruning. Precompute the segment's maximum (boosted) score once during scorer setup and return it, ignoring upTo. The segment-wide max is a valid upper bound for any upTo window, so pruning stays correct while becoming effective.

- Propagate Errors (and IOException/RuntimeException) unchanged from the rewrite catch via Utils.handleThrowable instead of wrapping everything in RuntimeException, matching the codebase idiom. - Widen numOrds to long before adding 63 in the bit-offset calc to avoid a potential int overflow for very large segments. - Replace the hand-rolled all-false Bits with Bits.MatchNoBits for the no-scorer (empty filter) case.

imotov · 2026-07-20T19:27:07Z

@jamxia155 I spent some time thinking more about the concerns I shared with you on Friday, and after reflecting on them over the weekend, I still feel quite uneasy about this cache introduction. While I can certainly see scenarios where it could deliver significant performance improvements, I can also think of several cases where it might actually have a negative impact. My main concern is that Lucene already has a query caching layer that has been carefully tuned over time and is further augmented by Elasticsearch at higher levels of the stack. As far as I know, Solr does not even use Lucene's query cache directly, instead relying on its own caching framework. In other words, downstream consumers of Lucene tend to be very opinionated about when query caches should be used and how they should be managed.

This PR makes several decisions on behalf of downstream consumers that they may not appreciate. Specifically:

The number of queries in the cache is hardcoded.
We don't keep track of total amount of memory that cache consumes and don't limit them.
Caching is mandatory for all queries, we don't provide a mechanism for downstream clients to control that.
There is no mechanism to disable the cache entirely.
There is no mechanism to clear the cache.

The cache implementation itself uses the full set of segments as the cache key, which introduces additional concerns. This works well for static indexes, but indexes are frequently updated in many real-world deployments. In those situations, the segment set changes continuously. Consider a simple scenario where documents are periodically added to the index and a single search is executed after each update. Over time, the cache would accumulate entries associated with obsolete segment configurations, resulting in a cache that entierly consists of stale results.

Lucene's own caching infrastructure, as well as Elasticsearch's cache implementations, are integrated with the segment lifecycle and merge process. They can therefore invalidate cache entries as soon as segment changes occur. Our implementation does not do that, which makes the cache substantially less useful for non-static indexes and increases the risk of wasting memory on entries that are unlikely to be reused.

More broadly, I am concerned that we are introducing a mandatory caching layer without providing the controls, invalidation mechanisms, and configurability that downstream consumers typically expect from Lucene. So, my recommendation would be to either spend a bit more time on addressing these issues or if we want to release it now make this cache configurable and disable it by default, enabling users who want it to work with it, but not force it on the rest of the users.

Adapt to the new cuvs-java per-partition multi-partition search API and fix the cache-staleness concern by keying per segment instead of per segment-set: - GPUKnnFloatVectorQuery builds one FilterBitsetHandle per segment (filter ∩ that segment's liveDocs, packed into its own long[]), or null for a segment with no explicit filter and no deletes, and passes the List<FilterBitsetHandle> to MultiPartitionCagraSearch.search. Every non-null handle is reference-released once per query. - FilterBitsetCache keys on (filter, single-segment reader key, field) instead of the whole segment-key list, so an index update only invalidates the changed segments' entries rather than the entire cache. - FilterCuVSProvider passthrough updated to the new SPI signature. - TestFilterBitsetCache updated for the per-segment acquire key.

Integrate the per-segment filter bitset cache with the segment lifecycle so obsolete entries don't linger until LRU eviction: - FilterBitsetCache.ensureCloseListener registers a one-time close listener per segment reader (tracked so each reader registers once); when the reader's core closes (merge, reopen with new deletes, or index close) its cache entries are evicted and their device allocations released. - invalidateReader(key) evicts and releases every entry for a segment reader key; the close listener calls it. - GPUKnnFloatVectorQuery registers the listener before caching each segment. This mirrors Lucene's LRUQueryCache (core key + closed listener). Adds TestFilterBitsetCache.invalidateReaderEvictsOnlyThatSegment.

Replace the fixed 16-entry filter bitset cache bound with a byte budget and add the client controls the reviewer asked for: - Byte-cap eviction: track total cached bytes and evict LRU entries until within budget (never the entry the in-flight query needs), plus a generous MAX_ENTRIES_GUARD against pathological tiny-entry streams. Host and device footprints match one-for-one, so one budget bounds both. - Controls: on by default; cuvs.lucene.filterBitsetCache.enabled and .maxBytes system properties plus setEnabled/setMaxBytes/clear(). Disabling routes every segment through the uncached, caller-owned build path. - Sizing: the budget defaults to min(16 x working set, 512 MiB), derived lazily from the first search's working set. An explicit budget smaller than that minimum working set fails loudly once; later index growth past the budget only warns and leans on eviction. A single segment larger than the whole budget is served uncached with a one-time warning. GPUKnnFloatVectorQuery computes each filtered segment's byte size in a pre-pass and reports the per-query working set before acquiring. Cache tests are reworked for the byte budget and reset the process-static cache in setup.

Default the cuVS filter device pool (com.nvidia.cuvs.filterBitsetPoolSize) to the resolved cache byte budget once it is known, before the first bitset upload, so the RMM pool backing filter device allocations tracks the cache size instead of churning cudaMalloc/cudaFree on every eviction/rebuild. An operator who sets the property explicitly overrides this default; the value is published only when unset. Adds filterPoolSizeDefaultsToBudget.

RMM's pool_memory_resource requires the initial pool size to be a multiple of 256 bytes. Round com.nvidia.cuvs.workspacePoolSize up before passing it to setWorkspacePool, and skip non-positive values so a misconfigured size falls back to the default workspace resource instead of throwing and disabling cuVS for the thread.

…earch

copy-pr-bot · 2026-07-22T17:03:21Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

jamxia155 · 2026-07-22T17:59:56Z

@jamxia155 I spent some time thinking more about the concerns I shared with you on Friday, and after reflecting on them over the weekend, I still feel quite uneasy about this cache introduction. While I can certainly see scenarios where it could deliver significant performance improvements, I can also think of several cases where it might actually have a negative impact. My main concern is that Lucene already has a query caching layer that has been carefully tuned over time and is further augmented by Elasticsearch at higher levels of the stack. As far as I know, Solr does not even use Lucene's query cache directly, instead relying on its own caching framework. In other words, downstream consumers of Lucene tend to be very opinionated about when query caches should be used and how they should be managed.

This PR makes several decisions on behalf of downstream consumers that they may not appreciate. Specifically:

The number of queries in the cache is hardcoded.

We don't keep track of total amount of memory that cache consumes and don't limit them.

Caching is mandatory for all queries, we don't provide a mechanism for downstream clients to control that.

There is no mechanism to disable the cache entirely.

There is no mechanism to clear the cache.

The cache implementation itself uses the full set of segments as the cache key, which introduces additional concerns. This works well for static indexes, but indexes are frequently updated in many real-world deployments. In those situations, the segment set changes continuously. Consider a simple scenario where documents are periodically added to the index and a single search is executed after each update. Over time, the cache would accumulate entries associated with obsolete segment configurations, resulting in a cache that entierly consists of stale results.

Lucene's own caching infrastructure, as well as Elasticsearch's cache implementations, are integrated with the segment lifecycle and merge process. They can therefore invalidate cache entries as soon as segment changes occur. Our implementation does not do that, which makes the cache substantially less useful for non-static indexes and increases the risk of wasting memory on entries that are unlikely to be reused.

More broadly, I am concerned that we are introducing a mandatory caching layer without providing the controls, invalidation mechanisms, and configurability that downstream consumers typically expect from Lucene. So, my recommendation would be to either spend a bit more time on addressing these issues or if we want to release it now make this cache configurable and disable it by default, enabling users who want it to work with it, but not force it on the rest of the users.

Thank you, @imotov, for the detailed comment and suggestions. I appreciate your leniency in offering the alternative of a short-term crutch, but I decided to try my hand at making the proper fix (we can always fall back to disabling by default if things are found wanting). These changes sit on top of the recently added reference-counting implementation prompted by your other comments, so an eviction concurrent with an in-flight search can never free a handle still in use.

Core concerns and resolutions:

Staleness / full-segment-set keying: The core problem of the cache going stale on any update came from the multi-partition CAGRA filter API taking a single concatenated bitset over all segments, which forced the cache to key on the full segment set (so one changed segment invalidated everything). I have updated the cuVS API (C++/C/JNI) to take per-partition device bitsets, so each segment's filter is independent. The cache now keys on (filter, per-segment reader key, field): an index update only invalidates entries for the changed segments, and unchanged segments keep hitting.
Integration with the segment lifecycle: Each segment registers a one-time close listener on its reader's CacheHelper. When that reader's core closes (merge, reopen-with-deletes, index close), its entries are evicted and their device allocations released immediately rather than aging out via LRU. This mirrors Lucene's LRUQueryCache (core key + closed listener).
Lucene's caching infra/Elasticsearch's cache impl: our cached value is a filter bitset that has been uploaded to the GPU (a device allocation, refcounted, freed on the CUDA memory resource that allocated it). It exists to avoid re-uploading that bitset to the device on every query. Lucene's LRUQueryCache and Elasticsearch's IndicesQueryCache, which builds on it, caches host-side DocIdSet/Bits; it has no notion of device memory, CUDA-resource-bound lifetimes, or refcounted GPU allocations, so it can't store our value. We still use Lucene for the filter evaluation itself (Weight/Scorer), and the fix deliberately modeled the cache's lifecycle on LRUQueryCache so the semantics are familiar, but the storage and GPU lifetime had to be our own.

Other items as in your comment:

Hard-coded cache size: Replaced the fixed entry count (16) with a byte budget which exactly bounds GPU memory footprint. Configurable via cuvs.lucene.filterBitsetCache.maxBytes / setMaxBytes; default min(16 * working-set, 512 MB), derived lazily from the first search. A generous secondary entry-count guard remains for pathological many-tiny-entry cases.
We don't track or limit total memory consumed: The cache now tracks total cached bytes and evicts LRU entries until within the byte budget (host and device footprints match 1:1, so one budget bounds both). An explicit budget smaller than the minimum working set fails loudly once; later growth past the budget degrades gracefully (evict/skip + warn) rather than failing a live search; a single entry larger than the whole budget is served uncached. Device allocations are additionally backed by a pre-warmed, growable RMM pool sized to the budget, so eviction/rebuild reuse device slabs instead of churning cudaMalloc/cudaFree.
Caching is mandatory; no mechanism for downstream clients to control it: The controls are now exposed as client-settable configuration: system properties read at startup plus static setters (setEnabled, setMaxBytes). Downstream can turn caching off (item 4) or bound its size (items 1–2) without code changes.
No mechanism to disable the cache entirely: setEnabled(false) / cuvs.lucene.filterBitsetCache.enabled=false disables it fully, and every segment routes through the uncached build path. It's on by default, defensible now that it has full invalidation, a hard memory cap, and this disable switch (as LRUQueryCache also ships on by default).
No mechanism to clear the cache: Added clear(), which evicts and frees every entry (releasing both the host arrays and the device allocations).

imotov · 2026-07-23T02:32:13Z

+        // Prefer this over the neighbor-index sentinel: the index sentinel is not uniform across
+        // CAGRA search algorithms (single-CTA emits 0x7FFFFFFF, multi-CTA 0xFFFFFFFF), so the
+        // distance is the reliable, algorithm-independent signal for an empty slot.
+        if (score == Float.MAX_VALUE) {


I think that check alone should be enough, but I would still keep ord < 0 as an extra check just to harden it a bit.

imotov · 2026-07-23T05:36:22Z

+  // ---- Configuration (guarded by FilterBitsetCache.class) ----
+
+  private static boolean enabled =
+      Boolean.parseBoolean(System.getProperty(PROP_ENABLED, "true"));


It is very uncommon for Lucene to rely on System.getProperty(), except when obtaining OS configuration parameters or in tests, benchmarks, and build scripts. Instead, I think this cache should be encapsulated within the codec and configured through codec parameters. When loaded via SPI, the codec can use reasonable defaults, while applications that need different behavior can explicitly instantiate and configure the codec with the desired settings and then provide it for SPI with their desired defaults. This keeps the configuration local to the codec and avoids introducing global JVM-level state.

jamxia155 added 5 commits April 22, 2026 16:05

Expose CAGRA SearchAlgo parameter

7a8b26b

jamxia155 mentioned this pull request Apr 22, 2026

Multi partition cagra search NVIDIA/cuvs#2035

Open

jamxia155 and others added 15 commits May 8, 2026 17:54

Merge remote-tracking branch 'origin/main' into multi-segment-cagra-s…

19773fa

…earch

Modify example to demo prefilter.

9abe58a

Revert "Allow rewrite() to use multi-partition path for k beyond 1024"

1b0355a

This reverts commit 7f394de.

Merge remote-tracking branch 'origin/main' into multi-segment-cagra-s…

32e07bf

…earch

Updates to pom to buidl snapshot

51e03b1

Bump cuvs-lucene to 26.08.0-SNAPSHOT

7f99c97

Merge remote-tracking branch 'origin/main' into multi-segment-cagra-s…

c840848

…earch

jamxia155 marked this pull request as ready for review July 1, 2026 00:35

jamxia155 requested review from a team as code owners July 1, 2026 00:35

Remove unnecessary formatting updates

62f6302

jamxia155 force-pushed the multi-segment-cagra-search branch from 97ef8e7 to 62f6302 Compare July 1, 2026 01:33

KyleFromNVIDIA approved these changes Jul 1, 2026

View reviewed changes

imotov reviewed Jul 15, 2026

View reviewed changes

jamxia155 added 10 commits July 16, 2026 15:08

Fix NPE when advance() precedes nextDoc() in docAndScoreQuery iterator

ce0237e

jamxia155 mentioned this pull request Jul 20, 2026

Propagate Error subclasses unchanged instead of indiscriminately wrapping in RuntimeException #178

Open

jamxia155 added 7 commits July 21, 2026 18:09

Update cuvs version

361daf3

Merge remote-tracking branch 'origin/main' into multi-segment-cagra-s…

222d9a9

…earch

imotov reviewed Jul 23, 2026

View reviewed changes

Uh oh!

Conversation

jamxia155 commented Apr 22, 2026

Uh oh!

copy-pr-bot Bot commented Apr 22, 2026

Uh oh!

jameslamb commented Jul 1, 2026

Uh oh!

imotov left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

imotov commented Jul 20, 2026

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

jamxia155 commented Jul 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants