Skip to content

topk: add an 8192-page tier and make the overflow check actionable - #4

Open
GeoffreyWang1117 wants to merge 1 commit into
Infini-AI-Lab:v1from
GeoffreyWang1117:pr/topk-8192-tier
Open

topk: add an 8192-page tier and make the overflow check actionable#4
GeoffreyWang1117 wants to merge 1 commit into
Infini-AI-Lab:v1from
GeoffreyWang1117:pr/topk-8192-tier

Conversation

@GeoffreyWang1117

Copy link
Copy Markdown

Problem

topk_output in csrc/topk.cu dispatches on max_num_pages through a ladder of (NUM_THREADS, ITEM_PER_THREAD) template pairs:

max_num_pages kernel
128 <128, 1>
256 <128, 2>
512 <128, 4>
1024 <256, 4>
2048 <256, 8>
4096 <512, 8>
TORCH_CHECK(false)

Two problems past 4096 pages:

  1. The ladder stops there. With page_size=16 that is a 64k-token context; with page_size=1 (which the token-level algorithms need) it is 4096 tokens.
  2. The failure is silent. 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

  1. Add a <1024, 8> tier for max_num_pages <= 8192, continuing the existing pattern (NUM_THREADS × ITEM_PER_THREAD = tier bound).
  2. Give the final TORCH_CHECK a 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 kernel

__shared__ union {
    typename BLF::TempStorage  lf;
    typename BLI::TempStorage  li;
    typename BSI::TempStorage  si;
    typename Sort::TempStorage sort;
} temp;

is 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:

tier static smem
<512, 8> (current top) 18,528 B
<1024, 8> (this PR) 37,056 B ✅
<1024, 16> 67,600 B ❌
ptxas error : Entry function '..._KernelILi1024ELi16E...'
              uses too much shared data (0x10810 bytes, 0xc000 max)

Reproduced with both ptxas 12.8 and 13.3, on sm_86, sm_89 and sm_90. Going past 8192 pages needs dynamic shared memory with an explicit cudaFuncSetAttribute opt-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 the TORCH_CHECK.

Verification

All seven tiers compile with nvcc 13.3 for sm_86, sm_89 and sm_90:

sm_86   all 7 tiers compile OK
sm_89   all 7 tiers compile OK
sm_90   all 7 tiers compile OK

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.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant