Skip to content

Optimize RaBitQ fastscan query setup + 80% qps#5396

Open
lyang24 wants to merge 1 commit into
facebookresearch:mainfrom
lyang24:rabitq-fastscan-lut-rounding-sgemv
Open

Optimize RaBitQ fastscan query setup + 80% qps#5396
lyang24 wants to merge 1 commit into
facebookresearch:mainfrom
lyang24:rabitq-fastscan-lut-rounding-sgemv

Conversation

@lyang24

@lyang24 lyang24 commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

What

Two hot-path fixes for RaBitQ FastScan search, from profiling Cohere 1M (d=1024, nprobe=64):

  1. LUT/query quantization: compute_LUT_uint8 and compute_query_factors call
    std::roundf per element (~330k libm calls per query, 16.9% of cycles in
    __roundf+PLT alone — half-away-from-zero rounding has no x86 instruction, so it
    never inlines). Replaced with an inline x + 0.5f + truncate helper, valid since
    these inputs are non-negative and range-bounded by construction. (recall impact was small enough to be unnoticed)

  2. Single-query pretransform: LinearTransform::apply_noalloc always calls sgemm,
    so for n=1 OpenBLAS repacks the whole rotation matrix every query
    (sgemm_incopy was 12.2% of cycles vs 2.7% for the actual kernel). Added an
    sgemv path for n=1.

Note: rounding is not bit-exact with roundf for inputs within 1-2 ulp of a
k + 0.5 boundary (±1 quantization code); ties agree. gemv also differs from gemm
in accumulation order. Neither affects recall (below), 4/10000 result ids changed.

image

Results

Cohere 1M, IVF_RABITQ_FASTSCAN, k=100, nprobe=64, single thread:

  • QPS 299.5 → 546.4 (median of 5)
  • perf stat: cycles -36.7%, instructions -36.4%, L1D misses -62.2%, IPC unchanged

Recall vs exact ground truth (recomputed by brute force; the HDF5-provided
neighbors are invalid for IP/COSINE): recall@1/10/100 = 0.95 / 0.924 / 0.6554,
identical before and after.

@meta-cla meta-cla Bot added the CLA Signed label Jul 10, 2026
@lyang24 lyang24 marked this pull request as ready for review July 10, 2026 04:28
@lyang24 lyang24 changed the title Optimize RaBitQ fastscan query setup Optimize RaBitQ fastscan query setup + 80% qps Jul 10, 2026
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.

1 participant