mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Michael Norris
2839c39f20
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5682 This diff adds an option, `deterministic_tie_break`, that makes a GPU IVF search return the same neighbours on every run. It is off by default. **CHANGES THAT APPLY BY DEFAULT. No option turns these off. The id tie break itself is the new option `deterministic_tie_break`, which is off by default.** 1. **Tie order:** the interleaved IVF scan (`GpuIndexIVFFlat`, `GpuIndexIVFScalarQuantizer`) orders an exact distance tie by the position in the list, for `k` up to 1024. Before, the order followed the arrival order in the sort network. Only the choice among exactly equal distances can change. REASON: it blew up binary size to make it not like this. see below. 2. **Launch shape:** the first pass of the `k <= 32` interleaved scan runs with 64 threads in a block instead of 128. Measured on one A100 with `deterministic_tie_break` off: 1.8% faster than master. REASON: perf improvement. 3. **Padding:** pass 2 of the interleaved scan rejects the padding slots of a list shorter than `k`. On master a pad cannot win, because `Comparator` is strict. The tie order of item 1 lets a pad win a tie with the heap sentinel, so this check goes with item 1. REASON: used for deterministic, but should have no impact on non deterministic flow. 4. **`copy_subset_to` with `SUBSET_TYPE_ID_MOD`:** a vector with a negative id goes to the subset that its id selects. Before, it was dropped from every subset, and a multi-GPU clone with `shard_type` 1 failed its `ntotal` check. REASON: bug fix. 5. **`merge_knn_results`** (used by `IndexShards` and `IndexShardsIVF`): only `-1` marks a missing result. Before, every negative id did, so a sharded search with signed ids lost results. REASON: bug fix. 6. **Shared k-selection code** (`Select.cuh`, the merge networks): every comparison receives the value as well as the key. The default comparator ignores the value, so the results do not change. REASON: used for deterministic, but should have no impact on non deterministic flow. Below changes do not change the default results: - The RAFT `TieBreak` template parameter defaults to false, and generates the same code as upstream. - The new cuVS `search_params::select_algo` defaults to `kAuto`. **Behaviour before:** - A query whose candidates tie on distance can get a different neighbour on a second run. - `copy_subset_to` with `SUBSET_TYPE_ID_MOD` drops every vector with a negative id from every subset. - `merge_knn_results`, which `IndexShards` and `IndexShardsIVF` use, treats every negative id as a missing result. It drops a shard's results from its first negative id onward. - The `k <= 32` IVF scan runs its first pass with 128 threads in a block. **Behaviour after:** - With `deterministic_tie_break` on, a repeated search returns the same ids and the same distances. - With `deterministic_tie_break` off, the search behaves as before. Only the choice among exactly equal distances can differ from master, as it can between any two layouts. - `SUBSET_TYPE_ID_MOD` keeps a vector with a negative id, in the subset that its id selects. - `merge_knn_results` treats only `-1` as a missing result, so a sharded search returns negative ids. - The `k <= 32` IVF scan runs its first pass with 64 threads in a block. **Gating:** - `GpuIndexIVFConfig::deterministic_tie_break` and `GpuClonerOptions::deterministic_tie_break` default to false. The cloner copies `GpuClonerOptions::deterministic_tie_break` into the IVF config. The caller sets it when it creates or clones the index. - These indexes honour `deterministic_tie_break`: - `GpuIndexIVFFlat` and `GpuIndexIVFScalarQuantizer` with the interleaved layout, which is the default, for `k` up to 1024. This path needs no cuVS change. - The cuVS `GpuIndexIVFPQ` and the cuVS `GpuIndexIVFScalarQuantizer`, for `k` and `nprobe` up to 256, in a build that defines `FAISS_CUVS_HAS_STABLE_SELECT`. - The classic `GpuIndexIVFPQ` and cuVS IVF-Flat ignore `deterministic_tie_break`. - The interleaved scan ignores `deterministic_tie_break` above `k` = 1024. Its kernels for a larger `k` keep `Comparator`, because a tie break in their large merge networks costs much code size. - The cuVS indexes ignore `deterministic_tie_break` above 256, and in a build without the `FAISS_CUVS_HAS_STABLE_SELECT` macro. **Why the `FAISS_CUVS_HAS_STABLE_SELECT` macro exists:** - The cuVS path sets `search_params::select_algo` to `raft::matrix::SelectAlgo::kWarpDistributedShmStable`. - That field and that enum value exist only in a cuVS and RAFT tree that has the cuVS and RAFT changes in this diff. They are not in cuVS 26.02, and they are not upstream yet. - Code that names them does not compile against such a tree. The macro tells the compiler that the symbols exist. It does not turn `deterministic_tie_break` on. - The internal Buck build defines the macro for its patched cuVS 26.06 tree. The CMake build does not define it. - When upstream cuVS and RAFT have the symbols, a cuVS version check can replace the macro. **Why each change is needed:** - The GPU k-selection ordered on the distance alone, so an exact tie followed the order in which the data reached the sort network. - With `deterministic_tie_break` on, `TieBreakComparator` orders a tie by the user id, which the IVF pass-1 kernel carries as the value. Pass 1 loads an id only when its candidate can still enter the top `k`. - With `deterministic_tie_break` off, pass 1 keeps the list offset as before, and pass 2 translates the offset to the user id as before. - IVF pass 2 rejects the padding that pass 1 leaves in a short list. A padded slot could otherwise win a tie and reach the output. - `raft/util/bitonic_sort.cuh` and `raft/matrix/detail/select_warpsort.cuh` add a tie break on the payload index as a template parameter, `TieBreak`, false by default. With false, the generated code is the same as upstream. - The new RAFT selector `kWarpDistributedShmStable` uses the new class `warp_sort_distributed_ext_stable`, which sets `TieBreak` to true. It supports `float` and `half` keys, and fails for other key types, so their `select_k` instantiations do not grow. - cuVS `ivf_pq.hpp` and `ivf_sq.hpp` add `search_params::select_algo`, with the default `kAuto`. The IVF-PQ and IVF-SQ searches pass it to their `select_k` calls. With `deterministic_tie_break` on, faiss asks for `kWarpDistributedShmStable`. - A sharded index with signed ids lost about half of its neighbours. On one A100, 4 shards of an `IVF64,SQ8` index with 20,000 signed random ids agreed with the unsharded index on 28% to 29% of the neighbours, with the id split and with the list split. With the fix, they agree on 100%. - `copy_subset_to` tested `id % n == i`. C++ gives a negative remainder for a negative id, so the test never matched. The fold `((id % n) + n) % n` maps every id into `[0, n)`. - The first pass of the `k <= 32` scan is faster with 64 threads. The measurements follow. **A caller that wants a repeatable search must also do two things itself:** - Train the codec once and reuse it, because k-means is not repeatable. - Sort each inverted list by id after the build. cuVS IVF-PQ, cuVS IVF-SQ, and an index with `INDICES_CPU` or `INDICES_IVF`, break a tie on the position in the list, and position follows insertion order. **Build size:** - The RAFT tie break is a template parameter, so only the stable selector instantiates the tie-break code. The other `select_k` instantiations generate the same code as upstream. - The scan kernels for `k` above 1024 keep `Comparator`, because a tie break in their large merge networks costs much code size. - A local dev ASAN build of `libfaiss_gpu_faiss.so` is 829.3 MiB. The `.nv_fatbin` section is 265.1 MiB. - CI reports no build size or build speed regression signals. **Measurements** Every number comes from one A100 80GB unless otherwise noted, and an opt build: - Base: 2,000,000 vectors, d = 768, float32, uniform in `[0, 1)` from `np.random.RandomState(1234).rand`. - Queries: 20,000 vectors from the same generator. - Metric: L2. - Coarse quantizer: `IndexFlatL2` with `nlist` = 519, so a list holds about 3,850 vectors. Training uses the first 200,000 base vectors. - Timing: one `search` call over all 20,000 queries. QPS is 20,000 divided by the best wall time. - Noise: six builds of the unchanged code differ by 0.33%. *Classic GPU IVF scan.* `IndexIVFScalarQuantizer` with `QT_8bit`, cloned with `use_cuvs = False`. `k` = 10, `nprobe` = 16, 61,657 candidates for each query. Best of 3 searches, two runs for each build: | build | QPS | | --- | --- | | master | 7,740 / 7,732 | | this diff, `deterministic_tie_break` off | 7,862 / 7,878 | | this diff, `deterministic_tie_break` on | 7,857 / 7,861 | - `deterministic_tie_break` costs about 0.1%, which is inside the noise. - With `deterministic_tie_break` off, this diff is 1.8% faster than master. The gain comes from the 64-thread first pass. First-pass threads in a block, same data, `deterministic_tie_break` on: | threads | QPS | | --- | --- | | 32 | 7,240 | | 64 | 7,877 | | 128 (the master launch shape) | 7,493 | | 256 | 6,428 | 512 threads does not build: `FinalBlockMerge` has no specialisation for 16 warps. The 64-thread gain over 128 holds across shapes on A100, H100, GB300. The H100 and GB300 rows below ran with `deterministic_tie_break` off: ``` ┌───────┬─────┬─────┬────────┬────────────────────────────┬─────────────────────────────────────┬─────────────────┐ │ GPU │ d │ k │ nprobe │ 64 threads (run 1 / run 2) │ 128 threads, master (run 1 / run 2) │ gain │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ H100 │ 768 │ 10 │ 16 │ 15,726 / 15,691 │ 14,316 / 14,258 │ +9.9% / +10.1% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ H100 │ 768 │ 32 │ 16 │ 15,668 / 15,645 │ 14,273 / 14,228 │ +9.8% / +10.0% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ H100 │ 128 │ 10 │ 16 │ 144,849 / 145,020 │ 125,843 / 125,592 │ +15.1% / +15.5% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ H100 │ 768 │ 10 │ 64 │ 3,946 / 3,946 │ 3,585 / 3,580 │ +10.1% / +10.2% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ GB300 │ 768 │ 10 │ 16 │ 16,627 / 16,629 │ 8,693 / 8,694 │ +91.3% / +91.3% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ GB300 │ 768 │ 32 │ 16 │ 16,576 / 16,570 │ 8,670 / 8,672 │ +91.2% / +91.1% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ GB300 │ 128 │ 10 │ 16 │ 157,620 / 157,534 │ 79,062 / 79,100 │ +99.4% / +99.2% │ ├───────┼─────┼─────┼────────┼────────────────────────────┼─────────────────────────────────────┼─────────────────┤ │ GB300 │ 768 │ 10 │ 64 │ 4,116 / 4,116 │ 2,167 / 2,167 │ +89.9% / +89.9% │ └───────┴─────┴─────┴────────┴────────────────────────────┴─────────────────────────────────────┴─────────────────┘ ``` A thread queue of 3 was 2.5% faster than 2 on the A100 and level on an H100, so this diff keeps 2. *cuVS IVF-PQ.* `GpuIndexIVFPQ` with `use_cuvs = True`, `M` = 384, 8 bits for each code, trained and filled on the GPU, copied to the CPU, then cloned twice: once with `deterministic_tie_break` off (`kAuto`) and once with it on (`kWarpDistributedShmStable`). Timed searches alternate between the two clones. Best of 5, two runs: | k | nprobe | `deterministic_tie_break` off | `deterministic_tie_break` on | cost | | --- | --- | --- | --- | --- | | 10 | 16 | 16,078 / 15,984 | 16,043 / 15,964 | +0.22% / +0.13% | | 100 | 64 | 3,967 / 3,965 | 3,988 / 3,960 | -0.52% / +0.11% | - The stable selector costs about 0.1% to 0.2%, which is inside the noise. - The two selectors returned different ids for 0.0015% to 0.045% of the results. These are exact distance ties: PQ sums table entries, so equal distances are more frequent than with raw floats. *cuVS IVF-SQ.* `GpuIndexIVFScalarQuantizer` with `QT_8bit` and `use_cuvs = True`, `INDICES_64_BIT`, built on the CPU, then cloned twice: once with `deterministic_tie_break` off (`kAuto`) and once with it on (`kWarpDistributedShmStable`). Timed searches alternate between the two clones. All 5 searches of both runs, in QPS: | run | k | nprobe | `deterministic_tie_break` off | `deterministic_tie_break` on | best off / best on | | --- | --- | --- | --- | --- | --- | | 1 | 10 | 16 | 14,189 / 14,290 / 14,359 / 14,311 / 14,321 | 14,252 / 14,266 / 14,337 / 14,340 / 14,245 | 14,359 / 14,340 | | 1 | 100 | 64 | 2,923 / 2,945 / 2,950 / 2,952 / 2,951 | 2,932 / 2,950 / 2,943 / 2,952 / 2,936 | 2,952 / 2,952 | | 2 | 10 | 16 | 14,232 / 14,290 / 14,288 / 14,259 / 14,143 | 14,219 / 14,280 / 14,278 / 14,184 / 14,301 | 14,290 / 14,301 | | 2 | 100 | 64 | 2,934 / 2,941 / 2,949 / 2,947 / 2,945 | 2,934 / 2,944 / 2,941 / 2,941 / 2,936 | 2,949 / 2,944 | - The best-of-5 difference is between -0.17% and +0.08%, which is inside the noise. - The two selectors returned different ids for 0.0090% (k = 10) and 0.0578% (k = 100) of the results, in both runs. These are exact distance ties. Reviewed By: alibeklfc Differential Revision: D121990149 fbshipit-source-id: ad68d32289da529356fc77a0737d3c637600f976