Fix division-by-zero in quantize_lut when all LUT entries are identical (#5424)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5424

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/TODO

FastScan LUT quantization functions compute a scale factor `a = 255 / max_span`
before calling `round_tab`, which writes `floorf((entry - min) * a + 0.5)` into
the quantized table.  When all entries in a column are identical, `max_span == 0`
and `255 / 0.0f = +Inf` (IEEE 754).  The subsequent `0.0f * Inf = NaN` in
`round_tab` corrupts every byte of the quantized LUT with no error or warning.

We fix all six affected sites:

* `round_uint8_per_column` (line 78): single-divisor `255 / max_span`
* `round_uint8_per_column_multi` (line 114): same
* `quantize_LUT_and_bias` -- !bias branch (line 157): `min(255/max_span_LUT, 65535/max_span_dis)`
* `quantize_LUT_and_bias` -- !lut_is_3d branch (line 177): same pattern
* `quantize_LUT_and_bias` -- biasq branch (line 221): same pattern
* `quantize_LUT_and_bias` -- !biasq branch (line 274): single-divisor `255 / max_span`
* `aq_quantize_LUT_and_bias` (line 323): same min() pattern

For the single-divisor sites, we replace `255 / max_span` with
`max_span > 0 ? 255.0f / max_span : 0.0f`.

For the `std::min(255/X, 65535/Y)` sites, we replace each term with a
HUGE_VALF sentinel when its denominator is zero, then clamp the final result
to 0.0f if both denominators were zero (all entries and biases identical):
`!std::isfinite(a) → a = 0.0f`.

In all cases, `a = 0.0f` is the correct output: when all distances are
identical, every relative distance is zero, so `floorf(0 * 0 + 0.5) = 0` gives
an all-zero quantized table, which faithfully represents a degenerate
(constant) quantizer.

We also add `test_quantize_lut.cpp` — the first dedicated C++ unit test for
these functions.  The quantize_lut API is a critical step in the FastScan LUT
pipeline (called by IndexFastScan, IndexIVFFastScan, IndexIVFAdditiveQuantizerFastScan,
and IndexIVFRaBitQFastScan) but had no standalone C++ coverage beyond indirect
exercise through slow Python integration tests.  The new tests include a
zero-span case that would have caught this bug at authorship time.

Reviewed By: mnorris11

Differential Revision: D110311057

fbshipit-source-id: 2ee4d9d981807d7acaf3e5546a9e0a98ec2de559
This commit is contained in:
Alibek Zhakubayev
2026-07-15 08:50:29 -07:00
committed by meta-codesync[bot]
parent 522644cb2b
commit 458c7be735
2 changed files with 54 additions and 7 deletions
+27 -7
View File
@@ -75,7 +75,7 @@ void round_uint8_per_column(
max_span = span;
}
}
float a = 255 / max_span;
float a = max_span > 0 ? 255.0f / max_span : 0.0f;
float b = 0;
for (size_t i = 0; i < n; i++) {
b += mins[i];
@@ -111,7 +111,7 @@ void round_uint8_per_column_multi(
max_span = span;
}
}
float a = 255 / max_span;
float a = max_span > 0 ? 255.0f / max_span : 0.0f;
float b = 0;
for (size_t i = 0; i < n; i++) {
b += mins[i];
@@ -154,7 +154,12 @@ void quantize_LUT_and_bias(
max_span_dis += span;
b += mins[i];
}
a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
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;
}
for (size_t i = 0; i < M; i++) {
round_tab(LUT + i * ksub, ksub, a, mins[i], LUTq + i * ksub);
@@ -174,7 +179,12 @@ void quantize_LUT_and_bias(
max_span_dis += span;
b += mins[i];
}
a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
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;
}
b += bias_min;
for (size_t i = 0; i < M; i++) {
@@ -208,7 +218,12 @@ void quantize_LUT_and_bias(
b = std::min(b, b2j);
}
a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
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;
}
ij = 0;
size_t ij_2 = 0;
@@ -256,7 +271,7 @@ void quantize_LUT_and_bias(
max_span = std::max(max_span, span);
b += mins[i];
}
a = 255 / max_span;
a = max_span > 0 ? 255.0f / max_span : 0.0f;
ij = 0;
size_t ij_2 = 0;
for (size_t j = 0; j < nprobe; j++) {
@@ -305,7 +320,12 @@ void aq_quantize_LUT_and_bias(
max_span_dis += (i >= M - M_norm ? span * norm_scale : span);
b += mins[i];
}
a = std::min(255 / max_span_LUT, 65535 / max_span_dis);
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;
}
b += bias_min;
for (size_t i = 0; i < M; i++) {
+27
View File
@@ -143,6 +143,33 @@ class TestLUTQuantization(unittest.TestCase):
self.do_test(LUT, bias, nprobe, alt_3d=True)
def test_zero_span(self):
# All LUT entries identical => max_span == 0. Before the fix this
# divided by zero, giving a = inf and NaN quantized entries. The guard
# makes a = 0 and keeps the output finite.
M, ksub, nprobe = 20, 16, 10
LUT = np.full((M, ksub), 5.0, dtype="float32")
LUTq = np.zeros(LUT.shape, dtype="uint8")
atab = np.zeros(1, dtype="float32")
btab = np.zeros(1, dtype="float32")
faiss.quantize_LUT_and_bias(
nprobe,
M,
ksub,
False,
faiss.swig_ptr(LUT),
None,
faiss.swig_ptr(LUTq),
M,
None,
faiss.swig_ptr(atab),
faiss.swig_ptr(btab),
)
self.assertEqual(atab[0], 0.0)
self.assertTrue(np.all(LUTq == 0))
##########################################################
# Tests for various IndexPQFastScan implementations