Lucene: Investigate performance of concurrent CAGRA search with batch size of 1

[jamxia155]

Given the lack of support in Lucene for batched queries, it should still be possible to improve GPU utilization by running search queries from multiple host threads. After a concurrency bug fix in cuVS (#1771), we can now run multiple CAGRA search queries concurrently. However, the impact of any overhead in cuvs-lucene remains to be assessed.

Proposed evaluation grid:

• Workload: 5M vectors, 64 dimesions, FP16
• Hardware: AWS g6e.16xlarge (1xL40S GPU, 64 vCPU, 512 GB RAM)
• Software: cuVS C++ (baseline) vs. cuvs-lucene (test)
• Search config: k = 1000, batch_size = 1, n_threads up to 64
• Recall targets: 80%, 90%, 95%, 99%

[jamxia155]

Initial testing with smaller dataset (1M vectors) resulted in 14 CAGRA index segments being built in cuvs-lucene. This is expected due to the fundamental design differences between Lucene and cuVS.

Profiling of cuvs-lucene search shows 1 invocation of cuVS single-CTA search_kernel per index segment per query, all processed sequentially in the same CUDA stream. The dominant runtime contribution appears to be some CPU processing in between kernel invocations.

cuvs-lucene search implementation reveals the following opportunities for optimization:

1. Currently, each local top-k result from the CagraIndex.search call is added to the global knnCollector sequentially right after each index segment is searched. We can instead write all the local top-k result to an intermediate results buffer and use GPU-accelerated select_k algorithm to compute the global top-k at once. This requires cuVS select_k to be exposed in cuvs-java.

2. Currently, each index segment is searched in the same CUDA stream sequentially. Given the low GPU utilization of the search kernel (only 1 CTA), we can create a pool of CUDA streams and launch these kernels in parallel.

3. Currently, cuvs-lucene is oversampling by a factor of 10 due to lack of pre-filter support for CagraQuery. Adding pre-filter support should remove the need for oversampling.

4. Nsys profile also revealed use of synchronous cudaMalloc/Free APIs. These allocations should be updated to RMM-managed containers or at least the async APIs.

[jamxia155]

Findings (smaller test case for convenience) starting to make sense:

• With single CPU thread, cuVS-Lucene latency is about 20x higher than cuVS.
• Notably, if we simulate a Lucene index segment in cuVS by reducing the CAGRA index size by 14 (the number of index segments that cuVS-Lucene builds), the cuVS latency remains constant. This is the major contributor for the 20x higher latency when cuVS-Lucene searches over 14 index segments instead of one singular index. Lucene/Java overhead also contributes.
• Moreover, cuVS performance scales up to about 8 CPU threads without increasing latency appreciably (scaling believed to be limited by memory system). cuVS-Lucene currently does not scale with CPU multithreading, due to unintended device synchronizations.
• Ongoing optimization effort for cuVS-Lucene exposes parallelism by searching multiple index segments concurrently on separate CUDA streams. This lowers latency to 1.8X of cuVS. However, the exposed parallelism likely saturates the system in a similar way to CPU multithreading so further scaling with CPU multithreading is not observed.

[jamxia155]

Experiments with persistent kernel mode in CAGRA search did not prove to be promising. Since Lucene queries can use arbitrary search parameters, to be completely general we can only expect the cost of persistent kernel initialization and teardown (~12 ms as measured) to be amortized over the number of index segments within each query, typically O(10). The amortized cost is still dominant over actual search latency which is on the order of a half ms.

[jamxia155]

Prototyped a new approach that internally batch the multiple per-segment searches for one query into a single search kernel.

(image credit: @singhmanas1)

In the new approach, for each query, we increase GPU utilization by:

1. Launching a search kernel for each query with as many CTAs as there are index segments
2. Having concurrent CPU threads each servicing a separate query.

The first level of parallelism is limited by the number of index segments, which is typically O(10). The second level of parallelism is limited by CUDA_DEVICE_MAX_CONNECTIONS, which has a maximum value of 32, and by contention for any global resources such as the CUDA device context lock, CPU side mutexes, memory bandwidth, etc. Practical scaling limit observed so far for the second level of parallelism is also O(10). Taken together, we are able to parallelize concurrent queries over hundreds of CTAs, enough to saturate a typical GPU.