Skip to content

Guard ScalarQuantizer compute_codes/decode OMP region for small n#5423

Closed
alibeklfc wants to merge 1 commit into
facebookresearch:mainfrom
alibeklfc:export-D111693838
Closed

Guard ScalarQuantizer compute_codes/decode OMP region for small n#5423
alibeklfc wants to merge 1 commit into
facebookresearch:mainfrom
alibeklfc:export-D111693838

Conversation

@alibeklfc

Copy link
Copy Markdown
Contributor

Summary:
ScalarQuantizer::compute_codes and ScalarQuantizer::decode unconditionally
open an #pragma omp parallel for region regardless of n, unlike every
sibling quantizer in the same family: ProductQuantizer::decode guards with
if (n > 100), AdditiveQuantizer::decode/decode_unpacked guard with
if (n > 100)/if (n > 1000), and RaBitQuantizer::compute_codes_core/
decode_core guard with if (n > 1000).

IndexIVFScalarQuantizer::reconstruct_from_offset calls sq.decode(sc.get(), recons, 1) once per vector, so this hot reconstruction path (used by
reconstruct(), reconstruct_n(), and reranking/IndexRefine flows on top
of IVF+SQ indexes) pays full OMP thread-team fork/join overhead for a single
scalar-encode/decode of one vector. IndexScalarQuantizer::sa_encode/
sa_decode (the public codec API, also reachable from Python) hit the same
cost for any single-vector call.

Fix: add if (n > 100) to both pragmas, matching the threshold already used
by ProductQuantizer::decode and AdditiveQuantizer::decode. No behavior
change for n > 100 (parallel region still opens); for small n the loop now
runs on the calling thread directly instead of forking an OMP team.

Benchmarked on this devserver (56 OMP threads available), calling
ScalarQuantizer::compute_codes/decode with n=1, 20000 reps, best-of-run
median:

qtype d compute_codes(n=1) before after decode(n=1) before after
QT_8bit 128 105.7us 2.7us 104.5us 2.4us
QT_8bit 768 112.7us 4.1us 114.8us 2.6us
QT_fp16 128 117.3us 2.3us 113.0us 2.4us
QT_fp16 768 108.3us 2.9us 109.9us 2.6us

~40-50x reduction in per-call latency for the n=1 case across both qtypes
and dimensions tested. n > 100 behavior is unaffected (same code path).

Differential Revision: D111693838

Summary:
`ScalarQuantizer::compute_codes` and `ScalarQuantizer::decode` unconditionally
open an `#pragma omp parallel for` region regardless of `n`, unlike every
sibling quantizer in the same family: `ProductQuantizer::decode` guards with
`if (n > 100)`, `AdditiveQuantizer::decode`/`decode_unpacked` guard with
`if (n > 100)`/`if (n > 1000)`, and `RaBitQuantizer::compute_codes_core`/
`decode_core` guard with `if (n > 1000)`.

`IndexIVFScalarQuantizer::reconstruct_from_offset` calls `sq.decode(sc.get(),
recons, 1)` once per vector, so this hot reconstruction path (used by
`reconstruct()`, `reconstruct_n()`, and reranking/`IndexRefine` flows on top
of IVF+SQ indexes) pays full OMP thread-team fork/join overhead for a single
scalar-encode/decode of one vector. `IndexScalarQuantizer::sa_encode`/
`sa_decode` (the public codec API, also reachable from Python) hit the same
cost for any single-vector call.

Fix: add `if (n > 100)` to both pragmas, matching the threshold already used
by `ProductQuantizer::decode` and `AdditiveQuantizer::decode`. No behavior
change for n > 100 (parallel region still opens); for small n the loop now
runs on the calling thread directly instead of forking an OMP team.

Benchmarked on this devserver (56 OMP threads available), calling
`ScalarQuantizer::compute_codes`/`decode` with n=1, 20000 reps, best-of-run
median:

| qtype   | d   | compute_codes(n=1) before | after   | decode(n=1) before | after  |
|---------|-----|---------------------------|---------|---------------------|--------|
| QT_8bit | 128 | 105.7us                   | 2.7us   | 104.5us             | 2.4us  |
| QT_8bit | 768 | 112.7us                   | 4.1us   | 114.8us             | 2.6us  |
| QT_fp16 | 128 | 117.3us                   | 2.3us   | 113.0us             | 2.4us  |
| QT_fp16 | 768 | 108.3us                   | 2.9us   | 109.9us             | 2.6us  |

~40-50x reduction in per-call latency for the n=1 case across both qtypes
and dimensions tested. n > 100 behavior is unaffected (same code path).

Differential Revision: D111693838
@meta-cla meta-cla Bot added the CLA Signed label Jul 14, 2026
@meta-codesync

meta-codesync Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

@alibeklfc has exported this pull request. If you are a Meta employee, you can view the originating Diff in D111693838.

@meta-codesync

meta-codesync Bot commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

This pull request has been merged in 764b86f.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant