From c73fa8842a48ef4e6194dde44caacea3c7c7d50a Mon Sep 17 00:00:00 2001
From: Erik Schultheis <erik.schultheis@aalto.fi>
Date: Mon, 15 Apr 2024 03:41:09 +0300
Subject: [PATCH] fixed numerics for online kernel, and an existing race
 condition in kernel 2

---
 dev/cuda/softmax_forward.cu | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dev/cuda/softmax_forward.cu b/dev/cuda/softmax_forward.cu
index 1742d2a..3e742b2 100644
--- a/dev/cuda/softmax_forward.cu
+++ b/dev/cuda/softmax_forward.cu
@@ -143,8 +143,8 @@ __global__ void softmax_forward_kernel2(float* out, const float* inp, int N, int
             shared[tid] = fmaxf(shared[tid], shared[tid + stride]);
         }
     }
-    float offset = shared[0];
     __syncthreads();
+    float offset = shared[0];
     // compute expf and write the result to global memory
     for (int i = tid; i < C; i += block_size) {
         out[idx * C + i] = expf(x[i] - offset);
@@ -327,7 +327,7 @@ __global__ void softmax_forward_online_kernel1(float* out, const float* inp, int
         float* out_row = out + i * C;
 
         float maxval = -INFINITY;
-        float sum = 0.0f;
+        double sum = 0.0f;
         for (int j = 0; j < C; j++) {
             float maxval_prev = maxval;
 			if (inp_row[j] > maxval) {