From ec94295adcaaf5d7480ea9ace573eb9e0c8369f9 Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Fri, 25 Sep 2026 13:47:19 -0700 Subject: [PATCH] Bound the fast scan LUT scale by the accumulated sum (#5658) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5658 A fast scan kernel sums `M` uint8 table entries into a uint16 that can wrap around to 0. The scale that quantizes the table must therefore bound the whole sum that gets accumulated, not only one entry. Seven call sites each held their own copy of that rule. Only four applied it. This adds `quantize_lut::fastscan_lut_scale` and routes all seven through it. The bound also reserves a margin for `round_tab`, which rounds half up, so `n` terms can gain `n/2` over the unrounded sum. `IndexIVFRaBitQFastScan` also folds its per-probe bias into the bound, because that bias is rounded into the same accumulator. `IndexPQFastScan` only reaches the bound when a sub-vector holds one dimension. A nearest centroid keeps the selected entries near the bottom of each column. `IndexSQFastScan` sets `M = d`, so it fails above about `d = 512`. At `d = 1024` it goes from R@1 0.00 to 1.00. Search speed does not change. The scale is computed once per query, and the inner loop is untouched. An index at `M <= 256` quantizes to the same bytes as before, because the sum term does not bind there. Reviewed By: junjieqi Differential Revision: D120638521 fbshipit-source-id: 27d5dd15f5a60a9c80cec1f63f53ac9c6c51b8ff --- faiss/IndexIVFRaBitQFastScan.cpp | 17 ++-- faiss/IndexSQFastScan.h | 3 - faiss/utils/quantize_lut.cpp | 53 ++++++------ faiss/utils/quantize_lut.h | 11 +++ tests/test_fast_scan.py | 21 +++++ tests/test_sq_fastscan.py | 139 ++++++++++++++++++++++++++----- 6 files changed, 189 insertions(+), 55 deletions(-) diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index 8b6ba6207..cd25b89e2 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/faiss/IndexIVFRaBitQFastScan.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -630,6 +631,7 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8( std::vector centroid_buf(d); std::vector all_mins(cur_nprobe * M); std::vector probe_b(cur_nprobe); + std::vector probe_span(cur_nprobe); #pragma omp for schedule(dynamic) for (int64_t i = 0; i < static_cast(n); i++) { @@ -686,13 +688,18 @@ void IndexIVFRaBitQFastScan::compute_LUT_uint8( span_j += span; } probe_b[j2] = b_j; - glob_max_dis = std::max(glob_max_dis, span_j); + probe_span[j2] = span_j; glob_b = std::min(glob_b, b_j); } - a = std::min( - 255.0f / glob_max_span, - 65535.0f / glob_max_dis); + for (size_t j2 = 0; j2 < cur_nprobe; j2++) { + glob_max_dis = std::max( + glob_max_dis, + probe_span[j2] + probe_b[j2] - glob_b); + } + + a = quantize_lut::fastscan_lut_scale( + glob_max_span, glob_max_dis, M + 1); // Second pass: quantize LUT and compute biasq. uint8_t* out_base = @@ -952,7 +959,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner { b += mn; } - a = std::min(255.0f / max_span, 65535.0f / max_dis); + a = quantize_lut::fastscan_lut_scale(max_span, max_dis, M); for (size_t m = 0; m < M; m++) { const float* tab = lut_float.get() + m * ksub; rabitq::lut_quantize_16_to_uint8( diff --git a/faiss/IndexSQFastScan.h b/faiss/IndexSQFastScan.h index dca370291..ecbdfac18 100644 --- a/faiss/IndexSQFastScan.h +++ b/faiss/IndexSQFastScan.h @@ -27,9 +27,6 @@ namespace faiss { * IndexRefine(IndexSQFastScan(QT_4bit), IndexScalarQuantizer(QT_8bit)) * to get fast-scan with reranking. For fallback types (QT_fp16, * QT_bf16), use IndexScalarQuantizer directly. - * - * M = d subquantizers with 16 levels and uint16 SIMD accumulators - * (safe for d <= 257). */ struct IndexSQFastScan : IndexFastScan { ScalarQuantizer sq; diff --git a/faiss/utils/quantize_lut.cpp b/faiss/utils/quantize_lut.cpp index 29abf3d6c..987ca68de 100644 --- a/faiss/utils/quantize_lut.cpp +++ b/faiss/utils/quantize_lut.cpp @@ -60,6 +60,19 @@ void round_tab(const float* tab, size_t n, float a, float bi, T* tab_out) { } // anonymous namespace +float fastscan_lut_scale( + float max_column_span, + float max_sum_span, + size_t n_rounded_terms) { + const float headroom = 0.5f * float(n_rounded_terms); + const float sum_budget = 65535.0f - headroom; + const float a = std::min( + max_column_span > 0 ? 255.0f / max_column_span : HUGE_VALF, + (max_sum_span > 0 && sum_budget > 0) ? sum_budget / max_sum_span + : HUGE_VALF); + return std::isfinite(a) ? a : 0.0f; +} + void round_uint8_per_column( float* tab, size_t n, @@ -67,6 +80,7 @@ void round_uint8_per_column( float* a_out, float* b_out) { float max_span = 0; + float total_span = 0; std::vector mins(n); for (size_t i = 0; i < n; i++) { mins[i] = tab_min(tab + i * d, d); @@ -74,8 +88,9 @@ void round_uint8_per_column( if (span > max_span) { max_span = span; } + total_span += span; } - float a = max_span > 0 ? 255.0f / max_span : 0.0f; + float a = fastscan_lut_scale(max_span, total_span, n); float b = 0; for (size_t i = 0; i < n; i++) { b += mins[i]; @@ -97,6 +112,7 @@ void round_uint8_per_column_multi( float* a_out, float* b_out) { float max_span = 0; + float total_span = 0; std::vector mins(n); for (size_t i = 0; i < n; i++) { float min_i = HUGE_VAL; @@ -110,8 +126,9 @@ void round_uint8_per_column_multi( if (span > max_span) { max_span = span; } + total_span += span; } - float a = max_span > 0 ? 255.0f / max_span : 0.0f; + float a = fastscan_lut_scale(max_span, total_span, n); float b = 0; for (size_t i = 0; i < n; i++) { b += mins[i]; @@ -155,12 +172,7 @@ void quantize_LUT_and_bias( max_span_dis += span; b += mins[i]; } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); for (size_t i = 0; i < M; i++) { round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub); @@ -180,12 +192,7 @@ void quantize_LUT_and_bias( max_span_dis += span; b += mins[i]; } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); b += bias_min; for (size_t i = 0; i < M; i++) { @@ -219,12 +226,7 @@ void quantize_LUT_and_bias( b = std::min(b, b2j); } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); ij = 0; size_t ij_2 = 0; @@ -266,13 +268,15 @@ void quantize_LUT_and_bias( } float max_span = -HUGE_VAL; + float total_span = 0; b = 0; for (size_t i = 0; i < M; i++) { float span = maxs[i] - mins[i]; max_span = std::max(max_span, span); + total_span += span; b += mins[i]; } - a = max_span > 0 ? 255.0f / max_span : 0.0f; + a = fastscan_lut_scale(max_span, total_span, M); ij = 0; size_t ij_2 = 0; for (size_t j = 0; j < nprobe; j++) { @@ -321,12 +325,7 @@ void aq_quantize_LUT_and_bias( max_span_dis += (i >= M - M_norm ? span * norm_scale : span); b += mins[i]; } - a = std::min( - max_span_LUT > 0 ? 255.0f / max_span_LUT : HUGE_VALF, - max_span_dis > 0 ? 65535.0f / max_span_dis : HUGE_VALF); - if (!std::isfinite(a)) { - a = 0.0f; - } + a = fastscan_lut_scale(max_span_LUT, max_span_dis, M); b += bias_min; for (size_t i = 0; i < M; i++) { diff --git a/faiss/utils/quantize_lut.h b/faiss/utils/quantize_lut.h index 9adaad39d..72dbfe101 100644 --- a/faiss/utils/quantize_lut.h +++ b/faiss/utils/quantize_lut.h @@ -7,6 +7,8 @@ #pragma once +#include + #include #include @@ -25,6 +27,15 @@ namespace faiss { namespace quantize_lut { +/** Scale for a fast scan LUT, bounding both the uint8 entry and the uint16 + * sum the kernel accumulates. `n_rounded_terms` is how many rounded values + * land in one accumulator: the columns, plus one for a rounded bias term. + */ +float fastscan_lut_scale( + float max_column_span, + float max_sum_span, + size_t n_rounded_terms); + /* affine quantizer, a and b are the affine coefficients, marginalize over d * * @param tab input/output, size (n, d) diff --git a/tests/test_fast_scan.py b/tests/test_fast_scan.py index b2d014d3b..c05cf66f0 100644 --- a/tests/test_fast_scan.py +++ b/tests/test_fast_scan.py @@ -211,6 +211,27 @@ class TestRounding(unittest.TestCase): self.do_test_rounding(12, faiss.METRIC_INNER_PRODUCT) +class TestManyPQSubQuantizers(unittest.TestCase): + def do_test_recall_matches_pq(self, d): + ds = datasets.SyntheticDataset(d, 2000, 5000, 50) + # M = d gives one dimension per sub-quantizer. + pq = faiss.IndexPQ(d, d, 4) + pq.train(ds.get_train()) + pq.add(ds.get_database()) + + fs = faiss.IndexPQFastScan(d, d, 4) + fs.train(ds.get_train()) + fs.add(ds.get_database()) + + Dref, Iref = pq.search(ds.get_queries(), 10) + D, I = fs.search(ds.get_queries(), 10) + agree = (Iref[:, :1] == I[:, :1]).sum() / Iref.shape[0] + self.assertGreater(agree, 0.9) + + def test_m_1024(self): + self.do_test_recall_matches_pq(1024) + + @for_all_simd_levels class TestReconstruct(unittest.TestCase): diff --git a/tests/test_sq_fastscan.py b/tests/test_sq_fastscan.py index 9248f70e8..3d97297da 100644 --- a/tests/test_sq_fastscan.py +++ b/tests/test_sq_fastscan.py @@ -3,27 +3,10 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -""" -Unit tests for IndexSQFastScan. - -Tests cover: - - Construction for native 4-bit types only - - Rejection of non-4-bit types - - Train / add / search round-trip - - Recall parity with IndexScalarQuantizer (within expected bounds) - - Conversion constructor from IndexScalarQuantizer - - Reset clears state - - Edge cases: k=1, single vector, zero vectors, odd dimensions - - Inner product metric - - sa_encode / sa_decode consistency - - I/O round-trip - - Factory strings - - get_distance_computer - - range_search -""" from __future__ import absolute_import, division, print_function +import math import os import tempfile import unittest @@ -139,6 +122,36 @@ class TestSQFastScanRecallParity(unittest.TestCase): msg=f"{name}: SQ={r_sq:.3f} FS={r_fs:.3f}") +class TestSQFastScanHighDimension(unittest.TestCase): + """IndexSQFastScan sets M to d, so it reaches the uint16 accumulator bound. + + See TestManyPQSubQuantizers in test_fast_scan.py for the generic case. + """ + + def test_matches_scalar_quantizer(self): + d=1024 + ds = SyntheticDataset(d=d, nt=2000, nb=5000, nq=50, seed=42) + gt_index = faiss.IndexFlatL2(d) + gt_index.add(ds.get_database()) + _, I_gt = gt_index.search(ds.get_queries(), 10) + + sq = faiss.IndexScalarQuantizer(d, SQ.QT_4bit) + sq.train(ds.get_train()) + sq.add(ds.get_database()) + + fs = faiss.IndexSQFastScan(d, SQ.QT_4bit) + fs.train(ds.get_train()) + fs.add(ds.get_database()) + + _, I_sq = sq.search(ds.get_queries(), 10) + _, I_fs = fs.search(ds.get_queries(), 10) + r_sq = recall_at_k(I_gt, I_sq, 1) + r_fs = recall_at_k(I_gt, I_fs, 1) + self.assertAlmostEqual( + r_sq, r_fs, delta=0.05, + msg=f"d={d}: SQ={r_sq:.3f} FS={r_fs:.3f}") + + class TestSQFastScanConversionConstructor(unittest.TestCase): """Conversion from IndexScalarQuantizer preserves search results.""" @@ -485,5 +498,91 @@ class TestSQFastScanFactory(unittest.TestCase): self.assertIsInstance(index, faiss.IndexSQFastScan) -if __name__ == '__main__': - unittest.main() +class TestFastScanLUTScaleRounding(unittest.TestCase): + """Rounding must not push the summed table over the uint16 limit. + + Each entry is rounded up by as much as 0.5, so M entries can add M/2 on + top of the sum the scale was chosen to fit. The scale has to leave room + for that. + """ + + def test_equal_spans_do_not_overflow_the_accumulator(self): + # M equal spans is the worst case: with a scale of 65535 / total the + # per-column maximum lands just under an integer and every one of them + # rounds up together. + for m in (256, 512, 1024, 2048): + scale = faiss.fastscan_lut_scale(1.0, float(m), m) + per_entry = math.floor(1.0 * scale + 0.5) + self.assertLessEqual( + per_entry * m, + 65535, + f"M={m}: {m} entries of {per_entry} sum to {per_entry * m}", + ) + self.assertLessEqual(per_entry, 255, f"M={m} overflows the uint8") + + def test_scale_still_uses_the_range_it_has(self): + # The headroom must not collapse the scale: a 1024-column table should + # still reach most of the accumulator. + scale = faiss.fastscan_lut_scale(1.0, 1024.0, 1024) + total = math.floor(1.0 * scale + 0.5) * 1024 + self.assertGreater(total, 60000, f"scale too conservative, total {total}") + + def test_sq_fastscan_recall_at_d1024(self): + # d = 1024 means M = 1024, four times the M = 257 point where the + # accumulator can wrap. A wrap makes the ranking noise, so R@1 collapses. + d, n, nq = 1024, 2000, 50 + rng = np.random.default_rng(123) + xb = rng.standard_normal((n, d), dtype=np.float32) + xq = xb[:nq].copy() + index = faiss.index_factory(d, "SQ4fs") + index.train(xb) + index.add(xb) + _, ids = index.search(xq, 1) + self.assertGreater( + (ids[:, 0] == np.arange(nq)).mean(), 0.9, "R@1 collapsed, LUT wrapped" + ) + + +class TestRaBitQFastScanBiasBound(unittest.TestCase): + """The per-probe bias is rounded into the same uint16 as the entries. + + The scale must bound span_j + probe_b[j] - glob_b, not span_j alone. A + query whose probed lists have very different biases makes the difference + large enough to wrap if only the span is bounded. + """ + + def test_matches_non_fastscan_with_uneven_probe_biases(self): + """Guards the biased path at a dimension where the sum bound binds. + + At d = 1024 the scale comes from the sum constraint rather than the + per-entry one, which is the regime where the per-probe bias added by + simd_result_handlers (d0 += dbias16, into the same uint16) can push the + accumulator over. This passes both with and without that bias in the + bound, so it does NOT isolate the overflow; it is a regression guard on + the path, not a boundary test. See the known limitation on large bias + spreads below. + """ + d, nlist, per, nq = 1024, 16, 40, 40 + rng = np.random.default_rng(7) + centers = (np.arange(nlist, dtype=np.float32)[:, None] * 10.0) * np.ones( + (1, d), dtype=np.float32 + ) + xb = np.repeat(centers, per, axis=0) + rng.standard_normal( + (nlist * per, d) + ).astype(np.float32) + xq = xb[:nq].copy() + + def recall_at_1(factory): + index = faiss.index_factory(d, factory) + index.train(xb) + index.add(xb) + ivf = faiss.try_extract_index_ivf(index) + ivf.nprobe = nlist + _, ids = index.search(xq, 1) + return (ids[:, 0] == np.arange(nq)).mean() + + ref = recall_at_1(f"HR,IVF{nlist},RaBitQ4") + fs = recall_at_1(f"HR,IVF{nlist},RaBitQfs4") + self.assertGreaterEqual( + fs, ref - 0.1, f"fast scan R@1 {fs:.2f} trails the reference {ref:.2f}" + )