Multi-GPU CAGRA#

Multi-GPU CAGRA extends the graph-based CAGRA algorithm to work across multiple GPUs, providing improved scalability and performance for large-scale vector search. It supports both replicated and sharded distribution modes.

Note

IMPORTANT: Multi-GPU CAGRA requires all data (datasets, queries, output arrays) to be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array) before use.

Index build parameters#

class cuvs.neighbors.mg.cagra.IndexParams(distribution_mode='sharded', *, **kwargs)#

Parameters to build multi-GPU CAGRA index for efficient search. Extends single-GPU IndexParams with multi-GPU specific parameters.

Parameters:

distribution_modestr, default = “sharded”: Distribution mode for multi-GPU setup. Valid values: [“replicated”, “sharded”]
**kwargsAdditional parameters passed to single-GPU IndexParams

Attributes:

distribution_mode

Methods

get_handle(self)

get_handle(self)[source]#

Index search parameters#

class cuvs.neighbors.mg.cagra.SearchParams(

search_mode='load_balancer',

*,

merge_mode='merge_on_root_rank',

n_rows_per_batch=1000,

**kwargs,

)#

Parameters to search multi-GPU CAGRA index.

Attributes:

merge_mode: Get the merge mode for multi-GPU search.
n_rows_per_batch: Get the number of rows per batch for multi-GPU search.
search_mode: Get the search mode for multi-GPU search.

Methods

get_handle(self)

get_handle(self)[source]#

merge_mode#: Get the merge mode for multi-GPU search.

n_rows_per_batch#: Get the number of rows per batch for multi-GPU search.

search_mode#: Get the search mode for multi-GPU search.

Index#

class cuvs.neighbors.mg.cagra.Index#

Multi-GPU CAGRA index object. Stores the trained multi-GPU CAGRA index state which can be used to perform nearest neighbors searches across multiple GPUs.

Attributes:

trained

Index build#

cuvs.neighbors.mg.cagra.build(IndexParams index_params, dataset, resources=None)[source]#

Build the multi-GPU CAGRA index from the dataset for efficient search.

Parameters:

index_paramscuvs.neighbors.cagra.IndexParams
datasetArray interface compliant matrix shape (n_samples, dim): Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU CAGRA, the dataset MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).
resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.: If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

Returns:

index: py:class:cuvs.neighbors.cagra.Index

Examples

>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> distances, neighbors = cagra.search(cagra.SearchParams(),
...                                         index, dataset, k)
>>> # Results are already in host memory (NumPy arrays)

Index search#

cuvs.neighbors.mg.cagra.search( SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None, )[source]#

Search the multi-GPU CAGRA index for the k-nearest neighbors of each query.

Parameters:

search_paramscuvs.neighbors.cagra.SearchParams
indexcuvs.neighbors.cagra.Index
queriesArray interface compliant matrix shape (n_queries, dim): Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU CAGRA, queries MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).
kint: The number of neighbors to search for each query.
neighborsArray interface compliant matrix shape (n_queries, k), optional: If provided, this array will be filled with the indices of the k-nearest neighbors. If not provided, a new host array will be allocated. IMPORTANT: Must be in host memory (CPU) for multi-GPU CAGRA. Expected dtype: int64
distancesArray interface compliant matrix shape (n_queries, k), optional: If provided, this array will be filled with the distances to the k-nearest neighbors. If not provided, a new host array will be allocated. IMPORTANT: Must be in host memory (CPU) for multi-GPU CAGRA.
resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.: If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

Returns:

distancesnumpy.ndarray: The distances to the k-nearest neighbors for each query (in host memory).
neighborsnumpy.ndarray: The indices of the k-nearest neighbors for each query (in host memory).

Examples

>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> queries = np.random.random_sample((n_queries, n_features)).astype(
...     np.float32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> distances, neighbors = cagra.search(cagra.SearchParams(),
...                                         index, queries, k)
>>> # Results are already in host memory (NumPy arrays)

Index save#

cuvs.neighbors.mg.cagra.save(Index index, filename, resources=None)[source]#

Serialize the multi-GPU CAGRA index to a file.

Parameters:

indexcuvs.neighbors.cagra.Index
filenamestr: The filename to serialize the index to.
resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.: If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

Examples

>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> cagra.save(index, "index.bin")

Index load#

cuvs.neighbors.mg.cagra.load(filename, resources=None)[source]#

Deserialize the multi-GPU CAGRA index from a file.

Parameters:

filenamestr: The filename to deserialize the index from.
resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.: If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

Returns:

indexIndex: The deserialized index.

Examples

>>> from cuvs.neighbors.mg import cagra
>>> index = cagra.load("index.bin")

Index distribute#

cuvs.neighbors.mg.cagra.distribute(filename, resources=None)[source]#

Distribute a single-GPU CAGRA index across multiple GPUs from a file.

Parameters:

filenamestr: The filename to distribute the index from.
resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.: If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

Returns:

indexIndex: The distributed index.

Examples

>>> from cuvs.neighbors.mg import cagra
>>> index = cagra.distribute("single_gpu_index.bin")