fix(index): widen the RaBitQ FastScan accumulator above 1024 dims - #8842
fix(index): widen the RaBitQ FastScan accumulator above 1024 dims#8842wombatu-kun wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The change fixes the demonstrated overflow at the correct boundary: it preserves the narrow SIMD path through 1024 dimensions, chunks and widens longer sums without changing the LUT or storage format, and verifies boundary, ragged-chunk, multi-batch, and distance-calculation behavior. This is preferable to shrinking the LUT range because it retains quantization resolution while reusing the established kernels.
|
That test builds an 8-bit IVF_PQ index, which shares no code with this change. This is the flake I added to #8789. The same test failed three times on This branch is based on |
Closes #7157
quantize_dist_table_intomaps the RaBitQ FastScan LUT onto the full0..=u8::MAXrange, and a row's binary sum takes2 * code_lenlookups from it. Past 128 code bytes (rotated dim above 1024) that sum no longer fits theu16accumulator insum_4bit_dist_table: the scalar path saturates and the AVX2/AVX-512/NEON kernels wrap, so a distance becomestrue_sum % 65536and the ranking collapses. At rotated dim 4096 a full-range sum reaches 261120, four times the ceiling.The bound is already known in this file.
sum_4bit_hacc_dist_table_avx2chunked its codes at a localSAFE_CODE_LEN = 128for exactly this reason, andquantize_ex_fastscan_dist_table_intocaps the ex-code LUT range so its own sum fits; only the binaryu8path was left with a narrow accumulator. This addssum_4bit_dist_table_u32, which sums through the sameu16kernels in chunks of the now-sharedSAFE_U16_CODE_LENand widens between chunks, and routesbinary_distances_with_scratchto it whencode_lenexceeds that length. The LUT is unchanged, so the wide sum is exactly what the narrow one would have produced without overflow, and nothing at or below rotated dim 1024 changes path. Theu32scratch was already threaded through the calculator forApproxMode::Accurate, so no signature moves, and the added work is one 32-lane widening pass per 128 code bytes.Measured on one synthetic partition, 4096 rows, 100 queries,
num_bits = 1, recall@10 against exact L2, medians of three repeats.ApproxMode::Accurateis the in-run control, since the two modes differ in exactly the accumulator width:Random selection would score 0.0024, so at 4096 the index was barely above chance, and
Accuratenever moved. The dip is not monotone in dimension because the kernels wrap rather than saturate: at 3072 a typical sum wraps once and keeps some of its order, while at 2048 it lands on the boundary and only some rows wrap.test_simd_matches_scalar_varied_dimensionsscaled its LUT down bycode_lenso the sum would fit, which is why the kernels' own tests never saw this. That scaling is gone: theu16entry point is now tested with full-range tables over the lengths it is contracted for, a new test models the ex-code caller's cap so the long lengths keep their coverage, and the wide entry point is checked againstsum_4bit_dist_table_u16_scalar.test_binary_distances_match_exact_at_high_dimcovers rotated dims 1024, 1536 and 4096, asserting which accumulator ran, that its sums are exact, and that the 4096 case really does overflow au16.AVX-512 and NEON were not executed here, since the machine is AVX2 only. The chunked call feeds every kernel the same
codesanddist_tablesub-slices in lockstep, which is the property their correctness rests on.