From 45fe51a28e43c02a8aa7060a0b4ff06379926540 Mon Sep 17 00:00:00 2001 From: Yifan Cui Date: Sat, 31 Jan 2026 13:42:21 +0800 Subject: [PATCH] Reduce topk kernel shared memory from 128KB to 32KB for better occupancy (#17747) Co-authored-by: Claude --- sgl-kernel/csrc/elementwise/topk.cu | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sgl-kernel/csrc/elementwise/topk.cu b/sgl-kernel/csrc/elementwise/topk.cu index 8bd0f3dcf..04a2b73be 100644 --- a/sgl-kernel/csrc/elementwise/topk.cu +++ b/sgl-kernel/csrc/elementwise/topk.cu @@ -32,7 +32,10 @@ constexpr size_t kSmem = static_cast(SGL_TOPK_DYNAMIC_SMEM_BYTES); constexpr size_t kSmem = 48 * 1024; // bytes #endif #else -constexpr size_t kSmem = 32 * 1024 * sizeof(uint32_t); // 128KB (bytes) +// Reduced from 128KB to 32KB to improve occupancy. +// Each radix pass needs at most ~TopK candidates in the threshold bin, +// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient. +constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t); // 32KB (bytes) #endif struct FastTopKParams {