[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 5e6db22

[copy-pr-bot]

/ok to test 5e6db22

@tarang-jain, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[tarang-jain]

/ok to test 5e6db22

[tfeher]

Thanks Tarang for addressing my questions. The PR looks good to me.

[KyleFromNVIDIA]

JIT+LTO changes look reasonable to me

[tarang-jain]

Notes: after a lot of discussion offline, we have brought back design 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 tarang-jain#1 . See how this will look merged onto main #2864

[tarang-jain]

/ok to test ad67e8b

[copy-pr-bot]

/ok to test ad67e8b

@tarang-jain, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[tarang-jain]

/ok to test ad67e8b

[tarang-jain]

/ok to test 6c393a7

[tarang-jain]

/ok to test 8dae48a

Co-authored parts of PR

[tarang-jain]

/merge