aryansingh/gpt-2-inference-engine
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ECE 408 sem project

Browse Source
This commit is contained in:
Aryan Singh committed 2026-06-07 13:43:20 -05:00
1 parent 1d79ccd2e3
commit fa3efcaf07
139 files changed
+11634 -1
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels_req_0/attention.cuh
+202
View File
@@ -0,0 +1,202 @@
#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_];
}
}
#define PREATT_T 128
#define PREATT_U 8
#define PREATT_S 16
__global__ void preattention_scores_kernel(float* q, float* k, int B, int T, int NH, int HS, float* preatt){
// each block handles a (b,nh) head and a tile of (t1, t2)
int head = blockIdx.z; //whcih pair
int b = head / NH;
int nh = head % NH;
int t1 = blockIdx.x * PREATT_T + threadIdx.x; // row of output query positin
int t2 = blockIdx.y * PREATT_U; //start column key position
//head pointers
const float* q_head = q + (b * NH + nh) * T * HS; //[T,HS]
const float* k_head = k + (b * NH + nh) * T * HS; //[T,HS]
float* preatt_head = preatt + (b * NH +nh) * T * T; //[T,T]
//reg tile for accumulation
float accumulation[PREATT_U] = {0.0f};
__shared__ float k_tile[PREATT_S][PREATT_U];
for(int tile_start = 0; tile_start < HS; tile_start+=PREATT_S){
// shared loading
// K is [T, HS], K[t2..t2+U, tile_start..tile_start+S]
int s_l = threadIdx.x / PREATT_U;
int u_l = threadIdx.x % PREATT_U;
int hs_idx = tile_start + s_l;
int t2_idx = t2 + u_l;
if (hs_idx < HS && t2_idx < T) {
k_tile[s_l][u_l] = k_head[t2_idx * HS + hs_idx];
} else {
k_tile[s_l][u_l] = 0.0f;
}
__syncthreads();
//each threads now loads its row of queries from global memory into registers and accumulates
if (t1 < T){
#pragma unroll
for(int s = 0; s < PREATT_S; s++){
int c = tile_start + s;
float q_val = (c < HS) ? q_head[t1 * HS +c] : 0.0f;
#pragma unroll
for(int u = 0; u < PREATT_U; u++){
accumulation[u] += q_val * k_tile[s][u];
}
}
}
__syncthreads();
}
if(t1 < T){
#pragma unroll
for(int u = 0; u < PREATT_U; u ++){
int t2_u = t2 + u;
if(t2_u < T){
preatt_head[t1 * T + t2_u] = accumulation[u];
}
}
}
}
#define ATTV_T 128
#define ATTV_U 8
#define ATTV_S 16
__global__ void attention_v_kernel(float* vaccum, float* att, float* v, int B, int T, int NH, int HS){
int head = blockIdx.z;
int b = head / NH;
int nh = head % NH;
int t = blockIdx.x * ATTV_T + threadIdx.x; // row of output query positin
int hs_base = blockIdx.y * ATTV_U; //start column key position
const float* att_head = att + (b * NH + nh) * T * T;
const float* v_head = v + (b * NH + nh) * T * HS;
float* out_head = vaccum + (b * NH + nh) * T * HS;
float accumulation[ATTV_U] = {0.0f};
__shared__ float v_tile[ATTV_S][ATTV_U];
for(int tile_start = 0; tile_start < T; tile_start += ATTV_S){
int s_l = threadIdx.x / ATTV_U;
int u_l = threadIdx.x % ATTV_U;
int t2_idx = tile_start + s_l;
int hs_idx = hs_base + u_l;
if (t2_idx < T && hs_idx < HS) {
v_tile[s_l][u_l] = v_head[t2_idx * HS + hs_idx];
} else {
v_tile[s_l][u_l] = 0.0f;
}
__syncthreads();
// Each thread loads att values from registers and accumulates
if (t < T) {
#pragma unroll
for (int s = 0; s < ATTV_S; s++) {
int t2 = tile_start + s;
float att_val = (t2 < T) ? att_head[t * T + t2] : 0.0f;
#pragma unroll
for (int u = 0; u < ATTV_U; u++) {
accumulation[u] += att_val * v_tile[s][u];
}
}
}
__syncthreads();
}
if (t < T) {
#pragma unroll
for (int u = 0; u < ATTV_U; u++) {
int hs = hs_base + u;
if (hs < HS) {
out_head[t * HS + hs] = accumulation[u];
}
}
}
}
// 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 block_preatt(PREATT_T);
dim3 grid_preatt((T + PREATT_T - 1) / PREATT_T,(T+PREATT_U-1)/PREATT_U,B*NH);
preattention_scores_kernel<<<grid_preatt, block_preatt>>>(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;
dim3 block_attv(PREATT_T);
dim3 grid_attv((T + ATTV_T - 1) / ATTV_T,(HS+ATTV_U-1)/ATTV_U,B*NH);
attention_v_kernel<<<grid_attv,block_attv>>>(vaccum, att, v, B, T, NH, HS);
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
}
#endif // __ATTENTION_CUH__
kernels_req_0/matmul.cuh
+90
View File
@@ -0,0 +1,90 @@
#ifndef __MATMUL_KERNEL_CUH__
#define __MATMUL_KERNEL_CUH__
#include <cuda_runtime.h>
#include "../utils/cuda_utils.cuh"
#define TILE_T 128
#define TILE_U 8
#define TILE_S 16
__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
const int BT = B * T;
// getting which pair this thread is going to process
const int row = blockIdx.x * TILE_T + threadIdx.x;
// OC
const int col = blockIdx.y * TILE_U;
//register tile
float accumilation[TILE_U] = {0.0f};
//init the tile with bais
if(bias != NULL && row < BT){
#pragma unroll
for(int i = 0; i < TILE_U; i++){
if(col + i < OC){
accumilation[i] = bias[col + i];
}
}
}
__shared__ float N_tile[TILE_S][TILE_U];
for(int tile_start = 0; tile_start < C; tile_start += TILE_S){
const int s_l = threadIdx.x / TILE_U;
const int u_l = threadIdx.x % TILE_U;
const int c_g = tile_start + s_l;
const int oc_g = col + u_l;
if(c_g < C && oc_g < OC){
N_tile[s_l][u_l] = weight[oc_g * C + c_g];
}
else{
N_tile[s_l][u_l] = 0.0f;
}
__syncthreads();
if(row < BT){
#pragma unroll
for(int s = 0 ; s < TILE_S; s++){
const int c = tile_start + s;
float m = 0.0f;
if(c < C){
m = inp[row * C + c] ;
}
#pragma unroll
for(int u = 0; u < TILE_U; u++){
accumilation[u] +=m *N_tile[s][u];
}
}
}
__syncthreads();
}
if(row < BT){
#pragma unroll
for(int u = 0; u < TILE_U; u++){
if(col + u < OC){
out[row*OC+col+u] = accumilation[u];
}
}
}
}
// 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
const int BT = B*T;
dim3 block(TILE_T);
dim3 grid((BT+TILE_T-1)/TILE_T,(OC+TILE_U-1)/TILE_U);
matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);
}
#endif // __MATMUL_KERNEL_CUH__