warpspeed run_to_run deterministic scan for SM90 using atomic global counter#9565
Open
srinivasyadav18 wants to merge 8 commits into
Open
warpspeed run_to_run deterministic scan for SM90 using atomic global counter#9565srinivasyadav18 wants to merge 8 commits into
srinivasyadav18 wants to merge 8 commits into
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR extends lookahead scan execution for SM90 by adding atomic tile scheduling, passing an atomic counter through dispatch and kernels, tightening tile-count bounds in warpspeed lookahead helpers, and adding non-SM100 masked bulk-store fallbacks. ChangesSM90 lookahead deterministic scan
Assessment against linked issues
Out-of-scope changes
Suggested reviewers: Comment |
This comment has been minimized.
This comment has been minimized.
srinivasyadav18
force-pushed
the
warpspeed_atomic_sm90_rebased
branch
from
1a42062 to
e2bb592
Compare
This comment has been minimized.
This comment has been minimized.
miscco
reviewed
Jun 24, 2026
miscco
left a comment
miscco
left a comment
Contributor
There was a problem hiding this comment.
I am really not sure whether this is the right approach as it greatly affects the performance of warpspeed for all architectures.
I
Comment on lines
+34
to
+39
| template <typename AccumT> | ||
| struct lookahead_tile_state_arg_t | ||
| { | ||
| warpspeed::tile_state_t<AccumT>* tile_states; | ||
| ::cuda::std::uint32_t* atomic_counter; | ||
| }; |
Contributor
There was a problem hiding this comment.
This is going to massively impact performance.
The issue is that with 64 bit types we barely fit into the 16 byte atomic window with the tile state and the perf difference is really significant. With this we are effectively never getting the fast atomics
Contributor
Author
There was a problem hiding this comment.
struct lookahead_tile_state_arg_t
We do not actually modify tile-state but rather lookahead_tile_state_arg_t, which is helper struct to pass arguments to kernel.
I have posted complete performance report below.
Contributor
There was a problem hiding this comment.
@srinivasyadav18 is right, this is just a helper struct for the union that is passed to the tile state setup kernel, so we can use the same kernel symbol between the old and new scan implementation. The struct is only used as a kernel argument. Should be fine.
Contributor
Author
Performance ReportH100_80GB_HBM3bench.exclusive_scan.deterministic (avg %Diff: +61.27%) mainly on large problem sizesH200bench.exclusive_scan.deterministic (avg %Diff: +40.93%) mainly on large problem sizesRTX_PRO_6000_Blackwell_Server_Editionbench.exclusive_scan.deterministic (avg %Diff: -0.36%)bench.exclusive_scan.sum (avg %Diff: +0.17%)bench.exclusive_scan.look_ahead.sum (avg %Diff: +0.68%)bench.exclusive_scan.custom (avg %Diff: +0.12%)B200bench.exclusive_scan.deterministic (avg %Diff: -0.18%)bench.exclusive_scan.sum (avg %Diff: +0.14%)bench.exclusive_scan.look_ahead.sum (avg %Diff: +0.32%)bench.exclusive_scan.custom (avg %Diff: +0.24%) |
This comment has been minimized.
This comment has been minimized.
srinivasyadav18
commented
Jun 26, 2026
| return error; | ||
| } | ||
|
|
||
| const int scan_grid_dim = atomic_scheduling ? ::cuda::std::min(sm_count, num_tiles) : num_tiles; |
Contributor
Author
There was a problem hiding this comment.
suggestion: try using occupancy or higher CTA count per SM if it improves performance.
Comment on lines
+34
to
+39
| template <typename AccumT> | ||
| struct lookahead_tile_state_arg_t | ||
| { | ||
| warpspeed::tile_state_t<AccumT>* tile_states; | ||
| ::cuda::std::uint32_t* atomic_counter; | ||
| }; |
Contributor
There was a problem hiding this comment.
@srinivasyadav18 is right, this is just a helper struct for the union that is passed to the tile state setup kernel, so we can use the same kernel symbol between the old and new scan implementation. The struct is only used as a kernel argument. Should be fine.
bernhardmgruber
left a comment
bernhardmgruber
left a comment
Contributor
There was a problem hiding this comment.
This PR looks like it is not only adding a deterministic scan for SM90, but also the non-deterministic version. Please point this out clearly in the PR description and update the PR title.
Otherwise, this is looking pretty good already! Please make sure that the Blackwell kernels are not impacts (no SASS changes on the scan benchmarks for SM100)
Comment on lines
+285
to
+298
| NV_IF_ELSE_TARGET( | ||
| NV_PROVIDES_SM_100, | ||
| (if (::cuda::ptx::elect_sync(~0)) { | ||
| ::cuda::ptx::cp_async_bulk_cp_mask( | ||
| ::cuda::ptx::space_global, | ||
| ::cuda::ptx::space_shared, | ||
| cpAsyncOobInfo.ptrGmemStartAlignDown, | ||
| srcSmem, | ||
| /*size*/ 16, | ||
| byteMaskStart); | ||
| }), | ||
| (const int rank = squad.threadRank(); if (rank < 16 && ((byteMaskStart >> rank) & 1u)) { | ||
| cpAsyncOobInfo.ptrGmemStartAlignDown[rank] = srcSmem[rank]; | ||
| })); |
Contributor
There was a problem hiding this comment.
Important question: Since you introduced code paths that access SMEM using regular loads and stores, is the line above:
::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared);
still valid? We need to acquire SMEM in the async proxy for the bulk copy, but now we are mixing bulk copies and regular loads from SMEM. Is this still legal?
@ahendriksen may be able to help here.
Contributor
Author
There was a problem hiding this comment.
Yes, I think its still needed, as we still do bulk async copies for the middle part (aligned portion of data).
|
|
||
| # pragma unroll 1 | ||
| while (true) | ||
| while (idxTile < numTiles) |
Contributor
There was a problem hiding this comment.
Question: Are we 100% certain this does not change SASS for any kernel on SM100? Why do we even need this, since we have nextIdxTileValid = static_cast<int>(regNextBlockIdx.x) < numTiles; later to exit the loop?
Contributor
Author
There was a problem hiding this comment.
while(true) produces bad codegen on SM90, leading to deadlock. Yes, there is some SASS difference, so I re-measured performance on GPUs >= SM90, and there is NO regression because of this.
Contributor
There was a problem hiding this comment.
The status quo increases the SASS of the kernels on SM100 by a lot. Here is a summary from claude:
+96 to +112 instructions each (~4-5% growth)
The opcode histogram reveals the nature of the change:
- +16 IMAD.WIDE.U32, +5 IMAD.WIDE.U32.X, +5 LEA.HI.X, +5 SHF.R.U64 — these are 64-bit address arithmetic instructions
- +5 IADD3.X, +5 IMAD.X — extended-precision integer ops
- +9 ISETP.LE.AND, +5 ISETP.LT.U32.AND/.EX — additional bound checks or comparisons
- +4 EXIT, -5 BRA — more exit points, fewer branches (code is being duplicated/specialized rather than branching)
- +3 LDC.64 — more 64-bit constant loads
Even if the measured perf on a benchmark does not show a runtime diff, it's a significant cost in code size (binary size and I-cache pressure).
Contributor
There was a problem hiding this comment.
Important: The SASS changes go completely away with:
Suggested change
| while (idxTile < numTiles) | |
| while ([&] { | |
| NV_IF_ELSE_TARGET(NV_PROVIDES_SM_100, { return true; }, { return idxTile < numTiles; }) | |
| }()) |
Please apply this change.
Contributor
Contributor
Author
|
/ok to test fa0e6e7 |
This comment has been minimized.
This comment has been minimized.
Contributor
Author
|
/ok to test 669019e |
This comment has been minimized.
This comment has been minimized.
Contributor
Author
|
/ok to test 1c1ae98 |
This comment has been minimized.
This comment has been minimized.
Contributor
Author
|
/ok to test 67e93ef |
Comment on lines
+299
to
+300
| ::cuda::std::conditional_t<(sizeof(AccumT) <= 4), scan_resources_t, scan_resources_t&> res; // this is the only shared | ||
| // mutable state |
Contributor
There was a problem hiding this comment.
Suggestion: I tried this quickly locally and it seems all kernels in cub.bench.scan.exclusive.sum.lookahead.base get shorter if we store scan_resources_t by value. Can you please send this diff:
Suggested change
| ::cuda::std::conditional_t<(sizeof(AccumT) <= 4), scan_resources_t, scan_resources_t&> res; // this is the only shared | |
| // mutable state | |
| scan_resources_t res; // this is the only shared mutable state |
as a separate PR with a benchmark? Thx!
Comment on lines
+798
to
+804
| const bool nextIdxTileValid = [&] { | ||
| bool valid = false; | ||
| NV_IF_ELSE_TARGET(NV_PROVIDES_SM_100, | ||
| (valid = ::cuda::ptx::clusterlaunchcontrol_query_cancel_is_canceled(regNextBlockIdx);), | ||
| (valid = static_cast<int>(regNextBlockIdx.x) < numTiles;)); | ||
| return valid; | ||
| }(); |
Contributor
There was a problem hiding this comment.
Suggested change
| const bool nextIdxTileValid = [&] { | |
| bool valid = false; | |
| NV_IF_ELSE_TARGET(NV_PROVIDES_SM_100, | |
| (valid = ::cuda::ptx::clusterlaunchcontrol_query_cancel_is_canceled(regNextBlockIdx);), | |
| (valid = static_cast<int>(regNextBlockIdx.x) < numTiles;)); | |
| return valid; | |
| }(); | |
| const bool nextIdxTileValid = [&] { | |
| NV_IF_ELSE_TARGET(NV_PROVIDES_SM_100, | |
| (return ::cuda::ptx::clusterlaunchcontrol_query_cancel_is_canceled(regNextBlockIdx);), | |
| (return static_cast<int>(regNextBlockIdx.x) < numTiles;)); | |
| }(); |
| { | ||
| device_scan_init_lookahead_body(tile_state.lookahead, num_tiles); | ||
| device_scan_init_lookahead_body(tile_state.lookahead.tile_states, num_tiles); | ||
| if (tile_state.lookahead.atomic_counter != nullptr && blockIdx.x == 0 && threadIdx.x == 0) |
Contributor
There was a problem hiding this comment.
Remark: I would reorder the condition so only one thread needs to load the value of atomic_counter and check it for nullptr.
Suggested change
| if (tile_state.lookahead.atomic_counter != nullptr && blockIdx.x == 0 && threadIdx.x == 0) | |
| if (blockIdx.x == 0 && threadIdx.x == 0 && tile_state.lookahead.atomic_counter != nullptr) |
Comment on lines
+81
to
+85
| device_scan_init_lookahead_body(tile_state.lookahead.tile_states, num_tiles); | ||
| if (tile_state.lookahead.atomic_counter != nullptr && blockIdx.x == 0 && threadIdx.x == 0) | ||
| { | ||
| *tile_state.lookahead.atomic_counter = 0; | ||
| } |
Contributor
There was a problem hiding this comment.
Suggestion: it's easy to miss the setup if atomic_counter when only looking at device_scan_init_lookahead_body. I would move the code into device_scan_init_lookahead_body
Comment on lines
+82
to
+85
| if (tile_state.lookahead.atomic_counter != nullptr && blockIdx.x == 0 && threadIdx.x == 0) | ||
| { | ||
| *tile_state.lookahead.atomic_counter = 0; | ||
| } |
Contributor
There was a problem hiding this comment.
Important: this is only needed when we are not on SM100+ or when StableReductionOrder == true. Please add a StableReductionOrder parameter to DeviceScanInitKernel and add an NV_IF_ELSE_TARGET here.
Contributor
🥳 CI Workflow Results🟩 Finished in 2h 33m: Pass: 100%/291 | Total: 12d 06h | Max: 2h 32m | Hits: 18%/1046130See results here. |
2 tasks
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.
Description
closes #9334
Checklist