IVF-PQ: low-precision coarse search

[achirkin]

Enable low-precision (half / int8) element type for use in the cuBLAS GEMM performed during coarse search (select clusters to probe). This makes cuBLAS use tensor cores and thus speeds up the coarse search.

Also propagate kMaxQueries compile time constant to a runtime search parameter: this allows to improve GPU utilization in extremely large batch size case, such as using IVF-PQ for constructing a nearest-neighbor graph for the whole dataset.

[tfeher]

Hi Artem, thanks for the PR! Could you add tests for the new options?

[achirkin]

Sure, thanks for pointing this out!

It's worth mentioning the int8 coarse search often gives the garbage recall and that is rather unavoidable. The problem is that we keep cluster norms as a part of cluster vectors and compute GEMM of the whole thing for the L2 case. But the norms are not normalized, so they grow very fast with the number of dimensions, which makes int8 representation impossible. I slightly improved the situation by encoding the norms into several int8 slots, but even that didn't help in many cases.

[tfeher]

Thank you Artem for the PR! It looks good overall, but I have a few questions.

[tfeher]

• Is the normalization requirement documented elsewhere?
• Can't we use our quantization API to set proper normalizaiton constants?
• Can we have larger min_recall by increasing nprobes?

[achirkin]

To make sure the int8_t coarse search works correctly, we need even stricter requirements that all elements are smaller than one. I'm also not sure it makes sense to require L2 normalization (that all norms are smaller than 2m), because that means reduce precision (if both the norm and the components are divided by the same constant).
I think increasing the nprobes won't help a lot, because if the norm is out of range we basically get the random selection.
All in all I doubt the int8_t variant will be useful, but we may reuse the code later by changing it to fp8 (and we can estimate the performance by running int8_t now). Therefore I suppose there's a value in having it as an experimental feature and not invest too much in documentation and testing.

[cjnolet]

Can we add the half/int8 precision directly to balanced kmeans so that we can reuse the solution across other algorithms which use that?

[achirkin]

Hi @cjnolet, the coarse search bits in IVF-PQ are rather not portable as it does cluster search + query type mapping at the same time and relies on IVF-PQ-specific representation (cluster center norms stored alongside the vectors), so unfortunately there's no code to share between the two.

[tfeher]

Thanks Artem for the update! The PR looks good to me.

I am not completely convinced that we need to add the int8 option, but you did add a clear explanation on its limitation, therefore I am fine with it.