Multi-GPU IVF-Flat

Multi-GPU IVF-Flat#

2026-03-06

8 min read time

Applies to Linux

Multi-GPU IVF-Flat extends the IVF-Flat 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 IVF-Flat 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.ivf_flat.IndexParams#

Bases: cuvs.neighbors.ivf_flat.ivf_flat.IndexParams

IndexParams(distribution_mode=u’sharded’, *, **kwargs)

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

distribution_modestr, default = “sharded”

Distribution mode for multi-GPU setup. Valid values: [“replicated”, “sharded”]

**kwargs : Additional parameters passed to single-GPU IndexParams

get_handle()#

IndexParams.get_handle(self)

Index search parameters#

class cuvs.neighbors.mg.ivf_flat.SearchParams#

Bases: cuvs.neighbors.ivf_flat.ivf_flat.SearchParams

SearchParams(n_probes=1, *, search_mode=u’load_balancer’, merge_mode=u’merge_on_root_rank’, n_rows_per_batch=1000, **kwargs)

Parameters to search multi-GPU IVF-Flat index.

get_handle()#

SearchParams.get_handle(self)

Index#

class cuvs.neighbors.mg.ivf_flat.Index#

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

Index build#

cuvs.neighbors.mg.ivf_flat.build(*args, resources=None, **kwargs)#

build(IndexParams index_params, dataset, resources=None)

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

index_params : cuvs.neighbors.ivf_flat.IndexParams dataset : Array interface compliant matrix shape (n_samples, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU IVF-Flat, 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.

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

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

Index extend#

cuvs.neighbors.mg.ivf_flat.extend(*args, resources=None, **kwargs)#

extend(Index index, new_vectors, new_indices=None, resources=None)

Extend the multi-GPU IVF-Flat index with new vectors.

index : cuvs.neighbors.ivf_flat.Index new_vectors : Array interface compliant matrix shape (n_new_vectors, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU IVF-Flat, new_vectors MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

new_indicesArray interface compliant matrix shape (n_new_vectors,)

, optional If provided, these indices will be used for the new vectors. If not provided, indices will be automatically assigned. IMPORTANT: Must be in host memory (CPU) for multi-GPU IVF-Flat.

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.

>>> import numpy as np
>>> from cuvs.neighbors.mg import ivf_flat
>>> n_samples = 50000
>>> n_features = 50
>>> n_new_vectors = 1000
>>> # For multi-GPU IVF-Flat, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> new_vectors = np.random.random_sample(
...     (n_new_vectors, n_features)).astype(np.float32)
>>> new_indices = np.arange(n_samples, n_new_vectors, dtype=np.int64)
>>> build_params = ivf_flat.IndexParams(metric="sqeuclidean")
>>> index = ivf_flat.build(build_params, dataset)
>>> ivf_flat.extend(index, new_vectors, new_indices)

Index save#

cuvs.neighbors.mg.ivf_flat.save(*args, resources=None, **kwargs)#

save(Index index, filename, resources=None)

Serialize the multi-GPU IVF-Flat index to a file.

index : cuvs.neighbors.ivf_flat.Index filename : str

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.

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

Index load#

cuvs.neighbors.mg.ivf_flat.load(*args, resources=None, **kwargs)#

load(filename, resources=None)

Deserialize the multi-GPU IVF-Flat index from a file.

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.

indexIndex

The deserialized index.

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

Index distribute#

cuvs.neighbors.mg.ivf_flat.distribute(*args, resources=None, **kwargs)#

distribute(filename, resources=None)

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

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.

indexIndex

The distributed index.

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