Repository navigation
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
bdice
added
bug
|
/ok to test c799bd9 |
There was a problem hiding this comment.
Notes for reviewers:
- Problem: this isn't a CI flake, it's a deterministically incorrect result that seemingly isn't caught by tests. I found it while reviewing workstream G (#2798) of the flaky-test sweep in #2783 (workstream K; the reproducer PR is #2800). When a query sits on a list centroid, for example a query equal to a row that is duplicated across a whole list, IVF-RaBitQ search returns nothing from that list: each such query returned 10 copies of the nearest other distinct row. It is deterministic. In the baseline run (38078120302, unmodified libcuvs), 18 of the 24 new test cases failed in all 20 repeats on both H100 and RTX PRO 6000.
- Root cause: the probe distance
||q - c||^2is computed in expanded form, and rounding can make it slightly negative. The block-sort search kernels takesqrtfof it for the error bound, which gives NaN, solow_dist < thresholdis false for every vector in that list. This affectsQUANT4/QUANT8at anybits_per_dim, andLUT16/LUT32withbits_per_dim > 1. The no-ex LUT kernels never compute the error bound, which is why the 6LUT16/LUT32cases withbits_per_dim = 1passed. - Fix: clamp the probe distances to
>= 0right after the GEMM. With the fix (run 38079761584), all 24 cases pass in every repeat, and the fullNEIGHBORS_ANN_IVF_RABITQ_TESTpasses. Mean recall over the 230 existing cases changed from 0.7479 to 0.7482 on RTX PRO 6000 and from 0.7440 to 0.7441 on H100, which is within run-to-run noise (details). - Not involved: the
max_topk_dist > 0check indevice_functions.cuhonly skips a threshold update, which is conservative. Open PRs #2291, #2373, #2377 and #2362 don't fix this. - Interactions: this is independent of #2798 and merges cleanly with it; the test is in its own file. #2291 changes the same spot in
PrepareClusterSearchInputs. Whichever PR lands second should apply the clamp after #2291's inner-productbinaryOp, and only tocentroid_distances, not tog_add, which can legitimately be negative for inner product.
| // Rounding in the expanded form can make ||q - c||^2 slightly negative when the query is on (or | ||
| // very near) a centroid. The search kernels take sqrt(||q - c||^2) for the error bound, which is | ||
| // NaN for a negative value, and a NaN lower bound drops every vector of that list, so clamp to 0. | ||
| raft::linalg::unaryOp(searcher.get_centroid_distances(), | ||
| searcher.get_centroid_distances(), | ||
| batch_size * num_centroids, | ||
| raft::plug_const_op<float, raft::max_op>{0.f}, | ||
| searcher_stream.get()); |
There was a problem hiding this comment.
The GEMM above accumulates -2 q.c onto ||q||^2 + ||c||^2 (L990-L1020). When q == c, cancellation can leave a tiny negative value. The search kernels read that value as q_g_add and take sqrtf(q_g_add) for the error bound: bitwise block sort L56-L58, LUT16 L52-L57, LUT emit-top-k L55. The NaN then propagates into low_dist, and low_dist < threshold is false for every vector in the list (bitwise L92-L96).
I put the clamp here because every search mode goes through this function: BatchClusterSearch, BatchClusterSearchLUT16 and BatchClusterSearchQuantizeQuery all call it (L1072, L1099, L1127). The same clamped values feed select_k for probe selection and est_dist in the kernels, so probe selection and the distance estimate use the same, valid distance.
Question for reviewers: this adds one elementwise pass over batch_size * n_lists floats per search batch. The alternative is sqrtf(fmaxf(q_g_add, 0.f)) in each kernel, as searcher_gpu_quantize_query.cu L117 already does for norm_quan. That has no extra pass, but every new kernel would have to remember the clamp, and the negative value would still reach select_k and est_dist. Which do you prefer?
| // Every row has kCopies > kK copies, so all kK neighbors must be exact copies of the query. | ||
| int64_t n_bad = 0; | ||
| for (int64_t q = 0; q < kDistinct; q++) { | ||
| bool bad = false; | ||
| for (int64_t j = 0; j < kK; j++) { | ||
| auto neighbor = h_neighbors[q * kK + j]; | ||
| auto dist = h_distances[q * kK + j]; | ||
| if (neighbor < 0 || neighbor >= kRows || neighbor % kDistinct != q || | ||
| !(std::abs(dist) <= 0.25f * gap[q])) { |
There was a problem hiding this comment.
Each list holds the 128 copies of one distinct row, so each query (one of the distinct rows) sits on its own list's centroid, and its probe distance to that list is zero up to rounding. Every list is probed, so the old failure isn't about probe selection: the query's own list is probed but contributes no candidates.
IVF-RaBitQ returns distances estimated from the quantized codes, so the check doesn't require exactly 0. It requires a distance within a quarter of the gap to the nearest other distinct row. A failing query returned exactly that gap; for example, query 0 at dim = 64 returned row 24 at 7.8915 (#2783 comment). The !(... <= ...) form also catches a NaN distance.
| ConfigureTest( | ||
| NAME NEIGHBORS_ANN_IVF_RABITQ_TEST | ||
| PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu | ||
| PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu neighbors/ann_ivf_rabitq/test_exact_match.cu |
There was a problem hiding this comment.
The new test is a separate translation unit in the existing NEIGHBORS_ANN_IVF_RABITQ_TEST binary. That keeps it out of the hunks #2798 adds to ann_ivf_rabitq.cuh and test_float_int64_t.cu, so the two PRs can merge in either order.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Summary
Merge Risk: ⚪ Minimal · up to The centroid-distance clamp and exact-match coverage appear ready to merge after normal checks. Pre-merge checks |
  |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. (1 skipped: 1… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
✅ Passed checks (4 passed)
-
Check name Status Explanation Description check The description clearly explains the negative probe-distance issue, the clamp fix, and the regression test. It directly matches the changeset. Title check The title clearly identifies the main change: fixing IVF-RaBitQ search for queries that match a list centroid. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. (1 skipped: 1 unsupported.)
- Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
-
🛠️ Fix failing CI checks 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
1 participant
IVF-RaBitQ computes the probe distance
||q - c||^2in expanded form (||q||^2 + ||c||^2 - 2 q.c). When a query sits on a list centroid, for example a query equal to a row that is duplicated across a whole list, rounding can make that value slightly negative. The block-sort search kernels takesqrtfof it for the error bound, which gives NaN. The NaN lower bound then failslow_dist < thresholdfor every vector in that list, so the query's own list contributes no candidates and all k results come from other lists. This affectsQUANT4/QUANT8(the default) at anybits_per_dim, andLUT16/LUT32withbits_per_dim > 1.This PR clamps the probe distances to non-negative values after the GEMM. It also adds a regression test where each list holds copies of one row and the queries are those rows. Without the fix, 18 of the 24 cases fail deterministically on H100 and RTX PRO 6000.
See #2800 and #2783.