Thread qb/centered search params through FastScan LUT and handler (#5095)

Summary:
Pull Request resolved: #5095

D100399519 added IVFRaBitQSearchParameters support to the FastScan scanner
but only patched the distance_to_code fallback path. The main search path
(LUT construction and SIMD distance correction in handle()) still read
qb/centered from the index, ignoring the search params override.

This diff completes the fix by:
1. Adding qb/centered fields to FastScanDistancePostProcessing context
2. Threading them through compute_LUT → compute_residual_LUT
3. Reading them from context in the handler's handle() method
4. Extracting them from IVFRaBitQSearchParameters in search_preassigned

Reviewed By: mnorris11

Differential Revision: D100674751

fbshipit-source-id: 72139ed1311f9422f8701933bf0464c411764d05

Author: alibeklfc

faiss/IndexIVFRaBitQFastScan.cpp

@@ -276,8 +276,10 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
         const float* residual,
         QueryFactorsData& query_factors,
         float* lut_out,
+        uint8_t qb_param,
+        bool centered_param,
         const float* original_query) const {
-    FAISS_THROW_IF_NOT(qb > 0 && qb <= 8);
+    FAISS_THROW_IF_NOT(qb_param > 0 && qb_param <= 8);
 
     std::vector<float> rotated_q(d);
     std::vector<uint8_t> rotated_qq(d);
@@ -287,8 +289,8 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
             residual,
             d,
             nullptr,
-            qb,
-            centered,
+            qb_param,
+            centered_param,
             metric_type,
             rotated_q,
             rotated_qq);
@@ -305,8 +307,8 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
         query_factors.rotated_q = rotated_q;
     }
 
-    if (centered) {
-        const float max_code_value = (1 << qb) - 1;
+    if (centered_param) {
+        const float max_code_value = (1 << qb_param) - 1;
 
         for (size_t m = 0; m < M; m++) {
             const size_t dim_start = m * 4;
@@ -372,15 +374,24 @@ void IndexIVFRaBitQFastScan::search_preassigned(
     FAISS_THROW_IF_NOT_MSG(!stats, "stats not supported for this index");
 
     size_t cur_nprobe = this->nprobe;
+    uint8_t used_qb = qb;
+    bool used_centered = centered;
     if (params) {
         FAISS_THROW_IF_NOT(params->max_codes == 0);
         cur_nprobe = params->nprobe;
+        if (auto rparams =
+                    dynamic_cast<const IVFRaBitQSearchParameters*>(params)) {
+            used_qb = rparams->qb;
+            used_centered = rparams->centered;
+        }
     }
 
     std::vector<QueryFactorsData> query_factors_storage(n * cur_nprobe);
     FastScanDistancePostProcessing context;
     context.query_factors = query_factors_storage.data();
     context.nprobe = cur_nprobe;
+    context.qb = used_qb;
+    context.centered = used_centered;
 
     const CoarseQuantized cq = {cur_nprobe, centroid_dis, assign};
     search_dispatch_implem(n, x, k, distances, labels, cq, context, params);
@@ -396,6 +407,10 @@ void IndexIVFRaBitQFastScan::compute_LUT(
     FAISS_THROW_IF_NOT(is_trained);
     FAISS_THROW_IF_NOT(by_residual);
 
+    // Use overridden qb/centered from context if provided, else index defaults
+    const uint8_t used_qb = context.qb > 0 ? context.qb : qb;
+    const bool used_centered = context.qb > 0 ? context.centered : centered;
+
     size_t cq_nprobe = cq.nprobe;
 
     size_t dim12 = 16 * M;
@@ -424,6 +439,8 @@ void IndexIVFRaBitQFastScan::compute_LUT(
                     xij,
                     query_factors_data,
                     dis_tables.get() + ij * dim12,
+                    used_qb,
+                    used_centered,
                     x + i * d);
 
             // Store query factors using compact indexing (ij directly)
@@ -624,6 +641,8 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
         context = FastScanDistancePostProcessing{};
         context.query_factors = &query_factors;
         context.nprobe = 1;
+        context.qb = qb;
+        context.centered = centered;
 
         index.compute_LUT_uint8(
                 1, xi, cq, dis_tables, biases, &normalizers[0], context);

faiss/IndexIVFRaBitQFastScan.h

@@ -119,6 +119,8 @@ struct IndexIVFRaBitQFastScan : IndexIVFFastScan {
             const float* residual,
             QueryFactorsData& query_factors,
             float* lut_out,
+            uint8_t qb_param,
+            bool centered_param,
             const float* original_query = nullptr) const;
 
     /// Decode FastScan code to RaBitQ residual vector with explicit
@@ -265,8 +267,9 @@ void IVFRaBitQHeapHandler<C, SL>::handle(
             (idx_base / index->bbs) * full_block_size + packed_block_size;
 
     // Cache index fields used in the inner loop.
-    const bool centered = index->centered;
-    const size_t qb = index->qb;
+    // Use overridden qb/centered from context if provided, else index defaults.
+    const bool centered = context->qb > 0 ? context->centered : index->centered;
+    const size_t qb = context->qb > 0 ? context->qb : index->qb;
     const size_t d = index->d;
 
 #ifndef NDEBUG

faiss/impl/fast_scan/FastScanDistancePostProcessing.h

@@ -35,6 +35,14 @@ struct FastScanDistancePostProcessing {
     /// Set to 0 to use index->nprobe as fallback.
     size_t nprobe = 0;
 
+    /// RaBitQ query quantization bits override.
+    /// Set to 0 to use the index default (index->qb).
+    uint8_t qb = 0;
+
+    /// RaBitQ centered scalar quantizer override.
+    /// Only used when qb > 0 (i.e., when params are overridden).
+    bool centered = false;
+
     /// Default constructor - no processing
     FastScanDistancePostProcessing() = default;
 

tests/test_rabitq_fastscan.py

@@ -1070,6 +1070,59 @@ def test_ivf_factory_with_batch_size(self):
         self.assertEqual(D.shape, (ds.nq, 5))
 
 
+class TestRaBitQFastScanSearchParams(unittest.TestCase):
+    """Test that IVFRaBitQSearchParameters qb/centered are respected."""
+
+    def test_higher_qb_improves_recall(self):
+        """Search with qb=4 should give better recall than qb=1."""
+        d = 64
+        nlist = 16
+        nprobe = 4
+        k = 10
+        ds = datasets.SyntheticDataset(d, 5000, 5000, 50)
+
+        # Ground truth with flat index
+        index_flat = faiss.IndexFlatL2(d)
+        index_flat.add(ds.get_database())
+        _, I_gt = index_flat.search(ds.get_queries(), k)
+
+        # Build IVF RaBitQ FastScan index with default qb=8
+        quantizer = faiss.IndexFlat(d, faiss.METRIC_L2)
+        index = faiss.IndexIVFRaBitQFastScan(
+            quantizer, d, nlist, faiss.METRIC_L2, 32, True
+        )
+        index.nprobe = nprobe
+        index.train(ds.get_train())
+        index.add(ds.get_database())
+
+        # Search with qb=1 (coarse quantization)
+        params_qb1 = faiss.IVFRaBitQSearchParameters()
+        params_qb1.nprobe = nprobe
+        params_qb1.qb = 1
+        _, I_qb1 = index.search(ds.get_queries(), k, params=params_qb1)
+
+        # Search with qb=4 (finer quantization)
+        params_qb4 = faiss.IVFRaBitQSearchParameters()
+        params_qb4.nprobe = nprobe
+        params_qb4.qb = 4
+        _, I_qb4 = index.search(ds.get_queries(), k, params=params_qb4)
+
+        # Compute recall@k
+        recall_qb1 = np.mean([
+            len(np.intersect1d(I_qb1[i], I_gt[i])) / k
+            for i in range(ds.nq)
+        ])
+        recall_qb4 = np.mean([
+            len(np.intersect1d(I_qb4[i], I_gt[i])) / k
+            for i in range(ds.nq)
+        ])
+
+        self.assertGreater(
+            recall_qb4, recall_qb1,
+            f"qb=4 recall ({recall_qb4:.3f}) should be higher "
+            f"than qb=1 recall ({recall_qb1:.3f})"
+        )
+
 
 if __name__ == "__main__":
     unittest.main()