Neighbors IVF RaBitQ | NVIDIA cuVS

Source header: cuvs/neighbors/ivf_rabitq.hpp

IVF-RaBitQ index build parameters

neighbors::ivf_rabitq::index_params

IVF-RaBitQ index build parameters

1 struct index_params : cuvs::neighbors::index_params {
2   uint32_t n_lists;
3   uint32_t bits_per_dim;
4   uint32_t kmeans_n_iters;
5   uint32_t max_train_points_per_cluster;
6   bool fast_quantize_flag;
7   size_t streaming_batch_size;
8   bool force_streaming;
9 };

Fields

Name	Type	Description
`n_lists`	`uint32_t`	The number of inverted lists (clusters) Hint: Increasing this parameter may alleviate shared memory pressure.
`bits_per_dim`	`uint32_t`	The total number of bits per dimension (single bit required for the binary RaBitQ algorithm + additional bits for extended RaBitQ). Supported values: [1, 2, 3, 4, 5, 6, 7, 8, 9]. Hint: the smaller the ‘bits_per_dim’, the smaller the index size and the better the search performance, but the lower the recall.
`kmeans_n_iters`	`uint32_t`	The number of iterations searching for kmeans centers (index building).
`max_train_points_per_cluster`	`uint32_t`	The number of data vectors (per cluster) to use during iterative kmeans building.
`fast_quantize_flag`	`bool`	Flag for using the fast quantize method
`streaming_batch_size`	`size_t`	Maximum number of vectors per batch when using streaming construction. This parameter controls the batch size during streaming construction from host memory. Batches contain complete clusters only (no partial clusters across batch boundaries). Note: Streaming construction is automatically used when the dataset doesn’t fit comfortably in GPU memory (determined by available workspace and kTolerableRatio).
`force_streaming`	`bool`	Force streaming construction regardless of dataset size. When set to true, streaming construction will be used even if the dataset would fit in GPU memory. This is useful for testing or when you want explicit control over the construction method. Note: This parameter only applies when the input dataset is in host memory. If the dataset is already in device memory, streaming construction is not applicable and this parameter has no effect. Default: false (auto-detect based on available memory)

IVF-RaBitQ index search parameters

neighbors::ivf_rabitq::search_mode

A type for specifying the mode for searching the RaBitQ index.

1 enum class search_mode {
2   LUT16 = 0,
3   LUT32 = 1,
4   QUANT4 = 2,
5   QUANT8 = 3
6 };

Values

Name	Value
`LUT16`	`0`
`LUT32`	`1`
`QUANT4`	`2`
`QUANT8`	`3`

IVF-RaBitQ index

neighbors::ivf_rabitq::index

IVF-RaBitQ index.

1 template <typename IdxT>
2 struct index;

neighbors::ivf_rabitq::index::index

Construct an empty index yet to be populated.

1 index(raft::resources const& handle);

Parameters

Name	Direction	Type	Description
`handle`		`raft::resources const&`

Returns

void

Additional overload: neighbors::ivf_rabitq::index::index

Construct an empty index yet to be populated.

1 index(raft::resources const& handle,
2 size_t n_rows,
3 uint32_t dim,
4 uint32_t n_lists,
5 uint32_t bits_per_dim);

Parameters

Name	Direction	Type	Description
`handle`		`raft::resources const&`
`n_rows`		`size_t`
`dim`		`uint32_t`
`n_lists`		`uint32_t`
`bits_per_dim`		`uint32_t`

Returns

void

neighbors::ivf_rabitq::index::dim

Dimensionality of the input data.

1 uint32_t dim() const noexcept;

Returns

uint32_t

neighbors::ivf_rabitq::index::size

Total length of the index.

1 IdxT size() const noexcept;

Returns

IdxT

neighbors::ivf_rabitq::index::rabitq_index

Accessor for underlying RaBitQ index

1 detail::IVFGPU& rabitq_index() noexcept;

Returns

detail::IVFGPU&

IVF-RaBitQ index build

neighbors::ivf_rabitq::build

Build the index from the dataset for efficient search.

1 auto build(raft::resources const& handle,
2 const cuvs::neighbors::ivf_rabitq::index_params& index_params,
3 raft::device_matrix_view<const float, int64_t, raft::row_major> dataset)
4 -> cuvs::neighbors::ivf_rabitq::index<int64_t>;

Usage example:

Parameters

Name	Direction	Type	Description
`handle`	in	`raft::resources const&`
`index_params`	in	`const cuvs::neighbors::ivf_rabitq::index_params&`	configure the index building
`dataset`	in	`raft::device_matrix_view<const float, int64_t, raft::row_major>`	a device_matrix_view to a row-major matrix [n_rows, dim]

Returns

cuvs::neighbors::ivf_rabitq::index<int64_t>

Additional overload: neighbors::ivf_rabitq::build

Build the index from the dataset for efficient search.

1 auto build(raft::resources const& handle,
2 const cuvs::neighbors::ivf_rabitq::index_params& index_params,
3 raft::host_matrix_view<const float, int64_t, raft::row_major> dataset)
4 -> cuvs::neighbors::ivf_rabitq::index<int64_t>;

Usage example:

Parameters

Name	Direction	Type	Description
`handle`	in	`raft::resources const&`
`index_params`	in	`const cuvs::neighbors::ivf_rabitq::index_params&`	configure the index building
`dataset`	in	`raft::host_matrix_view<const float, int64_t, raft::row_major>`	a host_matrix_view to a row-major matrix [n_rows, dim]

Returns

cuvs::neighbors::ivf_rabitq::index<int64_t>

IVF-RaBitQ index search

neighbors::ivf_rabitq::search

Search ANN using the constructed index.

1 void search(raft::resources const& handle,
2 const cuvs::neighbors::ivf_rabitq::search_params& search_params,
3 cuvs::neighbors::ivf_rabitq::index<int64_t>& index,
4 raft::device_matrix_view<const float, int64_t, raft::row_major> queries,
5 raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
6 raft::device_matrix_view<float, int64_t, raft::row_major> distances);

Parameters

Name	Direction	Type	Description
`handle`	in	`raft::resources const&`
`search_params`	in	`const cuvs::neighbors::ivf_rabitq::search_params&`	configure the search
`index`	in	`cuvs::neighbors::ivf_rabitq::index<int64_t>&`	ivf-rabitq constructed index
`queries`	in	`raft::device_matrix_view<const float, int64_t, raft::row_major>`	a device matrix view to a row-major matrix [n_queries, index->dim()]
`neighbors`	out	`raft::device_matrix_view<int64_t, int64_t, raft::row_major>`	a device matrix view to the indices of the neighbors in the source dataset [n_queries, k]
`distances`	out	`raft::device_matrix_view<float, int64_t, raft::row_major>`	a device matrix view to the distances to the selected neighbors [n_queries, k]

Returns

void

IVF-RaBitQ index serialize

neighbors::ivf_rabitq::serialize

Save the index to file.

1 void serialize(raft::resources const& handle,
2 const std::string& filename,
3 cuvs::neighbors::ivf_rabitq::index<int64_t>& index);

Parameters

Name	Direction	Type	Description
`handle`	in	`raft::resources const&`	the raft handle
`filename`	in	`const std::string&`	the file name for saving the index
`index`	in	`cuvs::neighbors::ivf_rabitq::index<int64_t>&`	IVF-RaBitQ index

Returns

void

neighbors::ivf_rabitq::deserialize

Load index from file.

1 void deserialize(raft::resources const& handle,
2 const std::string& filename,
3 cuvs::neighbors::ivf_rabitq::index<int64_t>* index);

Parameters

Name	Direction	Type	Description
`handle`	in	`raft::resources const&`	the raft handle
`filename`	in	`const std::string&`	the name of the file that stores the index
`index`	out	`cuvs::neighbors::ivf_rabitq::index<int64_t>*`	IVF-PQ index

Returns

void