202 lines
6.1 KiB
Plaintext
202 lines
6.1 KiB
Plaintext
#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__ |