From d09a08aeaebc6a9645c80f5b60da2a0598141772 Mon Sep 17 00:00:00 2001 From: santhoshini <89771081+pranathi000@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:53:08 +0530 Subject: [PATCH 1/2] Refactor softmaxKernel for numerical stability --- src/kernels.cu | 90 ++++++++++++++++++++------------------------------ 1 file changed, 35 insertions(+), 55 deletions(-) diff --git a/src/kernels.cu b/src/kernels.cu index c55fef3..4b60057 100644 --- a/src/kernels.cu +++ b/src/kernels.cu @@ -158,46 +158,37 @@ void causalMask(__nv_bfloat16 *input, int num_tokens) __global__ void softmaxKernel(__nv_bfloat16 *input, int num_tokens) { - // softmaxxing per head - // might waste a lot of memory by hardcoding the size here but can't use num_tokens directly - __shared__ float row[1024]; // row[0] will contain max value after the loop - __shared__ float max_val; - // find max of the row to subtract it for numerical stability + __shared__ float m[1024]; // running max per tree node + __shared__ float d[1024]; // running denominator (sum of exp) per tree node + int workIndex = blockIdx.x * num_tokens + threadIdx.x; - __nv_bfloat16 token = input[workIndex]; - row[threadIdx.x] = (float)token; + float token = (float)input[workIndex]; + + // leaf: this thread owns a single element + m[threadIdx.x] = token; + d[threadIdx.x] = 1.0f; __syncthreads(); + // one reduction: running max AND running sum, before the same __syncthreads() for (int i = 1; i < num_tokens; i = i * 2) { if (threadIdx.x % (i * 2) == 0 && threadIdx.x + i < num_tokens) { - row[threadIdx.x] = fmaxf(row[threadIdx.x], row[threadIdx.x + i]); - } - __syncthreads(); - } - if (threadIdx.x == 0) - { - max_val = row[0]; // so I don't need to allocate another shared value for max_val - } - __syncthreads(); + float m_a = m[threadIdx.x]; + float d_a = d[threadIdx.x]; + float m_b = m[threadIdx.x + i]; + float d_b = d[threadIdx.x + i]; - // turn into exp - row[threadIdx.x] = expf((float)token - max_val); - __syncthreads(); + float m_new = fmaxf(m_a, m_b); + float d_new = d_a * expf(m_a - m_new) + d_b * expf(m_b - m_new); - // now I can compute the numerical stable sum, similar pattern - tree reduction - // reusing row memory - for (int i = 1; i < num_tokens; i = i * 2) - { - if (threadIdx.x % (i * 2) == 0 && threadIdx.x + i < num_tokens) - { - row[threadIdx.x] = row[threadIdx.x] + row[threadIdx.x + i]; + m[threadIdx.x] = m_new; + d[threadIdx.x] = d_new; } __syncthreads(); } - input[workIndex] = (__nv_bfloat16)(expf((float)token - max_val) / row[0]); + input[workIndex] = (__nv_bfloat16)(expf(token - m[0]) / d[0]); } // input are masked attention scores (NUM_Q_HEADS, num_tok, num_tok) @@ -318,46 +309,35 @@ void ropeDecode(__nv_bfloat16 *input, int position_in_sequence, int proj_dim) // seq_len increases by 1 with every new token __global__ void softmaxKernelDecode(__nv_bfloat16 *input, int seq_len) { - // softmaxxing per head - // might waste a lot of memory by hardcoding the size here but can't use num_tokens directly - __shared__ float row[1024]; // row[0] will contain max value after the loop - __shared__ float max_val; - // find max of the row to subtract it for numerical stability + __shared__ float m[1024]; + __shared__ float d[1024]; + int workIndex = blockIdx.x * MAX_SEQ_LEN + threadIdx.x; - __nv_bfloat16 token = input[workIndex]; - row[threadIdx.x] = (float)token; + float token = (float)input[workIndex]; + + m[threadIdx.x] = token; + d[threadIdx.x] = 1.0f; __syncthreads(); for (int i = 1; i < seq_len; i = i * 2) { if (threadIdx.x % (i * 2) == 0 && threadIdx.x + i < seq_len) { - row[threadIdx.x] = fmaxf(row[threadIdx.x], row[threadIdx.x + i]); - } - __syncthreads(); - } - if (threadIdx.x == 0) - { - max_val = row[0]; // so I don't need to allocate another shared value for max_val - } - __syncthreads(); + float m_a = m[threadIdx.x]; + float d_a = d[threadIdx.x]; + float m_b = m[threadIdx.x + i]; + float d_b = d[threadIdx.x + i]; - // turn into exp - row[threadIdx.x] = expf((float)token - max_val); - __syncthreads(); + float m_new = fmaxf(m_a, m_b); + float d_new = d_a * expf(m_a - m_new) + d_b * expf(m_b - m_new); - // now I can compute the numerical stable sum, similar pattern - tree reduction - // reusing row memory - for (int i = 1; i < seq_len; i = i * 2) - { - if (threadIdx.x % (i * 2) == 0 && threadIdx.x + i < seq_len) - { - row[threadIdx.x] = row[threadIdx.x] + row[threadIdx.x + i]; + m[threadIdx.x] = m_new; + d[threadIdx.x] = d_new; } __syncthreads(); } - input[workIndex] = (__nv_bfloat16)(expf((float)token - max_val) / row[0]); + input[workIndex] = (__nv_bfloat16)(expf(token - m[0]) / d[0]); } // input are masked attention scores (NUM_Q_HEADS, seq_len) @@ -447,4 +427,4 @@ __global__ void pagedAttentionKernel(int layer, int num_active_slots, __nv_bfloa void pagedAttention(int layer, int num_active_slots, __nv_bfloat16 *q_proj, __nv_bfloat16 *kv_cache, int *block_table_gpu, int *gpu_seq_lens, int *gpu_active_slots, __nv_bfloat16 *output) { pagedAttentionKernel<<>>(layer, num_active_slots, q_proj, kv_cache, block_table_gpu, gpu_seq_lens, gpu_active_slots, output); -} \ No newline at end of file +} From 4ae160fbfc859331a11ac330f838ed2a4f470ee0 Mon Sep 17 00:00:00 2001 From: santhoshini <89771081+pranathi000@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:54:59 +0530 Subject: [PATCH 2/2] Fuse online softmax: compute running max and sum in one tree reduction --- src/kernels.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/kernels.cu b/src/kernels.cu index 4b60057..19ad83b 100644 --- a/src/kernels.cu +++ b/src/kernels.cu @@ -1,7 +1,7 @@ #include "cuda_to_hip.h" #include "kernels.cuh" #include - +// changes made // TODO perhaps share these between main.cpp and kernels.cu to not duplicate them? constexpr int N_LAYERS = 16; // TODO: hardcoded for llama 3.2 1B, just like any other value for now