[FEA] Binary IVF Flat Index

[tarang-jain]

Depends on NVIDIA/raft#2770

Implementation of binary ivf flat index (bitwise hamming metric for the IVF Flat index)

Key Features

1. Binary Index Structure

• Added binary_centers_ field to store cluster centers as packed uint8_t arrays for binary data
• Index automatically detects BitwiseHamming metric and configures itself for binary operation
• Only support uint8_t inputs with BitwiseHamming and add only single instantiations of newly added kernels

2. K-means Clustering for Binary Data

The clustering approach for binary data required special handling:

• Expanded Space Clustering: Binary data (uint8_t) is expanded to signed representation (int8_t) where each bit becomes ±1

  • 0 → -1, 1 → +1 transformation enables meaningful centroid computation
  • Clustering performed using L2 distance in the expanded dimensional space

• Centroid Quantization: After computing float centroids in expanded space, they are converted back to binary format:

  • Centroids are stored as packed uint8_t arrays
  • KMeans (coarse) prediction is done on these quantized centroids with the BitwiseHamming distance.

3. Distance Kernels

Coarse Search (Cluster Selection)

• Implemented specialized bitwise_hamming_distance_op for query-to-centroid distances in order to compute PairwiseDistances

Fine-Grained Search (Within Clusters)

Extended the interleaved scan kernel (ivf_flat_interleaved_scan.cuh) with specialized templates for BitwiseHamming:

• Veclen-based optimization: Different code paths based on vectorization width

  • Veclen=16,8,4: Load data as uint32_t, use __popc(x ^ y) for 4-byte Hamming distance
  • Veclen=1,2: Byte-wise XOR and population count

• Efficient memory access patterns:

  • Maintains interleaved data layout for coalesced memory access
  • Specialized loadAndComputeDist templates for uint8_t that leverage vectorized loads

cuvs-bench Benchmarks (latency):

[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.

[tarang-jain]

/ok to test 1aa722d

[tarang-jain]

/ok to test 64d103b

[tfeher]

Hi Tarang, I have looked through the PR again. Overall the ivf_flat::search() implementation looks good.

My main concern is that the balanced kmeans method is inefficient because it expands the bits. But I think it is better to treat that in a separate PR.

A few additional issues:

• the benchmark wrapper should be updated to have a binary option for IVF-Flat.
• Ideally we should use the dataset API to represent the binary quantized input.

[tfeher]

We don't train temp_quantizer because the centers are guaranteed to be +/-1?

[tarang-jain]

we must use the same method we used to expand the data, right? temp_quantizer need not be trained since we must use zero as the binary threshold.

[tarang-jain]

Ideally we should use the dataset API to represent the binary quantized input.

This index does not yet use the datasets API, no? So this would mean a much bigger change. Shall we do that in a follow-up for IVF-Flat as a whole?

[tarang-jain]

the benchmark wrapper should be updated to have a binary option for IVF-Flat.

I will do this in a follow-up. I want to keep this PR limited to the index itself since it is easier to review.

[tfeher]

This index does not yet use the datasets API, no? So this would mean a much bigger change. Shall we do that in a follow-up for IVF-Flat as a whole?

I am fine with that. But I would like to hear @achirkin's opinion, because he had some issues in his previous review regarding how the binary quantized data is represented.

[achirkin]

I've dug into the matter and this case turns out to be more interesting than I initially thought. I believe we have two ways to integrate this into the dataset API. In both cases I assume the spec-based dataset struct from #2858

1. Implement a new non-quantized dataset type bitmap_dataset. It's backed by raft::core::bitmap and raft::core::bitmap_view instead of mdarray/mdspan. Then we keep the current semantics of IVF-Flat being unaware of the encoding and computing the exact hamming distance. See my refactoring suggestion bitmap dataset tarang-jain/cuvs#2
2. Implement a new quantized dataset type binary_dataset, similar to SQ/VPQ datasets. It's backed by the uint8_t mdarray/mdspan plus the binary dictionary. This means the index is not flat anymore, rather IVF-Binary. See my refactoring suggestion Binary dataset tarang-jain/cuvs#1 . See how this will look merged onto main [DO NOT MERGE] IVF-Binary #2864

I think both cases make sense, but it's not clear which is better, because I don't know the target use case. Do we want to compute symmetric Hamming distance (1) or have an option to quantize the query during search or compute asymetric distance (similar to IVF-PQ) (2).

[tarang-jain]

@achirkin this PR currently does symmetric distances (so the query is already pre-quantized). But I agree that we can have asymmetric too. The only important piece would be to use the exact same quantizer that was used to quantize the training data.

[achirkin]

Thanks @tarang-jain , I've noticed it does symmetric distance now. I'm asking what is the original use case for this? Because I think the answer to that should determine whether we want to go route 1 or 2 from my message above. Anyway, let me know if you consider any of the two mergeable or your concerns otherwise.

[tarang-jain]

@achirkin I have studied your PRs and I am more inclined to use the new IVF Binary index. It seems much cleaner and simpler and having a new index type for a new quantization regime is along the same lines as IVF-PQ and IVF-RaBitQ. It will also avoid having to keep two getters: binary_centers() and centers() inside IVF Flat. About the search API, currently I would say we can stick to SDC. The user can binary quantize the queries themselves using our binary quantizer. And just using SDC should not deter us from introducing the new index type for binary datasets. Note however that there might be some inconsistency wrt cagra and nn descent which support bitwise hamming metric, but do not explicitly call their datasets "binary".

[achirkin]

/ok to test 54b8d65

[achirkin]

Yes, my main concern about binary-as-quantized-dataset path is that the primary API for all other indices is ADC search, and this one becomes SDC. However, I think we can bolt ADC later for API completeness.
Also couples suggestions below to fix small inconsistencies in the changes that I rushed in in my refactor.

[achirkin]

Would it make sense to rename it to something like storage_dim to make it clear that the member has to do with the underlying bytes rather than binary encoding?

[achirkin]

Shall we keep it IdxT template defaulted to int64_t for symmetry with other indices?