From 0236390f2fd8e63ff5e4c6d4ca981b7ba032cd17 Mon Sep 17 00:00:00 2001 From: lyang24 Date: Thu, 16 Jul 2026 19:50:35 -0700 Subject: [PATCH] Optimize RaBitQ fastscan query setup + 80% qps (#5396) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: ## 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. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5396 Reviewed By: mnorris11 Differential Revision: D112153740 Pulled By: alibeklfc fbshipit-source-id: d4eae49d7dd751b542a12c4c1de94c12dc5cbf6f --- faiss/IndexIVFRaBitQFastScan.cpp | 13 +-- faiss/VectorTransform.cpp | 32 ++++++++ faiss/impl/RaBitQUtils.cpp | 4 +- faiss/impl/RaBitQUtils.h | 45 +++++++++++ tests/CMakeLists.txt | 1 + tests/test_rabitq_utils.cpp | 131 +++++++++++++++++++++++++++++++ 6 files changed, 218 insertions(+), 8 deletions(-) create mode 100644 tests/test_rabitq_utils.cpp diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index 02bbdb484..cbea114c8 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/faiss/IndexIVFRaBitQFastScan.cpp @@ -28,6 +28,8 @@ namespace faiss { // Import shared utilities from RaBitQUtils using rabitq_utils::ExtraBitsFactors; using rabitq_utils::QueryFactorsData; +using rabitq_utils::round_nonnegative_to_uint16; +using rabitq_utils::round_nonnegative_to_uint8; using rabitq_utils::SignBitFactors; using rabitq_utils::SignBitFactorsWithError; @@ -603,13 +605,12 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8( float mn = all_mins[j2 * M + m]; uint8_t* out = out_base + j2 * dim12_2 + m * ksub; for (size_t s = 0; s < ksub; s++) { - out[s] = static_cast( - std::roundf(a * (tab[s] - mn))); + out[s] = round_nonnegative_to_uint8(a * (tab[s] - mn)); } } memset(out_base + j2 * dim12_2 + M * ksub, 0, (M2 - M) * ksub); - bq[j2] = static_cast( - std::roundf(a * (probe_b[j2] - glob_b))); + bq[j2] = + round_nonnegative_to_uint16(a * (probe_b[j2] - glob_b)); } normalizers[2 * i] = a; normalizers[2 * i + 1] = glob_b; @@ -851,8 +852,8 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner { for (size_t m = 0; m < M; m++) { const float* tab = lut_float.get() + m * ksub; for (size_t s = 0; s < ksub; s++) { - out[m * ksub + s] = static_cast( - std::roundf(a * (tab[s] - mins[m]))); + out[m * ksub + s] = + round_nonnegative_to_uint8(a * (tab[s] - mins[m])); } } memset(out + M * ksub, 0, (M2 - M) * ksub); diff --git a/faiss/VectorTransform.cpp b/faiss/VectorTransform.cpp index c3c0b9458..a0550db8f 100644 --- a/faiss/VectorTransform.cpp +++ b/faiss/VectorTransform.cpp @@ -47,6 +47,19 @@ int sgemm_( float* c, FINTEGER* ldc); +int sgemv_( + const char* trans, + FINTEGER* m, + FINTEGER* n, + float* alpha, + const float* a, + FINTEGER* lda, + const float* x, + FINTEGER* incx, + float* beta, + float* y, + FINTEGER* incy); + int dgemm_( const char* transa, const char* transb, @@ -192,6 +205,25 @@ void LinearTransform::apply_noalloc(idx_t n, const float* x, float* xt) const { float one = 1; FINTEGER nbiti = d_out, ni = static_cast(n), di = d_in; + if (n == 1) { + FINTEGER onei = 1; + // Avoid GEMM packing overhead for single-vector transforms. GEMV is + // mathematically equivalent but may not be bit-exact with GEMM because + // BLAS implementations can use different accumulation orders. + sgemv_("Transposed", + &di, + &nbiti, + &one, + A.data(), + &di, + x, + &onei, + &c_factor, + xt, + &onei); + return; + } + sgemm_("Transposed", "Not transposed", &nbiti, diff --git a/faiss/impl/RaBitQUtils.cpp b/faiss/impl/RaBitQUtils.cpp index 2d39de69b..8f226cfa1 100644 --- a/faiss/impl/RaBitQUtils.cpp +++ b/faiss/impl/RaBitQUtils.cpp @@ -224,8 +224,8 @@ QueryFactorsData compute_query_factors( uint8_t* rqq = rotated_qq.data(); for (size_t i = 0; i < d; i++) { const float v_q = rq[i]; - const uint8_t v_qq = std::clamp( - std::round((v_q - v_min) * inv_delta), 0, max_code); + const uint8_t v_qq = + round_clamped_to_uint8((v_q - v_min) * inv_delta, max_code); rqq[i] = v_qq; sum_qq += v_qq; diff --git a/faiss/impl/RaBitQUtils.h b/faiss/impl/RaBitQUtils.h index e3cc4b9d3..b008f583a 100644 --- a/faiss/impl/RaBitQUtils.h +++ b/faiss/impl/RaBitQUtils.h @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -86,6 +87,50 @@ struct QueryFactorsData { */ FAISS_API extern const float Z_MAX_BY_QB[8]; +/** Fast half-away rounding for non-negative RaBitQ quantization values. + * + * Contract: x must be finite, non-NaN, and in [0, 255.5). This is not a + * bit-exact replacement for roundf within a few ulps below k + 0.5, where the + * addition may round across the tie. It is intended only for RaBitQ LUT + * quantization paths where a +/- 1 code difference at those boundaries is + * acceptable. + */ +inline uint8_t round_nonnegative_to_uint8(float x) { + assert(x == x); + assert(x >= 0.0f); + assert(x < 255.5f); + return static_cast(static_cast(x + 0.5f)); +} + +/** Same as round_nonnegative_to_uint8 for uint16 RaBitQ bias values. */ +inline uint16_t round_nonnegative_to_uint16(float x) { + assert(x == x); + assert(x >= 0.0f); + assert(x < 65535.5f); + return static_cast(static_cast(x + 0.5f)); +} + +/** Fast clamped rounding for query-byte quantization. + * + * Contract: x must be non-NaN. Values outside [0, max_code] are clamped before + * applying the same non-bit-exact rounding used by + * round_nonnegative_to_uint8(). + */ +inline uint8_t round_clamped_to_uint8(float x, uint8_t max_code) { + assert(x == x); + + if (x <= 0.0f) { + return 0; + } + + const float max_code_f = static_cast(max_code); + if (x >= max_code_f) { + return max_code; + } + + return static_cast(static_cast(x + 0.5f)); +} + /** Compute factors for a single database vector using RaBitQ algorithm. * This function consolidates the mathematical logic that was duplicated * between IndexRaBitQ and IndexRaBitQFastScan. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 282aaf770..4aa2ca40d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -33,6 +33,7 @@ set(FAISS_TEST_SRC test_ivf_early_termination.cpp test_callback.cpp test_utils.cpp + test_rabitq_utils.cpp test_hamming.cpp test_zerocopy.cpp test_scalar_quantizer.cpp diff --git a/tests/test_rabitq_utils.cpp b/tests/test_rabitq_utils.cpp new file mode 100644 index 000000000..69b8f93e5 --- /dev/null +++ b/tests/test_rabitq_utils.cpp @@ -0,0 +1,131 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include + +namespace { + +uint8_t round_clamped_ref(float x, uint8_t max_code) { + float y = std::round(x); + y = std::max(0.0f, std::min(y, static_cast(max_code))); + return static_cast(y); +} + +void check_linear_transform_batch_match(bool have_bias) { + constexpr int d_in = 7; + constexpr int d_out = 5; + faiss::LinearTransform transform(d_in, d_out, have_bias); + transform.is_trained = true; + + transform.A.resize(d_in * d_out); + for (size_t i = 0; i < transform.A.size(); i++) { + transform.A[i] = std::sin(static_cast(i) * 0.37f); + } + + if (have_bias) { + transform.b.resize(d_out); + for (int i = 0; i < d_out; i++) { + transform.b[i] = 0.25f * i - 0.5f; + } + } + + std::vector x(d_in); + for (int i = 0; i < d_in; i++) { + x[i] = std::cos(static_cast(i) * 0.19f); + } + + std::vector one_out(d_out); + transform.apply_noalloc(1, x.data(), one_out.data()); + + std::vector batch_x(4 * d_in); + for (int i = 0; i < 4; i++) { + std::copy(x.begin(), x.end(), batch_x.begin() + i * d_in); + } + + std::vector batch_out(4 * d_out); + transform.apply_noalloc(4, batch_x.data(), batch_out.data()); + + for (int i = 0; i < d_out; i++) { + EXPECT_NEAR(one_out[i], batch_out[i], 1e-5f); + } +} + +} // namespace + +TEST(RaBitQUtils, RoundNonnegativeToUint8) { + using faiss::rabitq_utils::round_nonnegative_to_uint8; + + EXPECT_EQ(round_nonnegative_to_uint8(0.0f), 0); + EXPECT_EQ(round_nonnegative_to_uint8(0.49f), 0); + EXPECT_EQ(round_nonnegative_to_uint8(0.5f), 1); + EXPECT_EQ(round_nonnegative_to_uint8(1.5f), 2); + EXPECT_EQ(round_nonnegative_to_uint8(254.5f), 255); + EXPECT_EQ(round_nonnegative_to_uint8(255.0f), 255); + + const float below_half = std::nextafter(0.5f, 0.0f); + EXPECT_LE( + std::abs( + static_cast(round_nonnegative_to_uint8(below_half)) - + static_cast(std::round(below_half))), + 1); + + std::mt19937 rng(123); + std::uniform_real_distribution dist(0.0f, 255.49f); + for (int i = 0; i < 10000; i++) { + const float x = dist(rng); + EXPECT_LE( + std::abs( + static_cast(round_nonnegative_to_uint8(x)) - + static_cast(std::round(x))), + 1); + } +} + +TEST(RaBitQUtils, RoundNonnegativeToUint16) { + using faiss::rabitq_utils::round_nonnegative_to_uint16; + + EXPECT_EQ(round_nonnegative_to_uint16(0.0f), 0); + EXPECT_EQ(round_nonnegative_to_uint16(0.5f), 1); + EXPECT_EQ(round_nonnegative_to_uint16(65534.5f), 65535); + EXPECT_EQ(round_nonnegative_to_uint16(65535.0f), 65535); +} + +TEST(RaBitQUtils, RoundClampedToUint8) { + using faiss::rabitq_utils::round_clamped_to_uint8; + + EXPECT_EQ(round_clamped_to_uint8(-1.0f, 15), 0); + EXPECT_EQ(round_clamped_to_uint8(0.49f, 15), 0); + EXPECT_EQ(round_clamped_to_uint8(0.5f, 15), 1); + EXPECT_EQ(round_clamped_to_uint8(14.5f, 15), 15); + EXPECT_EQ(round_clamped_to_uint8(15.4f, 15), 15); + EXPECT_EQ(round_clamped_to_uint8(16.0f, 15), 15); + + std::mt19937 rng(456); + std::uniform_real_distribution dist(-4.0f, 260.0f); + for (int i = 0; i < 10000; i++) { + const float x = dist(rng); + EXPECT_LE( + std::abs( + static_cast(round_clamped_to_uint8(x, 255)) - + static_cast(round_clamped_ref(x, 255))), + 1); + } +} + +TEST(LinearTransform, SingleVectorMatchesBatchPath) { + check_linear_transform_batch_match(false); + check_linear_transform_batch_match(true); +}