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

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.

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

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

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.

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.

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.

Updated as suggested. Thanks.

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.

Thanks for the comment. I refactored the filter-bitset cache so its configuration and state are local to the codec/format rather than JVM-global.

FilterBitsetCache is now an instance class with no mutable static state or System.getProperty() configuration. An immutable FilterBitsetCacheConfig carries enabled and maxBytes; SPI-created codecs use DEFAULT (enabled=true, maxBytes=0, where zero lazily derives the budget), while explicitly instantiated codecs can provide a different configuration.

Each CuVS2510GPUVectorsFormat owns one cache and passes it to every reader it creates. The query obtains the cache from each reader and groups working-set accounting by cache identity, so the accounting also remains correct if readers backed by different format instances are encountered.

The default SPI codec still has one cache shared by the readers created from that SPI singleton, which is expected under Lucene’s SPI lifecycle. The change is that the cache configuration is now encapsulated in that codec/format instance rather than controlled through JVM properties or mutable cache statics.

I also removed the cache's automatic write to com.nvidia.cuvs.filterBitsetPoolSize. Device-pool policy is now owned independently by cuvs-java (updated in the cuVS PR), which supplies its own conservative 4 MB growable default; applications can explicitly size it or set it to zero to opt out.

imotov

Left a couple of small comments. Other than than LGTM.

imotov · 2026-07-23T22:38:01Z

+ *
+ * @param enabled whether filter bitsets are cached between searches
+ * @param maxBytes maximum bytes cached by this format; {@code 0} derives a budget from the first
+ *     filtered search


Sorry for not catching this earlier. I think this could lead to some fairly strange scenarios. Imagine a search engine cluster with many nodes where a query with a very small working set is routed to one subset of nodes, while a much heavier query is routed to another subset of nodes. In that case, the first set of nodes could become inexplicably less responsive than the second, leaving users scratching their heads trying to understand why. The nodes are the same, the mixture of queries that they send now are the same, but some nodes are very fast and some nodes are constantly threshing the cache and slow to respond. I think reacting to whatever happened to be the first query may not be the best strategy here. I'd lean toward either a fixed value or a conservative percentage of the available resources. That would provide more predictable behavior and make the system easier to reason about and troubleshoot.

imotov · 2026-07-23T23:17:25Z

+      CuVSResources resources = CuVSResources.create();
+      String poolSizeProp = System.getProperty(WORKSPACE_POOL_SIZE_PROPERTY);
+      if (poolSizeProp != null) {
+        long poolBytes = Long.parseLong(poolSizeProp);


I would add a separate handling for NumberFormatException here in case somebody sets that to "1024M" or something causing already allocated resources to silently leak.

Added NumberFormatException and additional handling (negative values, values that could overflow) in separate helper. For even more confidence, deferred CuVSResources creation till after poolBytes has been validated.

jamxia155 · 2026-07-24T18:20:48Z

/ok to test 9c55185

imotov · 2026-07-24T18:26:25Z

/ok to test 9c55185

imotov · 2026-07-27T22:37:45Z

/ok to test 083887e

jamxia155 · 2026-07-28T01:44:54Z

/ok to test 1a34e4d

jamxia155 · 2026-07-28T12:56:19Z

/ok to test a7b3760

jamxia155 · 2026-07-28T22:44:48Z

/ok to test 14afe00

jamxia155 · 2026-07-28T23:57:30Z

/ok to test 6b87d53

bdice

Approving CI changes.

This PR adds a CAGRA search API and associated implementation for batching workloads involving a single query and multiple indexes (e.g. index segments built from one dataset by cuvs-lucene). This promotes GPU utilization for non-batched queries by parallelizing work over multiple CTAs, while minimizing the number of concurrent kernels for potential host-side parallelism. Addresses cuvs-lucene [issue 124](NVIDIA/cuvs-lucene#124), [issue 158](NVIDIA/cuvs-lucene#158), [issue 159](NVIDIA/cuvs-lucene#159), and cuvs [issue 2166](#2166). Associated cuvs-lucene PR: NVIDIA/cuvs-lucene#133. Authors: - James Xia (https://github.com/jamxia155) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Kyle Edwards (https://github.com/KyleFromNVIDIA) URL: #2035

jamxia155 · 2026-07-30T02:37:33Z

/merge

jamxia155 mentioned this pull request Apr 22, 2026

Multi partition cagra search NVIDIA/cuvs#2035

Merged

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

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 mentioned this pull request Jul 20, 2026

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

Open

imotov reviewed Jul 23, 2026

View reviewed changes

jamxia155 requested a review from a team as a code owner July 28, 2026 01:44

jamxia155 requested a review from bdice July 28, 2026 01:44

jamxia155 added feature request New feature or request non-breaking Introduces a non-breaking change labels Jul 28, 2026

jamxia155 changed the base branch from main to release/26.08 July 29, 2026 14:54

Add multi-segment CAGRA search

1fad50f

jamxia155 force-pushed the multi-segment-cagra-search branch from e9b45a8 to 1fad50f Compare July 29, 2026 18:29

cjnolet approved these changes Jul 29, 2026

View reviewed changes

jamxia155 self-assigned this Jul 29, 2026

bdice approved these changes Jul 29, 2026

View reviewed changes

Restore RAPIDS_BRANCH

20038da

rapids-bot Bot merged commit ff7e4ac into NVIDIA:release/26.08 Jul 30, 2026
37 of 49 checks passed

imotov mentioned this pull request Jul 31, 2026

Fix panama headers when RAPIDS_VERSION points to a PR branch #185

Closed

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!

Choose a reason for hiding this comment

Uh oh!

jamxia155 Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

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!

jamxia155 commented Jul 24, 2026

Uh oh!

imotov commented Jul 24, 2026

Uh oh!

jamxia155 Jul 23, 2026 •

edited

Loading