From 08ef70e8b4690b7bf7535c7dbe9926a163067b05 Mon Sep 17 00:00:00 2001 From: Zhaohui Wang <173976389+GeoffreyWang1117@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:22:44 -0700 Subject: [PATCH] topk: add an 8192-page tier and make the overflow check actionable 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. --- csrc/topk.cu | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/csrc/topk.cu b/csrc/topk.cu index 62d747eb..df32f7c4 100644 --- a/csrc/topk.cu +++ b/csrc/topk.cu @@ -196,8 +196,26 @@ const int64_t max_num_pages reserved_bos, reserved_eos ); + } else if (max_num_pages <= 8192){ + TopKOutput_BF16_Kernel<1024, 8><<>>( + reinterpret_cast<__nv_bfloat16*>(x.data_ptr()), + dense_kv_indptr.data_ptr(), + sparse_kv_indptr.data_ptr(), + dense_kv_indices.data_ptr(), + sparse_kv_indices.data_ptr(), + topk_val, + reserved_bos, + reserved_eos + ); } else { - TORCH_CHECK(false); + // The next tier up (<1024, 16>) would need ~66 KB of *static* shared + // memory for the CUB block-sort scratch, above the 48 KB per-block + // limit ptxas enforces, so 8192 is the highest tier this dispatch + // shape can reach without switching to dynamic shared memory. + TORCH_CHECK(false, "topk_output: max_num_pages (", max_num_pages, + ") exceeds the maximum supported (8192). Reduce the " + "context length / raise page_size so that " + "ceil(max_seq_len / page_size) <= 8192."); } }