comp files

kernels/attention.cuh

@@ -2,10 +2,159 @@
 #define __ATTENTION_CUH__
 
 #include <cuda_runtime.h>
+#include <math.h>
+#include <float.h>
 #include "../utils/cuda_utils.cuh"
-#include "softmax.cuh"
 
-__global__ void permute_kernel(float* q, float* k, float* v, const float* inp, int B, int N, int NH, int d) {
+
+//added templates so the compiler can unroll loops and optimize better. 
+// op11
+// without swilling the QK^T inner loops read smem_q [query_row][i] where each thread query_row in teh warp holds a different row but reads teh same column i. as HS is 64 every thread bank index reduces to (i % 32) giving 32 way back conflict per warp on the Q reads. col -> col ^ row to swizzle the accesses and reduce bank conflicts. so each 32 read to 32 different banks instead of 32 reads to the same bank.
+template <int BR, int BC, int HS>
+__global__ void flash_attention_forward_kernel(const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ out, int B, int T, int NH, float scale) {
+    //following line by like psudo code from flash attention paper.
+    // https://tridao.me/publications/flash2/flash2.pdf
+
+    // step 1 map thread blocks to (batch , head , query_tile)
+    int b = blockIdx.x;
+    int nh = blockIdx.y;
+    int query_tile = blockIdx.z; 
+
+    int query_row = threadIdx.x; //load index
+    int query_global = query_tile * BR + query_row; 
+
+    //step 2 calculate all the base pointers for this block. We will be loading queries from Q and keys/values from K and V, and writing output to out.
+    const float* __restrict__ query_base = Q + b * NH * T * HS + nh * T * HS;
+    const float* __restrict__ key_base = K + b * NH * T * HS + nh * T * HS;
+    const float* __restrict__ value_base = V + b * NH * T * HS + nh * T * HS;
+    float* out_base = out + b * NH * T * HS + nh * T * HS;
+
+    // declare shared memory for the tile of queries and keys/values. We will load a tile of queries and then iterate over tiles of keys/values.
+    __shared__ float smem_q[BR * HS];
+    __shared__ float smem_k[BC * HS];
+    __shared__ float smem_v[BC * HS];
+
+    // we will compute the output for a tile of queries at a time.
+    int query_tile_start = query_tile * BR; // global index of the first query in this tile. This is used to load the correct queries into shared memory and also to know which output rows we are responsible for computing.
+    bool q_valid = (query_global < T); // check if this thread's query is within the sequence length. If not, it will do nothing.
+
+
+    //define swizzling function using the XOR on col and row index to reduce bank conflicts on smem_q reads in the inner loop.
+
+    #define SWIZZLE(col, row) ((row) * HS + ((col) ^ ((row) & (HS - 1))))
+
+    // step 4 load query tiles from global memory to shared memory.
+    // coalesced load of the queries for this tile into shared memory. a linear load of BR*HS elements, with each thread loading HS elements for one query.
+    #pragma unroll
+    for (int i = 0; i < HS; i++) {
+        smem_q[SWIZZLE(i, query_row)] = (query_tile_start + query_row < T) ? query_base[(query_tile_start + query_row) * HS + i] : 0.0f; // if the query is out of bounds, we can load zeros since they won't contribute to the output.
+    }
+
+    // step 5 init accumulators for the output of this tile of queries.
+    float out_reg[HS];
+    #pragma unroll  
+    for (int i = 0; i < HS; i++) {
+        out_reg[i] = 0.0f;
+    }
+    float max_score = -FLT_MAX;
+    float sum_exp = 0.0f;
+
+
+    // step 6 determine KV tile range for casual masking
+    int query_tile_end = (query_tile + 1) * BR;
+    int n_kv_tiles = (query_tile_end < T) ? (T + BC - 1) / BC : (query_tile_end + BC - 1) / BC; // if the end of the query tile is within the sequence length, we only need to iterate over the KV tiles that cover the sequence length. If the end of the query tile exceeds the sequence length, we need to iterate over all KV tiles that cover the end of the query tile to properly apply casual masking.
+
+    __syncthreads();
+
+    //step 7 main loop over kv values 
+    for(int j = 0; j < n_kv_tiles; j++){
+        __syncthreads();
+
+        int k_tile_start = j * BC;
+
+        // step 8 load k and v tiles from global memory
+        #pragma unroll
+        for (int i = 0; i < HS; i++) {
+            int l = i * BC + query_row;
+            int row = l / HS;
+            int col = l % HS;
+            int global_row = k_tile_start + row;
+            bool valid = (global_row < T) && (row < BC);
+
+            if (valid) {
+                smem_k[SWIZZLE(col, row)] = key_base[global_row * HS + col];
+                smem_v[SWIZZLE(col, row)] = value_base[global_row * HS + col];
+            } else {
+                smem_k[SWIZZLE(col, row)] = 0.0f; // pad with zeros if out of bounds
+                smem_v[SWIZZLE(col, row)] = 0.0f; // pad with zeros if out of bounds
+            }
+        }
+        __syncthreads();
+
+        // step 9 compute one row of S = Qi * Kj^T and apply casual masking
+        float score_reg[BC];
+        float row_max = -FLT_MAX;
+
+        #pragma unroll
+        for (int c = 0; c < BC; c++) {
+            int k_global = j * BC + c;
+            float dot; // compute the dot product of the query and key for this position
+
+            if (!q_valid || k_global > query_global) { // apply casual masking. If the key position is greater than the query position, we mask it out by setting the score to a very large negative number.
+                dot = -FLT_MAX;
+            } else {
+                float acc = 0.0f;
+                #pragma unroll
+                for (int h = 0; h < HS; h++) {
+                    acc += smem_q[SWIZZLE(h, query_row)] * smem_k[SWIZZLE(h, c)];
+                }
+                dot = acc * scale;
+            }
+            score_reg[c] = dot;
+            row_max = fmaxf(row_max, dot);  // local max for this row
+        }
+
+        // step 10 online softmax rescaling 
+        float max_new = fmaxf(max_score, row_max);
+        if (max_new == -FLT_MAX) { // if both are -FLT_MAX, we can skip the rescaling because all scores are -FLT_MAX and the output will be zero anyway.
+            continue; 
+        }
+        float alpha = (max_score == -FLT_MAX) ? 0.0f : expf(max_score - max_new); // if max_score is -FLT_MAX, we set alpha to 0 to avoid exp overflow. Otherwise, we compute the scaling factor for the old scores.
+        float local_tile = 0.0f; // this will hold the sum of exp for the scores in this tile after rescaling.
+
+        #pragma unroll
+        for (int c = 0; c < BC; c++) {
+            float p = (score_reg[c] == -FLT_MAX) ? 0.0f : expf(score_reg[c] - max_new); // compute the exp for this score after rescaling. If the score is -FLT_MAX, we set it to 0 to avoid exp underflow.
+            local_tile += p;
+            score_reg[c] = p; // store the rescaled exp in score_reg for later use in the weighted sum.
+        }
+
+        // update output acc 
+        #pragma unroll
+        for (int h = 0; h < HS; ++h) {
+            float pv = 0.0f;
+            #pragma unroll
+            for (int c = 0; c < BC; ++c) {
+                pv += score_reg[c] * smem_v[SWIZZLE(h, c)];
+            }
+            out_reg[h] = alpha * out_reg[h] + pv;
+        }
+        sum_exp = alpha * sum_exp + local_tile;  // update the local tile sum with the contribution from this tile and the previous tiles, scaled by alpha.
+        max_score = max_new; // update the max score for the next iteration
+
+    }
+    if (q_valid) {
+        float inv_l = (sum_exp > 0.0f) ? (1.0f / sum_exp) : 0.0f;
+        #pragma unroll
+        for (int h = 0; h < HS; ++h) {
+            out_base[query_global * HS + h] = out_reg[h] * inv_l;
+        }
+    }
+}
+
+
+
+__global__ void permute_kernel(float* __restrict__ q, float* __restrict__ k, float*  __restrict__ v, const float* __restrict__ inp, int B, int N, int NH, int d) {
     // Implement this
 
     int x = blockDim.x * blockIdx.x + threadIdx.x;
@@ -24,7 +173,7 @@ __global__ void permute_kernel(float* q, float* k, float* v, const float* inp, i
 }
 
 
-__global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, int d) {
+__global__ void unpermute_kernel(const float* __restrict__ inp, float* __restrict__ out, int B, int N, int NH, int d) {
     // Implement this
     int x = blockDim.x * blockIdx.x + threadIdx.x;
 
@@ -39,39 +188,42 @@ __global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, i
 
 }
 
-__global__ void preattention_scores_kernel(float* q, float* k, int B, int T, int NH, int HS, float* preatt){
-    int x = blockDim.x * blockIdx.x + threadIdx.x;
+// __global__ void preattention_scores_kernel(const float* __restrict__ q, const float* __restrict__ k, int B, int T, int NH, int HS, float* __restrict__ preatt){
+//     int x = blockDim.x * blockIdx.x + threadIdx.x;
 
-    if(x < B* T* NH* T){
-        int b = (x / T / T / NH);
-        int nh = (x / T / T) % NH;
-        int t1 = (x / T) % T;
-        int t2 = x % T;
+//     if(x < B* T* NH* T){
+//         int b = (x / T / T / NH);
+//         int nh = (x / T / T) % NH;
+//         int t1 = (x / T) % T;
+//         int t2 = x % T;
 
-        float sum = 0.0f;
-        for (int i = 0; i < HS; i++) {
-            sum += q[b*NH*T*HS + nh*T*HS + t1*HS + i] * k[b*NH*T*HS + nh*T*HS + t2*HS + i];
-        }
-        preatt[b * NH * T * T + nh * T * T + t1 * T + t2] = sum;
-    }   
-}
+//         float sum = 0.0f;
+//         const float* __restrict__ q_ptr = q + b * NH * T * HS + nh * T * HS + t1 * HS;
+//         const float* __restrict__ k_ptr = k + b * NH * T * HS + nh * T * HS + t2 * HS;
+//         for (int i = 0; i < HS; i++) {
+//             sum += q_ptr[i] * k_ptr[i];
+//         }
+//         preatt[b * NH * T * T + nh * T * T + t1 * T + t2] = sum;
+//     }   
+// }
 
-__global__ void attention_v_kernel(float* vaccum, float* att, float* v, int B, int T, int NH, int HS){
-    int x = blockDim.x * blockIdx.x + threadIdx.x;
+// __global__ void attention_v_kernel(float* __restrict__ vaccum, const float* __restrict__ att, const float* __restrict__ v, int B, int T, int NH, int HS){
+//     int x = blockDim.x * blockIdx.x + threadIdx.x;
 
-    if(x < B*NH*T*HS){
-        int b   = (x / HS / T / NH);
-        int nh  = (x / HS / T) % NH;
-        int t   = (x / HS) % T;
-        int hs  = x % HS;
-        float sum = 0.0f;
-        for (int t2 = 0; t2 < T; t2++) {
-            sum += att[b*NH*T*T + nh*T*T + t*T + t2] * v[b*NH*T*HS + nh*T*HS + t2*HS + hs];
-        }
+//     if(x < B*NH*T*HS){
+//         int b   = (x / HS / T / NH);
+//         int nh  = (x / HS / T) % NH;
+//         int t   = (x / HS) % T;
+//         int hs  = x % HS;
+//         float sum = 0.0f;
+//         const float* __restrict__ att_ptr = att + b * NH * T * T + nh * T * T + t * T;
+//         for (int t2 = 0; t2 < T; t2++) {
+//             sum += att_ptr[t2] * v[b*NH*T*HS + nh*T*HS + t2*HS + hs];
+//         }
 
-        vaccum[b * NH * T * HS + nh * T * HS + t * HS + hs] = sum;
-    }
-}
+//         vaccum[b * NH * T * HS + nh * T * HS + t * HS + hs] = sum;
+//     }
+// }
 
 // Launch all kernels related to attention here 
 void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, int T, int C, int NH) {
@@ -83,21 +235,23 @@ void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, i
     k = qkvr + 1 * B * T * C;
     v = qkvr + 2 * B * T * C;
 
-    dim3 block(32,1,1);
+    dim3 block(256,1,1);
     dim3 grid((B * NH * T * HS + block.x-1) / block.x,1,1);
 
     permute_kernel<<<grid, block>>>(q, k, v, inp, B, T, NH, HS);
 
-    float* preatt = inp;
-    dim3 grid_preatt((B * NH * T * T + block.x-1) / block.x,1,1);
-    preattention_scores_kernel<<<grid_preatt, block>>>(q, k, B, T, NH, HS, preatt);
+    //fa2 — op_7 sweep winner: BR=32, BC=32 was 26% faster than BR=64,BC=64
+    // at T=64 (smaller tiles fit two blocks per SM, doubling occupancy).
+    constexpr int BR = 32; // tile size for queries
+    constexpr int BC = 32; // tile size for keys/values
+    float scale = 1.0f / sqrtf((float)HS);
 
-    float scale = 1.0 / sqrtf(HS);
-    dim3 grid_softmax (( B * NH * T + block.x -1 )/ block.x,1,1);
-    softmax_forward_kernel<<<grid_softmax, block>>>(att, scale, preatt, B * NH, T);
+    int number_of_q_tiles = (T + BR - 1) / BR;
+    dim3 fa_grid(B, NH, number_of_q_tiles);      // one block per (batch, head, Q tile)
+    dim3 fa_block(BR, 1, 1); // one warp per query in the tile
 
     float* vaccum = inp;
-    attention_v_kernel<<<grid_preatt,block>>>(vaccum, att, v, B, T, NH, HS);
+    flash_attention_forward_kernel<BR, BC, 64><<<fa_grid, fa_block>>>(q, k, v, vaccum, B, T, NH, scale);
 
     unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
 }

kernels/attention_kv_cache_template.cuh

@@ -1,86 +0,0 @@
-#ifndef __ATTENTION_KV_CACHE_TEMPLATE_CUH__
-#define __ATTENTION_KV_CACHE_TEMPLATE_CUH__
-
-#include <cuda_runtime.h>
-#include "../utils/cuda_utils.cuh"
-
-extern bool g_enable_kv_cache; // Toggles KV-caching logic
-extern float* g_kv_cache;      // Pointer to the GPU-allocated KV cache buffer
-extern int g_layer_idx;        // Tracks the current transformer layer, useful for knowing when to switch to decode
-extern bool g_is_prefill;      // True during the initial prompt, False during token generation
-extern int g_current_pos;      // Tracks the absolute sequence position during Decode, helps index into the KV cache
-
-// --- Suggested Kernel Placeholders (TODO: Add your regular attention kernel here, and any other kernels you may need) ---
-
-/**
- * Kernel to store the newly computed Key and Value into the global KV cache. 
- * CAN BE FUSED WITH QKV MATMUL
- * In Prefill: Stores T tokens for the current layer.
- * In Decode: Stores 1 token at g_current_pos for the current layer.
- */
-__global__ void store_kv_kernel(/* args */) {
-    // Implement logic to copy K and V from the input QKV buffer into g_kv_cache
-}
-
-/**
- * Kernel to compute the attention for the Decode phase (T=1).
- * This involves:
- * 1. Extracting the current Query (Q)
- * 2. Gathering all previous Keys (K) and Values (V) from the cache
- * 3. Computing attention using the new Q and cached KV values.
- */
-__global__ void decode_attention_kernel(/* args */) {
-    // Implement T=1 attention logic using the KV cache
-}
-
-void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, int T, int C, int NH) {
-    // 1. Check if KV-caching is enabled, used by the verification script
-    if (!g_enable_kv_cache) {
-        // TODO: Call your standard Baseline Attention here (Milestone 1/2)
-        return;
-    }
-
-    // 2. Initialize KV Cache if necessary
-    if (g_kv_cache == nullptr) {
-        // TODO: Allocate the KV cache, based accordingly to the model parameters
-        // Option 1: Allocate a very big kv cache to store all the tokens
-        // Option 2: Dynamically resize the kv cache repeatedly for each new token
-    }
-
-    if (g_is_prefill) {
-        // --- PREFILL PHASE ---
-        // During prefill, we process the entire input sequence (T > 1).
-        
-        // A. Store the K and V results from this layer into the cache.
-        //    The projections are available in 'inp' (the QKV matmul output).
-        
-        // B. Run standard baseline attention for the initial prompt tokens.
-
-        // C. Update the global layer index to track progress through the model.
-        g_layer_idx++;
-        if (g_layer_idx == 12) {
-            g_layer_idx = 0;
-            g_is_prefill = false; // Transition to Decode phase for the next token
-        }
-    } else {
-        // --- DECODE PHASE ---
-        // During decode, we only process the single newest token (T = 1).
-        
-        // A. Store the single new Key and Value for this step into the cache at g_current_pos.
-        
-        // B. Run specialized Decode Attention.
-        //    This requires reading the history of Keys and Values from the cache
-        //    to compute the output for the current token.
-        //    TODO: Don't forget to scale your attention scores
-        //    TODO: Ensure the output shapes match the next layer
-        
-        // C. Update the global state.
-        g_layer_idx++;
-        if (g_layer_idx == 12) {
-            g_layer_idx = 0;
-            g_current_pos++; // Advance the sequence position
-        }
-    }
-}
-
-#endif // __ATTENTION_KV_CACHE_TEMPLATE_CUH__

kernels/encoder.cuh

@@ -4,27 +4,31 @@
 #include <cuda_runtime.h>
 #include "../utils/cuda_utils.cuh"
 
-__global__ void encoder_forward_kernel(float* out, const int* inp, const float* wte, const float* wpe,
+__global__ void encoder_forward_kernel(float* __restrict__ out, const int* __restrict__ inp, const float* __restrict__ wte, const float* __restrict__ wpe,
                                        int B, int T, int C) {
     // Implement this
-    int x = blockDim.x * blockIdx.x + threadIdx.x;
-    int y = blockDim.y * blockIdx.y + threadIdx.y;
-
-    if(x < B && y < T){
-        int token = inp[x*T + y];
-        for(int i = 0; i < C; i ++){
-            out[x*T*C + y*C + i] = wte[token*C+i] + wpe[y*C+i];
-        }
+    int row = blockIdx.x * blockDim.y + threadIdx.y;
+    if (row >= B * T){
+        return;
+    }
+    int b = row / T;
+    int t = row % T;
+    int token = inp[b * T + t];
+    const float* __restrict__ wte_ptr = wte + token * C;
+    const float* __restrict__ wpe_ptr = wpe + t * C;
+    float* __restrict__ out_ptr = out + row * C;
+    for (int i = threadIdx.x; i < C; i += 32) {
+        out_ptr[i] = wte_ptr[i] + wpe_ptr[i];
     }
 }
 
 // Launch kernel here
 void encoder_forward(float* out, const int* inp, const float* wte, const float* wpe, int B, int T, int C) {
     // Implement this
-    dim3 block(16,16,1);
-    dim3 grid((B + block.x -1)/ block.x,(T+block.y-1)/block.y);
-
-    encoder_forward_kernel<<<grid, block>>>(out, inp, wte, wpe,B,T,C);
+    dim3 block(32, 8, 1);
+    int total_rows = B * T;
+    dim3 grid((total_rows + 8 - 1) / 8, 1, 1);
+    encoder_forward_kernel<<<grid, block>>>(out, inp, wte, wpe, B, T, C);
 }
 
 

kernels/gelu.cuh

@@ -6,7 +6,7 @@
 
 #define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
 
-__global__ void gelu_forward_kernel(float* out, const float* inp, int N) {
+__global__ void gelu_forward_kernel(float* __restrict__ out, const float* __restrict__ inp, int N) {
     // Implement this
 
     int x = blockDim.x * blockIdx.x + threadIdx.x;
@@ -22,7 +22,7 @@ __global__ void gelu_forward_kernel(float* out, const float* inp, int N) {
 // Launch kernel here
 void gelu_forward(float* out, const float* inp, int N) {
     // Implement this
-    dim3 block(32,1,1);
+    dim3 block(256,1,1);
     dim3 grid((N + block.x -1) / block.x);
 
     gelu_forward_kernel<<<grid, block>>>(out, inp, N);

kernels/gpt2.cuh

@@ -0,0 +1,420 @@
+//-----------------------------------------------------------------------------------------------
+// Copyright (c) 2024 Andrej Karpathy
+// Licensed under the MIT License. See the LICENSE file for details.
+//
+// Modifications Copyright (c) 2025 Hanwen Liu, Hrishi Shah, Kelin Zeng, Charles Pei, and Vijay Daita, ALL RIGHTS RESERVED.
+//-----------------------------------------------------------------------------------------------
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <math.h>
+#include <time.h>
+#include <assert.h>
+#include <float.h>
+#include <string.h>
+#include <unistd.h>
+
+// GPU / CUDA related
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+#include "utils/utils.h"
+#include "utils/tokenizer.h"
+#include "utils/cuda_utils.cuh"
+
+// kernels 
+#include "kernels/attention.cuh"
+#include "kernels/encoder.cuh"
+#include "kernels/gelu.cuh"
+#include "kernels/layernorm.cuh"
+#include "kernels/matmul.cuh"
+#include "kernels/residual.cuh"
+
+// ----------------------------------------------------------------------------
+// GPT-2 model definition
+
+typedef struct {
+    int max_seq_len;
+    int vocab_size;
+    int padded_vocab_size;  // padded to multiple of 128
+    int num_layers;
+    int num_heads;
+    int channels;
+} GPT2Config;
+
+// the parameters of the model
+#define NUM_PARAMETER_TENSORS 16
+typedef struct {
+    float* wte; // (Vp, C)
+    float* wpe; // (maxT, C)
+    float* ln1w; // (L, C)
+    float* ln1b; // (L, C)
+    float* qkvw; // (L, 3*C, C)
+    float* qkvb; // (L, 3*C)
+    float* attprojw; // (L, C, C)
+    float* attprojb; // (L, C)
+    float* ln2w; // (L, C)
+    float* ln2b; // (L, C)
+    float* fcw; // (L, 4*C, C)
+    float* fcb; // (L, 4*C)
+    float* fcprojw; // (L, C, 4*C)
+    float* fcprojb; // (L, C)
+    float* lnfw; // (C)
+    float* lnfb; // (C)
+} ParameterTensors;
+
+void fill_in_parameter_sizes(size_t* param_sizes, GPT2Config config) {
+    int Vp = config.padded_vocab_size;
+    int C = config.channels;
+    int maxT = config.max_seq_len;
+    int L = config.num_layers;
+    param_sizes[0] = Vp * C; // wte
+    param_sizes[1] = maxT * C; // wpe
+    param_sizes[2] = L * C; // ln1w
+    param_sizes[3] = L * C; // ln1b
+    param_sizes[4] = L * (3 * C) * C; // qkvw
+    param_sizes[5] = L * (3 * C); // qkvb
+    param_sizes[6] = L * C * C; // attprojw
+    param_sizes[7] = L * C; // attprojb
+    param_sizes[8] = L * C; // ln2w
+    param_sizes[9] = L * C; // ln2b
+    param_sizes[10] = L * (4 * C) * C; // fcw
+    param_sizes[11] = L * (4 * C); // fcb
+    param_sizes[12] = L * C * (4 * C); // fcprojw
+    param_sizes[13] = L * C; // fcprojb
+    param_sizes[14] = C; // lnfw
+    param_sizes[15] = C; // lnfb
+}
+
+// allocate memory for the parameters and point the individual tensors to the right places
+float* malloc_and_point_parameters(ParameterTensors* params, size_t* param_sizes, int on_device) {
+    // on_device: 0 = CPU, 1 = GPU
+    // calculate the number of parameters
+    size_t num_parameters = 0;
+    for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) {
+        num_parameters += param_sizes[i];
+    }
+    // malloc all parameters all at once on the device
+    float* params_memory;
+    if (on_device) {
+        cudaCheck(cudaMalloc((void**)&params_memory, num_parameters * sizeof(float)));
+    } else {
+        params_memory = (float*)mallocCheck(num_parameters * sizeof(float));
+    }
+    // assign all the tensors their place in the array
+    float** ptrs[] = {
+        &params->wte, &params->wpe, &params->ln1w, &params->ln1b, &params->qkvw, &params->qkvb,
+        &params->attprojw, &params->attprojb, &params->ln2w, &params->ln2b, &params->fcw, &params->fcb,
+        &params->fcprojw, &params->fcprojb, &params->lnfw, &params->lnfb
+    };
+    float* params_memory_iterator = params_memory;
+    for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) {
+        *(ptrs[i]) = params_memory_iterator;
+        params_memory_iterator += param_sizes[i];
+    }
+    return params_memory;
+}
+
+#define NUM_ACTIVATION_TENSORS 21
+typedef struct {
+    float* encoded; // (B, T, C)
+    float* ln1; // (L, B, T, C)
+    float* ln1_mean; // (L, B, T)
+    float* ln1_rstd; // (L, B, T)
+    float* atty; // (L, B, T, C)
+    float* att; // (L, B, NH, T, T)
+    float* attproj; // (L, B, T, C)
+    float* residual2; // (L, B, T, C)
+    float* ln2; // (L, B, T, C)
+    float* ln2_mean; // (L, B, T)
+    float* ln2_rstd; // (L, B, T)
+    float* fch; // (L, B, T, 4*C)
+    float* fch_gelu; // (L, B, T, 4*C)
+    float* fcproj; // (L, B, T, C)
+    float* residual3; // (L, B, T, C)
+    float* lnf; // (B, T, C)
+    float* lnf_mean; // (B, T)
+    float* lnf_rstd; // (B, T)
+    float* losses; // (B, T)
+    float* qkvr; // (L, B, T, 3*C)
+    // This buffer will store the logits
+    // during the processing of transformer blocks, we will also use this as a
+    // general scratchpad buffer. Allocation is made large enough to hold (B, T, 3C),
+    // (B, NH, T, T), and (B, T, V) shaped tensors.
+    float* output;
+} ActivationTensors;
+
+void fill_in_activation_sizes(size_t* act_sizes, int B, int T, GPT2Config config) {
+    size_t Vp = config.padded_vocab_size;
+    size_t L = config.num_layers;
+    size_t NH = config.num_heads;
+    size_t C = config.channels;
+    act_sizes[0] = B * T * C; // encoded
+    act_sizes[1] = L * B * T * C; // ln1
+    //to align making 1 to 4
+    act_sizes[2] = 1; // ln1_mean, reserved for future use
+    act_sizes[3] = 1; // ln1_rstd, reserved for future use
+    act_sizes[4] = L * B * T * C; // atty
+    act_sizes[5] = L * B * NH * T * T; // att
+    act_sizes[6] = L * B * T * C; // attproj
+    act_sizes[7] = L * B * T * C; // residual2
+    act_sizes[8] = L * B * T * C; // ln2
+    act_sizes[9] = 1; // ln2_mean, reserved for future use
+    act_sizes[10] = 1; // ln2_rstd, reserved for future use
+    act_sizes[11] = 1; // fch
+    act_sizes[12] = L * B * T * 4*C; // fch_gelu
+    act_sizes[13] = L * B * T * C; // fcproj
+    act_sizes[14] = L * B * T * C; // residual3
+    act_sizes[15] = B * T * C; // lnf
+    act_sizes[16] = 1; // lnf_mean, reserved for future use
+    act_sizes[17] = 1; // lnf_rstd, reserved for future use
+    act_sizes[18] = B * T; // losses
+    act_sizes[19] = L * B * T * 3*C; // qkvr
+    act_sizes[20] = B * T * max(3*C, max(NH*T, Vp)); // output / scratch
+}
+
+// Reserved for future use
+#define NUM_BACKWARD_TENSORS 3
+typedef struct {
+    float* bt4c; // (B, T, 4*C)
+    float* preatt; // (B, NH, T, T)
+    float* residual3; // (B, T, C)
+} GradActTensors;
+
+// op_13: 256-byte aligned activations.
+// Each activation tensor's stride is rounded up to a multiple of 64 floats (256 bytes).
+#define align_to_256_bytes(x) \
+    (x + (128/sizeof(float) - (x % (128/sizeof(float)))) % (128/sizeof(float)))
+
+float* malloc_and_point(float** targets[], const size_t* act_sizes, int n) {
+    size_t num_activations = 0;
+    for (size_t i = 0; i < n; i++) {
+        num_activations += align_to_256_bytes(act_sizes[i]);
+    }
+    float* acts_memory;
+    cudaCheck(cudaMalloc((void**)&acts_memory, num_activations * sizeof(float)));
+    cudaCheck(cudaMemset(acts_memory, 0, num_activations * sizeof(float)));
+    float* acts_memory_iterator = acts_memory;
+    for (size_t i = 0; i < n; i++) {
+        *(targets[i]) = acts_memory_iterator;
+        acts_memory_iterator += align_to_256_bytes(act_sizes[i]);
+    }
+    return acts_memory;
+}
+
+float* malloc_and_point_activations(ActivationTensors* acts, const size_t* act_sizes) {
+    float** ptrs[] = {
+        &acts->encoded, &acts->ln1, &acts->ln1_mean, &acts->ln1_rstd, &acts->atty,
+        &acts->att, &acts->attproj, &acts->residual2, &acts->ln2, &acts->ln2_mean,
+        &acts->ln2_rstd, &acts->fch, &acts->fch_gelu, &acts->fcproj, &acts->residual3, &acts->lnf,
+        &acts->lnf_mean, &acts->lnf_rstd, &acts->losses, &acts->qkvr, &acts->output
+    };
+    return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS);
+}
+
+typedef struct {
+    GPT2Config config;
+    // the weights of the model, and their sizes
+    ParameterTensors params;
+    size_t param_sizes[NUM_PARAMETER_TENSORS];
+    float* params_memory;
+    size_t num_parameters;
+    ParameterTensors grads;
+    float* grads_memory;
+    float* m_memory;
+    float* v_memory;
+    // the activations of the model, and their sizes
+    ActivationTensors acts;
+    size_t act_sizes[NUM_ACTIVATION_TENSORS];
+    float* acts_memory;
+    size_t num_activations;
+    GradActTensors grads_acts;
+    size_t num_grad_acts;
+    float* grads_acts_memory;
+    // other run state configuration
+    int batch_size; // the batch size (B) of current forward pass
+    int seq_len; // the sequence length (T) of current forward pass
+    int* inputs; // the input tokens for the current forward pass
+    int* targets; // the target tokens for the current forward pass
+    float mean_loss;
+    float* cpu_losses;
+} GPT2;
+
+void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
+
+    // read in model from a checkpoint file
+    FILE *model_file = fopenCheck(checkpoint_path, "rb");
+    int model_header[256];
+    freadCheck(model_header, sizeof(int), 256, model_file);
+    if (model_header[0] != 20240326) { fprintf(stderr, "GPT2 - Bad magic model file\n"); exit(EXIT_FAILURE); }
+    if (model_header[1] != 3) {
+        // was bumped from 1 -> 3 to incorporate the padded vocab size
+        fprintf(stderr, "GPT2 - Bad version in model file\n");
+        exit(EXIT_FAILURE);
+    }
+
+    // read in hyperparameters
+    model->config.max_seq_len = model_header[2];
+    model->config.vocab_size = model_header[3];
+    model->config.num_layers = model_header[4];
+    model->config.num_heads = model_header[5];
+    model->config.channels = model_header[6];
+    model->config.padded_vocab_size = model_header[7];
+
+    // allocate space for all the parameters and read them in
+    fill_in_parameter_sizes(model->param_sizes, model->config);
+
+    // count the number of parameters
+    size_t num_parameters = 0;
+    for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) {
+        num_parameters += model->param_sizes[i];
+    }
+    model->num_parameters = num_parameters;
+
+    // create memory for model parameters on the device
+    model->params_memory = malloc_and_point_parameters(&model->params, model->param_sizes, 1);
+
+    // read in all the parameters from file and copy them to device
+    float* params_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float));
+    freadCheck(params_memory_cpu, sizeof(float), num_parameters, model_file);
+    cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, num_parameters * sizeof(float), cudaMemcpyHostToDevice));
+    free(params_memory_cpu);
+    fcloseCheck(model_file);
+
+    // other inits
+    model->acts_memory = NULL;
+    model->grads_memory = NULL;
+    model->m_memory = NULL;
+    model->v_memory = NULL;
+    model->grads_acts_memory = NULL;
+    model->inputs = NULL;
+    model->targets = NULL;
+    model->cpu_losses = NULL;
+    model->batch_size = 0;
+    model->seq_len = 0;
+    model->mean_loss = -1.0f; // -1.0f will designate no loss
+}
+
+void gpt2_forward(GPT2 *model, int* inputs, int B, int T) {
+
+    // ensure the model was initialized or error out
+    if (model->params_memory == NULL) {
+        printf("GPT2 - Error: model was not initialized properly.\n");
+        exit(EXIT_FAILURE);
+    }
+
+    // convenience parameters
+    int V = model->config.vocab_size;
+    int Vp = model->config.padded_vocab_size;
+    int L = model->config.num_layers;
+    int NH = model->config.num_heads;
+    int C = model->config.channels;
+
+    // validate inputs, all indices must be in the range [0, V)
+    for(int i = 0; i < B * T; i++) {
+        assert(0 <= inputs[i] && inputs[i] < V);
+    }
+
+    // allocate space for all the activations if needed (done here, lazily)
+    if(model->acts_memory == NULL || B != model->batch_size || T != model->seq_len) {
+        // record the current B,T as well
+        model->batch_size = B;
+        model->seq_len = T;
+        // and now allocate the space
+        fill_in_activation_sizes(model->act_sizes, B, T, model->config);
+        size_t num_activations = 0;
+        for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) {
+            num_activations += model->act_sizes[i];
+        }
+        model->num_activations = num_activations;
+        if (model->acts_memory != NULL) {
+            printf("GPT2 - Reallocating activations for B=%d T=%d\n", B, T);
+            cudaFree(model->acts_memory);
+            cudaFree(model->inputs);
+        }
+        model->acts_memory = malloc_and_point_activations(&model->acts, model->act_sizes);
+        printf("GPT2 - Allocated %zu MiB for activations\n", (num_activations * sizeof(float)) >> 20); // >> 20 is /(1024*1024)
+        cudaCheck(cudaMalloc((void**)&model->inputs, B * T * sizeof(int)));
+        cudaCheck(cudaMallocHost((void**)&model->cpu_losses, B * T * sizeof(float)));
+    }
+
+    // copy inputs to the model
+    cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice));
+
+    // forward pass
+    ParameterTensors params = model->params; // for brevity
+    ActivationTensors acts = model->acts;
+    float* residual;
+    encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0]
+
+    for (int l = 0; l < L; l++) {
+
+        residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C;
+
+        // get the pointers of the weights for this layer
+        float* l_ln1w = params.ln1w + l * C;
+        float* l_ln1b = params.ln1b + l * C;
+        float* l_qkvw = params.qkvw + l * 3*C * C;
+        float* l_qkvb = params.qkvb + l * 3*C;
+        float* l_attprojw = params.attprojw + l * C * C;
+        float* l_attprojb = params.attprojb + l * C;
+        float* l_ln2w = params.ln2w + l * C;
+        float* l_ln2b = params.ln2b + l * C;
+        float* l_fcw = params.fcw + l * 4*C * C;
+        float* l_fcb = params.fcb + l * 4*C;
+        float* l_fcprojw = params.fcprojw + l * C * 4*C;
+        float* l_fcprojb = params.fcprojb + l * C;
+
+        // get the pointers of the activations for this layer
+        float* l_ln1 = acts.ln1 + l * B * T * C;
+        float* l_ln1_mean = acts.ln1_mean;
+        float* l_ln1_rstd = acts.ln1_rstd;
+        float* l_qkvr = acts.qkvr + l * B * T * 3*C;
+        float* l_atty = acts.atty + l * B * T * C;
+        float* l_att = acts.att + l * B * NH * T * T;
+        float* l_attproj = acts.attproj + l * B * T * C;
+        float* l_residual2 = acts.residual2 + l * B * T * C;
+        float* l_ln2 = acts.ln2 + l * B * T * C;
+        float* l_ln2_mean = acts.ln2_mean;
+        float* l_ln2_rstd = acts.ln2_rstd;
+        float* l_fch = acts.fch + l * B * T * 4*C;
+        float* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C;
+        float* l_fcproj = acts.fcproj + l * B * T * C;
+        float* l_residual3 = acts.residual3 + l * B * T * C;
+        // these are only needed as scratchpads for the forward pass
+        float* scratch = acts.output;
+
+        // now do the forward pass
+        layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C);
+        matmul_forward(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C);
+        attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH);
+        matmul_forward(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C);
+        residual_forward(l_residual2, residual, l_attproj, B*T*C);
+        layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C);
+        //op16 matmul + bias + GeLU fusion
+        // matmul_forward(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C);
+        // gelu_forward(l_fch_gelu, l_fch, B*T*4*C);
+        matmul_forward_fused_gelu(l_fch_gelu, l_ln2, l_fcw, l_fcb, B, T, C, 4*C);
+        matmul_forward(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C);
+        residual_forward(l_residual3, l_residual2, l_fcproj, B*T*C);
+    }
+
+    residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3
+    layernorm_forward(acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C);
+    matmul_forward(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp);
+
+    model->mean_loss = -1.0f;
+}
+
+void gpt2_free(GPT2 *model) {
+    cudaCheck(cudaFree(model->params_memory));
+    cudaCheck(cudaFree(model->grads_memory));
+    cudaCheck(cudaFree(model->m_memory));
+    cudaCheck(cudaFree(model->v_memory));
+    cudaCheck(cudaFree(model->acts_memory));
+    cudaCheck(cudaFree(model->grads_acts_memory));
+    cudaCheck(cudaFree(model->inputs));
+    cudaCheck(cudaFree(model->targets));
+    cudaFreeHost(model->cpu_losses);
+}
+
+#define GPT2_EOT 50256

kernels/layernorm.cuh

@@ -6,32 +6,88 @@
 #include <float.h>
 #include "../utils/cuda_utils.cuh"
 
+#define BLOCK_DIM 256
+#define WARP_SIZE 32
+
+//lecture 15 slide ~32 
+__device__ __inline__ float warp_reduce(float val) {
+    float partialSum = val;
+    for(unsigned int stride = WARP_SIZE/2; stride > 0; stride /= 2) {
+        partialSum += __shfl_down_sync(0xffffffff, partialSum, stride);
+    }
+    return partialSum;
+}
+
 __global__ void layernorm_forward_kernel(float* out, float* mean, float* rstd, const float* inp, const float* weight,
                                          const float* bias, int B, int T, int C) {
     // Implement this
-    int b = blockDim.x * blockIdx.x + threadIdx.x;
-    int t = blockDim.y * blockIdx.y + threadIdx.y;
-
-    if (b >= B || t >= T) {
+    int bt = blockIdx.x;
+    if(bt >= B*T){
         return;
     }
 
-    float sum = 0.0f;
-    for (int c = 0; c < C; ++c) {
-        sum += inp[b * T * C + t * C + c];
-    }
-    float m = sum / C;
+    const float* x = inp + bt * C;
+    float* o = out + bt * C;
 
-    sum = 0.0f;
-    for (int c = 0; c < C; ++c) {
-        float diff = inp[b * T * C + t * C + c] - m;
-        sum += diff * diff;
+    // computer mean
+    float partial_sum = 0.0f;
+    for(int c = threadIdx.x ; c < C ; c += BLOCK_DIM){
+        partial_sum += x[c];
     }
-    float s = rsqrtf(sum / C + 1e-5f);
-    
-    for (int c = 0; c < C; ++c) {
-        float n = s * (inp[b * T * C + t * C + c] - m);
-        out[b * T * C + t * C + c] = n * weight[c] + bias[c];
+
+    partial_sum = warp_reduce(partial_sum);
+
+    __shared__ float partial_sum_s[BLOCK_DIM / WARP_SIZE];
+    if(threadIdx.x % WARP_SIZE == 0){
+        partial_sum_s[threadIdx.x / WARP_SIZE] = partial_sum;
+    }
+
+    __syncthreads();
+
+    if(threadIdx.x < WARP_SIZE){
+        partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sum_s[threadIdx.x] : 0.0f;
+        partial_sum = warp_reduce(partial_sum);
+    }
+
+    //thread 0 will have the final sum
+    __shared__ float s_mean;
+    if (threadIdx.x == 0) {
+        s_mean = partial_sum / C;
+    }
+
+    __syncthreads();
+    float m = s_mean;
+
+    //compute variance
+    partial_sum = 0.0f;
+    for (int c = threadIdx.x; c < C; c += BLOCK_DIM) {
+        float diff = x[c] - m;
+        partial_sum += diff * diff;
+    }
+
+    partial_sum = warp_reduce(partial_sum);
+
+    if (threadIdx.x % WARP_SIZE == 0) {
+        partial_sum_s[threadIdx.x / WARP_SIZE] = partial_sum;
+    }
+    __syncthreads();
+
+    if (threadIdx.x < WARP_SIZE) {
+        partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sum_s[threadIdx.x] : 0.0f;
+        partial_sum = warp_reduce(partial_sum);
+    }
+
+    __shared__ float s_rstd;
+    if (threadIdx.x == 0) {
+        s_rstd = rsqrtf(partial_sum / C + 1e-5f);
+    }
+    __syncthreads();
+    float s = s_rstd;
+
+    //normalize
+    for (int c = threadIdx.x; c < C; c += BLOCK_DIM) {
+        float n = s * (x[c] - m);
+        o[c] = n * weight[c] + bias[c];
     }
 }
 
@@ -39,10 +95,8 @@ __global__ void layernorm_forward_kernel(float* out, float* mean, float* rstd, c
 void layernorm_forward(float* out, float* mean, float* rstd, float* inp, float* weight, float* bias,
                        int B, int T, int C) {
     // Implement this
-    dim3 block(16,16);
-    dim3 grid((B+block.x-1)/block.x, (T+block.y-1)/block.y);
-
-    layernorm_forward_kernel<<<grid, block>>>(out,mean, rstd, inp, weight, bias,B,T, C);
+    int num_blocks = B*T;
+    layernorm_forward_kernel<<<num_blocks, BLOCK_DIM>>>(out,mean, rstd, inp, weight, bias,B,T, C);
 }
 
 #endif // __LAYERNORM_KERNEL_CUH__

kernels/matmul.cuh

@@ -3,35 +3,213 @@
 
 #include <cuda_runtime.h>
 #include "../utils/cuda_utils.cuh"
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/device/gemm.h"
+#include "/work/hdd/bche/hshah6/cutlass/include/cutlass/epilogue/thread/linear_combination_gelu.h"
 
-__global__ void matmul_forward_kernel(float* out, const float* inp, const float* weight, 
-                                      const float* bias,int B,int T,  int C, int OC) {
-    // Implement this
-    int b = blockDim.x*blockIdx.x + threadIdx.x; // B
-    int t = blockDim.y*blockIdx.y + threadIdx.y; // T
-    int oc = blockDim.z*blockIdx.z + threadIdx.z; // OC
+//referenced: https://github.com/NVIDIA/cutlass/blob/main/examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu
 
-    if (b >= B || t >= T || oc >= OC) {
-        return;
+//removing bias overhead and fusing with cutlass
+// __global__ void bias_add_filler_kernel(float* __restrict__ out, const float* __restrict__ bias, int BT, int OC) {
+//     int x = blockIdx.x * blockDim.x + threadIdx.x;
+//     if (x < BT * OC) out[x] = bias[x % OC];
+// }
+
+//using big cutlass for kernels for BT>=64
+using GemmDefault = cutlass::gemm::device::Gemm<
+    float, cutlass::layout::RowMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::RowMajor,
+    float,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm80,
+    cutlass::gemm::GemmShape<128,128,32>,
+    cutlass::gemm::GemmShape<64,64,32>,
+    cutlass::gemm::GemmShape<16,8,8>,
+    cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    3 // Stages
+>;
+
+//small cutlass for kernels for BT<64, we can use smaller threadblocks and warps to increase occupancy
+using GemmsmallDefault = cutlass::gemm::device::Gemm<
+    float, cutlass::layout::RowMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::RowMajor,
+    float,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm80,
+    cutlass::gemm::GemmShape<32,128,32>,
+    cutlass::gemm::GemmShape<32,64,32>,
+    cutlass::gemm::GemmShape<16,8,8>,
+    cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    2
+>;
+
+using GemmMediumDefault = cutlass::gemm::device::Gemm<
+    float, cutlass::layout::RowMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::RowMajor,
+    float,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm80,
+    cutlass::gemm::GemmShape<64,128,32>,
+    cutlass::gemm::GemmShape<32,64,32>,
+    cutlass::gemm::GemmShape<16,8,8>,
+    cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    3
+>;
+
+inline void matmul_forward(float* __restrict__ out, const float* __restrict__ inp, const float* __restrict__ weight, 
+                                      const float* __restrict__ bias,int B,int T,  int C, int OC) {
+    int BT = B*T;
+
+    if(bias != NULL){
+        if(BT <= 32){
+            GemmsmallDefault gemm_op;
+            typename GemmsmallDefault::Arguments args{ 
+                {BT, OC, C}, // problem size
+                {inp, C}, // A matrix
+                {weight, C}, // B matrix
+                {bias, 0}, // C matrix
+                {out, OC}, // D matrix (output)
+                {1.0f, 1.0f} // alpha, beta for the epilogue
+            };
+            gemm_op(args);
+        }
+        else if(BT <= 128){
+            GemmMediumDefault gemm_op;
+            typename GemmMediumDefault::Arguments args{ 
+                {BT, OC, C}, // problem size
+                {inp, C}, // A matrix
+                {weight, C}, // B matrix
+                {bias, 0}, // C matrix
+                {out, OC}, // D matrix (output)
+                {1.0f, 1.0f} // alpha, beta for the epilogue
+            };
+            gemm_op(args);
+        }
+        else {
+            GemmDefault gemm_op;
+            typename GemmDefault::Arguments args{ 
+                {BT, OC, C}, // problem size
+                {inp, C}, // A matrix
+                {weight, C}, // B matrix
+                {bias, 0}, // C matrix
+                {out, OC}, // D matrix (output)
+                {1.0f, 1.0f} // alpha, beta for the epilogue
+            };
+            gemm_op(args);
+        }
+    }
+    else{
+        if(BT <= 32){
+            GemmsmallDefault gemm_op;
+            typename GemmsmallDefault::Arguments args{ 
+                {BT, OC, C}, // problem size
+                {inp, C}, // A matrix
+                {weight, C}, // B matrix
+                {out, OC}, // C matrix (output)
+                {out, OC}, // D matrix (output)
+                {1.0f, 0.0f} // alpha, beta for the epilogue
+            };
+            gemm_op(args);
+        }
+        else if(BT <= 128){
+            GemmDefault gemm_op;
+            typename GemmDefault::Arguments args{ 
+                {BT, OC, C}, // problem size
+                {inp, C}, // A matrix
+                {weight, C}, // B matrix
+                {out, OC}, // C matrix (output)
+                {out, OC}, // D matrix (output)
+                {1.0f, 0.0f} // alpha, beta for the epilogue
+            };
+            gemm_op(args);
+        }
+        else{
+            GemmDefault gemm_op;
+            typename GemmDefault::Arguments args{ 
+                {BT, OC, C}, // problem size
+                {inp, C}, // A matrix
+                {weight, C}, // B matrix
+                {out, OC}, // C matrix (output)
+                {out, OC}, // D matrix (output)
+                {1.0f, 0.0f} // alpha, beta for the epilogue
+            };
+            gemm_op(args);
+        }
     }
 
-    float sum = (bias != NULL) ? bias[oc] : 0.0f;
-    const float* inp_token_ptr = inp + b * T * C + t * C;
-    for(int i = 0; i < C; i++){
-        sum += inp_token_ptr[i] * weight[oc*C + i] ;
-    }
-    out[b * T * OC + t * OC + oc] = sum;
+    
 }
 
+using Gemmgelu = cutlass::gemm::device::Gemm<
+    float, cutlass::layout::RowMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::RowMajor,
+    float,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm80,
+    cutlass::gemm::GemmShape<128, 128, 32>,
+    cutlass::gemm::GemmShape<64, 64, 32>,
+    cutlass::gemm::GemmShape<16, 8, 8>,
+    cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    3>;
+
+using GemmgeluSmall = cutlass::gemm::device::Gemm<
+    float, cutlass::layout::RowMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::RowMajor,
+    float,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm80,
+    cutlass::gemm::GemmShape<32, 128, 32>,
+    cutlass::gemm::GemmShape<32, 64, 32>,
+    cutlass::gemm::GemmShape<16, 8, 8>,
+    cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    2>;
+
+using GemmgeluMedium = cutlass::gemm::device::Gemm<
+    float, cutlass::layout::RowMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::RowMajor,
+    float,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm80,
+    cutlass::gemm::GemmShape<64, 128, 32>,
+    cutlass::gemm::GemmShape<32, 64, 32>,
+    cutlass::gemm::GemmShape<16, 8, 8>,
+    cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    3>;
+
 // Launch kernel here
-void matmul_forward(float* out, const float* inp, const float* weight, const float* bias,
+inline void matmul_forward_fused_gelu(float* out, const float* inp, const float* weight, const float* bias,
                     int B, int T, int C, int OC) {
     // Implement this
-    dim3 block(8,8,8);
-    dim3 grid((B+block.x-1)/block.x, (T+block.y-1)/block.y, (OC+block.z-1)/block.z);
-
-    matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);
+    int BT = B*T;
 
+    if (BT <= 32) {
+        GemmgeluSmall op;
+        typename GemmgeluSmall::Arguments args{
+            {BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
+        op(args);
+    } else if (BT <= 128) {
+        GemmgeluMedium op;
+        typename GemmgeluMedium::Arguments args{
+            {BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
+        op(args);
+    } else {
+        Gemmgelu op;
+        typename Gemmgelu::Arguments args{
+            {BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
+        op(args);
+    }
 }
 
 #endif // __MATMUL_KERNEL_CUH__

kernels/residual.cuh

@@ -4,7 +4,7 @@
 #include <cuda_runtime.h>
 #include "../utils/cuda_utils.cuh"
 
-__global__ void residual_forward_kernel(float* out, float* inp1, float* inp2, int N) {
+__global__ void residual_forward_kernel(float* __restrict__ out, const float* __restrict__ inp1, const float* __restrict__ inp2, int N) {
     // Implement this
     int x = blockDim.x * blockIdx.x + threadIdx.x;
 
@@ -16,7 +16,7 @@ __global__ void residual_forward_kernel(float* out, float* inp1, float* inp2, in
 // Launch kernel here
 void residual_forward(float* out, float* inp1, float* inp2, int N) {
     // Implement this
-    dim3 block(32,1,1);
+    dim3 block(256,1,1);
     dim3 grid((N + block.x - 1)/ block.x);
 
     residual_forward_kernel<<<grid, block>>>(out, inp1, inp2, N);

kernels/softmax.cuh

@@ -6,30 +6,93 @@
 #include <float.h>
 #include "../utils/cuda_utils.cuh"
 
+#define BLOCK_DIM 256
+#define WARP_SIZE 32
+
+// lecture 15 slide 32
+__device__ __inline__ float warpReduceSum(float partialSum) {
+    for (unsigned int stride = WARP_SIZE / 2; stride > 0; stride /= 2) {
+        partialSum += __shfl_down_sync(0xffffffff, partialSum, stride);
+    }
+    return partialSum;
+}
+
+// same pattern for max
+__device__ __inline__ float warpReduceMax(float partialMax) {
+    for (unsigned int stride = WARP_SIZE / 2; stride > 0; stride /= 2) {
+        partialMax = fmaxf(partialMax, __shfl_down_sync(0xffffffff, partialMax, stride));
+    }
+    return partialMax;
+}
+
 __global__ void softmax_forward_kernel(float* out, float inv_temperature, const float* inp, int N, int T) {
     // Implement this
-    int x = blockDim.x * blockIdx.x + threadIdx.x;
+    int row = blockIdx.x;
+    if(row >= N*T){
+        return;
+    }
 
-    if(x < N*T){
-        int own_pos = x % T;
-        const float* a = inp + x * T;
+    int own_pos = row % T;
+    const float* a = inp + row * T;
+    float* o = out + row * T;
 
-        float maxval = -FLT_MAX;
-        for (int i = 0; i <= own_pos; i++) {
-            maxval = fmaxf(maxval,a[i]);
-        }
+    //find max
+    //thread coarsening
+    float partial_max = -FLT_MAX;
+    for(int i = threadIdx.x; i <= own_pos; i+=BLOCK_DIM){
+        partial_max = fmaxf(partial_max,a[i]);
+    }
 
-        float sumval = 0.0f;
-        for (int i = 0; i <= own_pos; ++i) {
-            float ev = expf(inv_temperature * (a[i] - maxval));
-            sumval += ev;
-            out[x * T + i] = ev;
-        }
+    //lecture 15 slide 38
+    partial_max = warpReduceMax(partial_max);
 
-        float norm = 1.0f / sumval;
-        for (int i = 0; i <= own_pos; ++i) {
-            out[x * T + i] *= norm;
-        }
+    __shared__ float partial_max_s[BLOCK_DIM / WARP_SIZE];
+    if (threadIdx.x % WARP_SIZE == 0) {
+        partial_max_s[threadIdx.x / WARP_SIZE] = partial_max;
+    }
+    __syncthreads();
+    
+    if (threadIdx.x < WARP_SIZE) {
+        partial_max = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_max_s[threadIdx.x] : -FLT_MAX;
+        partial_max = warpReduceMax(partial_max);
+    }
+
+    __shared__ float s_max;
+    if (threadIdx.x == 0) {
+        s_max = partial_max;
+    }
+    __syncthreads();
+    float maxval = s_max;
+
+    float partial_sum = 0.0f;
+    for (int i = threadIdx.x; i <= own_pos; i += BLOCK_DIM) {
+        float ev = expf(inv_temperature * (a[i] - maxval));
+        o[i] = ev;
+        partial_sum += ev;
+    }
+
+    partial_sum = warpReduceSum(partial_sum);
+
+    __shared__ float partial_sums_s[BLOCK_DIM / WARP_SIZE];
+    if (threadIdx.x % WARP_SIZE == 0) {
+        partial_sums_s[threadIdx.x / WARP_SIZE] = partial_sum;
+    }
+    __syncthreads();
+
+    if (threadIdx.x < WARP_SIZE) {
+        partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sums_s[threadIdx.x] : 0.0f;
+        partial_sum = warpReduceSum(partial_sum);
+    }
+
+    __shared__ float s_sum;
+    if (threadIdx.x == 0) {
+        s_sum = partial_sum;
+    }
+    __syncthreads();
+    float norm = 1.0f / s_sum;
+
+    for (int i = threadIdx.x; i <= own_pos; i += BLOCK_DIM) {
+        o[i] *= norm;
     }
 }