CAGRA

View as Markdown

CAGRA is a GPU-optimized graph index for approximate nearest-neighbor search. Think of every vector as a point, and think of the graph as a map that connects each point to nearby points. During search, CAGRA follows that map toward better matches instead of checking every vector.

CAGRA works well when you want strong recall, high GPU throughput, and fast graph construction.

Example API Usage

C API | C++ API | Python API | Java API | Rust API | Go API

Building an index

1#include <cuvs/neighbors/cagra.h>
2
3cuvsResources_t res;
4cuvsCagraIndexParams_t index_params;
5cuvsCagraIndex_t index;
6DLManagedTensor *dataset;
7
8// populate tensor with data
9load_dataset(dataset);
10
11cuvsResourcesCreate(&res);
12cuvsCagraIndexParamsCreate(&index_params);
13cuvsCagraIndexCreate(&index);
14
15cuvsCagraBuild(res, index_params, dataset, index);
16
17cuvsCagraIndexDestroy(index);
18cuvsCagraIndexParamsDestroy(index_params);
19cuvsResourcesDestroy(res);

Searching an index

1#include <cuvs/neighbors/cagra.h>
2
3cuvsResources_t res;
4cuvsCagraSearchParams_t search_params;
5cuvsCagraIndex_t index;
6DLManagedTensor *queries;
7DLManagedTensor *neighbors;
8DLManagedTensor *distances;
9
10// populate tensor with data
11load_queries(queries);
12
13cuvsResourcesCreate(&res);
14cuvsCagraSearchParamsCreate(&search_params);
15
16// ... build or load index ...
17cuvsCagraSearch(res, search_params, index, queries, neighbors, distances);
18
19cuvsCagraSearchParamsDestroy(search_params);
20cuvsResourcesDestroy(res);

Saving and loading an index

Serialize a CAGRA index when you want to reuse the graph without rebuilding it. Include the dataset when the loaded index should be searchable immediately; omit it only when your workflow will attach or provide the dataset separately.

Go does not currently expose CAGRA save/load wrappers.

1#include <cuvs/neighbors/cagra.h>
2
3cuvsResources_t res;
4cuvsCagraIndex_t index;
5cuvsCagraIndex_t loaded_index;
6
7cuvsResourcesCreate(&res);
8cuvsCagraIndexCreate(&index);
9cuvsCagraIndexCreate(&loaded_index);
10
11// ... build index ...
12cuvsCagraSerialize(res, "/tmp/cuvs-cagra.bin", index, true);
13cuvsCagraDeserialize(res, "/tmp/cuvs-cagra.bin", loaded_index);
14
15cuvsCagraIndexDestroy(loaded_index);
16cuvsCagraIndexDestroy(index);
17cuvsResourcesDestroy(res);

How CAGRA works

CAGRA builds and searches a nearest-neighbor graph.

First, CAGRA builds an initial kNN graph. This is the first draft of the map: each vector is connected to vectors that look nearby. An exact brute-force build can create a very accurate initial graph, but it is usually too slow. In practice, the first graph does not need to be perfect because CAGRA improves it later. cuVS can build this initial graph with IVF-PQ or NN-Descent.

Second, CAGRA prunes the initial graph. This removes redundant paths and keeps the links that are most useful for search.

At search time, CAGRA starts from one or more graph vertices, follows links to better candidates, and keeps a working set of the best candidates it has seen so far.

When to use CAGRA

Use CAGRA when the index fits in GPU memory and you want fast approximate search.

Use CAGRA when build speed matters. CAGRA can build graphs quickly on the GPU.

Use CAGRA in hybrid environments where a GPU-built graph is converted to HNSW for CPU search.

Use brute-force instead when exact results are required or the dataset is small enough that a full scan is already fast enough.

Interoperability with HNSW

cuVS can convert a CAGRA graph to an HNSW graph. This lets the GPU build the graph while the CPU handles search later. This is useful when GPUs are available for indexing, but production search runs on CPUs.

If the graph is being serialized or converted to HNSW right after build, avoid keeping the dataset attached to the CAGRA index when the binding exposes that option. In C++, for example, set attach_dataset_on_build to false.

These examples cover the bindings that currently expose CAGRA and HNSW interoperability. Go supports CAGRA build and search, but does not currently expose HNSW conversion.

1#include <cuvs/neighbors/cagra.h>
2#include <cuvs/neighbors/hnsw.h>
3
4cuvsResources_t res;
5cuvsCagraIndexParams_t cagra_params;
6cuvsCagraIndex_t cagra_index;
7cuvsHnswIndexParams_t hnsw_params;
8cuvsHnswIndex_t hnsw_index;
9cuvsHnswSearchParams_t hnsw_search_params;
10DLManagedTensor *dataset;
11DLManagedTensor *queries;
12DLManagedTensor *neighbors;
13DLManagedTensor *distances;
14
15int64_t n_rows = 1000000;
16int64_t dim = 128;
17int M = 32;
18int ef_construction = 200;
19
20cuvsResourcesCreate(&res);
21cuvsCagraIndexParamsCreate(&cagra_params);
22cuvsCagraIndexCreate(&cagra_index);
23cuvsHnswIndexParamsCreate(&hnsw_params);
24cuvsHnswIndexCreate(&hnsw_index);
25cuvsHnswSearchParamsCreate(&hnsw_search_params);
26
27load_dataset(dataset);
28load_host_queries(queries);
29allocate_hnsw_outputs(neighbors, distances);
30
31cuvsCagraIndexParamsFromHnswParams(
32 cagra_params,
33 n_rows,
34 dim,
35 M,
36 ef_construction,
37 CUVS_CAGRA_HEURISTIC_SIMILAR_SEARCH_PERFORMANCE,
38 L2Expanded);
39
40hnsw_params->hierarchy = GPU;
41hnsw_search_params->ef = 200;
42hnsw_search_params->num_threads = 0;
43
44cuvsCagraBuild(res, cagra_params, dataset, cagra_index);
45cuvsHnswFromCagra(res, hnsw_params, cagra_index, hnsw_index);
46cuvsHnswSearch(res, hnsw_search_params, hnsw_index, queries, neighbors, distances);
47
48cuvsHnswSearchParamsDestroy(hnsw_search_params);
49cuvsHnswIndexDestroy(hnsw_index);
50cuvsHnswIndexParamsDestroy(hnsw_params);
51cuvsCagraIndexDestroy(cagra_index);
52cuvsCagraIndexParamsDestroy(cagra_params);
53cuvsResourcesDestroy(res);

Using Filters

CAGRA supports filtered search. A filter hides some vectors from the search result, so CAGRA may need to explore more of the graph to find enough valid neighbors.

CAGRA can adjust itopk_size internally based on the filtering rate. To disable this automatic adjustment, set filtering_rate to 0.0.

The examples below use a bitset filter. A bit value of 1 means a vector is allowed; a bit value of 0 means it is filtered out.

1#include <cuvs/neighbors/cagra.h>
2#include <cuvs/neighbors/common.h>
3
4cuvsResources_t res;
5cuvsCagraIndexParams_t index_params;
6cuvsCagraSearchParams_t search_params;
7cuvsCagraIndex_t index;
8DLManagedTensor *dataset;
9DLManagedTensor *queries;
10DLManagedTensor *neighbors;
11DLManagedTensor *distances;
12
13cuvsResourcesCreate(&res);
14cuvsCagraIndexParamsCreate(&index_params);
15cuvsCagraSearchParamsCreate(&search_params);
16cuvsCagraIndexCreate(&index);
17
18// Populate DLPack tensors with dataset, query, and output data.
19load_dataset(dataset);
20load_queries(queries);
21allocate_outputs(neighbors, distances);
22
23cuvsCagraBuild(res, index_params, dataset, index);
24
25// Create a device uint32 bitset with one bit per indexed vector. Bit 1 means
26// allowed; bit 0 means filtered out.
27DLManagedTensor *bitset = make_device_bitset(allowed_indices, n_vectors);
28
29cuvsFilter filter;
30filter.type = BITSET;
31filter.addr = (uintptr_t)bitset;
32
33cuvsCagraSearch(res, search_params, index, queries, neighbors, distances, filter);
34
35cuvsCagraIndexDestroy(index);
36cuvsCagraSearchParamsDestroy(search_params);
37cuvsCagraIndexParamsDestroy(index_params);
38cuvsResourcesDestroy(res);

Configuration parameters

Build parameters

NameDefaultDescription
metricL2Expanded / sqeuclideanDistance metric used to build and search the graph.
metric_arg2.0Extra argument for metrics that need one, such as Minkowski distance.
intermediate_graph_degree128Number of neighbors kept in the initial graph before pruning. Larger values can improve the final graph, but increase build time and memory use.
graph_degree64Number of neighbors kept for each vertex in the final graph. Larger values can improve recall, but use more memory and search work.
compressionNoneOptional vector product quantization parameters. When set, the compressed dataset is attached to the index and attach_dataset_on_build is effectively enabled.
graph_build_paramsstd::monostateParameters for the initial graph builder. The default lets cuVS choose a heuristic; explicit options include IVF-PQ, NN-Descent, ACE, and iterative-search graph build parameters.
guarantee_connectivityFalseUses a degree-constrained minimum spanning tree to guarantee the initial kNN graph is connected. This can improve recall on some datasets.
attach_dataset_on_buildTrueKeeps the dataset attached to the index after build. Set to False when serializing or converting to another graph format right after build.

Search parameters

NameDefaultDescription
max_queries0Maximum number of queries searched concurrently. 0 lets cuVS choose automatically.
itopk_size64Number of intermediate search results kept during search. This must be at least k and is the main search tuning knob.
max_iterations0Maximum number of search iterations. 0 lets cuVS choose automatically.
algoAUTOSearch implementation. Options include SINGLE_CTA, MULTI_CTA, MULTI_KERNEL, and AUTO.
team_size0Number of CUDA threads used to calculate each distance. Valid values are 4, 8, 16, or 32. 0 lets cuVS choose automatically.
search_width1Number of vertices selected as starting points for each search iteration.
min_iterations0Minimum number of search iterations.
thread_block_size0CUDA thread block size. Supported values include 64, 128, 256, 512, and 1024. 0 lets cuVS choose automatically.
hashmap_modeAUTOHash map implementation used during search. Options include HASH, SMALL, and AUTO.
hashmap_min_bitlen0Lower limit for the hash map bit length. 0 lets cuVS choose automatically.
hashmap_max_fill_rate0.5Maximum hash map fill rate. Valid values are greater than 0.1 and less than 0.9.
num_random_samplings1Number of initial random seed-node selection iterations.
rand_xor_mask0x128394Bit mask used for initial random seed-node selection.
persistentFalseUses the persistent search kernel where supported. Currently this applies only to SINGLE_CTA.
persistent_lifetime2.0Seconds before a persistent kernel stops when no requests are received.
persistent_device_usage1.0Fraction of the maximum grid size used by the persistent kernel. Lower values can leave GPU capacity for other work.
filtering_rate-1.0Expected fraction of nodes filtered out during filtered search. Negative values let cuVS estimate it automatically.

Tuning

The three parameters most often tuned are itopk_size, graph_degree, and intermediate_graph_degree.

Start with itopk_size. Increasing it usually improves recall, but lowers throughput because CAGRA keeps more candidates during search.

If search-time tuning is not enough, increase graph_degree. This gives each vertex more links to follow, but uses more memory and search work.

If the final graph quality is still too low, increase intermediate_graph_degree. This gives pruning more choices, but makes build more expensive.

Persistent search can improve throughput in services that run many concurrent CAGRA searches. Instead of launching a new search kernel for each request, cuVS can keep a persistent search kernel resident on the GPU and feed it incoming work. This reduces launch overhead and can help high-volume search services keep the GPU busy.

Enable it with the persistent search parameter. Persistent search currently applies to the SINGLE_CTA search implementation, so set algo to SINGLE_CTA when you want to force this path instead of relying on AUTO.

Use persistent_lifetime to control how long the persistent kernel waits for new work before stopping. Use persistent_device_usage to reserve less than the full GPU for the persistent kernel when the same GPU also needs to run other kernels. Keeping other GPU work active alongside a persistent kernel can be fragile, so tune this setting carefully and validate it under the same concurrency pattern used in production.

Memory footprint

CAGRA memory has two main parts: the dataset and the graph. During build, the dataset must be in GPU memory. After build, the dataset can be detached if it is not needed for search, for example when immediately converting the graph to HNSW.

To keep the formulas readable, this section uses short symbols. All estimates are in bytes. The examples convert bytes to MiB by dividing by 1024 * 1024.

  • N: Number of database vectors, or rows in the dataset being indexed.
  • D: Vector dimension, or number of values in each vector.
  • B: Bytes stored for each vector value. Use 4 for fp32, 2 for fp16, or the byte width of the attached dataset representation.
  • G: Final graph degree. This is the graph_degree build parameter, and each vector keeps G neighbor IDs after pruning.
  • I: Intermediate graph degree. This is the intermediate_graph_degree build parameter, and CAGRA uses this larger graph before pruning down to G.
  • C: Number of IVF-PQ coarse clusters/lists. This is the IVF-PQ n_lists value used by the graph build parameters.
  • R: IVF-PQ training-set ratio. This is train_set_ratio; R = 10 means training uses roughly N / 10 vectors.
  • Q: Query batch size, or number of query vectors processed together.
  • K: Search result count, or the requested k/topk nearest neighbors per query.
  • S_idx: Bytes per graph neighbor ID. This is sizeof(IdxT), usually 4 for int32_t or uint32_t.

The named terms in the formulas are also memory sizes:

  • dataset_size: Device memory used by the attached dataset vectors.
  • graph_size: Host memory used by the CAGRA graph neighbor IDs.
  • *_peak: Temporary peak memory for one build phase. Sequential phases are not added together.
  • query_size: Device memory for the current query batch.
  • result_size: Device memory for neighbor IDs and distances returned for the current query batch.
  • workspace_size: Query and result memory used during search.

Scratch and maximum vectors

Most CAGRA formulas below are linear in N once build parameters are fixed. The named temporary peaks are the main scratch terms for build phases, but real runs can also include allocator padding, CUDA library workspaces, memory-resource pools, and small implementation buffers. Reserve a headroom factor H = 0.20 for IVF-PQ graph builds and H = 0.30 for NN-Descent or iterative-search graph builds. If you can measure a representative smaller run, use:

Hmeasured=observed_peakformula_without_scratchformula_without_scratchH_{\text{measured}} = \frac{\text{observed\_peak} - \text{formula\_without\_scratch}} {\text{formula\_without\_scratch}}

Then set:

Musable=(MfreeMother)(1H)M_{\text{usable}} = (M_{\text{free}} - M_{\text{other}}) \cdot (1 - H)

The capacity variables in this subsection are:

  • M_free: Free memory in the relevant memory space before the operation starts. Use device memory for GPU-resident formulas and host memory for formulas explicitly marked as host memory.
  • M_other: Memory reserved for arrays, memory pools, concurrent work, or application buffers that are not included in the formula.
  • H: Scratch headroom fraction reserved for temporary buffers and allocator overhead.
  • M_usable: Memory budget left for the formula after subtracting M_other and reserving headroom.
  • observed_peak: Peak memory observed during a smaller representative run.
  • formula_without_scratch: Value of the selected peak formula with explicit scratch terms removed and without applying headroom.
  • peak_without_scratch(count): The selected peak formula rewritten as a function of the count being estimated, excluding scratch and headroom. The count is usually N for rows or vectors and B for K-selection batch rows.
  • B_per_row / B_per_vector: Bytes added by one more row or vector in the selected formula. For linear formulas, add the coefficients of the count being estimated after fixed values such as D, K, Q, and L are substituted.
  • B_fixed: Bytes in the selected formula that do not change with the estimated count, such as codebooks, centroids, fixed query batches, capped training buffers, or metadata.
  • N_max / B_max: Estimated largest row, vector, or batch-row count that fits in M_usable.

Choose the build or search formula that matches the operation, remove the explicit scratch/headroom from it, and rewrite it as:

peak_without_scratch(N)=NBper_vector+Bfixed\text{peak\_without\_scratch}(N) = N \cdot B_{\text{per\_vector}} + B_{\text{fixed}}

Then estimate:

Nmax=MusableBfixedBper_vectorN_{\max} = \left\lfloor \frac{M_{\text{usable}} - B_{\text{fixed}}} {B_{\text{per\_vector}}} \right\rfloor

For out-of-core IVF-PQ graph build, Q, C, and R can make several terms fixed or sublinear for a fixed configuration. Solve the full max(...) expression if the largest phase changes as N changes.

Baseline memory after build

The baseline memory footprint after index construction is:

dataset_size (device)=N×D×B\begin{aligned} \text{dataset\_size (device)} &= N \times D \times B \end{aligned} graph_size (host)=N×G×Sidx\begin{aligned} \text{graph\_size (host)} &= N \times G \times S_{\text{idx}} \end{aligned}

The dataset must be in GPU memory during index build, but can be detached afterward if it is not needed for search.

Example (1,000,000 vectors, dim = 1024, fp32, graph_degree = 64, IdxT = int32):

  • dataset_size = 4,096,000,000 B = 3906.25 MB
  • graph_size = 256,000,000 B = 244.14 MB

Build peak memory usage

Index build has two phases: construct an initial kNN graph, then optimize it by pruning redundant paths. These steps run sequentially, so their peak memory use is not additive. The overall peak depends on the configured RMM memory resource.

The initial graph can be built with IVF-PQ, NN-Descent, or the experimental iterative CAGRA-search builder. IVF-PQ can build in batches, which allows CAGRA to train on datasets larger than available GPU memory. The iterative builder requires the aligned dataset to fit in GPU memory because it repeatedly searches the partially built CAGRA graph.

Initial graph build using IVF-PQ

IVF-PQ builds the initial graph in two stages. First, it trains cluster centroids and PQ codebooks. Then it queries the IVF-PQ index in batches to form approximate nearest-neighbor lists.

IVF-PQ build peak:

Here, N / R is the IVF-PQ training sample size. The 4 byte factors are fp32 values for training vectors and cluster centroids. The uint32_t term stores one 32-bit ID per training vector.

IVFPQ_build_peak=NR×D×4+C×D×4+NR×sizeof(uint32_t)\begin{aligned} \text{IVFPQ\_build\_peak} &= \frac{N}{R} \times D \times 4 \\ &\quad + C \times D \times 4 \\ &\quad + \frac{N}{R} \times \operatorname{sizeof}(\mathrm{uint32\_t}) \end{aligned}

Example (N = 1e6, D = 1024, C = 1024, R = 10): 395.01 MB

IVF-PQ search peak:

Here, Q is the number of vectors in one search batch and I is the number of candidates kept per query while building the intermediate graph. The three terms estimate query vectors, candidate IDs, and candidate distances.

IVFPQ_search_peak=Q×D×4+Q×I×sizeof(uint32_t)+Q×I×4\begin{aligned} \text{IVFPQ\_search\_peak} &= Q \times D \times 4 \\ &\quad + Q \times I \times \operatorname{sizeof}(\mathrm{uint32\_t}) \\ &\quad + Q \times I \times 4 \end{aligned}

Example (Q = 1024, D = 1024, I = 128): 5.00 MB

Initial graph build using NN-Descent

Peak device memory:

The constants in the NN-Descent formulas are per-vector workspace estimates from the implementation. They are added to the vector storage terms before multiplying by N.

NND_device_peak=N×(D×2+276)\begin{aligned} \text{NND\_device\_peak} &= N \times (D \times 2 + 276) \end{aligned}
  • Data vectors are transferred to device and stored as fp16: D * 2 bytes per vector.
  • The small working graph, locks, and edge counters use 276 bytes per vector.
  • L2 metric adds 4 bytes per vector for precomputed norms.

Peak host memory:

NND_host_peak=N×(13×I+912)\begin{aligned} \text{NND\_host\_peak} &= N \times (13 \times I + 912) \end{aligned}
  • Full graph with distances: 1.3 * 8 * I bytes per vector.
  • Bloom filter for sampling: 1.3 * 2 * I bytes per vector.
  • 5 sample buffers with degree 32: 640 bytes per vector.
  • Graph update buffer with degree 32: 256 bytes per vector.
  • Edge counters: 16 bytes per vector.

The iterative builder starts with a small connected graph, then repeatedly uses CAGRA search to find neighbors for a larger prefix of the dataset. After each search pass, it optimizes the graph and doubles the active graph size until all rows are included.

This path is useful when the metric or data type is better served by CAGRA search itself, but it is not an out-of-core builder. The dataset is copied or aligned into GPU memory before the first iteration.

Variables used only in this subsection:

  • D_align: Aligned device stride used by CAGRA search. Use D when no padding is required.
  • Q_iter: Maximum query chunk size used by the iterative builder. The implementation currently uses min(N, 8192).
  • K_iter: Number of temporary neighbors kept per query during the last pass. Use I + 1.
  • G_iter: Largest graph degree used by the temporary searchable graph. Use G; early iterations use a smaller degree and the final iterations use G.
  • D_iter: Aligned device dataset memory.
  • G_tmp: Largest temporary device graph memory.
  • Q_tile: Query tile memory for one search chunk.
  • R_tile: Result tile memory for one search chunk.
  • W_iter: Temporary device workspace used by one iterative search pass.
  • H_iter: Host neighbors-list capacity in bytes after rounding up to a 2 MiB boundary. One MiB is 1024 * 1024 bytes.

The aligned device dataset is:

Diter=N×Dalign×B\begin{aligned} D_{\text{iter}} &= N \times D_{\text{align}} \times B \end{aligned}

The largest temporary device graph used during the search pass is:

Gtmp=N×Giter×Sidx\begin{aligned} G_{\text{tmp}} &= N \times G_{\text{iter}} \times S_{\text{idx}} \end{aligned}

Each search chunk needs query storage plus temporary neighbor IDs and distances:

Qtile=Qiter×Dalign×BRtile=Qiter×Kiter×(Sidx+4)\begin{aligned} Q_{\text{tile}} &= Q_{\text{iter}} \times D_{\text{align}} \times B \\ R_{\text{tile}} &= Q_{\text{iter}} \times K_{\text{iter}} \times (S_{\text{idx}} + 4) \end{aligned}

The host neighbors list stores the temporary neighbor candidates for all rows:

Hiter=round_up(N×Kiter×Sidx,2 MiB)\begin{aligned} H_{\text{iter}} &= \operatorname{round\_up} \big( N \times K_{\text{iter}} \times S_{\text{idx}}, 2\ \text{MiB} \big) \end{aligned}

The temporary device workspace for one search pass is:

Witer=Gtmp+Qtile+Rtile\begin{aligned} W_{\text{iter}} &= G_{\text{tmp}} \\ &\quad + Q_{\text{tile}} \\ &\quad + R_{\text{tile}} \end{aligned}

The practical device peak for the iterative graph build is:

iterative_device_peakDiter+max ⁣(Witer,optimize_peak)\begin{aligned} \text{iterative\_device\_peak} &\approx D_{\text{iter}} \\ &\quad + \max\!\big( W_{\text{iter}}, \text{optimize\_peak} \big) \end{aligned}

The practical host peak is:

iterative_host_peakHiter+N×G×Sidx\begin{aligned} \text{iterative\_host\_peak} &\approx H_{\text{iter}} + N \times G \times S_{\text{idx}} \end{aligned}

The final N * G * S_idx term is the host graph that remains after build. Check device and host memory separately. The usable N is the smaller value allowed by iterative_device_peak and iterative_host_peak.

Optimize phase

The optimize phase prunes and reorders the intermediate graph. Its peak memory scales linearly with the intermediate degree:

In this formula, the 4 byte term is per-vector bookkeeping. The (S_idx + 1) * I term stores I candidate neighbor IDs plus one byte of pruning state per candidate.

optimize_peak=N×(4+(Sidx+1)×I)\begin{aligned} \text{optimize\_peak} &= N \times \Big( 4 + (S_{\text{idx}} + 1) \times I \Big) \end{aligned}

Example (N = 1e6, I = 128, IdxT = int32): 614.17 MB

Out-of-core CAGRA build consists of IVF-PQ build, IVF-PQ search, and CAGRA optimization. These steps are sequential, so their temporary memory peaks are not added together.

Overall build peak memory usage

The overall device peak is the dataset size plus the largest temporary allocation from the sequential build steps.

Using IVF-PQ:

build_peak=dataset_size+max ⁣(IVFPQ_build_peak,IVFPQ_search_peak,optimize_peak)\begin{aligned} \text{build\_peak} &= \text{dataset\_size} \\ &\quad + \max\!\big( \text{IVFPQ\_build\_peak}, \\ &\qquad\qquad \text{IVFPQ\_search\_peak}, \\ &\qquad\qquad \text{optimize\_peak} \big) \end{aligned}

Example: 3906.25 + max(395.01, 5.00, 614.17) = 4520.42 MB

Using NN-Descent:

build_peak=dataset_size+max ⁣(NND_device_peak,optimize_peak)\begin{aligned} \text{build\_peak} &= \text{dataset\_size}^{*} \\ &\quad + \max\!\big( \text{NND\_device\_peak}, \\ &\qquad\qquad \text{optimize\_peak} \big) \end{aligned}

dataset_size* applies only when the user passes data that is already in device memory. NN-Descent internally copies the dataset to the device as fp16, so host-memory inputs do not add this term.

Using iterative CAGRA search:

Use iterative_device_peak for device memory and iterative_host_peak for host memory. These estimates already include the aligned dataset, temporary search chunks, temporary graph storage, optimization workspace, and final host graph.

Search peak memory usage

CAGRA search requires the dataset and graph to already be resident in GPU memory. When using CAGRA-Q, the original dataset can reside in host memory instead. Search also needs temporary workspace for query vectors and results.

If multiple batches run concurrently or overlap, each batch needs separate result buffers. The estimate below assumes one query batch at a time and reused buffers.

search_memory=dataset_size+graph_size+workspace_size\begin{aligned} \text{search\_memory} &= \text{dataset\_size} \\ &\quad + \text{graph\_size} \\ &\quad + \text{workspace\_size} \end{aligned}

The workspace contains query vectors and result storage:

In the query formula, sizeof(float) is 4 bytes because CAGRA search uses fp32 query storage here. In the result formula, each returned neighbor stores one graph ID of size S_idx and one fp32 distance.

query_size=Q×D×sizeof(float)\begin{aligned} \text{query\_size} &= Q \times D \times \operatorname{sizeof}(\mathrm{float}) \end{aligned} result_size=Q×K×(Sidx+sizeof(float))\begin{aligned} \text{result\_size} &= Q \times K \\ &\quad \times \big(S_{\text{idx}} + \operatorname{sizeof}(\mathrm{float})\big) \end{aligned} workspace_size=query_size+result_size\begin{aligned} \text{workspace\_size} &= \text{query\_size} + \text{result\_size} \end{aligned}

Example (D = 1024, Q = 100, K = 10, IdxT = int32):

  • query_size = 409,600 B = 0.39 MB
  • result_size = 8,000 B = 0.0076 MB
  • workspace_size = query_size + result_size = 0.40 MB
  • total search memory ~= 3906.25 + 244.14 + 0.40 = 4150.79 MB