[Fix][TOPI] Fuse GPU scan blocks to avoid CUDA gridDim.y overflow - #20108
Open
chenmiaoming wants to merge 1 commit into
Open
[Fix][TOPI] Fuse GPU scan blocks to avoid CUDA gridDim.y overflow#20108chenmiaoming wants to merge 1 commit into
chenmiaoming wants to merge 1 commit into
Conversation
Flatten the batch and scan-block axes into blockIdx.x for each scan stage so large batches do not exceed CUDA's grid-y limit.\n\nAdd a dynamic-shape CUDA cumprod regression covering 65,536 batches.\n\nFixes apache#20106
chenmiaoming
marked this pull request as ready for review
August 8, 2026 13:30
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.
Motivation
Fixes #20106.
The fallback TOPI GPU scan maps the per-row scan block to
blockIdx.xand thebatch to
blockIdx.y. CUDA limitsgridDim.yto 65,535, so a valid Relaxcumprodwith shape(65536, 1)fails at execution time withCUDA_ERROR_INVALID_VALUEandgrid=(1,65536,1).Changes
Fuse the virtual scan-block and batch dimensions into a single
blockIdx.xlaunch dimension, then recover both indices with integer division and modulo.
Apply this mapping to the initial copy, up-sweep, and down-sweep stages. Fusing
the dimensions fixes large batches without moving the same 65,535 limit onto
long scan axes.
Add a dynamic-shape CUDA
cumprodregression with an input of shape(65536, 3). The three-element scan axis exercises the up-sweep and down-sweepstages while retaining the batch size that previously caused the invalid
launch.
Testing
python -m pytest tests/python/relax/test_backend_dispatch_sort_scan.py -xvs(65536, 1)and a multi-block scanshape
(8, 2000)pre-commit run --files python/tvm/topi/gpu/scan.py tests/python/relax/test_backend_dispatch_sort_scan.py