Skip to content

Vectorize RaBitQ fast-scan handler adjustment (1 bit, no filter)#5385

Draft
lyang24 wants to merge 2 commits into
facebookresearch:mainfrom
lyang24:rabitq-fastscan-handler-simd
Draft

Vectorize RaBitQ fast-scan handler adjustment (1 bit, no filter)#5385
lyang24 wants to merge 2 commits into
facebookresearch:mainfrom
lyang24:rabitq-fastscan-handler-simd

Conversation

@lyang24

@lyang24 lyang24 commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

This patch adds an AVX512 fast path for the 1-bit IndexIVFRaBitQFastScan result handler.

The PQ4 fast-scan kernel already computes distances for 32 candidates at a time, but the RaBitQ handler was converting those 32 uint16 distances back into a scalar per-candidate loop to apply the 1-bit distance correction and heap threshold check.

For the common 1-bit, non-centered, no-selector, full-block path, this patch keeps the correction step vectorized:

  • converts 16 uint16 distances to float at a time
  • loads/deinterleaves the per-vector RaBitQ aux factors
  • computes adjusted distances with AVX512
  • compares all lanes against the current heap threshold
  • drains only surviving lanes into the scalar heap update path

The scalar fallback is unchanged for multibit, selectors, partial blocks, centered mode, and non-AVX512 builds.

Validation

Tested on AWS Sapphire Rapids with a 1M Cohere IVF_RABITQ_FASTSCAN index.

Config:

  • nprobe=64
  • k=100
  • nq=1000
  • 5 rounds
  • 8 search threads

Results:
baseline 312.87
patched 322.61

This is about a 3% end-to-end QPS improvement with unchanged results.

@meta-cla meta-cla Bot added the CLA Signed label Jul 7, 2026
@lyang24 lyang24 force-pushed the rabitq-fastscan-handler-simd branch from 71c4a2c to 4d6b2b4 Compare July 7, 2026 06:25
@mdouze

mdouze commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Please refer to https://github.com/facebookresearch/faiss/blob/main/faiss/docs/simd_dynamic_dispatch_migration.md on how to integrate SIMD code in Faiss.

@lyang24

lyang24 commented Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

Please refer to https://github.com/facebookresearch/faiss/blob/main/faiss/docs/simd_dynamic_dispatch_migration.md on how to integrate SIMD code in Faiss.

Thanks will read and get back later

@lyang24 lyang24 marked this pull request as draft July 8, 2026 08:06
@lyang24

lyang24 commented Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

Please refer to https://github.com/facebookresearch/faiss/blob/main/faiss/docs/simd_dynamic_dispatch_migration.md on how to integrate SIMD code in Faiss.

fixed but would like to wait for #5396 and then run benchmark again

@lyang24 lyang24 force-pushed the rabitq-fastscan-handler-simd branch from 93c6e94 to 577c483 Compare July 10, 2026 06:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants