Switch sparse_permute_1d to OverflowOnly cap (follow-up to D104903707) - #6057
Open
q10 wants to merge 2 commits into
Open
Switch sparse_permute_1d to OverflowOnly cap (follow-up to D104903707)#6057q10 wants to merge 2 commits into
q10 wants to merge 2 commits into
Conversation
Summary:
The Tier-2 fix in `sparse_batched_unary_embeddings.cu` capped only the
x-dim grid for `batched_unary_embeddings_forward_kernel`:
```
const auto blocks_x = std::min<uint32_t>(
blocks_x_uncapped,
utils::cuda::get_max_thread_blocks(stream));
```
That is sufficient when the y/z grid dims (`gridDim.y = T`,
`gridDim.z = N`) are small, but fails when `T * N` is large. The total
launch threads on ROCm are
`blocks_x * threads * T * N`; with `T = 65535, N = 16, threads = 512`,
this exceeds `2^32` even when `blocks_x` is at the
`MAX_THREAD_BLOCKS_FACTOR * SM_count` cap (≈ 19200 on MI350).
The new repro test `test_batched_unary_embeddings_backward_large_grid`
made this gap visible by exercising the forward path during backward
setup. On MI350 it failed at:
```
sparse_batched_unary_embeddings.hip(108): batched_unary_embeddings_forward_kernel
[grid 9 x 65535 x 16] [block 512 x 1 x 1]:
Total 4,831,764,480 > HIP 2^32
```
This diff extends the host-side cap to also factor in `T * N` so the
total launch thread count stays under `2^32` on ROCm:
```
const uint64_t threads_per_x_block =
threads * T * N;
const uint32_t total_cap_blocks_x = threads_per_x_block > 0
? (uint32_t)std::min<uint64_t>(
((uint64_t)max_uint32 - 1) / threads_per_x_block,
max_uint32)
: max_uint32;
const auto blocks_x = std::min<uint32_t>(
{blocks_x_uncapped, sm_cap, total_cap_blocks_x});
```
The kernel already grid-strides over `b`, so capping `blocks_x` (even
to a small value when `T * N` is huge) is correctness-preserving — the
kernel just iterates more times per (t, n) pair. NVIDIA codegen is
unchanged because the cap is gated entirely by `#ifdef USE_ROCM`.
Differential Revision: D105669555
Summary: Drops the explicit `BlockCapPolicy::Always` argument from the two `utils::cuda::determine_grid_blocks(...)` call sites in `src/sparse_ops/sparse_permute_1d.cu`, falling through to the default `BlockCapPolicy::OverflowOnly` policy. Net behaviour change vs the landed D104903707 + D106267802 state: - ROCm small/medium grid (`blocks * threads_per_block <= UINT32_MAX`): uncapped grid restored. Was: capped to `MAX_THREAD_BLOCKS_FACTOR * #SMs` unconditionally. Now: passes through unchanged (matches pre-D104903707 behaviour). - ROCm large grid (`blocks * threads_per_block > UINT32_MAX`): cap still applied. `permute_1D_lengths_kernel` uses `CUDA_KERNEL_LOOP` (already grid-strides), and `permute_1D_data_kernel_vec` grid-strides over `b_t`, so the cap remains correctness-preserving for the large-grid regime where the HIP 2^32 thread-per-launch limit would otherwise fire. - NVIDIA: bit-identical to today (the threshold check lives entirely under `#ifdef USE_ROCM` inside the helper). This is the new follow-up to D104903707 (landed) per the master plan. D104903707 itself cannot be amended; this diff at the stack tip retroactively switches the policy from `Always` (set by D106267802 / Diff X for behaviour preservation during the helper introduction) to `OverflowOnly` (the default — only cap when the unguarded launch would actually exceed 2^32 threads). Reviewed By: henrylhtsang Differential Revision: D106269670
Contributor
|
@q10 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D106269670. |
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:
Drops the explicit
BlockCapPolicy::Alwaysargument from the twoutils::cuda::determine_grid_blocks(...)call sites insrc/sparse_ops/sparse_permute_1d.cu, falling through to the defaultBlockCapPolicy::OverflowOnlypolicy.Net behaviour change vs the landed D104903707 + D106267802 state:
blocks * threads_per_block <= UINT32_MAX):uncapped grid restored. Was: capped to
MAX_THREAD_BLOCKS_FACTOR * #SMsunconditionally. Now: passesthrough unchanged (matches pre-D104903707 behaviour).
blocks * threads_per_block > UINT32_MAX):cap still applied.
permute_1D_lengths_kernelusesCUDA_KERNEL_LOOP(already grid-strides), andpermute_1D_data_kernel_vecgrid-strides overb_t, so the capremains correctness-preserving for the large-grid regime where
the HIP 2^32 thread-per-launch limit would otherwise fire.
under
#ifdef USE_ROCMinside the helper).This is the new follow-up to D104903707 (landed) per the master
plan. D104903707 itself cannot be amended; this diff at the stack
tip retroactively switches the policy from
Always(set byD106267802 / Diff X for behaviour preservation during the helper
introduction) to
OverflowOnly(the default — only cap when theunguarded launch would actually exceed 2^32 threads).
Reviewed By: henrylhtsang
Differential Revision: D106269670