bfloat16 (bf16) support in faiss

[mlomeli1]

Many LLMs are trained with bf16, if we want to use the hidden states of LLMs for retrieval, those vectors will be in bf16 dtype. It would be helpful to support bf16 in Faiss so that we can use LLMs as retriever or embedding model.

[mdouze]

Note that we cannot pass through the numpy wrapper because numpy does not support bf16.
Adapting gpu_knn code for pytorch should be easy

faiss/contrib/torch_utils.py

Line 497 in dc55e11

def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRIC_L2, device=-1, use_raft=False):

[mdouze]

• add bfloat16 support here

https://github.com/facebookresearch/faiss/blob/main/faiss/gpu/GpuDistance.h#L76

[mlomeli1]

Example of how we currently would use a PQ codec to encode/decode pytorch bf16 tensors:

torch.from_numpy( codec.sa_decode(codec.sa_encode(x.to(device='cpu', dtype=torch.float32).numpy()) )

this piece of code showcases all the cpu moves + up casting + converting to numpy array. Successively, at decoding, we need to convert back to a tensor. Ideally, we could avoid some of these if this was supported for pytorch tensors.

[alexanderguzhva]

just in case, there is a ScalarQuantizer implementation for bf16, maybe portions of it can be reused