faiss 0.5.0 → 0.5.1

data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/RaBitQUtils.h>
+#include <faiss/impl/RaBitQuantizerMultiBit.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/rabitq_simd.h>
 #include <algorithm>
@@ -20,15 +21,47 @@
 namespace faiss {
 
 // Import shared utilities from RaBitQUtils
-using rabitq_utils::FactorsData;
+using rabitq_utils::ExtraBitsFactors;
 using rabitq_utils::QueryFactorsData;
-
-static size_t get_code_size(const size_t d) {
-    return (d + 7) / 8 + sizeof(FactorsData);
+using rabitq_utils::SignBitFactors;
+using rabitq_utils::SignBitFactorsWithError;
+
+RaBitQuantizer::RaBitQuantizer(size_t d, MetricType metric, size_t nb_bits)
+        : Quantizer(d, 0), // code_size will be set below
+          metric_type{metric},
+          nb_bits{nb_bits} {
+    // Validate nb_bits range
+    FAISS_THROW_IF_NOT(nb_bits >= 1 && nb_bits <= 9);
+
+    // Set code_size using compute_code_size
+    code_size = compute_code_size(d, nb_bits);
 }
 
-RaBitQuantizer::RaBitQuantizer(size_t d, MetricType metric)
-        : Quantizer(d, get_code_size(d)), metric_type{metric} {}
+size_t RaBitQuantizer::compute_code_size(size_t d, size_t num_bits) const {
+    // Validate inputs
+    FAISS_THROW_IF_NOT(num_bits >= 1 && num_bits <= 9);
+
+    size_t ex_bits = num_bits - 1;
+
+    // Base: 1-bit codes + base factors
+    // Layout for 1-bit: [binary_code: (d+7)/8 bytes][SignBitFactors: 8 bytes]
+    //   base_factors = or_minus_c_l2sqr (4) + dp_multiplier (4)
+    // Layout for multi-bit: [binary_code: (d+7)/8
+    // bytes][SignBitFactorsWithError: 12 bytes]
+    //   factors = or_minus_c_l2sqr (4) + dp_multiplier (4) + f_error (4)
+    size_t base_size = (d + 7) / 8 +
+            (ex_bits == 0 ? sizeof(SignBitFactors)
+                          : sizeof(SignBitFactorsWithError));
+
+    // Extra: ex-bit codes + ex factors (only if ex_bits > 0)
+    // Layout: [ex_code: (d*ex_bits+7)/8 bytes][ex_factors: 8 bytes]
+    size_t ex_size = 0;
+    if (ex_bits > 0) {
+        ex_size = (d * ex_bits + 7) / 8 + sizeof(ExtraBitsFactors);
+    }
+
+    return base_size + ex_size;
+}
 
 void RaBitQuantizer::train(size_t n, const float* x) {
     // does nothing
@@ -54,23 +87,49 @@ void RaBitQuantizer::compute_codes_core(
         return;
     }
 
-    // compute codes
+    const size_t ex_bits = nb_bits - 1;
+
+    // Compute codes
 #pragma omp parallel for if (n > 1000)
     for (int64_t i = 0; i < n; i++) {
-        // the code
+        // Pointer to this vector's code
         uint8_t* code = codes + i * code_size;
-        FactorsData* fac = reinterpret_cast<FactorsData*>(code + (d + 7) / 8);
 
-        // cleanup it
-        if (code != nullptr) {
-            memset(code, 0, code_size);
-        }
+        // Clear code memory
+        memset(code, 0, code_size);
 
         const float* x_row = x + i * d;
 
+        // Pointer arithmetic for code layout:
+        // For 1-bit: [binary_code: (d+7)/8 bytes][SignBitFactors: 8 bytes]
+        // For multi-bit: [binary_code: (d+7)/8 bytes][SignBitFactorsWithError:
+        // 12 bytes]
+        //                [ex_code: (d*ex_bits+7)/8 bytes][ex_factors: 8 bytes]
+        uint8_t* binary_code = code;
+
+        // Step 1: Compute 1-bit quantization and base factors
+        // Store residual for potential ex-bits quantization
+        std::vector<float> residual(d);
+
         // Use shared utilities for computing factors
-        *fac = rabitq_utils::compute_vector_factors(
-                x_row, d, centroid_in, metric_type);
+        SignBitFactorsWithError factors_data =
+                rabitq_utils::compute_vector_factors(
+                        x_row, d, centroid_in, metric_type, ex_bits > 0);
+
+        // Write appropriate factors based on nb_bits
+        if (ex_bits == 0) {
+            // For 1-bit: write only SignBitFactors (8 bytes)
+            SignBitFactors* base_factors =
+                    reinterpret_cast<SignBitFactors*>(code + (d + 7) / 8);
+            base_factors->or_minus_c_l2sqr = factors_data.or_minus_c_l2sqr;
+            base_factors->dp_multiplier = factors_data.dp_multiplier;
+        } else {
+            // For multi-bit: write full SignBitFactorsWithError (12 bytes)
+            SignBitFactorsWithError* full_factors =
+                    reinterpret_cast<SignBitFactorsWithError*>(
+                            code + (d + 7) / 8);
+            *full_factors = factors_data;
+        }
 
         // Pack bits into standard RaBitQ format
         for (size_t j = 0; j < d; j++) {
@@ -78,13 +137,35 @@ void RaBitQuantizer::compute_codes_core(
             const float centroid_val =
                     (centroid_in == nullptr) ? 0.0f : centroid_in[j];
             const float or_minus_c = x_val - centroid_val;
+            residual[j] = or_minus_c;
+
             const bool xb = (or_minus_c > 0.0f);
 
-            // store the output data
-            if (code != nullptr && xb) {
-                rabitq_utils::set_bit_standard(code, j);
+            // Store the 1-bit sign code
+            if (xb) {
+                rabitq_utils::set_bit_standard(binary_code, j);
             }
         }
+
+        // Step 2: Compute ex-bits quantization (if nb_bits > 1)
+        if (ex_bits > 0) {
+            // Pointer to ex-bit code section
+            uint8_t* ex_code =
+                    code + (d + 7) / 8 + sizeof(SignBitFactorsWithError);
+            // Pointer to ex-factors section
+            ExtraBitsFactors* ex_factors = reinterpret_cast<ExtraBitsFactors*>(
+                    ex_code + (d * ex_bits + 7) / 8);
+
+            // Quantize residual to ex-bits (pass centroid for IP metric)
+            rabitq_multibit::quantize_ex_bits(
+                    residual.data(),
+                    d,
+                    nb_bits,
+                    ex_code,
+                    *ex_factors,
+                    metric_type,
+                    centroid_in);
+        }
     }
 }
 
@@ -101,6 +182,7 @@ void RaBitQuantizer::decode_core(
     FAISS_ASSERT(x != nullptr);
 
     const float inv_d_sqrt = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
+    const size_t ex_bits = nb_bits - 1;
 
 #pragma omp parallel for if (n > 1000)
     for (int64_t i = 0; i < n; i++) {
@@ -108,10 +190,19 @@ void RaBitQuantizer::decode_core(
 
         // split the code into parts
         const uint8_t* binary_data = code;
-        const FactorsData* fac =
-                reinterpret_cast<const FactorsData*>(code + (d + 7) / 8);
 
+        // Cast to appropriate type based on nb_bits
+        // For 1-bit: use SignBitFactors (8 bytes)
+        // For multi-bit: use SignBitFactorsWithError (12 bytes, but only first
+        // 8 bytes used for decode)
+        const SignBitFactors* fac = (ex_bits == 0)
+                ? reinterpret_cast<const SignBitFactors*>(code + (d + 7) / 8)
+                : reinterpret_cast<const SignBitFactorsWithError*>(
+                          code + (d + 7) / 8);
+
+        // this is the baseline code
         //
+        // compute <q,o> using floats
         for (size_t j = 0; j < d; j++) {
             // extract i-th bit
             const uint8_t masker = (1 << (j % 8));
@@ -124,51 +215,69 @@ void RaBitQuantizer::decode_core(
     }
 }
 
-struct RaBitDistanceComputer : FlatCodesDistanceComputer {
-    // dimensionality
-    size_t d = 0;
-    // a centroid to use
-    const float* centroid = nullptr;
+// Implementation of RaBitQDistanceComputer (declared in header)
 
-    // the metric
-    MetricType metric_type = MetricType::METRIC_L2;
+float RaBitQDistanceComputer::lower_bound_distance(const uint8_t* code) {
+    FAISS_ASSERT(code != nullptr);
 
-    RaBitDistanceComputer();
+    // Compute estimated distance using 1-bit codes
+    float est_distance = distance_to_code_1bit(code);
 
-    float symmetric_dis(idx_t i, idx_t j) override;
-};
+    // Extract f_error from the code
+    size_t size = (d + 7) / 8;
+    const SignBitFactorsWithError* base_fac =
+            reinterpret_cast<const SignBitFactorsWithError*>(code + size);
+    float f_error = base_fac->f_error;
 
-RaBitDistanceComputer::RaBitDistanceComputer() = default;
+    // Compute proper lower bound using RaBitQ error formula:
+    // lower_bound = est_distance - f_error * g_error
+    // This guarantees: lower_bound ≤ true_distance
+    float lower_bound = est_distance - (f_error * g_error);
 
-float RaBitDistanceComputer::symmetric_dis(idx_t i, idx_t j) {
-    FAISS_THROW_MSG("Not implemented");
+    // Distance cannot be negative
+    return std::max(0.0f, lower_bound);
 }
 
-struct RaBitDistanceComputerNotQ : RaBitDistanceComputer {
+namespace {
+
+struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
     // the rotated query (qr - c)
     std::vector<float> rotated_q;
     // some additional numbers for the query
     QueryFactorsData query_fac;
 
-    RaBitDistanceComputerNotQ();
+    RaBitQDistanceComputerNotQ();
 
-    float distance_to_code(const uint8_t* code) override;
+    // Compute distance using only 1-bit codes (fast)
+    float distance_to_code_1bit(const uint8_t* code) override;
+
+    // Compute full distance using 1-bit + ex-bits (accurate)
+    float distance_to_code_full(const uint8_t* code) override;
 
     void set_query(const float* x) override;
 };
 
-RaBitDistanceComputerNotQ::RaBitDistanceComputerNotQ() = default;
+RaBitQDistanceComputerNotQ::RaBitQDistanceComputerNotQ() = default;
 
-float RaBitDistanceComputerNotQ::distance_to_code(const uint8_t* code) {
+float RaBitQDistanceComputerNotQ::distance_to_code_1bit(const uint8_t* code) {
     FAISS_ASSERT(code != nullptr);
     FAISS_ASSERT(
             (metric_type == MetricType::METRIC_L2 ||
              metric_type == MetricType::METRIC_INNER_PRODUCT));
+    FAISS_ASSERT(rotated_q.size() == d);
 
     // split the code into parts
     const uint8_t* binary_data = code;
-    const FactorsData* fac =
-            reinterpret_cast<const FactorsData*>(code + (d + 7) / 8);
+
+    // Cast to appropriate type based on nb_bits
+    // For 1-bit: use SignBitFactors (8 bytes)
+    // For multi-bit: use SignBitFactorsWithError (12 bytes) which includes
+    // f_error
+    size_t ex_bits = nb_bits - 1;
+    const SignBitFactors* base_fac = (ex_bits == 0)
+            ? reinterpret_cast<const SignBitFactors*>(code + (d + 7) / 8)
+            : reinterpret_cast<const SignBitFactorsWithError*>(
+                      code + (d + 7) / 8);
 
     // this is the baseline code
     //
@@ -177,48 +286,70 @@ float RaBitDistanceComputerNotQ::distance_to_code(const uint8_t* code) {
     // It was a willful decision (after the discussion) to not to pre-cache
     //   the sum of all bits, just in order to reduce the overhead per vector.
     uint64_t sum_q = 0;
-    for (size_t i = 0; i < d; i++) {
-        // extract i-th bit
-        const uint8_t masker = (1 << (i % 8));
-        const bool b_bit = ((binary_data[i / 8] & masker) == masker);
 
+    for (size_t i = 0; i < d; i++) {
+        // Extract i-th bit
+        bool bit = rabitq_utils::extract_bit_standard(binary_data, i);
         // accumulate dp
-        dot_qo += (b_bit) ? rotated_q[i] : 0;
+        dot_qo += bit ? rotated_q[i] : 0;
         // accumulate sum-of-bits
-        sum_q += (b_bit) ? 1 : 0;
+        sum_q += bit ? 1 : 0;
     }
 
-    float final_dot = 0;
-    // dot-product itself
-    final_dot += query_fac.c1 * dot_qo;
-    // normalizer coefficients
-    final_dot += query_fac.c2 * sum_q;
-    // normalizer coefficients
-    final_dot -= query_fac.c34;
-
-    // this is ||or - c||^2 - (IP ? ||or||^2 : 0)
-    const float or_c_l2sqr = fac->or_minus_c_l2sqr;
+    // Apply query factors
+    float final_dot =
+            query_fac.c1 * dot_qo + query_fac.c2 * sum_q - query_fac.c34;
 
     // pre_dist = ||or - c||^2 + ||qr - c||^2 -
     //     2 * ||or - c|| * ||qr - c|| * <q,o> - (IP ? ||or||^2 : 0)
-    const float pre_dist = or_c_l2sqr + query_fac.qr_to_c_L2sqr -
-            2 * fac->dp_multiplier * final_dot;
+    float pre_dist = base_fac->or_minus_c_l2sqr + query_fac.qr_to_c_L2sqr -
+            2 * base_fac->dp_multiplier * final_dot;
 
     if (metric_type == MetricType::METRIC_L2) {
         // ||or - q||^ 2
         return pre_dist;
     } else {
         // metric == MetricType::METRIC_INNER_PRODUCT
+        return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr);
+    }
+}
 
-        // this is ||q||^2
-        const float query_norm_sqr = query_fac.qr_norm_L2sqr;
+float RaBitQDistanceComputerNotQ::distance_to_code_full(const uint8_t* code) {
+    FAISS_ASSERT(code != nullptr);
+    FAISS_ASSERT(
+            (metric_type == MetricType::METRIC_L2 ||
+             metric_type == MetricType::METRIC_INNER_PRODUCT));
+    FAISS_ASSERT(rotated_q.size() == d);
 
-        // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
-        return -0.5f * (pre_dist - query_norm_sqr);
+    size_t ex_bits = nb_bits - 1;
+
+    if (ex_bits == 0) {
+        // No ex-bits, just return 1-bit distance
+        return distance_to_code_1bit(code);
     }
+
+    // Extract pointers to code sections
+    const uint8_t* binary_data = code;
+    size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError);
+    const uint8_t* ex_code = code + offset;
+    const ExtraBitsFactors* ex_fac = reinterpret_cast<const ExtraBitsFactors*>(
+            ex_code + (d * ex_bits + 7) / 8);
+
+    // Call shared utility directly with rotated_q pointer
+    return rabitq_utils::compute_full_multibit_distance(
+            binary_data,
+            ex_code,
+            *ex_fac,
+            rotated_q.data(),
+            query_fac.qr_to_c_L2sqr,
+            query_fac.qr_norm_L2sqr,
+            d,
+            ex_bits,
+            metric_type);
 }
 
-void RaBitDistanceComputerNotQ::set_query(const float* x) {
+void RaBitQDistanceComputerNotQ::set_query(const float* x) {
+    q = x;
     FAISS_ASSERT(x != nullptr);
     FAISS_ASSERT(
             (metric_type == MetricType::METRIC_L2 ||
@@ -237,6 +368,10 @@ void RaBitDistanceComputerNotQ::set_query(const float* x) {
         rotated_q[i] = x[i] - ((centroid == nullptr) ? 0 : centroid[i]);
     }
 
+    // Compute g_error (query norm for lower bound computation)
+    // g_error = ||qr - c|| (L2 norm of rotated query)
+    g_error = std::sqrt(query_fac.qr_to_c_L2sqr);
+
     // compute some numbers
     const float inv_d = (d == 0) ? 1.0f : (1.0f / std::sqrt((float)d));
 
@@ -257,8 +392,10 @@ void RaBitDistanceComputerNotQ::set_query(const float* x) {
 }
 
 //
-struct RaBitDistanceComputerQ : RaBitDistanceComputer {
+struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
     // the rotated and quantized query (qr - c)
+    std::vector<float> rotated_q;
+    // the rotated and quantized query (qr - c) for fast 1-bit computation
     std::vector<uint8_t> rotated_qq;
     // we're using the proposed relayout-ed scheme from 3.3 that allows
     //    using popcounts for computing the distance.
@@ -272,16 +409,20 @@ struct RaBitDistanceComputerQ : RaBitDistanceComputer {
     // the smallest value divisible by 8 that is not smaller than dim
     size_t popcount_aligned_dim = 0;
 
-    RaBitDistanceComputerQ();
+    RaBitQDistanceComputerQ();
 
-    float distance_to_code(const uint8_t* code) override;
+    // Compute distance using only 1-bit codes (fast)
+    float distance_to_code_1bit(const uint8_t* code) override;
+
+    // Compute full distance using 1-bit + ex-bits (accurate)
+    float distance_to_code_full(const uint8_t* code) override;
 
     void set_query(const float* x) override;
 };
 
-RaBitDistanceComputerQ::RaBitDistanceComputerQ() = default;
+RaBitQDistanceComputerQ::RaBitQDistanceComputerQ() = default;
 
-float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
+float RaBitQDistanceComputerQ::distance_to_code_1bit(const uint8_t* code) {
     FAISS_ASSERT(code != nullptr);
     FAISS_ASSERT(
             (metric_type == MetricType::METRIC_L2 ||
@@ -290,21 +431,28 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
     // split the code into parts
     size_t size = (d + 7) / 8;
     const uint8_t* binary_data = code;
-    const FactorsData* fac = reinterpret_cast<const FactorsData*>(code + size);
+
+    // Cast to appropriate type based on nb_bits
+    // For 1-bit: use SignBitFactors (8 bytes)
+    // For multi-bit: use SignBitFactorsWithError (12 bytes) which includes
+    // f_error
+    size_t ex_bits = nb_bits - 1;
+    const SignBitFactors* base_fac = (ex_bits == 0)
+            ? reinterpret_cast<const SignBitFactors*>(code + size)
+            : reinterpret_cast<const SignBitFactorsWithError*>(code + size);
 
     // this is ||or - c||^2 - (IP ? ||or||^2 : 0)
     float final_dot = 0;
     if (centered) {
         int64_t int_dot = ((1 << qb) - 1) * d;
+        // See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
         int_dot -= 2 *
                 rabitq::bitwise_xor_dot_product(
                            rearranged_rotated_qq.data(), binary_data, size, qb);
         final_dot += int_dot * query_fac.int_dot_scale;
     } else {
-        // See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
         auto dot_qo = rabitq::bitwise_and_dot_product(
                 rearranged_rotated_qq.data(), binary_data, size, qb);
-
         // It was a willful decision (after the discussion) to not to pre-cache
         // the sum of all bits, just in order to reduce the overhead per vector.
         // process 64-bit popcounts
@@ -317,32 +465,60 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
         final_dot -= query_fac.c34;
     }
 
-    // this is ||or - c||^2 - (IP ? ||or||^2 : 0)
-    const float or_c_l2sqr = fac->or_minus_c_l2sqr;
-
     // pre_dist = ||or - c||^2 + ||qr - c||^2 -
     //     2 * ||or - c|| * ||qr - c|| * <q,o> - (IP ? ||or||^2 : 0)
-    const float pre_dist = or_c_l2sqr + query_fac.qr_to_c_L2sqr -
-            2 * fac->dp_multiplier * final_dot;
+    const float pre_dist = base_fac->or_minus_c_l2sqr +
+            query_fac.qr_to_c_L2sqr - 2 * base_fac->dp_multiplier * final_dot;
 
     if (metric_type == MetricType::METRIC_L2) {
         // ||or - q||^ 2
         return pre_dist;
     } else {
         // metric == MetricType::METRIC_INNER_PRODUCT
+        // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
+        return -0.5f * (pre_dist - query_fac.qr_norm_L2sqr);
+    }
+}
+
+float RaBitQDistanceComputerQ::distance_to_code_full(const uint8_t* code) {
+    FAISS_ASSERT(code != nullptr);
+    FAISS_ASSERT(
+            (metric_type == MetricType::METRIC_L2 ||
+             metric_type == MetricType::METRIC_INNER_PRODUCT));
+    FAISS_ASSERT(rotated_q.size() == d);
 
-        // this is ||q||^2
-        const float query_norm_sqr = query_fac.qr_norm_L2sqr;
+    size_t ex_bits = nb_bits - 1;
 
-        // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
-        return -0.5f * (pre_dist - query_norm_sqr);
+    if (ex_bits == 0) {
+        // No ex-bits, just return 1-bit distance
+        return distance_to_code_1bit(code);
     }
+
+    // Extract pointers to code sections
+    const uint8_t* binary_data = code;
+    size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError);
+    const uint8_t* ex_code = code + offset;
+    const ExtraBitsFactors* ex_fac = reinterpret_cast<const ExtraBitsFactors*>(
+            ex_code + (d * ex_bits + 7) / 8);
+
+    // Call shared utility directly with rotated_q pointer
+    return rabitq_utils::compute_full_multibit_distance(
+            binary_data,
+            ex_code,
+            *ex_fac,
+            rotated_q.data(),
+            query_fac.qr_to_c_L2sqr,
+            query_fac.qr_norm_L2sqr,
+            d,
+            ex_bits,
+            metric_type);
 }
 
 // Use shared constant from RaBitQUtils
 using rabitq_utils::Z_MAX_BY_QB;
 
-void RaBitDistanceComputerQ::set_query(const float* x) {
+void RaBitQDistanceComputerQ::set_query(const float* x) {
+    q = x;
     FAISS_ASSERT(x != nullptr);
     FAISS_ASSERT(
             (metric_type == MetricType::METRIC_L2 ||
@@ -351,10 +527,15 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
     FAISS_THROW_IF_NOT(qb > 0);
 
     // Use shared utilities for core query factor computation
-    std::vector<float> rotated_q;
+    // rotated_q is populated directly by compute_query_factors as an output
+    // parameter
     query_fac = rabitq_utils::compute_query_factors(
             x, d, centroid, qb, centered, metric_type, rotated_q, rotated_qq);
 
+    // Compute g_error (query norm for lower bound computation)
+    // g_error = ||qr - c|| (L2 norm of rotated query)
+    g_error = std::sqrt(query_fac.qr_to_c_L2sqr);
+
     // Rearrange the query vector for SIMD operations (RaBitQuantizer-specific)
     popcount_aligned_dim = ((d + 7) / 8) * 8;
     size_t offset = (d + 7) / 8;
@@ -371,24 +552,28 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
     }
 }
 
+} // anonymous namespace
+
 FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
         uint8_t qb,
         const float* centroid_in,
         bool centered) const {
     if (qb == 0) {
-        auto dc = std::make_unique<RaBitDistanceComputerNotQ>();
+        auto dc = std::make_unique<RaBitQDistanceComputerNotQ>();
         dc->metric_type = metric_type;
         dc->d = d;
         dc->centroid = centroid_in;
+        dc->nb_bits = nb_bits;
 
         return dc.release();
     } else {
-        auto dc = std::make_unique<RaBitDistanceComputerQ>();
+        auto dc = std::make_unique<RaBitQDistanceComputerQ>();
         dc->metric_type = metric_type;
         dc->d = d;
         dc->centroid = centroid_in;
         dc->qb = qb;
         dc->centered = centered;
+        dc->nb_bits = nb_bits;
 
         return dc.release();
     }

data/vendor/faiss/faiss/impl/RaBitQuantizer.h

@@ -37,11 +37,28 @@ struct RaBitQuantizer : Quantizer {
     //   possible. Thus, a quantizer has to introduce a metric.
     MetricType metric_type = MetricType::METRIC_L2;
 
-    RaBitQuantizer(size_t d = 0, MetricType metric = MetricType::METRIC_L2);
+    // Number of bits per dimension (1-9). Default is 1 for backward
+    // compatibility.
+    // - nb_bits = 1: standard 1-bit RaBitQ (sign bits only)
+    // - nb_bits = 2-9: multi-bit RaBitQ (1 sign bit + ex_bits extra bits)
+    size_t nb_bits = 1;
+
+    RaBitQuantizer(
+            size_t d = 0,
+            MetricType metric = MetricType::METRIC_L2,
+            size_t nb_bits = 1);
+
+    // Compute code size based on dimensionality and number of bits
+    // Returns: size in bytes for one encoded vector
+    // - nb_bits=1: (d+7)/8 + 8 bytes (1-bit codes + base factors)
+    // - nb_bits>1: (d+7)/8 + 8 + d*ex_bits/8 + 8 bytes
+    //              (1-bit codes + base factors + ex-bit codes + ex factors)
+    size_t compute_code_size(size_t d, size_t num_bits) const;
 
     void train(size_t n, const float* x) override;
 
-    // every vector is expected to take (d + 7) / 8 + sizeof(FactorsData) bytes,
+    // every vector is expected to take (d + 7) / 8 + sizeof(SignBitFactors)
+    // bytes,
     void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
 
     void compute_codes_core(
@@ -71,9 +88,59 @@ struct RaBitQuantizer : Quantizer {
     // specify qb = 0 to get an DC that does not quantize a query
     // specify qb > 0 to have SQ qb-bits query
     FlatCodesDistanceComputer* get_distance_computer(
-            uint8_t qb,
-            const float* centroid_in = nullptr,
+            uint8_t qb = 0,
+            const float* centroid = nullptr,
             bool centered = false) const;
 };
 
+// RaBitQDistanceComputer: Base class for RaBitQ distance computers
+//
+// This intermediate class exists to provide a unified interface for
+// two-stage multi-bit search. While most Faiss quantizers extend
+// FlatCodesDistanceComputer directly, RaBitQ requires this additional
+// abstraction layer due to its unique split encoding strategy
+// (1 sign bit + magnitude bits) which enables:
+//
+// 1. distance_to_code_1bit() - Fast 1-bit filtering using only sign bits
+// 2. distance_to_code_full() - Accurate multi-bit refinement using all bits
+// 3. lower_bound_distance() - Error-bounded adaptive filtering
+//                              (based on 1-bit estimator)
+//
+// These three methods implement RaBitQ's two-stage search pattern and are
+// shared between the quantized (Q) and non-quantized (NotQ) query variants.
+// The intermediate class allows two-stage search code to work with both
+// variants via a single dynamic_cast.
+struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
+    size_t d = 0;
+    const float* centroid = nullptr;
+    MetricType metric_type = MetricType::METRIC_L2;
+    size_t nb_bits = 1;
+
+    // Query norm for lower bound computation (g_error in rabitq-library)
+    // This is the L2 norm of the rotated query: ||query - centroid||
+    float g_error = 0.0f;
+
+    float symmetric_dis(idx_t /*i*/, idx_t /*j*/) override {
+        // Not used for RaBitQ
+        FAISS_THROW_MSG("Not implemented");
+    }
+
+    // Compute 1-bit distance estimate (fast)
+    virtual float distance_to_code_1bit(const uint8_t* code) = 0;
+
+    // Compute full multi-bit distance (accurate)
+    virtual float distance_to_code_full(const uint8_t* code) = 0;
+
+    // Compute lower bound of distance using error bounds
+    // Guarantees: actual_distance >= lower_bound_distance
+    // Used for adaptive filtering in two-stage search
+    virtual float lower_bound_distance(const uint8_t* code);
+
+    // Override from FlatCodesDistanceComputer
+    // Delegates to distance_to_code_full() for multi-bit distance computation
+    float distance_to_code(const uint8_t* code) final {
+        return distance_to_code_full(code);
+    }
+};
+
 } // namespace faiss