Merge remote-tracking branch 'karpathy/master' into linear16

README.md

@@ -2,13 +2,7 @@
 
 LLM training in simple, pure C/CUDA. There is no need for 245MB of PyTorch or 107MB of cPython. Training GPT-2 (CPU, fp32) is ~1,000 lines of clean code in the single file [train_gpt2.c](train_gpt2.c), and training it on GPU is ~2,000 lines (adds CUDA kernels) in [train_gpt2.cu](train_gpt2.cu). The code compiles and runs instantly, it exactly matches the PyTorch reference implementation, and it ~matches the speed of (compiled) PyTorch (fp32, no flash attention). I chose GPT-2 as the first working example because it is the grand-daddy of LLMs, the first time the modern stack was put together.
 
-Currently, we are working on:
-
-- optimize the CUDA implementation further to match/exceed PyTorch speed
-- lower the precision from fp32 to mixed precision training
-- add multi-gpu training, starting with DDP
-- reproduce the GPT-2 training run (add data, evals)
-- more modern architectures, Llama 2, Gemma, Mistral, etc.
+Our current goal is to reproduce GPT-2 with a multi-node, mixed-precision, efficient implementation. For an overview of current ongoing work, see the latest [State of the Union](https://github.com/karpathy/llm.c/discussions/224) post.
 
 I'd like this repo to only maintain C and CUDA code. Ports of this repo to other languages are very welcome, but should be done in separate repos, and then I am happy to link to them below in the "notable forks" section, just like I did in [llama2.c notable forks](https://github.com/karpathy/llama2.c/tree/master?tab=readme-ov-file#notable-forks).
 
@@ -255,10 +249,12 @@ Lastly, I will be a lot more sensitive to complexity in the root folder of the p
 
 - C#
   - [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project
- 
+
 - Rust
   -  [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project
 
+- Metal
+  - [llm.metal](https://github.com/regrettable-username/llm.metal) by @[regrettable-username](https://github.com/regrettable-username): LLM training in simple, raw C/Metal Shading Language
 ## discussions
 
 Ways of organizing development:

dev/cuda/common.h

@@ -86,8 +86,8 @@ void validate_result(T* device_result, const T* cpu_reference, const char* name,
         if (i < 5) {
             printf("%f %f\n", cpu_reference[i], out_gpu[i]);
         }
-        // ensure correctness for all elements
-        if (fabs(cpu_reference[i] - out_gpu[i]) > tolerance) {
+        // ensure correctness for all elements. We can set an "ignore" mask by writing NaN
+        if (fabs(cpu_reference[i] - out_gpu[i]) > tolerance && !isnan(cpu_reference[i])) {
             printf("Mismatch of %s at %d: CPU_ref: %f vs GPU: %f\n", name, i, cpu_reference[i], out_gpu[i]);
             nfaults ++;
             if (nfaults >= 10) {

dev/cuda/layernorm_backward.cu

@@ -6,6 +6,9 @@ nvcc -O3 --use_fast_math layernorm_backward.cu -o layernorm_backward
 
 version 1 is naive port from CPU code to kernel: parallelizes over B,T, loops over C
 ./layernorm_backward 1
+
+version 2 moves a lot of reduction to shared memory over global memory
+./layernorm_backward 2
 */
 
 #include <stdio.h>
@@ -152,19 +155,18 @@ __global__ void layernorm_backward_kernel1(float* dinp, float* dweight, float* d
     }
 }
 
-
-// super naive kernel that just parallelizes over B,T and loops over C
+// uses shared memory instead for the reduces
 __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* dbias,
                                            const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd,
                                            int B, int T, int C) {
+    extern __shared__ float shared[]; // size = 2 * C
+
     namespace cg = cooperative_groups;
     cg::thread_block block = cg::this_thread_block();
     cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
     int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
     int N = B * T;
-    if(idx >= N) {
-        return;
-    }
+    if(idx >= N) { return; } // thread guards
 
     int b = idx / T;
     int t = idx % T;
@@ -175,6 +177,18 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d
     const float mean_bt = mean[b * T + t];
     const float rstd_bt = rstd[b * T + t];
 
+    // the first half of shared memory is bias, second is weight
+    float* dbias_shared = shared;
+    float* dweight_shared = shared + C;
+
+    // init shared memory to zero
+    #pragma unroll
+	for(int i = threadIdx.x; i < C; i+= blockDim.x){
+       dbias_shared[i] = 0.0f;
+       dweight_shared[i] = 0.0f;
+    }
+    __syncthreads();
+
     // first: two reduce operations
     float dnorm_mean = 0.0f;
     float dnorm_norm_mean = 0.0f;
@@ -184,10 +198,8 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d
         dnorm_mean += dnorm_i;
         dnorm_norm_mean += dnorm_i * norm_bti;
     }
-
     dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
     dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
-
     dnorm_mean = dnorm_mean / C;
     dnorm_norm_mean = dnorm_norm_mean / C;
 
@@ -196,9 +208,9 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d
         float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
         float dnorm_i = weight[i] * dout_bt[i];
         // gradient contribution to bias
-        atomicAdd(&dbias[i], dout_bt[i]);
+        atomicAdd(&dbias_shared[i], dout_bt[i]);
         // gradient contribution to weight
-        atomicAdd(&dweight[i], norm_bti * dout_bt[i]);
+        atomicAdd(&dweight_shared[i], norm_bti * dout_bt[i]);
         // gradient contribution to input
         float dval = 0.0f;
         dval += dnorm_i; // term 1
@@ -207,6 +219,13 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d
         dval *= rstd_bt; // final scale
         dinp_bt[i] += dval;
     }
+    __syncthreads();
+
+    // write to global memory
+	for(int i = threadIdx.x; i < C; i+= blockDim.x){
+        atomicAdd(&dbias[i], dbias_shared[i]);
+        atomicAdd(&dweight[i], dweight_shared[i]);
+	}
 }
 
 // ----------------------------------------------------------------------------
@@ -225,7 +244,8 @@ void layernorm_backward2(float* dinp, float* dweight, float* dbias,
                         int B, int T, int C, const int block_size) {
     const int N = B * T;
     const int grid_size = ceil_div(32*N, block_size);
-    layernorm_backward_kernel2<<<grid_size, block_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
+    size_t shared_mem_size = 2 * C * sizeof(float);
+    layernorm_backward_kernel2<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
 }
 
 // kernel version dispatch
@@ -358,4 +378,4 @@ int main(int argc, char **argv) {
     cudaCheck(cudaFree(d_rstd));
 
     return 0;
-}
+}

dev/cuda/matmul_backward_bias.cu

@@ -7,6 +7,7 @@ nvcc -O3 matmul_backward_bias.cu -lineinfo -o matmul_backward_bias
 ./matmul_backward_bias 1
 ./matmul_backward_bias 2
 ./matmul_backward_bias 3
+./matmul_backward_bias 4
 
 ncu:
 sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1
@@ -14,6 +15,7 @@ sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <assert.h>
 #include <cublas_v2.h>
 #include <cuda_runtime.h>
 #include <omp.h>
@@ -124,6 +126,46 @@ __global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, in
     }
 }
 
+// this kernel performs a column-wise reduction over dout, in PyTorch equivalent to:
+// dbias = dout.sum((0,1))
+// the idea is to employ one block to reduce along several columns,
+// where each block has a width of 32 columns to ensure coalesced access.
+// at the end we accumulate the reductions performed by the warps in each block via shared memory
+__global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, int B, int T, int OC) {
+    // this kernel is launched with 1D grid_dim of OC/32
+    // for example let's say block_size is 128
+    extern __shared__ float smem[]; // of size block_size (128)
+    const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3
+    const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31
+    const int tl = blockIdx.x * warpSize; // pointer to the start column for this block
+    const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4
+
+    // pointer to the start of the column for one lane of threads
+    // so e.g. 4 threads (of the same lane_id) will reduce this one column
+    const float* dout_col = dout + tl + lane_id;
+
+    // column reductions by looping through the rows
+    // each of the 4 threads offsets by its warp_id and then skips by vstep
+    // together these 4 threads cover all B*T rows of this (lane_id) column
+    // importantly, consecutive threads (in threadId) are processing adjacent columns,
+    // leading to a coalesced memory access pattern
+    float dout_sum = 0.0f;
+    for (int row = warp_id; row < B * T; row += vstep) {
+        dout_sum += dout_col[row * OC];
+    }
+    smem[lane_id + warp_id * warpSize] = dout_sum;
+    __syncthreads();
+
+    // warp_id 0 reduces the shared memory column-wise, linearly
+    dout_sum = 0.0f;
+    if (warp_id == 0) {
+        for (int j = 0; j < vstep; j++) {
+            dout_sum += smem[lane_id + j * warpSize];
+        }
+        dbias[tl + lane_id] += dout_sum;
+    }
+}
+
 // ----------------------------------------------------------------------------
 // kernel launcher
 
@@ -152,6 +194,14 @@ void matmul_backward_bias3(float* dinp, float* dweight, float* dbias,
     matmul_backward_bias_kernel3<<<OC, block_size>>>(dbias, dout, B, T, OC);
 }
 
+void matmul_backward_bias4(float* dinp, float* dweight, float* dbias,
+                      float* dout, float* inp, float* weight, float* ones,
+                      int B, int T, int C, int OC, int block_size) {
+    assert(OC % 32 == 0); // OC must be divisible by 32 for this kernel
+    const int grid_size = OC / 32;
+    matmul_backward_bias_kernel4<<<grid_size, block_size, block_size * sizeof(float)>>>(dbias, dout, B, T, OC);
+}
+
 void matmul_backward_bias(int kernel_num,
                      float* dinp, float* dweight, float* dbias,
                      float* dout, float* inp, float* weight, float* ones,
@@ -166,6 +216,9 @@ void matmul_backward_bias(int kernel_num,
         case 3:
             matmul_backward_bias3(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
             break;
+        case 4:
+            matmul_backward_bias4(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
+            break;
         default:
             printf("Invalid kernel number\n");
             exit(1);
@@ -230,7 +283,7 @@ int main(int argc, char **argv) {
         matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, 128);
         // compare
         printf("Checking correctness...\n");
-        validate_result(d_dbias, dbias, "dbias", OC, 1e-3f);
+        validate_result(d_dbias, dbias, "dbias", OC, 5e-3f);
         printf("All results match for block_size=%d.\n\n", block_size);
     }
 

dev/cuda/trimat_forward.cu

@@ -0,0 +1,596 @@
+/*
+Triangular matrix multiplication as in autoregressive attention. A short story.
+by @ngc92
+
+Compile:
+nvcc -O3 --use_fast_math trimat_forward.cu -o trimat_forward -lcublas
+
+Run:
+
+cuBLAS baseline kernel
+./trimat_forward 0
+
+naive
+./trimat_forward 1
+
+registers
+./trimat_forward 2
+
+tri3
+./trimat_forward 3
+
+tri4
+./trimat_forward 4
+*/
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <assert.h>
+#include <float.h>
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include "common.h"
+
+static cublasHandle_t cublas_handle;
+static float* d_qkvr;   // scratch for the cublas kernel
+
+/*                    ** Chapter I - Introduction **
+ *
+ *  You are Trimul. You've always wanted to do fast matrix multiplication, but they said
+ *  "Don't bother, big dumb Cublas is much faster than you!"
+ *  "I don't need to be faster than Cublas", you replied, "I can be smarter. Cublas calculates
+ *  the entire matrix, but I need only half. If I'm more than half as fast as Cublas, I'm
+ *  going to win."
+ *
+ *  So to prove everyone wrong, you enter the TriMatlon, the most prestigious competition
+ *  for anyone paying Attention.
+ *
+ *  Before you start preparing, lets have a look at the players involved
+ *
+ *  First, there is the Referee (`trimul_cpu`), slow and ponderous, but producing results
+ *  beyond reproof.
+ *  Then, there is Cublas. Cublas' mind is so inflexible, it doesn't actually comprehend
+ *  what we are trying to do here, so Cublas has brought an assistant (`permute_kernel`)
+ *  that translates the competition into a task that it can solve. But once it recognizes
+ *  the problem, its muscle memory kicks in, and matrix products are produced faster than
+ *  the eye can see. Stuck in its routine, Cublas doesn't realize the task is already
+ *  finished with the lower triangle, though.
+ *
+ *  If you can do without an assistant, and can solve the right task, then that's your opportunity
+ *  to shine!
+ */
+
+
+// taken from then attention forward pass
+void trimul_cpu(float* out, const float* inp,
+                int B, int T, int C, int NH) {
+    int C3 = C*3;
+    int hs = C / NH; // head size
+    float scale = 1.0 / sqrtf(hs);
+
+    for (int b = 0; b < B; b++) {
+        for (int t = 0; t < T; t++) {
+            for (int h = 0; h < NH; h++) {
+                const float* query_t = inp + b * T * C3 + t * C3 + h * hs;
+                float* out_bth = out + b * NH * T * T + h * T * T + t * T;
+
+                // pass 1: calculate query dot key and maxval
+                for (int t2 = 0; t2 <= t; t2++) {
+                    const float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key
+
+                    // (query_t) dot (key_t2)
+                    float val = 0.0f;
+                    for (int i = 0; i < hs; i++) {
+                        val += query_t[i] * key_t2[i];
+                    }
+                    val *= scale;
+
+                    out_bth[t2] = val;
+                }
+                for(int t2 = t + 1; t2 < T; ++t2) {
+                    out_bth[t2] = NAN;
+                }
+            }
+        }
+    }
+}
+
+__global__ void permute_kernel(float* q, float* k, float* v,
+                               const float* inp,
+                               int B, int N, int NH, int d) {
+    // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d)
+    // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d)
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+    // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_]
+
+    if (idx < B * NH * N * d) {
+        int b = idx / (NH * N * d);
+        int rest = idx % (NH * N * d);
+        int nh_ = rest / (N * d);
+        rest = rest % (N * d);
+        int n = rest / d;
+        int d_ = rest % d;
+
+        int inp_idx = \
+            (b * N * 3 * NH * d)
+            +   (n * 3 * NH * d)
+            +       (0 * NH * d)
+            +          (nh_ * d)
+            +                d_;
+
+        q[idx] = inp[inp_idx];
+        k[idx] = inp[inp_idx + NH * d];
+        v[idx] = inp[inp_idx + 2 * (NH * d)];
+    }
+}
+
+
+void trimul_cublas(float* preatt,
+                   const float* inp,
+                   int B, int T, int C, int NH) {
+    int HS = C / NH; // head size
+
+    // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)
+    float* q, * k, * v;
+    q = d_qkvr + 0 * B * T * C;
+    k = d_qkvr + 1 * B * T * C;
+    v = d_qkvr + 2 * B * T * C;
+    int total_threads = B * NH * T * HS;
+    int num_blocks = ceil_div(total_threads, 256);
+    permute_kernel<<<num_blocks, 256>>>(q, k, v, inp, B, T, NH, HS);
+    cudaCheck(cudaGetLastError());
+
+    // batched matrix multiply with cuBLAS
+    const float alpha = 1.0f / sqrtf(HS);
+    const float beta = 0.0f;
+    cublasCheck(cublasSgemmStridedBatched(cublas_handle,
+                                          CUBLAS_OP_T, CUBLAS_OP_N,
+                                          T, T, HS,
+                                          &alpha,
+                                          k, HS, T * HS,
+                                          q, HS, T * HS,
+                                          &beta,
+                                          preatt, T, T * T,
+                                          B * NH));
+}
+
+/*                    ** Chapter II - Getting a Team **
+ *
+ *  OK, you've registered for the competition, now what to do. TriMatlon is a team competition, so first, you need
+ *  to figure out what kind of team you need, and how to organize it. The individual instances and heads of the
+ *  problem are completely independent, so you just can send separate teams to work there completely independently.
+ *
+ *  To figure out how to organize each team, you take out your spyglass (`Nsight Compute`) and look how the Cublas teams
+ *  are handling their work.
+ *  Turns out, you need 256 athletes in each group, and those handle 128 x 128 of the tasks. They work together in
+ *  a tight square formation, 16 wide and 16 deep.
+ *
+ *  So, you went out and got your 100 000 friends, and split them into groups (`trimul_launcher`). Each group gets
+ *  informed about where they should work (`trimul_global`) and goes off to do their thing (`matmul_tri_naive`).
+ *  Let's observe how we're doing.
+ */
+
+// using creates an alias for a function pointer
+using matmul_fn_ptr = void(*)(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha);
+
+template<matmul_fn_ptr matmul_tri>
+__global__ void __launch_bounds__(256, 2) trimul_global(float* out, const float* inp, int T, int C, int NH) {
+    // skip above the diagonal
+    if(blockIdx.y > blockIdx.x)
+        return;
+
+    // set up indices
+    int C3 = C*3;
+    int hs = C / NH; // head size
+    float scale = 1.0 / sqrtf(hs);
+
+    // we put the "batch x head" dimension into the z block index.
+    int h = blockIdx.z % NH;
+    int b = blockIdx.z / NH;
+
+    // Get the base address for the current batch and head
+    const float* q = inp + b * T * C3 + h * hs;
+    const float* k = inp + b * T * C3 + h * hs + C;
+    float* r = out + (b*NH + h)*T*T;
+
+    // start the multiplication
+    matmul_tri(r, T, q, C3, k, C3, T, hs, scale);
+}
+
+template<matmul_fn_ptr matmul_tri>
+void trimul_launcher(float* out, const float* inp, int B, int T, int C, int NH) {
+    // we assume nice shapes here. Let's not make the code a mess by supporting weird shapes that you
+    // wouldn't want to use anyway.
+    assert(T % 128 == 0);
+    // No need to ceil_div, if it's not a multiple of 128, we would get wrong results anyway.
+    trimul_global<matmul_tri><<<dim3(T / 128, T / 128, NH * B), dim3(16, 16)>>>(out, inp, T, C, NH);
+    cudaCheck(cudaGetLastError());
+}
+
+/*                     ** Chapter III - ... **
+ *
+ *  You go over to the playing field. On one end of the field, there is a huge pile of funnily shaped cookie cutters.
+ *  Some in the shape of animals, some in the shape of a landscape. Each group of workers has assigned some runners,
+ *  fetching the cookie cutters for them. The workers seem very relaxing, chatting with each other, lounging about.
+ *  You focus in on one of them.
+ *
+ *  He seems to be giving an instruction to a runner, and then turns back to reading a novel. The runner, meanwhile,
+ *  crosses the field and back, handing him an elephant shape. Then she's off again to pick up a savannah background.
+ *  Having received the two shapes, pressed them into the dough, and makes an elephant-in-the-savannah cookie. He hands
+ *  the cutters back to the runner. "Can you please fetch me an elephant and a jungle next?"
+ *  While she's on her way, he takes a sip off his cocktail.
+ *  This time, she's making only one trip, keeping the elephant in her pocket (_Cache_). Still, it seems to take forever.
+ *  You keep observing:
+ *  - Elephant and zoo
+ *  - Elephant and island
+ *  ...
+ *  - Lion and savannah
+ *  - Lion and jungle
+ *  - Lion and zoo
+ *  ...
+ *
+ *  The worker has his poor runner fetch the same things over and over again, looking like she's about to faint from exhaustion.
+ *  Even though she realizes this and always keeps one of them in her pocket, there is so much running,
+ *  and little actual work happening.
+ *
+ *  Clearly, this isn't going to be effective, so you call a team meeting.
+ */
+
+// baseline implementation: 20 ms
+__device__ void matmul_tri_naive(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) {
+    // get coordinates of our block
+    int i_base = 128 * blockIdx.x + 8 * threadIdx.x;
+    int j_base = 128 * blockIdx.y + 8 * threadIdx.y;
+
+    // one more check to skip the upper diagonal in blocks that are on the diagonal.
+    if(j_base > i_base)
+        return;
+
+    // Simple nested loop that calculates 8x8 results in one thread.
+    for(int io = 0; io < 8; ++io) {
+        int i = i_base + io;
+        for(int jo = 0; jo < 8; ++jo) {
+            int j = j_base + jo;
+            float val = 0;
+            for (int s = 0; s < hs; ++s) {
+                val += k[i * ks + s] * q[j * qs + s];
+            }
+            p[i * ps + j] = val * alpha;
+        }
+    }
+}
+
+/*                     ** Chapter IV - ... **
+ *
+ *  Each worker is producing 64 combined cookies from 8 animals and 8 landscapes. They send there runners of 64 times
+ *  to fetch the corresponding shapes. This is terribly inefficient; The runners need a minute or so for each trip,
+ *  but making a cookie can be done in just a second.
+ *
+ *  "Let's try something different tomorrow: Just get all 16 cookie cutters that you need, and do all 64 combinations
+ *  of them! See all this free space on your workbench (_registers_), you can keep them all there for easy access."
+ *
+ *  The next morning, you come back to the field for another practice session. Initially, there is bustling activity
+ *  with the runners, picking up 16 shapes for each worker. But then, the workers have to put down their newspapers
+ *  and start making cookies. Now there are 64 combinations, so it takes them a full minute.
+ *
+ *  Not all groups of workers are equally fast. When the first group finishes with all animal-landscape combinations,
+ *  they already start asking the runners for the next set of cookie cutters, combining plants and houses. Even though
+ *  the workers are much busier than before, they are still spending most of their time just waiting.
+ *
+ *  Still, instead of being busy for 20 hours, your team is now done with the task in just 3h 30 minutes; already, this
+ *  is five times faster.
+ *
+ *  You think to yourself: "Why should we stop at 8 x 8 combinations? Lets to 16 x 16, that's only twice the work for
+ *  the runners, but four times as much for the actual workers."
+ *  You head over to the baking area, and make that suggestion to one of your team leaders.
+ *  "In theory, that sounds great", she agrees, "but see, we only have limited space on our workbenches (_registers_).
+ *  There is still some room left, but we simply cannot bake 256 cookies at the same time, sorry."
+ *
+ *  A different strategy is needed, then.
+ */
+
+// reorganize loops to enable data reuse: 3.5 ms
+__device__ void matmul_tri_registers(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) {
+    int i_base = 128 * blockIdx.x + 8 * threadIdx.x;
+    int j_base = 128 * blockIdx.y + 8 * threadIdx.y;
+
+    if (j_base > i_base)
+        return;
+
+    // shift our pointers to the sub-block this thread is responsible for
+    k += i_base * ks;
+    q += j_base * qs;
+    p += i_base * ps + j_base;
+
+    float vals[8][8] = {};
+    for (int s = 0; s < hs; ++s) {
+        float lhs[8];
+        float rhs[8];
+        for (int u = 0; u < 8; ++u) {
+            lhs[u] = k[u * ks + s];
+            rhs[u] = q[u * qs + s];
+        }
+
+        for (int i = 0; i < 8; ++i) {
+            for (int j = 0; j < 8; ++j) {
+                vals[i][j] += lhs[i] * rhs[j];
+            }
+        }
+    }
+
+    for (int i = 0; i < 8; ++i) {
+        for (int j = 0; j < 8; ++j) {
+            p[i * ps + j] = vals[i][j] * alpha;
+        }
+    }
+}
+
+/*                     ** Chapter IV - By the Bucketload **
+ *
+ *  Despite the hectic activity, you pick out one of the runners. "Why are you always brining just one shape? Wouldn't
+ *  it be much more efficient if you took more than one?"
+ *  "Of course", the runner answers, "but they've asked me for an elephant, a lion, a zebra, and a goldfish. These
+ *  are all over the place, I can't just pick them up at one spot (_strided acccess_).
+ *  "But the lion is right next to the palm tree. You could bring those two together?", you confirm.
+ *  "Yes", he says, "if the just asked for the different categories at the same time, that would make things
+ *  so much easier. See, I have this bucket, I could carry lots of things in one go if I could just scoop them up
+ *  from the same place (_coalesced access_).
+ *
+ *  OK, then lets fetch the first animal, first plant, first vehicle, and first landmark shape in one go (_vectorized load_).
+ *  [Here, the metaphor breaks down a bit: Since we're accumulating all the results, getting more data at the same time
+ *  depth-wise doesn't require more space on the workbench. We're stacking the cookies!]
+ *
+ *  You also streamline the shape combination further. Instead of picking up all animals and landscapes at one, it is
+ *  more efficient, using less workbench space, to just pick up all animals. Then, you get one landscape, combine it
+ *  will all animals, get the next landscape, combine, and so on.
+ *
+ *  In this way, instead of 2 x 8 x 4 cookie cutters that take up space, you only need (8+1) x 4 at the same time.
+ *
+ *  With these optimizations, you are down to 100 minutes for this task. Still slower than Cublas, but not by much.
+ *
+ *  In the arena, each team also has access to a small storage hut, much closer to their workbenches than the piles of
+ *  cookie cutters on the other side. Cublas is using them heavily, so maybe you should, too.
+ */
+
+// convenient helper functions to make the code below more readable
+__device__ float4 ld_vec(const float* address) {
+    return *reinterpret_cast<const float4*>(address);
+}
+
+__device__ void st_vec(float* address, float4 val) {
+    *reinterpret_cast<float4*>(address) = val;
+}
+
+// vector instructions for coalesced memory access: 1.7 ms
+__device__ void matmul_tri3(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) {
+    int i_base = 128 * blockIdx.x + 8 * threadIdx.x;
+    int j_base = 128 * blockIdx.y + 8 * threadIdx.y;
+
+    if (j_base > i_base)
+        return;
+
+    // shift our pointers to the sub-block this thread is responsible for
+    k += i_base * ks;
+    q += j_base * qs;
+    p += i_base * ps + j_base;
+
+    float vals[8][8] = {};
+    for (int s = 0; s < hs; s += 4) {
+        // load in float4 to improve coalescing
+        float4 rhs[8];
+        for (int u = 0; u < 8; ++u) {
+            rhs[u] = ld_vec(q + u * qs + s);
+        }
+
+        for (int i = 0; i < 8; ++i) {
+            // no need to keep lhs around for the i loop, its only reused in the j loop anyway.
+            float4 lhs = ld_vec(k + i * ks + s);
+            for (int j = 0; j < 8; ++j) {
+                vals[i][j] += lhs.x * rhs[j].x;
+                vals[i][j] += lhs.y * rhs[j].y;
+                vals[i][j] += lhs.z * rhs[j].z;
+                vals[i][j] += lhs.w * rhs[j].w;
+            }
+        }
+    }
+
+    for (int i = 0; i < 8; ++i) {
+        for (int j = 0; j < 8; j += 4) {
+            float4 result;
+            result.x = vals[i][j + 0] * alpha;
+            result.y = vals[i][j + 1] * alpha;
+            result.z = vals[i][j + 2] * alpha;
+            result.w = vals[i][j + 3] * alpha;
+            st_vec(p + i * ps + j, result);
+        }
+    }
+}
+
+/*                     ** Chapter V - Sharing is Caring **
+ *
+ *  You take a look around the shed, and see that there are 32 shelves there. They are much larger than the workbenches,
+ *  giving you enough space for all the cookie cutters needed by the entire team.
+ *
+ *  Within the team, workers have banded together in groups of 32. They are always doing the same thing, reducing the
+ *  amount of effort required for coordination. However, that also means that if you send them all to pick up different
+ *  cookie cutters from the same shelf, they will have to wait and queue up (_shared memory bank conflict_).
+ *
+ *  In order to achieve maximum efficiency, we send the runners fetching cutters with the maximum bucket size: 32 different
+ *  categories at the same time.
+ *
+ *  [I'm having trouble getting the specifics into the story in a sensible way. For now, please read the code for more
+ *  details.]
+ *
+ */
+__device__ void matmul_tri4(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) {
+    int i_base = 128 * blockIdx.x + 8 * threadIdx.x;
+    int j_base = 128 * blockIdx.y + 8 * threadIdx.y;
+
+    // we need all threads for loading data, so none of them can chicken out early, even
+    // if they are not responsible for any useful result.
+    if (blockIdx.y > blockIdx.x)
+        return;
+
+    k += 128 * blockIdx.x * ks;
+    q += 128 * blockIdx.y * qs;
+
+    __shared__ float lhs_s[128][32];
+    __shared__ float rhs_s[128][32];
+
+    float vals[8][8] = {};
+    for (int so = 0; so < hs; so += 32) {
+        // Read a large slice of the input, worked on together by all threads.
+        // They are organized differently for this part. We want to ensure
+        // fully coalesced loads, so we let a single warp handle consecutive
+        // addresses, which means we need to combine two threadIdx.y values
+        // in one read operation.
+        // note: threads may read data here that they don't need themselves.
+        //       this really is a block-level operation.
+        __syncthreads();
+        for(int y = threadIdx.y / 2; y < 128; y += 8) {
+            int xo = (threadIdx.y % 2) * 16;
+            lhs_s[y][threadIdx.x + xo] = k[y * ks + so + threadIdx.x + xo];
+            rhs_s[y][threadIdx.x + xo] = q[y * qs + so + threadIdx.x + xo];
+        }
+        __syncthreads();
+
+        for (int si = 0; si < 32; ++si) {
+            float rhs[8];
+            for (int u = 0; u < 8; ++u) {
+                rhs[u] = rhs_s[u + 8 * threadIdx.y][(si + threadIdx.x) % 32];
+            }
+
+            for (int ii = 0; ii < 8; ++ii) {
+                float lhs = lhs_s[ii + 8 * threadIdx.x][(si + threadIdx.x) % 32];
+                for (int ji = 0; ji < 8; ++ji) {
+                    vals[ii][ji] += lhs * rhs[ji];
+                }
+            }
+        }
+    }
+
+    // don't write above the diagonal
+    if (j_base > i_base)
+        return;
+
+    for (int ii = 0; ii < 8; ++ii) {
+        for (int ji = 0; ji < 8; ji += 4) {
+            int i = i_base + ii;
+            int j = j_base + ji;
+            float4 result;
+            result.x = vals[ii][ji + 0] * alpha;
+            result.y = vals[ii][ji + 1] * alpha;
+            result.z = vals[ii][ji + 2] * alpha;
+            result.w = vals[ii][ji + 3] * alpha;
+            st_vec(p + i * ps + j, result);
+        }
+    }
+}
+
+/*                     ** Chapter VI - Competition Day **
+ *
+ * Finally, you feel ready to take on Cublas. You hand out tickets to the event for you friends to see.
+ *
+ *    ---------------------------------------------------------------------------------
+ *    |           CuBLAS vs TriMul - Fight of the Century                             |
+ *    |                                                                               |
+ *    |   Ticket code:                                                                |
+ *    |   > nvcc -O3 --use_fast_math trimat_forward.cu -o trimat_forward -lcublas     |
+ *    |   > ./trimat 4                                                                |
+ *    |                                                                               |
+ *    ---------------------------------------------------------------------------------
+ */
+
+void trimul_gpu(int kernel_num,
+                float* out,  const float* inp,
+                int B, int T, int C, int NH) {
+    switch (kernel_num) {
+        case 0:
+            trimul_cublas(out, inp, B, T, C, NH);
+            break;
+        case 1:
+            trimul_launcher<matmul_tri_naive>(out, inp, B, T, C, NH);
+            break;
+        case 2:
+            trimul_launcher<matmul_tri_registers>(out, inp, B, T, C, NH);
+            break;
+        case 3:
+            trimul_launcher<matmul_tri3>(out, inp, B, T, C, NH);
+            break;
+        case 4:
+            trimul_launcher<matmul_tri4>(out, inp, B, T, C, NH);
+            break;
+        default:
+            printf("Invalid kernel number\n");
+            exit(1);
+    }
+}
+
+
+
+int main(int argc, char **argv) {
+    srand(0);
+
+    int B = 8;
+    int T = 1024;
+    int C = 768;
+    int NH = 12;
+
+    int deviceIdx = 0;
+    cudaCheck(cudaSetDevice(deviceIdx));
+    cublasCreate(&cublas_handle);
+
+    // create host memory of random numbers
+    float* out = (float*)malloc(B * NH * T * T * sizeof(float));
+    float* inp = make_random_float(B * T * 3 * C);
+
+    // move to GPU
+    float* d_out;
+    float* d_inp;
+    cudaCheck(cudaMalloc(&d_out, B * NH * T * T * sizeof(float)));
+    cudaCheck(cudaMalloc(&d_inp, B * T * 3 * C * sizeof(float)));
+    cudaCheck(cudaMemcpy(d_inp, inp, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice));
+
+    // buffer for cublas
+    cudaCheck(cudaMalloc(&d_qkvr, B * T * 3 * C * sizeof(float)));
+
+    // read kernel_num from command line
+    int kernel_num = 1;
+    if (argc > 1) {
+        kernel_num = atoi(argv[1]);
+    }
+    printf("Using kernel %d\n", kernel_num);
+
+    // first check the correctness of the kernel
+    trimul_cpu(out, inp, B, T, C, NH);
+    trimul_gpu(kernel_num, d_out, d_inp, B, T, C, NH);
+    validate_result(d_out, out, "out", B * NH * T * T, 1e-4f);
+
+    printf("All results match. Starting benchmarks.\n\n");
+
+    // benchmark speed of the kernel
+    int repeat_times = 100;
+
+    float elapsed_time = benchmark_kernel(repeat_times, trimul_gpu,
+                                          kernel_num, d_out, d_inp,
+                                          B, T, C, NH);
+
+
+    float cublas_time = benchmark_kernel(repeat_times, trimul_gpu,
+                                         0, d_out, d_inp,
+                                         B, T, C, NH);
+
+    printf("time %.2f ms vs %.2f ms for CuBLAS\n", elapsed_time, cublas_time);
+
+    // free memory
+    free(out);
+    free(inp);
+    cudaCheck(cudaFree(d_out));
+    cudaCheck(cudaFree(d_inp));
+    cublasDestroy(cublas_handle);
+
+    return 0;
+}

train_gpt2.c

@@ -364,8 +364,10 @@ void gelu_forward(float* out, float* inp, int N) {
 }
 
 // we want to use -Ofast optimization, but sadly GeLU breaks, so disable this flag just for it (#168)
-#pragma float_control(precise, on, push) // On msvc /fp:fast is a lot faster, but the expf inside coshf breaks the model
-__attribute__((optimize("no-finite-math-only"))) // same for gcc -Ofast
+#pragma float_control(precise, on, push)
+#if defined(__GNUC__) && !defined(__clang__)
+__attribute__((optimize("no-finite-math-only")))
+#endif
 void gelu_backward(float* dinp, float* inp, float* dout, int N) {
     for (int i = 0; i < N; i++) {
         float x = inp[i];

train_gpt2.cu

@@ -664,51 +664,79 @@ __global__ void softmax_forward_kernel7(float* out, const float* inp, int N, int
     }
 }
 
-// cooperative groups solution, one warp per output channel
-template <typename Td>
-__global__ void matmul_backward_bias_kernel2(Td* dbias, const Td* dout, int B, int T, int OC) {
-    // dout is (B, T, OC), dbias is (OC)
-    // e.g. if block_size = 128, then we have 4 warps per block, each in charge of one output channel
-    namespace cg = cooperative_groups;
-    cg::thread_block block = cg::this_thread_block();
-    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
-    // meta_group_size is the number of warps in a block (e.g. 4), meta_group_rank is the warp index (0,1,2,3)
-    int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
-    if(idx >= OC) { return; }
-    int BT = B * T; // number of elements to reduce in total, per channel
-    // first, thread coarsening to sum reduce the problem size from B*T to 32
-    float sum = 0.0f;
-    for(int i = warp.thread_rank(); i < BT; i += warp.size()) {
-        sum += (float)dout[i * OC + idx];
+// this kernel performs a column-wise reduction over dout, in PyTorch equivalent to:
+// dbias = dout.sum((0,1))
+// the idea is to employ one block to reduce along several columns,
+// where each block has a width of 32 columns to ensure coalesced access.
+// at the end we accumulate the reductions performed by the warps in each block via shared memory
+__global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, int B, int T, int OC) {
+    // this kernel is launched with 1D grid_dim of OC/32
+    // for example let's say block_size is 128
+    extern __shared__ float smem[]; // of size block_size (128)
+    const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3
+    const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31
+    const int tl = blockIdx.x * warpSize; // pointer to the start column for this block
+    const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4
+
+    // pointer to the start of the column for one lane of threads
+    // so e.g. 4 threads (of the same lane_id) will reduce this one column
+    const float* dout_col = dout + tl + lane_id;
+
+    // column reductions by looping through the rows
+    // each of the 4 threads offsets by its warp_id and then skips by vstep
+    // together these 4 threads cover all B*T rows of this (lane_id) column
+    // importantly, consecutive threads (in threadId) are processing adjacent columns,
+    // leading to a coalesced memory access pattern
+    float dout_sum = 0.0f;
+    for (int row = warp_id; row < B * T; row += vstep) {
+        dout_sum += dout_col[row * OC];
     }
-    // now do a warp-level reduce to get the sum across the 32 threads in this warp
-    sum = cg::reduce(warp, sum, cg::plus<float>{});
-    // write the result to output (global memory)
-    if(warp.thread_rank() == 0) {
-        dbias[idx] = (Td)((float)dbias[idx] + sum);
+    smem[lane_id + warp_id * warpSize] = dout_sum;
+    __syncthreads();
+
+    // warp_id 0 reduces the shared memory column-wise, linearly
+    dout_sum = 0.0f;
+    if (warp_id == 0) {
+        for (int j = 0; j < vstep; j++) {
+            dout_sum += smem[lane_id + j * warpSize];
+        }
+        dbias[tl + lane_id] += dout_sum;
     }
 }
 
-template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
-__global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
-                        const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
-                        int B, int T, int C) {
+// uses shared memory instead for the reduces
+__global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* dbias,
+                                           const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd,
+                                           int B, int T, int C) {
+    extern __shared__ float shared[]; // size = 2 * C
+
+    namespace cg = cooperative_groups;
     cg::thread_block block = cg::this_thread_block();
     cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
     int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
     int N = B * T;
-    if(idx >= N) {
-        return;
-    }
+    if(idx >= N) { return; } // thread guards
 
     int b = idx / T;
     int t = idx % T;
 
-    const Tdout* dout_bt = dout + b * T * C + t * C;
-    const Trest* inp_bt = inp + b * T * C + t * C;
-    Tdinp* dinp_bt = dinp + b * T * C + t * C;
-    float mean_bt = (float)mean[b * T + t];
-    float rstd_bt = (float)rstd[b * T + t];
+    const float* dout_bt = dout + b * T * C + t * C;
+    const float* inp_bt = inp + b * T * C + t * C;
+    float* dinp_bt = dinp + b * T * C + t * C;
+    const float mean_bt = mean[b * T + t];
+    const float rstd_bt = rstd[b * T + t];
+
+    // the first half of shared memory is bias, second is weight
+    float* dbias_shared = shared;
+    float* dweight_shared = shared + C;
+
+    // init shared memory to zero
+    #pragma unroll
+	for(int i = threadIdx.x; i < C; i+= blockDim.x){
+       dbias_shared[i] = 0.0f;
+       dweight_shared[i] = 0.0f;
+    }
+    __syncthreads();
 
     // first: two reduce operations
     float dnorm_mean = 0.0f;
@@ -719,10 +747,8 @@ __global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams
         dnorm_mean += dnorm_i;
         dnorm_norm_mean += dnorm_i * norm_bti;
     }
-
     dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus<float>{});
     dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus<float>{});
-
     dnorm_mean = dnorm_mean / C;
     dnorm_norm_mean = dnorm_norm_mean / C;
 
@@ -731,9 +757,9 @@ __global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams
         float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt;
         float dnorm_i = (float)weight[i] * (float)dout_bt[i];
         // gradient contribution to bias
-        atomicAddX(&dbias[i], (Tparams)dout_bt[i]);
+        atomicAdd(&dbias_shared[i], dout_bt[i]);
         // gradient contribution to weight
-        atomicAddX(&dweight[i], (Tparams)(norm_bti * (float)dout_bt[i]));
+        atomicAdd(&dweight_shared[i], norm_bti * dout_bt[i]);
         // gradient contribution to input
         float dval = 0.0f;
         dval += dnorm_i; // term 1
@@ -742,6 +768,13 @@ __global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams
         dval *= rstd_bt; // final scale
         dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval);
     }
+    __syncthreads();
+
+    // write to global memory
+	for(int i = threadIdx.x; i < C; i+= blockDim.x){
+        atomicAdd(&dbias[i], dbias_shared[i]);
+        atomicAdd(&dweight[i], dweight_shared[i]);
+	}
 }
 
 
@@ -1164,9 +1197,9 @@ void matmul_backward_fp32(float* dinp, float* dweight, float* dbias,
     cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, C, dout, OC, &one, dweight, C));
     // backward to bias, if given, does a +=
     if (dbias != NULL) {
-        const int block_size = 512;
-        const int grid_size = CEIL_DIV(OC * 32, block_size);
-        matmul_backward_bias_kernel2<<<grid_size, block_size>>>(dbias, dout, B, T, OC);
+        const int block_size = 1024;
+        const int grid_size = OC / 32; // for now, OC must be divisible by 32 for this kernel to work
+        matmul_backward_bias_kernel4<<<grid_size, block_size, block_size * sizeof(float)>>>(dbias, dout, B, T, OC);
         cudaCheck(cudaGetLastError());
     }
 }
@@ -1193,11 +1226,11 @@ template <typename Tdinp, typename Tparams, typename Tdout, typename Trest>
 void layernorm_backward(Tdinp* dinp, Tparams* dweight, Tparams* dbias,
                         const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd,
                         int B, int T, int C) {
-    const int block_size = 256;
+    const int block_size = 512;
     const int N = B * T;
-    // one warp per token, so we need to divide by 32 here.
-    const int grid_size = CEIL_DIV(N, block_size / 32);
-    layernorm_backward_kernel<<<grid_size, block_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
+    const int grid_size = CEIL_DIV(32*N, block_size);
+    size_t shared_mem_size = 2 * C * sizeof(float);
+    layernorm_backward_kernel2<<<grid_size, block_size, shared_mem_size>>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C);
     cudaCheck(cudaGetLastError());
 }