[cuda] make BitonicArgSort_1024 / _2048 stable on tied values

[maxwbuckley]

Summary

BitonicArgSort_1024 and BitonicArgSort_2048 in include/LightGBM/cuda/cuda_algorithms.hpp used the comparator

if ((scores[indices[a]] > scores[indices[b]]) == ascending) { /* swap */ }

For descending sort (ascending == false) and equal scores a == b, this evaluates to false == false → true — i.e., ties get swapped. The output index permutation for tied elements then depends on the bitonic network structure rather than input order. CPU code uses std::stable_sort, which never swaps ties.

Symptom

LambdaRank round-1 has all scores = 0. The pair-iteration gradient computation depends on the sorted index order. With non-stable sort, CUDA's pair assignments diverge from CPU's, producing different gradients on the very first round.

Test: 10 queries × 20 items, lambdarank
Before this fix:  CPU/CUDA max|Δ| raw_score round 1 = 0.29
After this fix:   CPU/CUDA max|Δ| raw_score round 1 = 0.14

The remaining 0.14 is FP-precision in atomicAdd_block accumulation order — documented expected behavior per lightgbm-org#6055 — and is unaffected by this change.

Fix

Replace the comparator with a strict-direction form so equal elements never trigger a swap:

-          if ((scores[indices[threadIdx.x]] > scores[indices[index_to_compare]]) == ascending) {
+          const bool need_swap = ASCENDING
+              ? (scores[indices[threadIdx.x]] > scores[indices[index_to_compare]])
+              : (scores[indices[threadIdx.x]] < scores[indices[index_to_compare]]);
+          if (need_swap) {
             ...
           }

Same change applied to the analogous block in BitonicArgSort_2048.

Other call sites

BitonicArgSort_1024 is also used by the categorical split-finder kernel (cuda_best_split_finder.cu:640) for many-vs-many splits. The CPU/CUDA parity sweep was re-run after this fix and the categorical case (and every other previously-clean case) still matches at FP epsilon — see verification below.

Test plan

• Added test_cuda_lambdarank_round1_matches_cpu_within_fp_drift in tests/python_package_test/test_dual.py (gated on LIGHTGBM_TEST_CUDA=1). Asserts max|Δ| < 0.2 — strict enough to catch the regression, loose enough for the FP-precision residual.
• Re-ran the full scratch/cpu_cuda_parity.py sweep on RTX 5090 / CUDA 13.2 — no regression on any case (categorical, regression, binary, multiclass, etc. all still match at FP epsilon where they did before).

🤖 Generated with Claude Code

[BelixRogner]

Thanks Max, and Claude Code. The motivation is right — (a > b) == ascending swapping ties when ascending == false is a real bug, and lambdarank round-1 with all-zero scores is the perfect case to surface it.

But I want to flag what I think is a regression in the BitonicArgSort_1024 half of the fix before merging.

The two halves of the diff disagree on which variable to consult. In _2048, the comparator (correctly) reads lowercase ascending — the per-pass local. In _1024, the comparator was switched to uppercase ASCENDING — the template parameter. They should be the same variable.

Look at the upstream _1024 body around line 240:

template <..., bool ASCENDING>
__device__ ... BitonicArgSort_1024(...) {
  for (outer_depth ...) {
    ...
    const bool ascending = ASCENDING ? (outer_segment_index % 2 == 0) : (outer_segment_index % 2 > 0);
    //          ^^^^^^^^^^^^^^^^^^^^^^^ per-pass local, NOT the template parameter
    for (inner_depth ...) {
      if ((... > ...) == ascending) { swap }
      //                  ^^^^^^^^^ uses the per-pass local

That's textbook bitonic sort: outer phases alternate direction so the inner merge can pull a bitonic sequence into a monotone one. The comparator has to consult per-pass ascending, not the global ASCENDING.

PR #10's _1024 fix replaces the comparator with:

const bool need_swap = ASCENDING  // ← uppercase, template param
    ? (scores[...] > scores[...])
    : (scores[...] < scores[...]);

I believe that breaks the bitonic merge for non-tied inputs. The reason your lambdarank test passes despite the breakage is exactly the same reason the original bug only surfaced there: round-1 scores are all zero, so the strict comparator returns false regardless of direction — no swaps, no exposure of the sort being wrong.

The _2048 change in the same diff uses lowercase ascending and looks correct.

BitonicArgSort_1024 is also called by the categorical split-finder at cuda_best_split_finder.cu:640 with real, distinct scores. If _1024's comparator no longer alternates direction by pass, that sort should produce wrong orderings on inputs that aren't all-equal. Your parity sweep showed no regression on the categorical case — I'd guess the call happens to use ASCENDING == true and the relevant per-pass ascending is also true on the inputs you tested, i.e. they coincide by luck on small inputs.

Suggested fix for _1024 — symmetric with what you have in _2048:

const bool need_swap = ascending  // lowercase, the per-pass local
    ? (scores[indices[threadIdx.x]] > scores[indices[index_to_compare]])
    : (scores[indices[threadIdx.x]] < scores[indices[index_to_compare]]);

A test with non-trivial input — e.g. a categorical-split sort with all-distinct, well-separated scores, asserting CPU and CUDA agree on the resulting permutation — would also catch this case that lambdarank round-1 can't.

Happy to be wrong here if you can show me an input that exercises _1024's alternating-direction passes and validates the output is correct with the ASCENDING-only comparator. On read though I think the per-pass ascending is what we need.

Same env var convention note as #6/#7/#8/#9 — align to TASK=cuda while you're updating.

[maxwbuckley]

Good catch — you're right. Pushed f47e0499 with three changes:

1. BitonicArgSort_1024 comparator now reads the per-pass ascending local, not the template ASCENDING. The fix is now symmetric with _2048. Your read of why the lambdarank test passed despite the breakage was exactly correct: round-1 all-zero scores meant the strict comparator returned false either way, so the broken comparator never made a wrong swap on that input.

2. Added test_cuda_bitonic_argsort_1024_with_distinct_scores_matches_cpu in test_dual.py — trains a regression on a single categorical feature with distinct, well-separated per-category gradient/hessian sums, which exercises BitonicArgSort_1024 over non-tied scores via the categorical split-finder at cuda_best_split_finder.cu:640. Asserts CPU and CUDA agree on predictions after one round; the wrong-comparator case would diverge.

3. Env var aligned to TASK=cuda.

[BelixRogner]

Quick nudge — the chronic Python - latest versions failure that was hitting every open PR was fixed in #15 (merged earlier today) and the propagated update-branch on this PR confirms CI is otherwise green. Only thing keeping this one from merging now is the ruff/cpplint issue on test_dual.py (and on cuda_rank_objective.cpp for #14):

• PR [cuda] enforce max_depth on CUDA tree learner #7: PT006 — parametrize("max_depth,num_leaves", …) wants a tuple.
• PR [cuda] fix illegal-memory-access crash in weighted L1 / quantile training #8: I001 — imports unsorted.
• PR [cuda] make BitonicArgSort_1024 / _2048 stable on tied values #10: I001 + ruff format reformat.

Single pre-commit run --all-files && git commit --amend --no-edit && git push --force-with-lease should clear all three (and #14 if you do those together).

Let me know when pushed and I'll merge.

[BelixRogner]

Quick rebase nudge — #7 and #8 just landed on master and touched files this PR also modifies (cuda_algorithms.hpp / cuda_best_split_finder.cu / etc.), so this branch now shows a merge conflict on GitHub. One more git merge master && git push round should clear it; CI is otherwise green (the apparent failures yesterday were all environmental — dask socket flakes, a cancelled job rolling up, and a Boost-headers wheel-build issue, none of which were touching the actual PR content).

Ready to merge as soon as the conflict's resolved.