[BUG] Cagra index search segfaults for vectors with 1024+ dimensions

[Ngalstyan4]

Describe the bug
Cagra index creation via the CPP api (cagra::build) seems to succeed for vectors of dimension 1025 but search (cagra::search) results in a Segmentation Fault.

The issues exists both in the cuda API and in the runtime API of cagra.

Steps/Code to reproduce bug

Change this line in cagra example to int64_t n_dim = 1025;, compile and run. This should cause a segmentation fault.

To test on the runtime API, use raft::runtime::neighbors::cagra::{build, search} in place of raft::neighbors::cagra::{build, search}.

Expected behavior
Expected the index creation and search to succeed without any issues, as it does for smaller dimensional vectors.

Environment details (please complete the following information):

• Environment location: GCP, Debian 11, Cuda compilation tools, release 12.2, V12.2.140, Tesla V100-SXM2-16GB
• Method of RAFT install: install from source

Additional context
This is important because OpenAI embeddings are 1536 dimensional. There are many other models that produce 1024+ dimensional embeddings.

The fix might be as simple as adding an element to this array but I have not tried it.

[cjnolet]

@Ngalstyan4 CAGRA CUDA kernels are templated on the dataset dimension and we need to define a new template argument for larger dimensions.

Some of these RAG embeddings go up to extremely large numbers of dimensions. Any specific range you'd are targeting or just as high as we can go?

[Ngalstyan4]

1536 dimensional openai embeddings are probably most widely used these days, so having support for that out of the box would be the primary ask.

It would also be great to throw an informative error when the templates are used with unsupported vector sizes instead of segfaulting, as it happens now.

[cjnolet]

It would also be great to throw an informative error when the templates are used with unsupported vector sizes instead of segfaulting, as it happens now.

Totally agree witht his. We're going to prioritize this, and do a PoC to at least enable 1536 in the meantime. We can work our way up to even higher dimensions depending on the perf of enabling 1536. For example, it might take some significant improvements to the CUDA kernels to support 4096 dims, but we'll know once we dive in.

[narangvivek10]

While integrating CAGRA into Lucene, I ran into this issue as well. It seems that dimensions of more than 1024 segfaults while running search. Here is how we encountered this issue.