Exposes `DeviceBatchedTopK::{Min,Max}{Keys,Pairs}` for non-deterministic, unordered, and small segments-only by elstehle · Pull Request #9331 · NVIDIA/cccl

elstehle · 2026-06-09T14:58:32Z

Description

Rendered docs:
cub:DeviceBatchedTopK — CUDA Core Compute Libraries.pdf

This work depends on the following work that needs to be merged first:

[libcu++] Adds a cuda::execution::tie_break requirement #9238 (Introduce an option to specify a tie_breaking behavior in the requirements API #9255)
[libcu++] Make argument namespace and wrappers construction public #9251 (Expose argument annotation framework for public usage in cub::DeviceBatchedTopK #9254)

Closes #7616

copy-pr-bot · 2026-06-09T14:58:37Z

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.

elstehle · 2026-06-09T15:00:53Z

/ok to test 5e447db

elstehle · 2026-06-09T15:32:01Z

/ok to test 84707e0

coderabbitai · 2026-06-13T04:18:40Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds the public cub::DeviceBatchedTopK header exposing eight static methods (MaxKeys/MinKeys/MaxPairs/MinPairs × device-storage/env-based overloads) backed by a shared detail::dispatch_batched_topk dispatcher with compile-time requirement validation. Includes two example programs, two new Catch2 test files, migration of existing segmented-topk tests to the public API, an improved static_assert message, and updated execution-requirements documentation.

Changes

DeviceBatchedTopK Implementation and Coverage

Layer / File(s)	Summary
Shared dispatcher and API contract `cub/cub/device/device_batched_topk.cuh`	`detail::dispatch_batched_topk` enforces determinism/ordering/tie-break at compile time, validates `num_segments` annotation, resolves policy selector and stream from env, and forwards to `batched_topk::dispatch`; `DeviceBatchedTopK` struct with Doxygen/RST docs added.
MaxKeys/MinKeys/MaxPairs/MinPairs overloads `cub/cub/device/device_batched_topk.cuh`	Device-storage and env-based overloads for all four operations; keys-only variants pass null value iterators; pairs variants forward both iterators; all scope NVTX ranges and delegate to `dispatch_batched_topk` or `dispatch_with_env`.
Improved static_assert in policy selector `cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh`	Single-line `static_assert` failure in `find_smallest_covering_policy` replaced with multi-line message naming the one-block-per-segment constraint and instructing reduction of the encoded maximum segment size.
Execution requirements documentation `docs/cub/api_docs/device_topk_requirements.rst`	Intro updated to cover `DeviceBatchedTopK`; full requirements reference added covering the two-orthogonal-concerns model, constraint tables, worked example matrix, and goal-to-require mapping table.
MaxKeys and MaxPairs example programs `cub/examples/device/example_device_batched_topk_keys.cu`, `cub/examples/device/example_device_batched_topk_pairs.cu`	Keys example demonstrates `float` MaxKeys with compile-time constants and runtime `num_segments`; pairs example demonstrates MaxPairs with global-index values and host-side pair consistency verification.
Device-storage API Catch2 tests `cub/test/catch2_test_device_batched_topk_api.cu`	Covers temp-storage query+allocate+invoke for MaxKeys, MinKeys, MaxPairs, MinPairs on two fixed 8-key segments; keys tests sort output and assert values; pairs tests validate index-to-key consistency.
Environment-based API Catch2 tests `cub/test/catch2_test_device_batched_topk_env_api.cu`	Covers env-based (auto-temp-storage) overloads for all four operations using `cuda::std::execution::env` with stream and determinism/output_ordering requirements.
Migration of existing tests to public API `cub/test/catch2_test_device_segmented_topk_keys.cu`, `cub/test/catch2_test_device_segmented_topk_pairs.cu`	Replaces internal dispatch header with `device_batched_topk.cuh`; rewrites dispatch helpers to call `DeviceBatchedTopK` public methods via `if constexpr` on renamed `SelectDirection`; updates `DECLARE_TMPL_LAUNCH_WRAPPER` wiring.

Assessment against linked issues

Objective	Addressed	Explanation
Expose public-facing interface for `cub::DeviceSegmentedTopK` (now `cub::DeviceBatchedTopK`) [`#7616`]	✅
Add examples for fixed small segments and variable small segments [`#7616`]	✅
Add documentation and mention current constraints on small segment sizes [`#7616`]	✅
Make existing tests go through the public `cub::DeviceBatchedTopK` interface instead of internal dispatch [`#7616`]	✅

Possibly related PRs

NVIDIA/cccl#9238: Adds cuda::execution::tie_break, which detail::dispatch_batched_topk validates in conjunction with determinism via compile-time static_assert.
NVIDIA/cccl#9311: Modifies the same catch2_test_device_segmented_topk_keys.cu and catch2_test_device_segmented_topk_pairs.cu dispatch logic that this PR rewrites to use the public API.
NVIDIA/cccl#9251: Updates the cuda::args constant/immediate wrapping framework that DeviceBatchedTopK's dispatcher and examples rely on.

Suggested reviewers

pauleonix
davebayer
wmaxey
gevtushenko

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

coderabbitai

Actionable comments posted: 3

🧹 Nitpick comments (2)

cub/cub/device/device_batched_topk.cuh (1)
131-131: 💤 Low value

suggestion: Missing const on non-modified variable.

Per coding guidelines, all variables that are not modified must use const.
-  auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
+  const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
Source: Coding guidelines
cub/test/catch2_test_device_batched_topk_api.cu (1)

181-191: ⚡ Quick win

suggestion: strengthen pair tests to assert top-k/min-k key content per segment (like the keys-only tests) and add explicit bounds checks before indexing h_keys_in[h_values_out[...]]. Current checks only verify key/value association and can miss ranking regressions.
As per coding guidelines, cub/**/* reviews should focus on algorithm correctness and test coverage.

Also applies to: 249-257

Source: Coding guidelines

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 15d66aaf-5a41-4f97-8f96-2cc85967f600

📥 Commits

Reviewing files that changed from the base of the PR and between 8c0e6cb and bf86bed.

📒 Files selected for processing (8)

cub/cub/device/device_batched_topk.cuh
cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh
cub/examples/device/example_device_batched_topk_keys.cu
cub/examples/device/example_device_batched_topk_pairs.cu
cub/test/catch2_test_device_batched_topk_api.cu
cub/test/catch2_test_device_batched_topk_env_api.cu
cub/test/catch2_test_device_segmented_topk_keys.cu
cub/test/catch2_test_device_segmented_topk_pairs.cu

coderabbitai

Actionable comments posted: 2

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: eeb0ed08-dd40-4e36-a833-46526b2d846f

📥 Commits

Reviewing files that changed from the base of the PR and between bf86bed and 4685962.

📒 Files selected for processing (4)

cub/cub/device/device_batched_topk.cuh
cub/cub/device/device_topk.cuh
docs/cub/api.rst
docs/cub/api_docs/device_topk_requirements.rst

✅ Files skipped from review due to trivial changes (2)

docs/cub/api.rst
cub/cub/device/device_topk.cuh

🚧 Files skipped from review as they are similar to previous changes (1)

cub/cub/device/device_batched_topk.cuh

coderabbitai

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

cub/cub/device/device_batched_topk.cuh (2)

294-298: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

important: Remove the default {} environment from these public overloads, or stop documenting env as optional.

detail::dispatch_batched_topk hard-rejects any environment that does not explicitly require cuda::execution::determinism::not_guaranteed plus cuda::execution::output_ordering::unsorted at Lines 89-93, so every EnvT env = {} here is an unusable default. Right now the public signature and docs advertise an optional parameter that immediately fails at compile time on the default call path.

Also applies to: 305-315, 358-366, 413-423, 464-472, 538-550, 590-599, 649-661, 701-710

25-35: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

important: Add the direct include for ::cuda::std::int64_t.

This header forms default_policy_selector_t with ::cuda::std::int64_t at Line 117 but never includes the defining header. That leaves a public header depending on transitive includes and include order. As per coding guidelines, "Files must include all headers related to the symbols that they are using. No transitive header inclusions are allowed."

Also applies to: 114-121

Source: Coding guidelines

🧹 Nitpick comments (1)

cub/cub/device/device_batched_topk.cuh (1)

306-315: ⚡ Quick win

suggestion: Mark the device-storage overloads [[nodiscard]] too.

These four public entry points return cudaError_t, but only the env-managed overloads enforce status checking. That makes it easy to drop launch or temp-storage query failures silently on the new API surface. As per coding guidelines, "Most functions with a non-void return type shall use [[nodiscard]] attribute."

Also applies to: 414-423, 539-550, 650-661

Source: Coding guidelines

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: db52dab9-9922-4749-8a59-2b2958426960

📥 Commits

Reviewing files that changed from the base of the PR and between 4685962 and 074f048.

📒 Files selected for processing (1)

cub/cub/device/device_batched_topk.cuh

coderabbitai

♻️ Duplicate comments (1)

docs/cub/api_docs/device_topk_requirements.rst (1)

45-69: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

important: Reword the default contract and the GPU-to-GPU row.

This still describes the deterministic/stable-sorted contract as the current default, which conflicts with the compile-time checks in cub::DeviceTopK/cub::DeviceBatchedTopK that only accept determinism::not_guaranteed + output_ordering::unsorted. The “Bit-identical across GPUs” row also still needs output_ordering::stable_sorted if equal-key positions must be pinned.

Also applies to: 284-285

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2df71ae6-19f0-4660-9f51-f60bdeb0ba5c

📥 Commits

Reviewing files that changed from the base of the PR and between 074f048 and af679b6.

📒 Files selected for processing (10)

cub/cub/device/device_batched_topk.cuh
cub/cub/device/device_topk.cuh
cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh
cub/examples/device/example_device_batched_topk_keys.cu
cub/examples/device/example_device_batched_topk_pairs.cu
cub/test/catch2_test_device_batched_topk_api.cu
cub/test/catch2_test_device_batched_topk_env_api.cu
cub/test/catch2_test_device_segmented_topk_keys.cu
cub/test/catch2_test_device_segmented_topk_pairs.cu
docs/cub/api_docs/device_topk_requirements.rst

✅ Files skipped from review due to trivial changes (2)

cub/cub/device/device_topk.cuh
cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh

🚧 Files skipped from review as they are similar to previous changes (7)

cub/test/catch2_test_device_batched_topk_env_api.cu
cub/test/catch2_test_device_segmented_topk_keys.cu
cub/test/catch2_test_device_batched_topk_api.cu
cub/test/catch2_test_device_segmented_topk_pairs.cu
cub/examples/device/example_device_batched_topk_pairs.cu
cub/examples/device/example_device_batched_topk_keys.cu
cub/cub/device/device_batched_topk.cuh

…segmented-topk # Conflicts: # cub/test/catch2_test_device_segmented_topk_keys.cu # cub/test/catch2_test_device_segmented_topk_pairs.cu

…segmented-topk

Jacobfaib

LGTM generally speaking

Jacobfaib · 2026-06-26T15:04:13Z

+    "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A "
+    "raw pointer or iterator is not accepted.");
+
+  const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);


@pciolkosz IMO we should at least have an internal cuda::__null_stream_ref that does cuda::stream_ref{cudaStream_t{}}. I see this all the time (and use it myself many times as well).

Jacobfaib · 2026-06-26T15:06:13Z

+  template <typename KeyInputIteratorItT,
+            typename KeyOutputIteratorItT,
+            typename SegmentSizeParameterT,
+            typename KParameterT,


Do you validate anywhere that KParameterT is an expected type? So either a cuda::args or bare value or whatever.

Yup, detail::dispatch_batched_topk static_asserts that each of SegmentSizeParameterT, KParameterT, and NumSegmentsParameterT is either a cuda::args wrapper or a plain integral. Raw pointers/iterators are rejected. Every public overload funnels through detail::dispatch_batched_topk.

…segmented-topk

gevtushenko · 2026-06-30T05:56:23Z

suggestion: consider adding new header in cub.cuh and docs/cub/api_docs/device_wide.rst.

…segmented-topk

elstehle · 2026-06-30T07:43:19Z

suggestion: consider adding new header in cub.cuh and docs/cub/api_docs/device_wide.rst.

Good catch. Fixed.

…segmented-topk

github-actions · 2026-06-30T15:48:02Z

🥳 CI Workflow Results

🟩 Finished in 3h 06m: Pass: 100%/287 | Total: 3d 02h | Max: 1h 05m | Hits: 100%/200001

See results here.

github-project-automation Bot added this to CCCL Jun 9, 2026

github-project-automation Bot moved this to Todo in CCCL Jun 9, 2026

cccl-authenticator-app Bot moved this from Todo to In Progress in CCCL Jun 9, 2026

elstehle force-pushed the topk/expose-device-segmented-topk branch from 5e447db to 84707e0 Compare June 9, 2026 15:24

This comment has been minimized.

Sign in to view

elstehle force-pushed the topk/expose-device-segmented-topk branch 2 times, most recently from 5244000 to bf86bed Compare June 13, 2026 04:10

elstehle marked this pull request as ready for review June 13, 2026 04:11

elstehle requested review from a team as code owners June 13, 2026 04:11

elstehle requested review from alliepiper and pauleonix June 13, 2026 04:11

cccl-authenticator-app Bot moved this from In Progress to In Review in CCCL Jun 13, 2026

coderabbitai Bot reviewed Jun 13, 2026

View reviewed changes

Comment thread cub/examples/device/example_device_batched_topk_keys.cu Outdated

Comment thread cub/test/catch2_test_device_batched_topk_api.cu Outdated

Comment thread cub/test/catch2_test_device_batched_topk_env_api.cu

coderabbitai Bot reviewed Jun 13, 2026

View reviewed changes

Comment thread docs/cub/api_docs/device_topk_requirements.rst Outdated

Comment thread docs/cub/api_docs/device_topk_requirements.rst Outdated

This comment has been minimized.

Sign in to view

elstehle force-pushed the topk/expose-device-segmented-topk branch from 4685962 to 074f048 Compare June 14, 2026 20:53

coderabbitai Bot reviewed Jun 14, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

elstehle force-pushed the topk/expose-device-segmented-topk branch from 074f048 to af679b6 Compare June 15, 2026 09:53

coderabbitai Bot reviewed Jun 15, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

elstehle force-pushed the topk/expose-device-segmented-topk branch from af679b6 to a2c19d1 Compare June 15, 2026 15:38

This comment has been minimized.

Sign in to view

elstehle force-pushed the topk/expose-device-segmented-topk branch from a2c19d1 to 419fde2 Compare June 17, 2026 05:51

This comment has been minimized.

Sign in to view

elstehle added 5 commits June 22, 2026 21:12

introduces the both-or-nothing rule for determinism and tie-break

2f7f695

adds documentation on function parameters

d851e22

validates unwrapped parameters

c777859

clarifies support for plain integrals

6cb0864

fixes copydetails

b9ffd43

elstehle force-pushed the topk/expose-device-segmented-topk branch from acdc270 to b9ffd43 Compare June 23, 2026 16:16

Merge remote-tracking branch 'upstream/main' into topk/expose-device-…

ef9d2b6

…segmented-topk # Conflicts: # cub/test/catch2_test_device_segmented_topk_keys.cu # cub/test/catch2_test_device_segmented_topk_pairs.cu

This comment has been minimized.

Sign in to view

Merge remote-tracking branch 'upstream/main' into topk/expose-device-…

3625924

…segmented-topk

ORippler mentioned this pull request Jun 24, 2026

ggml : process data in smaller chunks in CUDA ggml_top_k() implementation to reduce temporary buffers memory usage ggml-org/llama.cpp#24776

Merged

Jacobfaib approved these changes Jun 26, 2026

View reviewed changes

elstehle added 3 commits June 27, 2026 06:33

Merge remote-tracking branch 'upstream/main' into topk/expose-device-…

c1226fb

…segmented-topk

adds api test for unwrapped integrals

4cbb929

addresses review comments

6b0e679

This comment has been minimized.

Sign in to view

gevtushenko approved these changes Jun 30, 2026

View reviewed changes

Comment thread cub/test/catch2_test_device_batched_topk_env_api.cu

Comment thread cub/test/catch2_test_device_batched_topk_api.cu

Comment thread cub/cub/device/device_batched_topk.cuh Outdated

Comment thread cub/cub/device/device_batched_topk.cuh Outdated

elstehle added 4 commits June 29, 2026 23:37

Merge remote-tracking branch 'upstream/main' into topk/expose-device-…

800ce18

…segmented-topk

adds to the umbrella cub.cuh and device-wide docs

840ebdd

const-rev env and docs on in/out overlapping

484c418

improves api examples; adds mis-aligned temp storage test

4c7a6e1

elstehle enabled auto-merge (squash) June 30, 2026 07:43

This comment has been minimized.

Sign in to view

Merge remote-tracking branch 'upstream/main' into topk/expose-device-…

57000c6

…segmented-topk

elstehle disabled auto-merge June 30, 2026 13:27

This comment has been minimized.

Sign in to view

elstehle merged commit 6f00ddc into NVIDIA:main Jul 1, 2026
606 of 612 checks passed

Conversation

elstehle commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

copy-pr-bot Bot commented Jun 9, 2026

Uh oh!

elstehle commented Jun 9, 2026

Uh oh!

elstehle commented Jun 9, 2026

Uh oh!

This comment has been minimized.

coderabbitai Bot commented Jun 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Assessment against linked issues

Possibly related PRs

Suggested reviewers

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment has been minimized.

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

Jacobfaib left a comment

Choose a reason for hiding this comment

Uh oh!

Jacobfaib Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Jacobfaib Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

elstehle Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment has been minimized.

This comment has been minimized.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gevtushenko commented Jun 30, 2026

Uh oh!

elstehle commented Jun 30, 2026

Uh oh!

This comment has been minimized.

This comment has been minimized.

github-actions Bot commented Jun 30, 2026

🥳 CI Workflow Results

🟩 Finished in 3h 06m: Pass: 100%/287 | Total: 3d 02h | Max: 1h 05m | Hits: 100%/200001

Uh oh!

Uh oh!

Reviewers

elstehle commented Jun 9, 2026 •

edited

Loading

coderabbitai Bot commented Jun 13, 2026 •

edited

Loading