ECE 408 sem project

2026-06-07 13:43:20 -05:00
1 parent 1d79ccd2e3
commit fa3efcaf07
@@ -0,0 +1,105 @@
+#ifndef __ATTENTION_CUH__
+#define __ATTENTION_CUH__
+
+#include <cuda_runtime.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) {
+    // Implement this
+
+    int x = blockDim.x * blockIdx.x + threadIdx.x;
+
+    if( x < B*NH*N*d){
+        int b = (x / d / N / NH);
+        int nh = (x / d / N) % NH;
+        int t = (x / d) % N;
+        int d_ = x % d;
+
+        int inp_idx = (b * N * 3 * NH * d) + (t * 3 * NH * d) + (0 * NH * d) + (nh * d) + d_;
+        q[b * NH * N * d + nh * N * d + t * d + d_] = inp[inp_idx];
+        k[b * NH * N * d + nh * N * d + t * d + d_] = inp[inp_idx + NH * d];
+        v[b * NH * N * d + nh * N * d + t * d + d_] = inp[inp_idx + 2 * (NH * d)];
+    }
+}
+
+
+__global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, int d) {
+    // Implement this
+    int x = blockDim.x * blockIdx.x + threadIdx.x;
+
+    if( x < B*NH*N*d){
+        int b = (x / d / N / NH);
+        int nh = (x / d / N) % NH;
+        int t = (x / d) % N;
+        int d_ = x % d;
+
+        out[b*N*NH*d + t*NH*d + nh*d + d_] = inp[b*NH*N*d + nh*N*d + t*d + d_];
+    }
+
+}
+
+__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;
+
+    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;
+    }   
+}
+
+__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;
+
+    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];
+        }
+
+        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) {
+    // Implement this
+    int HS = C / NH;
+
+    float *q, *k, *v;
+    q = qkvr + 0 * B * T * C;
+    k = qkvr + 1 * B * T * C;
+    v = qkvr + 2 * B * T * C;
+
+    dim3 block(32,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);
+
+    float scale = 1.0 / sqrtf(HS);
+    int softmax_rows = B * NH * T;
+    softmax_forward_kernel<<<softmax_rows, 256>>>(att, scale, preatt, B * NH, T);
+    
+    float* vaccum = inp;
+    attention_v_kernel<<<grid_preatt,block>>>(vaccum, att, v, B, T, NH, HS);
+
+    unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
+}
+
+#endif // __ATTENTION_CUH__
@@ -0,0 +1,102 @@
+#ifndef __LAYERNORM_KERNEL_CUH__
+#define __LAYERNORM_KERNEL_CUH__
+
+#include <cuda_runtime.h>
+#include <math.h>
+#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 bt = blockIdx.x;
+    if(bt >= B*T){
+        return;
+    }
+
+    const float* x = inp + bt * C;
+    float* o = out + bt * C;
+
+    // computer mean
+    float partial_sum = 0.0f;
+    for(int c = threadIdx.x ; c < C ; c += BLOCK_DIM){
+        partial_sum += x[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];
+    }
+}
+
+// Launch kernel here
+void layernorm_forward(float* out, float* mean, float* rstd, float* inp, float* weight, float* bias,
+                       int B, int T, int C) {
+    // Implement this
+    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__
@@ -0,0 +1,99 @@
+#ifndef __SOFTMAX_KERNEL_CUH__
+#define __SOFTMAX_KERNEL_CUH__
+
+#include <cuda_runtime.h>
+#include <math.h>
+#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 row = blockIdx.x;
+    if(row >= N*T){
+        return;
+    }
+
+    int own_pos = row % T;
+    const float* a = inp + row * T;
+    float* o = out + row * T;
+
+    //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]);
+    }
+
+    //lecture 15 slide 38
+    partial_max = warpReduceMax(partial_max);
+
+    __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;
+    }
+}
+
+#endif // __SOFTMAX_KERNEL_CUH__