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,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);
|
||||
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,86 @@
|
||||
#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__
|
||||
@@ -0,0 +1,31 @@
|
||||
#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* out, const int* inp, const float* wte, const float* 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];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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* out, const float* 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,48 @@
|
||||
#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* 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) {
|
||||
return;
|
||||
}
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int c = 0; c < C; ++c) {
|
||||
sum += inp[b * T * C + t * C + c];
|
||||
}
|
||||
float m = sum / C;
|
||||
|
||||
sum = 0.0f;
|
||||
for (int c = 0; c < C; ++c) {
|
||||
float diff = inp[b * T * C + t * C + c] - m;
|
||||
sum += diff * diff;
|
||||
}
|
||||
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];
|
||||
}
|
||||
}
|
||||
|
||||
// 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,37 @@
|
||||
#ifndef __MATMUL_KERNEL_CUH__
|
||||
#define __MATMUL_KERNEL_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
|
||||
__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
|
||||
|
||||
if (b >= B || t >= T || oc >= OC) {
|
||||
return;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
// 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* out, float* inp1, float* 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,36 @@
|
||||
#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* out, float inv_temperature, const float* 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* a = inp + 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;
|
||||
out[x * T + i] = ev;
|
||||
}
|
||||
|
||||
float norm = 1.0f / sumval;
|
||||
for (int i = 0; i <= own_pos; ++i) {
|
||||
out[x * T + i] *= norm;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#endif // __SOFTMAX_KERNEL_CUH__
|
||||
Reference in new issue
Block a user