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,191 @@
|
||||
#ifndef __ATTENTION_CUH__
|
||||
#define __ATTENTION_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
#include "softmax.cuh"
|
||||
#include <mma.h>
|
||||
|
||||
using namespace nvcuda;
|
||||
|
||||
#define WMMA_M 16
|
||||
#define WMMA_N 16
|
||||
#define WMMA_K 8
|
||||
|
||||
__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_WARPS_M 2
|
||||
#define PREATT_WARPS_N 4
|
||||
#define PREATT_BM (PREATT_WARPS_M * WMMA_M)
|
||||
#define PREATT_BN (PREATT_WARPS_N * WMMA_N)
|
||||
|
||||
__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;
|
||||
|
||||
//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]
|
||||
|
||||
int start_row = blockIdx.x * PREATT_BM;
|
||||
int start_col = blockIdx.y * PREATT_BN;
|
||||
|
||||
int warp_id = threadIdx.x / 32;
|
||||
int warp_row = warp_id / PREATT_WARPS_N;
|
||||
int warp_col = warp_id % PREATT_WARPS_N;
|
||||
|
||||
int warp_out_row = start_row + warp_row * WMMA_M;
|
||||
int warp_out_col = start_col + warp_col * WMMA_N;
|
||||
|
||||
// Q * K^T: A=Q (row_major), B=K (col_major gives K^T)
|
||||
wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> a_frag;
|
||||
wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::col_major> b_frag;
|
||||
wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;
|
||||
|
||||
wmma::fill_fragment(c_frag, 0.0f);
|
||||
|
||||
for(int kk = 0; kk < HS; kk += WMMA_K){
|
||||
if(warp_out_row < T && kk < HS){
|
||||
wmma::load_matrix_sync(a_frag, q_head + warp_out_row * HS + kk, HS);
|
||||
#pragma unroll
|
||||
for(int i = 0; i < a_frag.num_elements; i++){
|
||||
a_frag.x[i] = wmma::__float_to_tf32(a_frag.x[i]);
|
||||
}
|
||||
}
|
||||
|
||||
if(warp_out_col < T && kk < HS) {
|
||||
wmma::load_matrix_sync(b_frag, k_head + warp_out_col * HS + kk, HS);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < b_frag.num_elements; i++)
|
||||
b_frag.x[i] = wmma::__float_to_tf32(b_frag.x[i]);
|
||||
}
|
||||
|
||||
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
|
||||
}
|
||||
|
||||
//write to global mem
|
||||
if(warp_out_row < T && warp_out_col < T){
|
||||
wmma::store_matrix_sync(preatt_head + warp_out_row * T + warp_out_col, c_frag, T, wmma::mem_row_major);
|
||||
}
|
||||
}
|
||||
|
||||
#define ATTV_WARPS_M 2
|
||||
#define ATTV_WARPS_N 4
|
||||
#define ATTV_BM 32
|
||||
#define ATTV_BN 64
|
||||
|
||||
__global__ void attention_v_kernel(float* vaccum, float* att, float* v, int B, int T, int NH, int HS){
|
||||
int head = blockIdx.z; //which pair
|
||||
int b = head / NH;
|
||||
int nh = head % NH;
|
||||
|
||||
//head pointers
|
||||
const float* att_head = att + (b * NH + nh) * T * T; //[T,HS]
|
||||
const float* v_head = v + (b * NH + nh) * T * HS; //[T,HS]
|
||||
float* out_head = vaccum + (b * NH +nh) * T * HS; //[T,T]
|
||||
|
||||
int row_start = blockIdx.x * ATTV_BM;
|
||||
int col_start = blockIdx.y * ATTV_BN;
|
||||
|
||||
int warpId = threadIdx.x / 32;
|
||||
int warp_row = warpId / ATTV_WARPS_N;
|
||||
int warp_col = warpId % ATTV_WARPS_N;
|
||||
|
||||
int warp_out_row = row_start + warp_row * WMMA_M;
|
||||
int warp_out_col = col_start + warp_col * WMMA_N;
|
||||
|
||||
// att * V: both row-major; standard matmul
|
||||
wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> a_frag;
|
||||
wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> b_frag;
|
||||
wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;
|
||||
|
||||
wmma::fill_fragment(c_frag, 0.0f);
|
||||
|
||||
for(int kk = 0; kk < T; kk+=WMMA_K){
|
||||
if (warp_out_row < T && kk < T) {
|
||||
wmma::load_matrix_sync(a_frag, att_head + warp_out_row * T + kk, T);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < a_frag.num_elements; i++)
|
||||
a_frag.x[i] = wmma::__float_to_tf32(a_frag.x[i]);
|
||||
}
|
||||
|
||||
if (kk < T && warp_out_col < HS) {
|
||||
wmma::load_matrix_sync(b_frag, v_head + kk * HS + warp_out_col, HS);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < b_frag.num_elements; i++)
|
||||
b_frag.x[i] = wmma::__float_to_tf32(b_frag.x[i]);
|
||||
}
|
||||
|
||||
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
|
||||
}
|
||||
|
||||
if (warp_out_row < T && warp_out_col < HS) {
|
||||
wmma::store_matrix_sync(out_head + warp_out_row * HS + warp_out_col, c_frag, HS, wmma::mem_row_major);
|
||||
}
|
||||
}
|
||||
|
||||
// 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((T + PREATT_BM - 1) / PREATT_BM, (T + PREATT_BN - 1) / PREATT_BN, B * NH);
|
||||
preattention_scores_kernel<<<grid_preatt, 256>>>(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 grid_attv((T + ATTV_BM - 1) / ATTV_BM, (HS + ATTV_BN - 1) / ATTV_BN, B * NH);
|
||||
attention_v_kernel<<<grid_attv,256>>>(vaccum, att, v, B, T, NH, HS);
|
||||
|
||||
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
|
||||
}
|
||||
|
||||
#endif // __ATTENTION_CUH__
|
||||
Reference in new issue
Block a user