ECE 408 sem project
This commit is contained in:
Aryan Singh committed
1 parent
1d79ccd2e3
commit
fa3efcaf07
139 files changed
+11634
-1
No files matched your search
@@ -0,0 +1,108 @@
|
||||
#ifndef __ATTENTION_CUH__
|
||||
#define __ATTENTION_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
#include "softmax.cuh"
|
||||
|
||||
__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;
|
||||
|
||||
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(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;
|
||||
|
||||
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(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;
|
||||
|
||||
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* __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;
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// 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);
|
||||
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);
|
||||
|
||||
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,33 @@
|
||||
#ifndef ENCODER_FORWARD_KERNEL_CUH
|
||||
#define ENCODER_FORWARD_KERNEL_CUH
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
__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];
|
||||
const float* __restrict__ wte_ptr = wte + token * C;
|
||||
const float* __restrict__ wpe_ptr = wpe + y * C;
|
||||
for(int i = 0; i < C; i ++){
|
||||
out[x*T*C + y*C + 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);
|
||||
}
|
||||
|
||||
|
||||
#endif // ENCODER_FORWARD_KERNEL_CUH
|
||||
@@ -0,0 +1,32 @@
|
||||
#ifndef GELU_KERNEL_CUH_
|
||||
#define GELU_KERNEL_CUH_
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
|
||||
|
||||
__global__ void gelu_forward_kernel(float* __restrict__ out, const float* __restrict__ inp, int N) {
|
||||
// Implement this
|
||||
|
||||
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
|
||||
if(x < N){
|
||||
float val = inp[x];
|
||||
float cube = 0.044715f * val * val * val;
|
||||
out[x] = 0.5f * val * (1.0f + tanhf(GELU_SCALING_FACTOR * (val + cube)));
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// Launch kernel here
|
||||
void gelu_forward(float* out, const float* inp, int N) {
|
||||
// Implement this
|
||||
dim3 block(32,1,1);
|
||||
dim3 grid((N + block.x -1) / block.x);
|
||||
|
||||
gelu_forward_kernel<<<grid, block>>>(out, inp, N);
|
||||
|
||||
}
|
||||
|
||||
#endif // GELU_KERNEL_CUH_
|
||||
@@ -0,0 +1,51 @@
|
||||
#ifndef __LAYERNORM_KERNEL_CUH__
|
||||
#define __LAYERNORM_KERNEL_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <float.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
__global__ void layernorm_forward_kernel(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd, const float* __restrict__ inp, const float* __restrict__ weight,
|
||||
const float* __restrict__ 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) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float* __restrict__ x = inp + b * T * C + t * C;
|
||||
float* __restrict__ o = out + b * T * C + t * C;
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int c = 0; c < C; ++c) {
|
||||
sum += x[c];
|
||||
}
|
||||
float m = sum / C;
|
||||
|
||||
sum = 0.0f;
|
||||
for (int c = 0; c < C; ++c) {
|
||||
float diff = x[c] - m;
|
||||
sum += diff * diff;
|
||||
}
|
||||
float s = rsqrtf(sum / C + 1e-5f);
|
||||
|
||||
for (int c = 0; c < C; ++c) {
|
||||
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
|
||||
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);
|
||||
}
|
||||
|
||||
#endif // __LAYERNORM_KERNEL_CUH__
|
||||
@@ -0,0 +1,38 @@
|
||||
#ifndef __MATMUL_KERNEL_CUH__
|
||||
#define __MATMUL_KERNEL_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
__global__ void matmul_forward_kernel(float* __restrict__ out, const float* __restrict__ inp, const float* __restrict__ weight,
|
||||
const float* __restrict__ 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
|
||||
|
||||
if (b >= B || t >= T || oc >= OC) {
|
||||
return;
|
||||
}
|
||||
|
||||
float sum = (bias != NULL) ? bias[oc] : 0.0f;
|
||||
const float* __restrict__ inp_token_ptr = inp + b * T * C + t * C;
|
||||
const float* __restrict__ w_r = weight + oc * C;
|
||||
for(int i = 0; i < C; i++){
|
||||
sum += inp_token_ptr[i] * w_r[i] ;
|
||||
}
|
||||
out[b * T * OC + t * OC + oc] = sum;
|
||||
}
|
||||
|
||||
// Launch kernel here
|
||||
void matmul_forward(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);
|
||||
|
||||
}
|
||||
|
||||
#endif // __MATMUL_KERNEL_CUH__
|
||||
@@ -0,0 +1,25 @@
|
||||
#ifndef RESIDUAL_KERNEL_CUH_
|
||||
#define RESIDUAL_KERNEL_CUH_
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
__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;
|
||||
|
||||
if(x < N){
|
||||
out[x] = inp1[x] + inp2[x];
|
||||
}
|
||||
}
|
||||
|
||||
// Launch kernel here
|
||||
void residual_forward(float* out, float* inp1, float* inp2, int N) {
|
||||
// Implement this
|
||||
dim3 block(32,1,1);
|
||||
dim3 grid((N + block.x - 1)/ block.x);
|
||||
|
||||
residual_forward_kernel<<<grid, block>>>(out, inp1, inp2, N);
|
||||
}
|
||||
|
||||
#endif // RESIDUAL_KERNEL_CUH_
|
||||
@@ -0,0 +1,37 @@
|
||||
#ifndef __SOFTMAX_KERNEL_CUH__
|
||||
#define __SOFTMAX_KERNEL_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <float.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
__global__ void softmax_forward_kernel(float* __restrict__ out, float inv_temperature, const float* __restrict__ inp, int N, int T) {
|
||||
// Implement this
|
||||
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
|
||||
if(x < N*T){
|
||||
int own_pos = x % T;
|
||||
const float* __restrict__ a = inp + x * T;
|
||||
float* __restrict__ o = out + x * T;
|
||||
|
||||
float maxval = -FLT_MAX;
|
||||
for (int i = 0; i <= own_pos; i++) {
|
||||
maxval = fmaxf(maxval,a[i]);
|
||||
}
|
||||
|
||||
float sumval = 0.0f;
|
||||
for (int i = 0; i <= own_pos; ++i) {
|
||||
float ev = expf(inv_temperature * (a[i] - maxval));
|
||||
sumval += ev;
|
||||
o[i] = ev;
|
||||
}
|
||||
|
||||
float norm = 1.0f / sumval;
|
||||
for (int i = 0; i <= own_pos; ++i) {
|
||||
o[i] *= norm;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#endif // __SOFTMAX_KERNEL_CUH__
|
||||
Reference in new issue
Block a user