From 7f6bf10b1c83b1634c3cfd3367ba1e76907c79b9 Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Fri, 25 Sep 2026 13:47:19 -0700 Subject: [PATCH] Halve IndexSQFastScan rerank memory with a split-nibble layout (#5662) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5662 `IndexIVFSQFastScan` kept whole original codes in a parallel inverted list, so an 8-bit quantizer cost 12 bits per dimension against 8 for plain SQ8. It now keeps only the bits the scan does not read, which is the same split the flat index uses in the diff below. ## The split is bit-width generic An 8-bit code keeps 4 leftover bits. A 6-bit code keeps 2. `sq_fastscan_utils.h` packs the remainder with `BitstringWriter`, which lays bits out exactly as `Codec6bit` and `Codec8bit` do. The LUT reconstruction collapses to one formula: `step = 2^(bits-4)`, `offset = step / 2`, `scale = 2^bits - 1`. At 4 bits this reduces to the native `(v + 0.5) / 15`, so the native and split paths no longer need separate constants. This also corrects a small bias, because the 8-bit case previously used an offset of 7.5 where the true band midpoint is 8.0. ## `IndexIVFSQFastScan` splits as well The IVF form kept whole codes in a parallel inverted list. It now keeps only the leftover bits. ## Measured results, flat index, 200k vectors | index | before | after | | --- | --- | --- | | `SQ8fs` | 12 bits/dim | **8.00 bits/dim** | | `SQ6fs` | 10 bits/dim | **6.00 bits/dim** | The IVF form reaches 1.37x `SQ8` at d=64 and 1.10x at d=512, against 1.79x before. The remainder is the ids duplicated into the parallel list and the block padding. Both are per-vector, so the ratio falls as `d` grows. Recall does not drop. The split reaches 0.914 to 0.932 for `SQ6fs` and 0.969 to 0.980 for `SQ8fs` at `rerank_factor=4`, measured against exact neighbours restricted to the admitted set. Also fixes a double-append in `add()`. `IndexFastScan::add` blocks adds over 65536 and calls `add()` virtually. The override therefore re-enters once per chunk, and every side-table entry was written twice for a large index. ## Serialization The `ISfs` payload now carries the leftover bits instead of whole original codes. The fourcc does not change, because `ISfs` has not been released with rerank state and no stored index is affected. Reviewed By: junjieqi Differential Revision: D121092769 fbshipit-source-id: d6e8a4f541ce73057eba86215a83a0a80f7e106d --- faiss/IndexIVFSQFastScan.cpp | 164 +++++++++++++++++++++++------------ faiss/IndexIVFSQFastScan.h | 30 +++++-- faiss/impl/index_read.cpp | 8 +- faiss/impl/index_write.cpp | 6 +- 4 files changed, 138 insertions(+), 70 deletions(-) diff --git a/faiss/IndexIVFSQFastScan.cpp b/faiss/IndexIVFSQFastScan.cpp index 7ed7cfac9..14dc1b1ef 100644 --- a/faiss/IndexIVFSQFastScan.cpp +++ b/faiss/IndexIVFSQFastScan.cpp @@ -75,8 +75,8 @@ IndexIVFSQFastScan::IndexIVFSQFastScan( new BlockInvertedLists(nlist_in, get_CodePacker()), true); if (needs_rerank(qtype)) { - orig_codes_invlists = - new ArrayInvertedLists(nlist_in, sq.code_size); + lo_codes_invlists = + new ArrayInvertedLists(nlist_in, sq_rerank_size(d, qtype)); direct_map.set_type(DirectMap::Hashtable, invlists, 0); } } @@ -136,7 +136,7 @@ IndexIVFSQFastScan::IndexIVFSQFastScan( replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true); if (needs_rerank(sq.qtype)) { - orig_codes_invlists = new ArrayInvertedLists(nlist, sq.code_size); + lo_codes_invlists = new ArrayInvertedLists(nlist, lo_code_size()); direct_map.set_type(DirectMap::Hashtable, invlists, 0); } @@ -149,15 +149,7 @@ IndexIVFSQFastScan::IndexIVFSQFastScan( InvertedLists::ScopedCodes orig_codes(orig.invlists, list_no); InvertedLists::ScopedIds orig_ids(orig.invlists, list_no); - if (needs_rerank(sq.qtype)) { - orig_codes_invlists->add_entries( - list_no, list_size, orig_ids.get(), orig_codes.get()); - for (size_t j = 0; j < list_size; j++) { - direct_map.add_single_id(orig_ids.get()[j], list_no, j); - } - } - - // Decode to float and re-quantize to 4-bit. + // Decode to float and re-encode for the scan. std::vector recon(list_size * d); sq.decode(orig_codes.get(), recon.data(), list_size); @@ -165,8 +157,21 @@ IndexIVFSQFastScan::IndexIVFSQFastScan( if (is_native_4bit(sq.qtype)) { sq.compute_codes(recon.data(), flat4.data(), list_size); } else { - float_to_4bit_nibbles( - recon.data(), flat4.data(), list_size, d, M2, sq); + // Top 4 bits feed the scan, the rest are kept for rerank. + std::vector lo(list_size * lo_code_size(), 0); + float_to_split_codes( + recon.data(), + flat4.data(), + lo.data(), + list_size, + d, + M2, + sq); + lo_codes_invlists->add_entries( + list_no, list_size, orig_ids.get(), lo.data()); + for (size_t j = 0; j < list_size; j++) { + direct_map.add_single_id(orig_ids.get()[j], list_no, j); + } } BlockInvertedLists* bil = dynamic_cast(invlists); @@ -180,7 +185,7 @@ IndexIVFSQFastScan::IndexIVFSQFastScan( } IndexIVFSQFastScan::~IndexIVFSQFastScan() { - delete orig_codes_invlists; + delete lo_codes_invlists; } // ----------------------------------------------------------------------- @@ -229,7 +234,10 @@ void IndexIVFSQFastScan::encode_vectors( } else if (is_native_4bit(sq.qtype)) { sq.compute_codes(to_encode, codes_out, n); } else { - float_to_4bit_nibbles(to_encode, codes_out, n, d, M2, sq); + // add_core keeps the leftover bits in lo_codes_invlists. + std::vector lo_scratch(size_t(n) * lo_code_size()); + float_to_split_codes( + to_encode, codes_out, lo_scratch.data(), n, d, M2, sq); } if (include_listnos) { @@ -345,6 +353,34 @@ void IndexIVFSQFastScan::sa_decode(idx_t n, const uint8_t* bytes, float* x) } } +void IndexIVFSQFastScan::recombine_code( + int64_t list_no, + int64_t offset, + uint8_t* scratch, + int* values, + uint8_t* code_out) const { + // Interleaved for SIMD, so read one sub-quantizer at a time. + memset(scratch, 0, size_t(M2) / 2); + InvertedLists::ScopedCodes list_codes(invlists, list_no); + BitstringWriter bsw(scratch, M2 / 2); + for (size_t m = 0; m < M; m++) { + bsw.write( + pq4_get_packed_element(list_codes.get(), bbs, M2, offset, m), + 4); + } + InvertedLists::ScopedCodes lo(lo_codes_invlists, list_no); + split_to_codes( + scratch, lo.get() + offset * lo_code_size(), d, sq.qtype, values); + // BitstringWriter matches Codec6bit/Codec8bit, so distance_to_code reads + // it. + memset(code_out, 0, sq.code_size); + BitstringWriter out(code_out, sq.code_size); + const int bits = sq_bits(sq.qtype); + for (int m = 0; m < d; m++) { + out.write(values[m], bits); + } +} + void IndexIVFSQFastScan::reconstruct_from_offset( int64_t list_no, int64_t offset, @@ -353,10 +389,13 @@ void IndexIVFSQFastScan::reconstruct_from_offset( // Fallback: codes stored directly in invlists. InvertedLists::ScopedCodes codes(invlists, list_no); sq.decode(codes.get() + offset * sq.code_size, recons, 1); - } else if (orig_codes_invlists) { - // Rerank: use original full-precision codes for reconstruction. - InvertedLists::ScopedCodes codes(orig_codes_invlists, list_no); - sq.decode(codes.get() + offset * sq.code_size, recons, 1); + } else if (lo_codes_invlists) { + // Rerank: rebuild the original code from its two nibble halves. + std::vector scratch(M2 / 2), code(sq.code_size); + std::vector values(d); + recombine_code( + list_no, offset, scratch.data(), values.data(), code.data()); + sq.decode(code.data(), recons, 1); } else { // Native 4-bit: unpack from block inverted lists. std::vector code(M2 / 2, 0); @@ -416,6 +455,8 @@ void IndexIVFSQFastScan::search( { std::unique_ptr dc( sq.get_distance_computer(metric_type)); + std::vector scratch(M2 / 2), code(sq.code_size); + std::vector values(d); #pragma omp for for (idx_t i = 0; i < n; i++) { @@ -441,10 +482,13 @@ void IndexIVFSQFastScan::search( idx_t list_no = lo_listno(lo); idx_t offset = lo_offset(lo); - InvertedLists::ScopedCodes codes( - orig_codes_invlists, list_no); - float dis = dc->distance_to_code( - codes.get() + offset * sq.code_size); + recombine_code( + list_no, + offset, + scratch.data(), + values.data(), + code.data()); + float dis = dc->distance_to_code(code.data()); if (metric_type == METRIC_L2) { if (dis < heap_dis[0]) { @@ -516,6 +560,8 @@ void IndexIVFSQFastScan::search_preassigned( { std::unique_ptr dc( sq.get_distance_computer(metric_type)); + std::vector scratch(M2 / 2), code(sq.code_size); + std::vector values(d); #pragma omp for for (idx_t i = 0; i < n; i++) { @@ -541,10 +587,13 @@ void IndexIVFSQFastScan::search_preassigned( idx_t list_no = lo_listno(lo); idx_t offset = lo_offset(lo); - InvertedLists::ScopedCodes codes( - orig_codes_invlists, list_no); - float dis = dc->distance_to_code( - codes.get() + offset * sq.code_size); + recombine_code( + list_no, + offset, + scratch.data(), + values.data(), + code.data()); + float dis = dc->distance_to_code(code.data()); if (metric_type == METRIC_L2) { if (dis < heap_dis[0]) { @@ -611,6 +660,8 @@ void IndexIVFSQFastScan::range_search( std::unique_ptr dc( sq.get_distance_computer(metric_type)); + std::vector scratch(M2 / 2), code(sq.code_size); + std::vector values(d); #pragma omp for for (idx_t i = 0; i < n; i++) { @@ -622,20 +673,23 @@ void IndexIVFSQFastScan::range_search( if (list_no < 0) { continue; } - size_t list_size = orig_codes_invlists->list_size(list_no); + size_t list_size = lo_codes_invlists->list_size(list_no); if (list_size == 0) { continue; } - InvertedLists::ScopedCodes codes( - orig_codes_invlists, list_no); - InvertedLists::ScopedIds ids(orig_codes_invlists, list_no); + InvertedLists::ScopedIds ids(lo_codes_invlists, list_no); for (size_t j = 0; j < list_size; j++) { if (sel && !sel->is_member(ids[j])) { continue; } - float dis = dc->distance_to_code( - codes.get() + j * sq.code_size); + recombine_code( + list_no, + j, + scratch.data(), + values.data(), + code.data()); + float dis = dc->distance_to_code(code.data()); if (metric_type == METRIC_L2) { if (dis < radius) { qres.add(dis, ids[j]); @@ -701,28 +755,28 @@ void IndexIVFSQFastScan::add_with_ids( AlignedTable flat_codes(n * code_size); encode_vectors(n, x, idx.get(), flat_codes.get()); - // Store original SQ codes for rerank types - std::vector orig_flat_codes; + // Same quantisation as encode_vectors, or the halves disagree. + std::vector lo_flat_codes; if (needs_rerank(sq.qtype)) { - orig_flat_codes.resize(n * sq.code_size); + std::vector to_encode(size_t(n) * d); for (idx_t i = 0; i < n; i++) { - if (idx[i] < 0) { - continue; - } - if (by_residual) { - std::vector residual(d); - quantizer->compute_residual(x + i * d, residual.data(), idx[i]); - sq.compute_codes( - residual.data(), - orig_flat_codes.data() + i * sq.code_size, - 1); + if (by_residual && idx[i] >= 0) { + quantizer->compute_residual( + x + i * d, to_encode.data() + i * d, idx[i]); } else { - sq.compute_codes( - x + i * d, - orig_flat_codes.data() + i * sq.code_size, - 1); + memcpy(to_encode.data() + i * d, x + i * d, sizeof(float) * d); } } + lo_flat_codes.assign(size_t(n) * lo_code_size(), 0); + std::vector hi_discard(size_t(n) * (M2 / 2)); + float_to_split_codes( + to_encode.data(), + hi_discard.data(), + lo_flat_codes.data(), + n, + d, + M2, + sq); } // Sort by list assignment @@ -763,10 +817,10 @@ void IndexIVFSQFastScan::add_with_ids( code_size); if (needs_rerank(sq.qtype)) { - orig_codes_invlists->add_entry( + lo_codes_invlists->add_entry( list_no, id, - orig_flat_codes.data() + order[i] * sq.code_size); + lo_flat_codes.data() + order[i] * lo_code_size()); direct_map.add_single_id(id, list_no, ofs); } } @@ -790,9 +844,9 @@ void IndexIVFSQFastScan::add_with_ids( void IndexIVFSQFastScan::reset() { IndexIVFFastScan::reset(); - if (orig_codes_invlists) { + if (lo_codes_invlists) { for (size_t i = 0; i < nlist; i++) { - orig_codes_invlists->resize(i, 0); + lo_codes_invlists->resize(i, 0); } direct_map.clear(); } diff --git a/faiss/IndexIVFSQFastScan.h b/faiss/IndexIVFSQFastScan.h index a1a0b1397..c3c4ace16 100644 --- a/faiss/IndexIVFSQFastScan.h +++ b/faiss/IndexIVFSQFastScan.h @@ -10,6 +10,7 @@ #include #include #include +#include namespace faiss { @@ -29,10 +30,8 @@ namespace faiss { * ScalarQuantizer's own SIMD-optimised InvertedListScanner (same * behavior as IndexIVFScalarQuantizer, no fast-scan acceleration). * - * For reranked types, both the original full-precision codes and the - * packed 4-bit codes are stored in the inverted lists. The block - * inverted lists hold packed 4-bit codes; original codes are stored - * in a parallel ArrayInvertedLists (orig_codes_invlists). + * For reranked types the block inverted lists hold the packed high nibbles + * and a parallel ArrayInvertedLists (lo_codes_invlists) holds the low ones. */ struct IndexIVFSQFastScan : IndexIVFFastScan { ScalarQuantizer sq; @@ -40,10 +39,25 @@ struct IndexIVFSQFastScan : IndexIVFFastScan { /// Overselection ratio for reranking (default 2). float rerank_factor = 2; - /// Parallel inverted lists storing original SQ codes (for rerank - /// types). Owned by this index. nullptr for native 4-bit and - /// fallback types. - InvertedLists* orig_codes_invlists = nullptr; + /// Parallel inverted lists holding the leftover low bits of each code. + /// The top 4 bits of every code live in the packed scan blocks. + /// Owned by this index. nullptr for native 4-bit and fallback types. + InvertedLists* lo_codes_invlists = nullptr; + + /// Bytes per vector in lo_codes_invlists: the bits of each code left + /// over once the scan has taken the top 4. + size_t lo_code_size() const { + return sq_fastscan::sq_rerank_size(d, sq.qtype); + } + + /// Rebuild the code at (list_no, offset) from its two halves. `scratch` + /// needs M2 / 2 bytes, `code_out` needs sq.code_size. + void recombine_code( + int64_t list_no, + int64_t offset, + uint8_t* scratch, + int* values, + uint8_t* code_out) const; IndexIVFSQFastScan( Index* quantizer, diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index 822b52e45..cd3b33cad 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -2838,18 +2838,18 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { bool has_orig; READ1(has_orig); if (has_orig) { - ivfsqfs->orig_codes_invlists = read_InvertedLists(f, io_flags); - // Rebuild direct_map from orig_codes_invlists for reranking + ivfsqfs->lo_codes_invlists = read_InvertedLists(f, io_flags); + // Rebuild direct_map from lo_codes_invlists for reranking ivfsqfs->direct_map.set_type( DirectMap::Hashtable, ivfsqfs->invlists, 0); for (size_t list_no = 0; list_no < ivfsqfs->nlist; list_no++) { size_t list_size = - ivfsqfs->orig_codes_invlists->list_size(list_no); + ivfsqfs->lo_codes_invlists->list_size(list_no); if (list_size == 0) { continue; } InvertedLists::ScopedIds ids( - ivfsqfs->orig_codes_invlists, list_no); + ivfsqfs->lo_codes_invlists, list_no); for (size_t j = 0; j < list_size; j++) { ivfsqfs->direct_map.add_single_id(ids[j], list_no, j); } diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index f56f8eacf..29ca06f4b 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -995,11 +995,11 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(ivfsqfs->rerank_factor); write_ScalarQuantizer(&ivfsqfs->sq, f); write_InvertedLists(ivfsqfs->invlists, f); - // Write orig_codes_invlists if present - bool has_orig = (ivfsqfs->orig_codes_invlists != nullptr); + // Write lo_codes_invlists if present + bool has_orig = (ivfsqfs->lo_codes_invlists != nullptr); WRITE1(has_orig); if (has_orig) { - write_InvertedLists(ivfsqfs->orig_codes_invlists, f); + write_InvertedLists(ivfsqfs->lo_codes_invlists, f); } } else if ( const IndexIVFPQFastScan* ivpq_2 =