Segmented Topk tuning on B200 indexed.cluster#9801
Draft
gonidelis wants to merge 91 commits into
Draft
Conversation
Also avoid unnecessary grid sync in reduce-then-scan path
from gmem (L2)
Also remove ifdefs since perf is clearly worse with the alternative code paths.
Also fixes use of generic addressing due to spilling of smem addresses to local memory.
BlockLoadToShared seems to have some problems when used for pipelining. So for now it was replaced with direct PTX.
This might help on Hopper. Need to check that it does not pessimize Blackwell. Also some cleanup.
This is in preparation of implementing deterministic TopK (via counting scan)
Adds env-based dispatch and moves to new argument-framework.
It is too expensive to do it after every tile. Just try to avoid unnecessary global loads. Also avoid unnecessary global loads in the CTAs that don't straddle the candidate boundary. They do not have to care about iterating through keys in any special order.
Always put the unaligned part into the tail edge buffer and do not force the remaining laigned tail to be resident.
Use aligned cluster barrier instructions among other things.
If it fits a single chunk there isn't much latency hiding we can do.
This reverts commit 1a8e1bb.
Enable auto-tuning of the cluster backend.
Start ping-pong with the right direction from the start.
Mostly SASS neutral up to a SASS reduction from process_tiles dedup.
Slight SASS impact, need to benchmark.
Contributor
|
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. |
pauleonix
reviewed
Jul 11, 2026
| make_sm100_pairs_cluster_policy(::cuda::std::int64_t static_max_segment_size, ::cuda::std::int64_t max_k) | ||
| -> cluster_topk_policy | ||
| { | ||
| if (static_max_segment_size <= 512 && max_k <= 512) |
Contributor
There was a problem hiding this comment.
I think our benchmark only covers k = max_seg_size = 512 which always takes the fast-path and just copies gmem->gmem. So I don't know if anything in this tuning is actually good for k < max_seg_size <= 512. For that copy fast path the only important parameters are threads_per_block, min_blocks_per_sm and copy_items_per_thread I think. The rest might be arbitrary.
pauleonix
reviewed
Jul 11, 2026
| /*tie_break_items_per_thread=*/8, | ||
| /*copy_items_per_thread=*/9}; | ||
| } | ||
| if (static_max_segment_size <= 1024 && max_k <= 1024) |
Contributor
There was a problem hiding this comment.
https://github.com/NVIDIA/cccl/pull/9801/changes#r3562779477 Same applies here.
Member
Author
|
Final full results for |
gonidelis
changed the title
indexed.cluster
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
3 participants
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.
Adds tuning encodings for segmented TopK. It's based on top of https://github.com/NVIDIA/cccl/pull/9224/changes so I draft until #9224 gets merged.
I like the speedups, not sure if we expected more. This of course is using @bernhardmgruber's new (amazing) tuning API.
If anyone wants to review it pre-emptively see only the changes in this file
cub/cub/device/dispatch/tuning/tuning_batched_topk.cuhVerification benchmark below.