138 lines
4.3 KiB
Plaintext
138 lines
4.3 KiB
Plaintext
#ifndef __ATTENTION_CUH__
|
|
#define __ATTENTION_CUH__
|
|
|
|
// WINDOW = 129 because the CPU reference uses an inclusive range [q_t - 128, q_t],
|
|
// which covers 129 positions (128 look-back steps + the current token itself).
|
|
#define WINDOW 129
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
#include "local_softmax.cuh"
|
|
#include <math.h>
|
|
|
|
__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;
|
|
|
|
//Make int terms of WINDOW instead of full T, so its T * WINDOW
|
|
if(x < B* T* NH* WINDOW){
|
|
int b = (x / WINDOW / T / NH);
|
|
int nh = (x / WINDOW / T) % NH;
|
|
int t1 = (x / WINDOW) % T;
|
|
int w = x % WINDOW;
|
|
|
|
//Key start point based on offset
|
|
int t2 = t1 - w;
|
|
|
|
float sum = 0.0f;
|
|
|
|
//Need to check to not look too far back and stay in window range
|
|
//First check if t2 (Key Step) is behind or equal to t1 (Query Step)
|
|
// int lowest_margin = t1 - WINDOW;
|
|
if (t2 >= 0) {
|
|
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];
|
|
}
|
|
}
|
|
else {
|
|
sum = -FLT_MAX;
|
|
}
|
|
preatt[b * NH * T * WINDOW + nh * T * WINDOW + t1 * WINDOW + w] = sum;
|
|
//preatt[x] = 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 t1 = (x / HS) % T;
|
|
int hs = x % HS;
|
|
float sum = 0.0f;
|
|
for (int w = 0; w < WINDOW; w++) {
|
|
int t2 = t1 - w;
|
|
if (t2 >= 0) {
|
|
sum += att[b*NH*T*WINDOW + nh*T*WINDOW + t1*WINDOW + w] * v[b*NH*T*HS + nh*T*HS + t2*HS + hs];
|
|
}
|
|
}
|
|
|
|
vaccum[b * NH * T * HS + nh * T * HS + t1 * HS + hs] = sum;
|
|
//vaccum[x] = sum;
|
|
}
|
|
}
|
|
|
|
// Launch all kernels related to attention here
|
|
void local_attention_forward_gpu(float* out, float* qkvr, float* att, float* inp, int B, int T, int NH, int HS) {
|
|
// Implement this
|
|
|
|
float *q, *k, *v;
|
|
q = qkvr + 0 * B * T * NH * HS;
|
|
k = qkvr + 1 * B * T * NH * HS;
|
|
v = qkvr + 2 * B * T * NH * HS;
|
|
|
|
dim3 block(256,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;
|
|
cudaMalloc(&preatt, B * NH * T * WINDOW * sizeof(float));
|
|
|
|
dim3 grid_preatt((B * NH * T * WINDOW + block.x-1) / block.x,1,1);
|
|
preattention_scores_kernel<<<grid_preatt, block>>>(q, k, B, T, NH, HS, preatt);
|
|
|
|
float* packed_att;
|
|
cudaMalloc(&packed_att, B * NH * T * WINDOW * sizeof(float));
|
|
|
|
float scale = 1.0 / sqrtf(HS);
|
|
dim3 grid_softmax (( B * NH * T + block.x -1 )/ block.x,1,1);
|
|
local_softmax_forward_kernel<<<grid_softmax, block>>>(packed_att, scale, preatt, B * NH, T);
|
|
|
|
float* vaccum;
|
|
cudaMalloc(&vaccum, B * NH * T * HS * sizeof(float));
|
|
|
|
attention_v_kernel<<<grid,block>>>(vaccum, packed_att, v, B, T, NH, HS);
|
|
|
|
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
|
|
|
|
cudaFree(preatt);
|
|
cudaFree(vaccum);
|
|
cudaFree(packed_att);
|
|
}
|
|
|
|
#endif // __ATTENTION_CUH__ |