[CUB] Make ReduceByKey run-to-run deterministic

[ahmedtaha100]

@nanan-nvidia @guillaume-michel

Description

closes #9995

DeviceReduce::ReduceByKey can produce different FP64 result bits across identical calls when a segment spans multiple tiles. Cross-tile partial reductions are combined in an order that can depend on tile completion.

This change opts the current public ReduceByKey dispatch into the stable ordering already supported by TilePrefixCallbackOp. The shared kernel defaults to the previous ordering, so RLE and the deprecated dispatcher retain their current behavior. The virtual shared-memory sizing path uses the same stable policy as the kernel.

The regression reproduces the reported 65,536-item input with a 50,000-item first run and compares that aggregate bit-for-bit across 20 identical calls. It runs in one generated test shard. Existing parameterized tests cover the run count and complete output correctness.

Verification

• The checked-in regression fails on the unfixed base with 0xbfa59be1de5076d4 versus 0xbfa59be1de5076d8.
• The patched exact reproducer produces 0xbfa59be1de5076d8 in all 200 calls on an RTX 5090.
• All four ReduceByKey type shards, the ReduceByKey virtual-shared-memory target, and the RLE compatibility target pass: 8,212 assertions across 51 test cases.
• The reproducer compiles for sm75, sm90, and sm120. Runtime execution was available for sm120.
• Repository pre-commit hooks pass for all three changed files.

Performance

Measurements used CUDA 13.3 on the RTX 5090. The exact reported workload changed from 13.477 us to 13.997 us, a 3.86% increase based on the median of three order-balanced process pairs.

Selected U64/F64/I32 runs of the existing ReduceByKey NVBench benchmark:

Elements	Maximum segment	Base (us)	Patched (us)	Change	Maximum noise
2^16	2^8	14.410	15.415	+6.98%	8.246%
2^20	2^8	24.926	30.797	+23.55%	0.870%
2^24	2^1	296.125	297.260	+0.38%	0.214%
2^24	2^4	203.515	213.069	+4.69%	0.318%
2^24	2^8	193.411	207.811	+7.45%	0.063%

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 619b0dc6-8952-4847-b36a-5ea3b3b649bb

📥 Commits

Reviewing files that changed from the base of the PR and between 0720444 and 009b23f.

📒 Files selected for processing (1)

• cub/test/catch2_test_device_reduce_by_key.cu

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added an option for reduce-by-key operations to preserve a stable reduction order.
  • Enabled stable ordering in supported reduce-by-key execution paths by default.

• Bug Fixes

  • Improved run-to-run determinism for reductions using custom operations and floating-point data.

• Tests

  • Added coverage verifying that repeated reductions produce identical aggregate results.

Walkthrough

The reduce-by-key policy now carries a stable reduction-order flag. Dispatch enables the flag for both kernel paths. A regression test checks bit-identical fp64 results across repeated multi-tile reductions.

Changes

Reduce-by-key determinism

Layer / File(s)	Summary
Stable reduction-order propagation cub/cub/agent/agent_reduce_by_key.cuh, cub/cub/device/dispatch/dispatch_reduce_by_key.cuh	The policy exposes StableReductionOrder, passes it to TilePrefixCallbackOpT, and enables stable ordering in both dispatch paths.
Repeated fp64 determinism coverage cub/test/catch2_test_device_reduce_by_key.cu	The regression test repeats reductions with deterministic double inputs and compares the bit representation of the first aggregate across runs.

Assessment against linked issues

Objective	Addressed	Explanation
[#9995] Ensure run-to-run deterministic fp64 reductions, including equal-key runs spanning multiple tiles.	✅
[#9995] Preserve existing single-tile and exact integer-valued fp64 behavior across GPU architectures.	❓	The dispatch keeps the default policy unstable and enables stability in the production paths, but the provided regression test does not explicitly cover these cases or GPU architectures.

Suggested reviewers: srinivasyadav18, bernhardmgruber, davebayer

Merge Risk: ⚪ Minimal · up to 009b2

The change makes ReduceByKey results deterministic with targeted tests and reported validation; no actionable merge-blocking risk remains beyond normal checks and review.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 69b6565e-8ae2-4392-844b-369f6c27dc16

📥 Commits

Reviewing files that changed from the base of the PR and between f047516 and 0720444.

📒 Files selected for processing (3)

• cub/cub/agent/agent_reduce_by_key.cuh
• cub/cub/device/dispatch/dispatch_reduce_by_key.cuh
• cub/test/catch2_test_device_reduce_by_key.cu

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.