topk: add an 8192-page tier and make the overflow check actionable - #4
Open
GeoffreyWang1117 wants to merge 1 commit into
Open
topk: add an 8192-page tier and make the overflow check actionable#4GeoffreyWang1117 wants to merge 1 commit into
GeoffreyWang1117 wants to merge 1 commit into
Conversation
topk_output dispatches on max_num_pages through a ladder of
(NUM_THREADS, ITEM_PER_THREAD) pairs that tops out at <512, 8> = 4096
pages. Past that it hits a bare TORCH_CHECK(false), which aborts with no
message at all -- the user sees an assertion with no indication of what
was exceeded or what to change.
Two changes:
1. Add a <1024, 8> tier for max_num_pages <= 8192. 1024 is the CUDA
per-block thread limit, and 1024 x 8 keeps the CUB block-sort scratch
at ~36 KB, still under the 48 KB static shared-memory limit. Verified
to compile with nvcc 13.3 for sm_86, sm_89 and sm_90.
2. Give the final TORCH_CHECK a message that names the offending value,
the supported maximum, and what to change.
Note on why the ladder stops here: the natural next step, <1024, 16>,
needs ~66 KB of static shared memory for the same scratch union, and
ptxas rejects it ("uses too much shared data (0x10810 bytes, 0xc000
max)") on every arch. Going beyond 8192 pages needs dynamic shared
memory or a different top-k strategy, so the error message is the
honest fix for that range.
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
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.
Problem
topk_outputincsrc/topk.cudispatches onmax_num_pagesthrough a ladder of(NUM_THREADS, ITEM_PER_THREAD)template pairs:max_num_pages≤<128, 1><128, 2><128, 4><256, 4><256, 8><512, 8>TORCH_CHECK(false)Two problems past 4096 pages:
page_size=16that is a 64k-token context; withpage_size=1(which the token-level algorithms need) it is 4096 tokens.TORCH_CHECK(false)with no message aborts with a bare assertion. Nothing tells you which value was exceeded, what the limit is, or what to change. I spent a while bisecting config values to find out that this was the thing that tripped.Change
<1024, 8>tier formax_num_pages <= 8192, continuing the existing pattern (NUM_THREADS × ITEM_PER_THREAD= tier bound).TORCH_CHECKa message that names the offending value, the supported maximum, and what to change.Why the ladder stops at 8192 and not higher
The natural next rung,
<1024, 16>, does not build. The CUB scratch union in the kernelis dominated by the block-exchange buffer, which scales as
NUM_THREADS × ITEM_PER_THREAD × 4 B. That is static shared memory, capped at 48 KB per block:<512, 8>(current top)<1024, 8>(this PR)<1024, 16>Reproduced with both ptxas 12.8 and 13.3, on
sm_86,sm_89andsm_90. Going past 8192 pages needs dynamic shared memory with an explicitcudaFuncSetAttributeopt-in, or a different top-k strategy — either is a bigger change than this PR, so the honest thing for that range is an error message that says so. The reasoning is left as a comment at theTORCH_CHECK.Verification
All seven tiers compile with nvcc 13.3 for
sm_86,sm_89andsm_90:1024 is the CUDA per-block thread limit, and at 63 registers/thread the new tier stays inside the 64K register file (1024 × 63 = 64,512).
No change to any existing tier's behaviour — this only adds a rung and a message.