[Q] Do you use predict in fastkmeans in index lookup

[yichuan-w]

I noticed that you use fastkmeans in index build, but I am not that familiar with the codebase. Do you still use fast-kmeans Triton kernel in the centroid lookup stuff, or do you implement this on your own?

[raphaelsty]

We call fastkmeans once to get the centroids. Once we have the centroids we start building the index but we don't rely on kmeans anymore

[yichuan-w]

Yeah, you can use fast-k-means to do the prediction to find the nearest centroid. What's the reason not to do that?
I am curious about that. Or is the nearest centroid lookup not a bottleneck?

[raphaelsty]

This is a good question and a good idea. I would say we would benefit from having triton kernels

But within the search we are interested in the top-k centroids for each token embeddings unless you set n_ivf_probe to 1 (default is 8). FastKmeans predict method returns a single centroid.

The bottleneck in fast-plaid is the residual decompression and the last scoring step to get the final ranking itself. The main benefit would be a triton kernel like in the original PLAID for the final ranking. Not sure how to properly integrate it with rust but we might find a way

Saw your LEANN project, this is amazing, don't hesitate to ask questions, in order to spot pitfalls, willing to improve fast-plaid :D