Iterative CAGRA-Q

[irina-resh-nvda]

Build CAGRA on PQ datasets with Iterative CAGRA-Q

Iterative cagra graph construction using CAGRA-Q search.

This PR improves the iterative CAGRA build method by enabling PQ compression: the dataset is compressed before the iterative search starts, and CAGRA-Q is used to iteratively update the KNN graph.

This is the first time we are introducing building CAGRA on (PQ) quantized datasets directly.

This PR also adds support for C and Python APIs for creating a PQ dataset and building a cagra graph on it. We also add support to update_dataset from dense->PQ or PQ->dense. Serializing the index is also supported.

Example Runnable workflow

import cupy as cp

from cuvs.common import make_device_pq_dataset
from cuvs.neighbors import cagra, hnsw
from cuvs.preprocessing.quantize import pq

dataset = cp.random.random((5000, 64), dtype=cp.float32)
pq_dataset = make_device_pq_dataset(
    pq.PQDatasetParams(pq_dim=32), dataset
)

index = cagra.build(
    cagra.IndexParams(build_algo="iterative_cagra_search"), pq_dataset
)
cagra.save("index.bin", index, include_dataset=False)

loaded_index = cagra.Index()
cagra.load(loaded_index, "index.bin")
cagra.update_dataset(loaded_index, dataset)

queries = dataset[:100]
cagra.search(cagra.SearchParams(), loaded_index, queries, 10)

hnsw_index = hnsw.from_cagra(hnsw.IndexParams(), loaded_index)
hnsw.search(hnsw.SearchParams(), hnsw_index, cp.asnumpy(queries), 10)

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[aamijar]

/ok to test 100db4f

[irina-resh-nvda]

It looks good to me

[aamijar]

/ok to test 19f7397

[aamijar]

Wait - this should be removed. We are doing a D2H of the entire graph. And then if you look at the caller we are doing H2D back again. update_graph member function does a copy to device.

cuvs/cpp/src/neighbors/cagra.cuh

Lines 328 to 331 in df7730d

	auto cagra_graph = detail::iterative_build_graph<T, IdxT>(res, effective_params, dataset);
	index_type idx(res, effective_params.metric);
	idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view()));

cuvs/cpp/include/cuvs/neighbors/cagra.hpp

Lines 704 to 708 in 5191775

	void update_graph(
	raft::resources const& res,
	raft::host_matrix_view<const graph_index_type, int64_t, raft::row_major> knn_graph)
	{
	RAFT_LOG_DEBUG("Copying CAGRA knn graph from host to device");

[aamijar]

Addressed in ffd1a1d

[aamijar]

/ok to test ffd1a1d

[aamijar]

/ok to test abd796c

[aamijar]

/ok to test e218f9e

[tfeher]

Thanks everyone for pushing it through the finish line, the PR looks good to me!

[tfeher]

Nitpick, but we can construct the view first and use in two places

Suggested change

	reconstruct_vpq_queries<T, half, int64_t>(
	res,
	*vpq_queries,
	static_cast<uint64_t>(offset),
	static_cast<uint32_t>(batch_size),
	raft::make_device_matrix_view<T, int64_t>(
	reconstructed_batch_queries->data_handle(), batch_size, query_ld));
	auto batch_query_view = raft::make_device_matrix_view<const T, int64_t>(
	reconstructed_batch_queries->data_handle(), batch_size, query_ld);
	auto batch_query_view = raft::make_device_matrix_view<const T, int64_t>(
	reconstructed_batch_queries->data_handle(), batch_size, query_ld);
	reconstruct_vpq_queries<T, half, int64_t>(
	res,
	*vpq_queries,
	static_cast<uint64_t>(offset),
	static_cast<uint32_t>(batch_size),
	batch_query_view);

Most of Artem's comments were addressed, I have created issue #2610 to track the open question from Artem's review.