Batch the IVF ScalarQuantizer scan four codes at a time - #5456
Batch the IVF ScalarQuantizer scan four codes at a time#5456blazingphoenix7 wants to merge 1 commit into
Conversation
The SQ inverted-list scanners call distance_to_code one code at a time via run_scan_codes1, even though the SQ distance computers already expose query_to_codes_batch_4 (the batched kernel HNSW uses). Add run_scan_codes4, which distances four consecutive codes per step and applies the threshold and heap bookkeeping in id order so the heap-update sequence matches run_scan_codes1 exactly, then route the SQ scanners' no-selector path through it. The batched kernel reconstructs each code, which is bit-identical to the single-code path for every quantizer except the uniform ones, whose single-code path predecodes the query. A compile-time check keeps the uniform quantizers on the scalar scan, so search results are unchanged for every quantizer type. The selector path and all non-SQ scanners are unchanged.
|
@mnorris11 could I ask you to take a look at this one when you get a chance? I opened it the same day as #5457, which was imported and merged last week, so I think this one may have just been missed rather than held up on anything. It routes the no-selector IVF ScalarQuantizer scan through the four-at-a-time distance kernel the SQ computers already expose and HNSW already uses, applying the threshold compare and heap updates in the same id order as the scalar scan. Distances and ids come out bit-identical to that scan for QT_8bit, QT_8bit_uniform, QT_6bit, QT_4bit, QT_fp16 and QT_8bit_direct in both metrics, and end to end L2 search is about 1.5x faster at d=128 and 1.6x at d=256. Happy to rebase it or split it up if that would make review easier. |
1 participant
The IVF ScalarQuantizer inverted-list scan distances one code at a time through
run_scan_codes1, even though the SQ distance computers already expose a four-at-a-time kernel,query_to_codes_batch_4(the one HNSW uses throughdistances_batch_4). This addsrun_scan_codes4, a batched inner loop that distances four consecutive codes per step and then applies the same threshold compare,add_result, and threshold refresh asrun_scan_codes1, in id order, so the sequence of heap updates is identical. The two SQ scanners route their no-selector path through it; the selector path stays onrun_scan_codes1.Four independent accumulator chains let the out-of-order core overlap the latency-bound reduction that a single code cannot hide. Results are unchanged: the batched kernel reconstructs each code, which is bit-identical to the single-code path for every quantizer except the uniform ones, whose single-code path predecodes the query, so a compile-time check keeps those on the scalar scan. Verified bit-identical distances and ids against the scalar scan for QT_8bit, QT_8bit_uniform, QT_6bit, QT_4bit, QT_fp16, and QT_8bit_direct, in both metrics, at d in {64, 128, 256}.
Measured on AVX2, end-to-end
IndexIVFScalarQuantizersearch of QT_8bit, single thread, batched against the scalar scan:The L2 gain is the larger one, around 1.5-1.6x, and it holds when the index spills out of cache. Inner product gains less since its shorter chain is already partly overlapped by the core. When batching does not help it matches the scalar scan, so there is no regression.