Reduces the latency of air-based block-level top-k - #11060
Open
elstehle wants to merge 1 commit into
Open
Conversation
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughBlock TopK selection now operates over the full key width. Public APIs remove bit-range parameters. The implementation adds configurable radix-pass unrolling, double-buffered histograms, fused bucket selection, early termination, and optional original-key scattering. ChangesBlock TopK selection
Suggested reviewers: Merge Risk: 🟡 Moderate · up to The optimized block-level top-k path assumes blocked input, but striped input can produce incorrect top-k results when it reaches this implementation. Merge should wait until the input contract is enforced or clearly documented and affected callers are addressed. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
cub/cub/block/specializations/block_topk_air.cuh (1)
252-254: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: mark
filter_opanddigit_extractorasconst. Both are never modified.diff
- auto filter_op = compare_key_prefix_op<bit_ordered_type>{prefix_mask, kth_key_prefix}; - auto digit_extractor = + const auto filter_op = compare_key_prefix_op<bit_ordered_type>{prefix_mask, kth_key_prefix}; + const auto digit_extractor = traits::template digit_extractor<fundamental_digit_extractor_t>(pass_begin_bit, pass_bits, decomposer);As per coding guidelines: "All variables that are not modified must be declared
const, including cast results, function return values, and loop-invariant computations."Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5b06dd72-070b-4dd3-b910-5ab68916a0a5
📒 Files selected for processing (2)
cub/cub/block/block_topk.cuhcub/cub/block/specializations/block_topk_air.cuh
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Comment on lines
+381
to
+385
| // TODO (elstehle): Short-circuit if k is greater than the number of items in the tile | ||
| if ((!IsFullTile && k >= valid_items) || k >= tile_items) | ||
| { | ||
| return; | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🔎 Supported by static analysis
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f 'catch2_test_block_topk.cu' cub | while IFS= read -r f; do
rg -n -C 6 'blocked_input|LoadDirectStriped|num_valid|GENERATE_COPY|check_topk' "$f"
done
# Any other callers relying on compaction
rg -n -C 5 -e 'max_keys|min_keys|max_pairs|min_pairs' --iglob '!**/catch2_test_block_topk.cu' cubRepository: NVIDIA/cccl
Length of output: 13808
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- block_topk_air.cuh ---'
sed -n '300,520p' cub/cub/block/specializations/block_topk_air.cuh
printf '%s\n' '--- block_topk test ---'
fd -t f 'catch2_test_block_topk.cu' cub | xargs -r sed -n '1,240p'
printf '%s\n' '--- direct striped callers ---'
rg -n -C 8 'LoadDirectStriped|block_topk_t|block_topk<' cub/cub --glob '*.{cuh,cu,hpp,h}'Repository: NVIDIA/cccl
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- block_topk_air declarations and output writeback ---'
sed -n '1,180p' cub/cub/block/specializations/block_topk_air.cuh
sed -n '500,700p' cub/cub/block/specializations/block_topk_air.cuh
printf '%s\n' '--- batched top-k load policy and arrangement ---'
sed -n '80,180p' cub/cub/agent/agent_batched_topk.cuh
sed -n '220,330p' cub/cub/agent/agent_batched_topk.cuh
printf '%s\n' '--- block_topk references and documentation ---'
rg -n -C 5 'block_topk|select_topk|selected items|blocked arrangement|striped arrangement' cub/cub --glob '*.{cuh,cu,hpp,h}' | head -240Repository: NVIDIA/cccl
Length of output: 38034
important: Require blocked input for block_topk_air. select_topk uses blocked indices for validity checks and output gathering. LoadDirectStriped does not satisfy this arrangement, so partial striped input can produce incorrect top-k results. Document the blocked-input precondition instead of treating the early return as striped-input support.
Contributor
🔬 CUB benchmark SASS comparisonHow to request a benchmark run
Targets with a SASS change
|
Contributor
🥳 CI Workflow Results🟩 Finished in 2h 56m: Pass: 100%/284 | Total: 13d 12h | Max: 2h 55m | Hits: 17%/1179295See results here. |
Contributor
Author
|
segmented_topk.variable.keys`segmented_topk.variable.indexed` |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR makes a couple of performance improvements for our AIR-based BlockTopK specialization.
Perf Summary
cub.bench.segmented_topk.fixed.keys: 15% faster on average, up to 24% for small segments and small k, at least 5% faster in every one of the 76 configurations.Changes
begin_bit/end_bitparameters fromblock_topkto give us more flexibility in the implementation for the time being. I think we eventually want to add an overload that allows passing bit-ranges with support for argument-annotations (e.g., compile-time bit ranges).UnrollBitPassestemplate parameter.ScatterOriginalKeystemplate parameter(default true). Scattering a register copy of the original keys removes the untwiddling and
the -0.0 restoration and measured neutral (for integral types) to better (f16, f32, f64). The opt-out exists for register-constrained callers.
ShiftDigitExtractorinstead ofBFEDigitExtractorfor the histogram digit extraction so the bucket computation can be fused with atomic addressing.