From a67753f51632a517f2f0a81b644cd2aafeee25ae Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 1 Jul 2026 14:15:58 +0800 Subject: [PATCH] softmax: cap block size at 512 threads launch_softmax_{f32,bf16} clamped block to 1024 threads when cols was larger. Halving the ceiling to 512 keeps two blocks per SM resident on the large vocab kernels that dominate speculative verify workloads without changing rows/block indexing, and never exceeds cols. --- csrc/reduce/softmax.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/reduce/softmax.cu b/csrc/reduce/softmax.cu index 1be36de..c3fe584 100644 --- a/csrc/reduce/softmax.cu +++ b/csrc/reduce/softmax.cu @@ -90,7 +90,7 @@ __global__ void softmax_bf16( extern "C" { void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stream) { - int block = (cols < 1024) ? cols : 1024; + int block = (cols < 512) ? cols : 512; if (block < 32) block = 32; softmax_f32<<>>( (const float*)x, (float*)out, cols); @@ -98,7 +98,7 @@ void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stre } void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* stream) { - int block = (cols < 1024) ? cols : 1024; + int block = (cols < 512) ? cols : 512; if (block < 32) block = 32; softmax_bf16<<>>( (const __nv_bfloat16*)x, (__nv_bfloat16*)out, cols);