From 60cf281e46e0336808fa9d243c41dcc863f142a6 Mon Sep 17 00:00:00 2001
From: Erik Schultheis <erik.schultheis@aalto.fi>
Date: Sun, 14 Apr 2024 14:56:51 +0300
Subject: [PATCH] use V dimension in softmax tests adjusted tests to be more
 challenging

---
 dev/cuda/softmax_forward.cu | 46 ++++++++++++++++++++++++++++---------
 1 file changed, 35 insertions(+), 11 deletions(-)

diff --git a/dev/cuda/softmax_forward.cu b/dev/cuda/softmax_forward.cu
index d4d227d..e075256 100644
--- a/dev/cuda/softmax_forward.cu
+++ b/dev/cuda/softmax_forward.cu
@@ -29,6 +29,7 @@ version 7 is softmax optimized for very large C.
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <assert.h>
 #include <cuda_runtime.h>
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
@@ -50,13 +51,17 @@ void softmax_forward_cpu(float* out, const float* inp, int N, int C) {
                 maxval = inp_row[j];
             }
         }
-        float sum = 0.0f;
+        // Note: since we want to ensure that the CUDA-kernels are accurate,
+        // we do this accumulation in higher precision, so we can be assured
+        // that our ground-truth is of high quality.
+        double sum = 0.0;
         for (int j = 0; j < C; j++) {
             out_row[j] = expf(inp_row[j] - maxval);
             sum += out_row[j];
         }
+        float norm = 1.f / (float)sum;
         for (int j = 0; j < C; j++) {
-            out_row[j] /= sum;
+            out_row[j] *= norm;
         }
     }
 }
@@ -592,20 +597,31 @@ int main(int argc, char **argv) {
 
     int B = 8;
     int T = 1024;
+    int V = 50257;
 
     int deviceIdx = 0;
     cudaCheck(cudaSetDevice(deviceIdx));
 
     // create host memory of random numbers
-    float* out = (float*)malloc(B * T * T * sizeof(float));
-    float* inp = make_random_float(B * T * T);
+    float* out = (float*)malloc(B * T * V * sizeof(float));
+    float* inp = make_random_float(B * T * V);
+
+    // make the input less uniformly random: Otherwise, all probabilities will be basically zero,
+    // and the tests are not actually meaningful.
+    const int* outliers = make_random_int(B * T * 3, V);
+    for(int k = 0; k < 3; ++k) {
+        for(int j = 0; j < B * T; ++j) {
+            inp[j * V +  outliers[j*3 + k]] *= 20;
+        }
+    }
+
 
     // move to GPU
     float* d_out;
     float* d_inp;
-    cudaCheck(cudaMalloc(&d_out, B * T * T * sizeof(float)));
-    cudaCheck(cudaMalloc(&d_inp, B * T * T * sizeof(float)));
-    cudaCheck(cudaMemcpy(d_inp, inp, B * T * T * sizeof(float), cudaMemcpyHostToDevice));
+    cudaCheck(cudaMalloc(&d_out, B * T * V * sizeof(float)));
+    cudaCheck(cudaMalloc(&d_inp, B * T * V * sizeof(float)));
+    cudaCheck(cudaMemcpy(d_inp, inp, B * T * V * sizeof(float), cudaMemcpyHostToDevice));
 
     // read kernel_num from command line
     int kernel_num = 1;
@@ -616,14 +632,22 @@ int main(int argc, char **argv) {
 
     int block_sizes[] = {32, 64, 128, 256, 512, 1024};
 
-    softmax_forward_cpu(out, inp, B * T, T);
+    softmax_forward_cpu(out, inp, B * T, V);
+    {
+        float max_el = -INFINITY;
+        for(int i = 0; i <  B * T * V; ++i) {
+            max_el = max(max_el, out[i]);
+        }
+        assert(max_el > 1e-4);
+        printf("Largest output is: %f\n", max_el);
+    }
 
     // first check the correctness of the kernel
     for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
         int block_size = block_sizes[j];
         printf("Checking block size %d.\n", block_size);
-        softmax_forward(kernel_num, d_out, d_inp, B * T, T, block_size);
-        validate_result(d_out, out, "out", B * T * T, 1e-4f);
+        softmax_forward(kernel_num, d_out, d_inp, B * T, V, block_size);
+        validate_result(d_out, out, "out", B * T * V, 1e-4f);
     }
 
     printf("All results match. Starting benchmarks.\n\n");
@@ -634,7 +658,7 @@ int main(int argc, char **argv) {
 
         int repeat_times = 1000;
         float elapsed_time = benchmark_kernel(repeat_times, softmax_forward,
-                                              kernel_num, d_out, d_inp, B * T, T, block_size
+                                              kernel_num, d_out, d_inp, B * T, V, block_size
                                               );
 
         printf("block_size %4d | time %.4f ms | per token %.2f µs\n", block_size, elapsed_time, elapsed_time * 1'000 / (B*T));