Skip to content

fix(index): widen the RaBitQ FastScan accumulator above 1024 dims - #8842

Open
wombatu-kun wants to merge 1 commit into
lance-format:mainfrom
wombatu-kun:issue/7157-rq-dist-table-u16-overflow
Open

fix(index): widen the RaBitQ FastScan accumulator above 1024 dims#8842
wombatu-kun wants to merge 1 commit into
lance-format:mainfrom
wombatu-kun:issue/7157-rq-dist-table-u16-overflow

Conversation

@wombatu-kun

Copy link
Copy Markdown
Contributor

Closes #7157

quantize_dist_table_into maps the RaBitQ FastScan LUT onto the full 0..=u8::MAX range, and a row's binary sum takes 2 * code_len lookups from it. Past 128 code bytes (rotated dim above 1024) that sum no longer fits the u16 accumulator in sum_4bit_dist_table: the scalar path saturates and the AVX2/AVX-512/NEON kernels wrap, so a distance becomes true_sum % 65536 and 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_avx2 chunked its codes at a local SAFE_CODE_LEN = 128 for exactly this reason, and quantize_ex_fastscan_dist_table_into caps the ex-code LUT range so its own sum fits; only the binary u8 path was left with a narrow accumulator. This adds sum_4bit_dist_table_u32, which sums through the same u16 kernels in chunks of the now-shared SAFE_U16_CODE_LEN and widens between chunks, and routes binary_distances_with_scratch to it when code_len exceeds 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. The u32 scratch was already threaded through the calculator for ApproxMode::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::Accurate is the in-run control, since the two modes differ in exactly the accumulator width:

rotated dim code_len Normal before Normal after Accurate before / after
768 96 0.804 0.814 0.797 / 0.816
1024 128 0.824 0.819 0.811 / 0.817
2048 256 0.022 0.890 0.878 / 0.881
3072 384 0.129 0.889 0.896 / 0.895
4096 512 0.008 0.908 0.907 / 0.911

Random selection would score 0.0024, so at 4096 the index was barely above chance, and Accurate never 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_dimensions scaled its LUT down by code_len so the sum would fit, which is why the kernels' own tests never saw this. That scaling is gone: the u16 entry 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 against sum_4bit_dist_table_u16_scalar. test_binary_distances_match_exact_at_high_dim covers rotated dims 1024, 1536 and 4096, asserting which accumulator ran, that its sums are exact, and that the 4096 case really does overflow a u16.

AVX-512 and NEON were not executed here, since the machine is AVX2 only. The chunked call feeds every kernel the same codes and dist_table sub-slices in lockstep, which is the property their correctness rests on.

@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer bug Something isn't working labels Aug 28, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 28, 2026
@wombatu-kun

Copy link
Copy Markdown
Contributor Author

mac-build (stable) is the only red job here; the other 35 checks are green. It failed on index::vector::ivf::v2::tests::test_legacy_ivf_pq_cosine_multivec_smoke at recall: 0.49 against a 0.5 bar, one row of 100 short.

That test builds an 8-bit IVF_PQ index, which shares no code with this change. simd::dist_table is imported by exactly one production file in the tree, rust/lance-index/src/vector/bq/storage.rs:30, and everything it reaches runs behind RabitDistCalculator, which only an RQ index constructs. PQ has its own kernel, lance_linalg::simd::u8::u8x16 at rust/lance-index/src/vector/pq/distance.rs:10, and pq_matrix_params sets num_bits: 8, so compute_pq_distance never enters the 4-bit branch and accumulates in f32. The one edit outside RaBitQ replaced a function-local const SAFE_CODE_LEN = 128 in the private sum_4bit_hacc_dist_table_avx2 with the shared SAFE_U16_CODE_LEN, the same value, on a path that is not compiled for the ARM runner that failed.

This is the flake I added to #8789. The same test failed three times on main itself on 2026-08-27 at recall: 0.48: 33125150016, 33108104729, 33067105290. The fixture vectors are seeded, StdRng::from_seed([13; 32]) at rust/lance-testing/src/datagen.rs:244; the variance is unseeded k-means init at rust/lance-index/src/vector/kmeans.rs:931 and the OS-seeded training sampler at rust/lance/src/index/vector/utils.rs:1053. #8767's d692d5abd seeds both param structs in pq_matrix_params and widens the IVF sample budget to cover the whole multivector fixture, which targets exactly this test.

This branch is based on 29c4d594b, which is main's current head, so there is nothing to rebase onto, and I have no write access to re-run the job.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-index Vector index, linalg, tokenizer bug Something isn't working K-approved Latest Gatekeeper recommendation permits acceptance.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

IVF_RQ retrieval quality degrades when number of embedding dimensions increases?

1 participant