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. |
cbcfb91 to
9e8ca0d
Compare
The random donor selection took candidate points from a shared atomic counter, so which receiver got which donor depended on the order in which the warps ran. Derive each receiver's candidates from its index instead, with a stride that is coprime with the number of rows, so that every receiver still finds a donor whenever one exists. The pseudo-random donor offset rotated with a function-local static counter, so a fit depended on how many fits had run before it in the same process. Rotate it with the balancing step instead.
9e8ca0d to
ed9ac8b
Compare
…ally built libcuvs
bdice
changed the title
…ally built libcuvs
The per-cluster codebook training listed each cluster's rows in the order given by an atomic counter, so the per-cluster training subsets depended on scheduling. Group the row indices with a stable sort by label instead, so each cluster's rows are listed in increasing order.
ed9ac8b to
fe6517c
Compare
…ally built libcuvs
…ally built libcuvs
|
/ok to test fe6517c |
bdice
added
bug
bdice
left a comment
bdice
left a comment
There was a problem hiding this comment.
This PR doesn't fix one specific flaky test. It removes two scheduling-dependent choices from the balanced k-means path that the IVF-Flat/PQ/SQ/RaBitQ builds use, plus a process-history dependence. All three came up in workstream E of #2783 (comment). It is the narrowed follow-up to #2801 and independent of it.
- Problem: the random donor selection (used by IVF-PQ codebook training) drew candidates from a shared atomic counter, so the donor each receiver got depended on warp order. The donor offset rotated with a function-local
staticcounter, so a second build in the same process used different offsets. IVF-PQ per-cluster codebooks were trained on rows grouped by anatomicAddkernel, so the training subset depended on scheduling. - Fix: each receiver derives its candidate sequence from its own index, with a stride coprime with
n_rows. The offset rotates with the balancing step. The per-cluster grouping uses a stable sort. - Still not bitwise reproducible:
calc_centers_and_sizesstill sums centers with float atomics viareduce_rows_by_key(main), as the description says. I ask about a follow-up inline. - Perf (#2803, run 38154672010): IVF/CAGRA builds are within about ±3%. The exceptions are A100 IVF-SQ 1M (+6.3%) and IVF-RaBitQ 1M (+5.2%), whose per-sample ranges overlap main's. Recall is unchanged. One H100 microbenchmark of 64 tiny fits was +12–16% in 2 of 5 configs. Nsight Systems shows more EM steps from a different donor trajectory there, not slower kernels.
- CI on
fe6517c6(run 38168567842): one test job failed,conda-cpp-tests (3, 4) / 13.3.0 / amd64 / rtxpro6000(job). The failure wasScaNN/f32_i64.build_host_input/12(max reconstruction error 1.4321 vs. the 1.425 limit). That is the same test and threshold as the open flake #2549, which predates this PR. ScaNN's build does use balanced k-means (scann_build.cuh), though, so I'd like a re-run of that job before calling it unrelated. The other failing checks are Label Checker and thepr-builderaggregate.
| // The candidate sequence of each receiver is derived from its index (not from a shared atomic | ||
| // counter), so the donor points do not depend on the order in which the warps run. | ||
| // `stride >= n_clusters`, so different receivers use different counters (and thus, before | ||
| // wrapping around, different candidates); `stride` is coprime with `n_rows`, so each receiver | ||
| // visits every row within `n_rows` attempts and finds a donor whenever one exists. | ||
| IdxT attempt = 0; | ||
| do { | ||
| auto old = atomicAdd(search_count, IdxT{1}); | ||
| auto candidate = | ||
| static_cast<IdxT>((static_cast<int64_t>(seed) * static_cast<int64_t>(old + 1)) % | ||
| static_cast<int64_t>(n_rows)); | ||
| auto counter = static_cast<int64_t>(attempt) * static_cast<int64_t>(stride) + | ||
| static_cast<int64_t>(receiver_cluster); | ||
| auto candidate = static_cast<IdxT>((static_cast<int64_t>(seed) * (counter + 1)) % | ||
| static_cast<int64_t>(n_rows)); |
There was a problem hiding this comment.
Before this change, every receiver warp took its next candidate from one shared counter (main), so which receiver saw which candidate depended on the order in which the warps ran. Now receiver r at attempt a uses counter = a · stride + r:
stride ≥ n_clusterskeeps the counters of different receivers disjoint.strideis coprime withn_rows(computed on the host, L864-L868). The seedofstis a prime that doesn't dividen_rows(L852-L854). Soseed · (counter + 1) mod n_rowsvisits every row withinn_rowsattempts, and theattempt < n_rowsbound still guarantees that a donor is found whenever one exists.
Only balanced_donor_selection::Random uses this kernel, and only the IVF-PQ codebook training selects it (per-subspace, per-cluster). The candidate sequence differs from main's, so the codebooks change, but the IVF-PQ recall@10 medians in the description stay within main's min–max.
| // Rotate the pseudo-random offset with the balancing step rather than with a process-wide | ||
| // counter, so that the donor points do not depend on how many fits ran before in the process. | ||
| IdxT ofst; | ||
| size_t i_primes = balancing_step; | ||
| do { | ||
| i_primes = (i_primes + 1) % kPrimes.size(); | ||
| ofst = kPrimes[i_primes]; | ||
| ofst = kPrimes[i_primes++ % kPrimes.size()]; | ||
| } while (n_rows % ofst == 0); |
There was a problem hiding this comment.
This replaces the function-local static IdxT i_primes (main, used here). That counter carried over between fits, so the N-th build in a process picked different offsets from the first one. It was also an unsynchronized static shared by every host thread that builds concurrently. Now the offset index comes from the EM iteration (iter at L995).
This changes the donor trajectory, and with it the number of EM steps a fit takes. That explains the +12–16% in 2 of 5 configs of the H100 64-tiny-fit microbenchmark: Nsight Systems shows 18,525 vs. 15,217 distance kernels at the same per-kernel cost, and in that profiled run the total wall time was 2% lower than main. The IVF-PQ builds that run this per-subspace training are within ±2%.
| // Group the row indices by cluster with a stable sort, so that the order of the rows within a | ||
| // cluster (which decides the per-cluster training subsets) is the same in every build. | ||
| rmm::device_uvector<uint32_t> sorted_labels(n_rows, stream); | ||
| raft::copy(sorted_labels.data(), labels, n_rows, stream); | ||
| thrust::sequence(exec_policy, data_indices, data_indices + n_rows); | ||
| thrust::stable_sort_by_key(exec_policy, sorted_labels.begin(), sorted_labels.end(), data_indices); |
There was a problem hiding this comment.
Before this change, fill_indices_kernel placed each row with atomicAdd on its cluster's offset (main, launch), so the order of rows within a cluster depended on scheduling. That order matters because per-cluster codebook training gathers the cluster's residuals in that order and then keeps only the first pq_n_rows of them (main). The stable sort lists each cluster's rows in increasing row order. The per-cluster-codebook build is -2.2% (H100) / -0.4% (A100) vs. main in the description's table.
Question for reviewers: there's no direct test for this function, because it dereferences *thrust::max_element(...) on the host (L238-L239) and so only works with the managed cluster_sizes that build() passes. (A test with plain device memory segfaulted in the stress run.) Should I compute the max on the device and add a direct test here, or leave that for a separate PR?
| /* | ||
| * The balancing step must pick the same donor points for identical inputs: they may depend neither | ||
| * on the order in which the GPU runs the warps nor on how many fits ran before in the process. | ||
| * Most rows are packed in a few clusters and many clusters are empty, so that many receivers need a | ||
| * donor. | ||
| */ | ||
| void check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection donor) |
There was a problem hiding this comment.
This calls adjust_centers three times on identical inputs and requires bitwise-identical centers, for both donor strategies. On main, the static counter gives each call a different offset. The test covers the balancing step only, not a full fit, because the center sums in a fit are still float atomics. Note that n_rows = 10007 is prime, so the host-side gcd loop returns n_clusters immediately and its search isn't exercised here.
Question for reviewers: do you want a follow-up that makes the balanced center sums (calc_centers_and_sizes, main) deterministic with a fixed-order sum, at least for the small problems balanced k-means runs (such as the per-subspace / per-cluster PQ codebook fits)? It could reuse the sort-and-sum approach from #2801. Or is a non-bitwise-reproducible IVF build acceptable as it is?
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Summary
Merge Risk: ⚪ Minimal · up to No actionable merge-blocking issue is established; the change is ready for normal merge checks. Pre-merge checks |
  |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (2 skipped: 2… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
✅ Passed checks (4 passed)
-
Check name Status Explanation Description check The description directly explains the scheduling-independent donor selection and IVF-PQ grouping changes, including scope, testing, performance, and numerical limitations. Title check The title clearly and concisely identifies the two main changes: scheduling-independent balanced k-means donor selection and IVF-PQ codebook grouping. 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 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (2 skipped: 2 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
Balanced k-means (used by the IVF-Flat, IVF-PQ, IVF-SQ and IVF-RaBitQ builds) picks donor points for underfull clusters, and the IVF-PQ per-cluster codebook training picks the training rows of each cluster. Both choices depended on GPU scheduling, and the donor choice also depended on how many fits had already run in the process. This PR makes both depend only on the inputs:
staticcounter, so a second build in the same process used different offsets. It now rotates with the balancing step.This PR does not make balanced k-means bitwise reproducible:
calc_centers_and_sizesstill sums each cluster's rows with floating-point atomics (unchanged frommain), so centers can still differ in the last bits between identical builds. Fully deterministic balanced sums are left out on purpose. A possible follow-up is a fixed-order sort-and-sum kernel for the small problems that balanced k-means runs, without fixed-point arithmetic.Tests:
KmeansBalancedAdjustCenters.Reproducible{SizeSorted,RandomDonor}call the balancing step three times on identical inputs and require identical centers (with thestaticcounter onmain, each call used a different offset). There is no dedicated test for the IVF-PQ grouping.calculate_offsets_and_indicesdereferences a device pointer on the host (*thrust::max_element), so it only works with the managed memory thatbuild()passes. The existing IVF-PQ tests, which include per-cluster codebooks, cover it.Performance
Method:
main(24b0b17) andmain+ this PR are built identically in the same CI run (CUDA 13.3, driver 595.91). One benchmark binary is linked against eachlibcuvs.so, and the variants are interleaved in 5 rounds with rotating order, 1 warmup + 2 timed builds per round (10 samples). Data is Gaussian blobs, with 10K queries from the same distribution. A third build (#2801, which changes only regular k-means and so has code identical tomainon these paths) serves as an A/A control: it stays within -2.3%..+1.9% ofmainon every row below.IVF-Flat/PQ/SQ/RaBitQ and CAGRA builds are within -2.2%..+1.5% on H100. On A100 (arm64 host) they range from -3.2% to +6.3%. IVF-SQ 1M (+6.3%) and IVF-RaBitQ 1M (+5.2%) are the only builds above +3% (plus the int8 balanced fit, +4.9%), their per-sample ranges overlap
main's (SQ 228.6–249.9 vs. 223.5–237.2 ms), and the same builds at 5M are -1.6%. GH200, GB300, RTX PRO 6000 and L4 (IVF-Flat/IVF-PQ at 1M and 5M, CAGRA 1M, 5 samples each) are within -2.6%..+4.3%. The RTX PRO 6000 numbers were bimodal on that runner. Search is unchanged (QPS within ±1.7%, except one A100 config whosemainsamples alone span 1785–2611 ms).The cost of balanced k-means depends on how many EM iterations the balancing triggers, and that depends on which donor points are picked. This PR changes that choice (now a fixed function of the input), so individual fits can take more or fewer iterations than on
main. In a microbenchmark of 64 back-to-back IVF-PQ-codebook-sized fits (65536 x 2, k=256, size-sorted or random donors, same or distinct data per fit), A100 is within -2.0%..+1.5% in all 6 measurements. On H100, size-sorted donors refitting the same data are +16.0% and +16.5% (two runs), random donors on distinct data +12.6%, random donors on the same data +3.6% and -2.4%, and size-sorted donors on distinct data +0.6%. An Nsight Systems run of the +16.5% case (same data refit 64 times, size-sorted donors) shows the cause: 18,525 vs. 15,217 distance kernels, i.e. more E steps, with the same per-kernel cost. In that profiled run, total wall time was 2% lower thanmain. These fits are launch- and allocation-bound (about 5,000 kernel launches and 2,700cudaMalloc/cudaFreeper fit). The IVF-PQ builds that run this per-subspace training are within ±2% (table above).Numerical accuracy
The center sums are unchanged from
main(float atomics), so centers still differ in the last bits between identical builds. Build quality is unchanged. H100 recall@10 (median, min–max over 10 builds) formainvs. this PR:Of the other index configs, all 14 on A100 and 9 of 10 on H100 have their median inside
main's min–max. The exception, IVF-Flat 10M x 96, is 0.9769 vs. 0.9769–0.9773. Balanced-fit inertia, computed in float64 for 1M x 128, k=1024, is 5.11297e8 for both.C++ tests of this exact source (
CLUSTER_TEST*KmeansBalanced*,KmeansBalancedAdjustCenters.*x50, fullNEIGHBORS_ANN_IVF_PQ_TESTandNEIGHBORS_ANN_IVF_FLAT_TEST) pass on H100 and L4.Logs: https://github.com/NVIDIA/cuvs/actions/runs/38154672010 (H100 IVF job, A100 IVF job, tests H100/L4). 64-fit microbenchmarks and profiles: https://github.com/NVIDIA/cuvs/actions/runs/38160888485, https://github.com/NVIDIA/cuvs/actions/runs/38165062741. Benchmark code: #2803.
Part of #2783. Independent of #2801.