Add grid-stride + ROCm cap to ssd_generate/update_row_addrs kernels - #6102
Closed
q10 wants to merge 2 commits into
Closed
Add grid-stride + ROCm cap to ssd_generate/update_row_addrs kernels#6102q10 wants to merge 2 commits into
q10 wants to merge 2 commits into
Conversation
…nels Summary: linearize_index_kernel and linearize_index_index_select_kernel launch grid = div_round_up(total_B, kMaxThreads) with block = kMaxThreads (total_B = B * T), so on ROCm total threads can exceed the HIP 2^32 threads-per-launch limit for very large total_B. Both kernels are warp-collective (shfl_sync across all lanes), so a naive grid-stride would desync the shuffles. Cap both launches with utils::cuda::cap_grid_dim_x_from_workload(total_B, kMaxThreads, stream) and add a WARP-ALIGNED grid-stride loop whose bound uses (b_t - lane_id) so whole warps iterate together and every shfl_sync still has all lanes present. Per-lane `valid` / `t < T` masking is unchanged. The grid-stride loop is unconditional rather than #ifdef USE_ROCM. On CUDA the grid already spans the workload, so the body executes exactly once per thread and the kernel is behaviorally unchanged; keeping a single code path avoids braces that straddle a preprocessor conditional, and keeps CUDA correct if the grid is ever clamped. Loop bounds are compared in 64 bits (no uint32 narrowing), and transpose_embedding_input now TORCH_CHECKs that total_B fits in int32_t: b_t is a 32-bit thread index handed to FixedDivisor::DivMod(int32_t), so an oversized workload fails loudly instead of silently wrapping the grid-stride index. NOTE: warp-collective transform -- validate carefully on MI300 in the per-diff pass (shfl correctness under the strided loop). Reviewed By: henrylhtsang Differential Revision: D113351688
Summary: ssd_generate_row_addrs_kernel and ssd_update_row_addrs_kernel launch grid = div_round_up(numel, kNumWarps) with block = dim3(kWarpSize, kNumWarps), so total threads ~= numel * kWarpSize exceeds the HIP 2^32 threads-per-launch limit on ROCm for large lxu_cache_locations / ssd_row_addrs counts. Cap both launches with utils::cuda::cap_grid_dim_x(..., OverflowOnly) and add a ROCm grid-stride loop over the warp index (n / n_curr). In the update kernel the "row not used in both iterations" early-return (n_next < 0) becomes `continue` under ROCm. Added the cuda_utilities.cuh include. No-op on CUDA. Reviewed By: henrylhtsang Differential Revision: D113351690
Contributor
|
@q10 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113351690. |
Contributor
|
This pull request has been merged in f6c66ba. |
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.
Summary:
ssd_generate_row_addrs_kernel and ssd_update_row_addrs_kernel launch
grid = div_round_up(numel, kNumWarps) with block = dim3(kWarpSize, kNumWarps),
so total threads ~= numel * kWarpSize exceeds the HIP 2^32 threads-per-launch
limit on ROCm for large lxu_cache_locations / ssd_row_addrs counts.
Cap both launches with utils::cuda::cap_grid_dim_x(..., OverflowOnly) and add a
ROCm grid-stride loop over the warp index (n / n_curr). In the update kernel the
"row not used in both iterations" early-return (n_next < 0) becomes
continueunder ROCm. Added the cuda_utilities.cuh include. No-op on CUDA.
Reviewed By: henrylhtsang
Differential Revision: D113351690