comp files
This commit is contained in:
Aryan Singh committed
1 parent
fa3efcaf07
commit
1f77d0131d
27 files changed
+1238
-798
No files matched your search
+193
-39
@@ -2,10 +2,159 @@
|
||||
#define __ATTENTION_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <float.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) {
|
||||
|
||||
//added templates so the compiler can unroll loops and optimize better.
|
||||
// op11
|
||||
// without swilling the QK^T inner loops read smem_q [query_row][i] where each thread query_row in teh warp holds a different row but reads teh same column i. as HS is 64 every thread bank index reduces to (i % 32) giving 32 way back conflict per warp on the Q reads. col -> col ^ row to swizzle the accesses and reduce bank conflicts. so each 32 read to 32 different banks instead of 32 reads to the same bank.
|
||||
template <int BR, int BC, int HS>
|
||||
__global__ void flash_attention_forward_kernel(const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V, float* __restrict__ out, int B, int T, int NH, float scale) {
|
||||
//following line by like psudo code from flash attention paper.
|
||||
// https://tridao.me/publications/flash2/flash2.pdf
|
||||
|
||||
// step 1 map thread blocks to (batch , head , query_tile)
|
||||
int b = blockIdx.x;
|
||||
int nh = blockIdx.y;
|
||||
int query_tile = blockIdx.z;
|
||||
|
||||
int query_row = threadIdx.x; //load index
|
||||
int query_global = query_tile * BR + query_row;
|
||||
|
||||
//step 2 calculate all the base pointers for this block. We will be loading queries from Q and keys/values from K and V, and writing output to out.
|
||||
const float* __restrict__ query_base = Q + b * NH * T * HS + nh * T * HS;
|
||||
const float* __restrict__ key_base = K + b * NH * T * HS + nh * T * HS;
|
||||
const float* __restrict__ value_base = V + b * NH * T * HS + nh * T * HS;
|
||||
float* out_base = out + b * NH * T * HS + nh * T * HS;
|
||||
|
||||
// declare shared memory for the tile of queries and keys/values. We will load a tile of queries and then iterate over tiles of keys/values.
|
||||
__shared__ float smem_q[BR * HS];
|
||||
__shared__ float smem_k[BC * HS];
|
||||
__shared__ float smem_v[BC * HS];
|
||||
|
||||
// we will compute the output for a tile of queries at a time.
|
||||
int query_tile_start = query_tile * BR; // global index of the first query in this tile. This is used to load the correct queries into shared memory and also to know which output rows we are responsible for computing.
|
||||
bool q_valid = (query_global < T); // check if this thread's query is within the sequence length. If not, it will do nothing.
|
||||
|
||||
|
||||
//define swizzling function using the XOR on col and row index to reduce bank conflicts on smem_q reads in the inner loop.
|
||||
|
||||
#define SWIZZLE(col, row) ((row) * HS + ((col) ^ ((row) & (HS - 1))))
|
||||
|
||||
// step 4 load query tiles from global memory to shared memory.
|
||||
// coalesced load of the queries for this tile into shared memory. a linear load of BR*HS elements, with each thread loading HS elements for one query.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < HS; i++) {
|
||||
smem_q[SWIZZLE(i, query_row)] = (query_tile_start + query_row < T) ? query_base[(query_tile_start + query_row) * HS + i] : 0.0f; // if the query is out of bounds, we can load zeros since they won't contribute to the output.
|
||||
}
|
||||
|
||||
// step 5 init accumulators for the output of this tile of queries.
|
||||
float out_reg[HS];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < HS; i++) {
|
||||
out_reg[i] = 0.0f;
|
||||
}
|
||||
float max_score = -FLT_MAX;
|
||||
float sum_exp = 0.0f;
|
||||
|
||||
|
||||
// step 6 determine KV tile range for casual masking
|
||||
int query_tile_end = (query_tile + 1) * BR;
|
||||
int n_kv_tiles = (query_tile_end < T) ? (T + BC - 1) / BC : (query_tile_end + BC - 1) / BC; // if the end of the query tile is within the sequence length, we only need to iterate over the KV tiles that cover the sequence length. If the end of the query tile exceeds the sequence length, we need to iterate over all KV tiles that cover the end of the query tile to properly apply casual masking.
|
||||
|
||||
__syncthreads();
|
||||
|
||||
//step 7 main loop over kv values
|
||||
for(int j = 0; j < n_kv_tiles; j++){
|
||||
__syncthreads();
|
||||
|
||||
int k_tile_start = j * BC;
|
||||
|
||||
// step 8 load k and v tiles from global memory
|
||||
#pragma unroll
|
||||
for (int i = 0; i < HS; i++) {
|
||||
int l = i * BC + query_row;
|
||||
int row = l / HS;
|
||||
int col = l % HS;
|
||||
int global_row = k_tile_start + row;
|
||||
bool valid = (global_row < T) && (row < BC);
|
||||
|
||||
if (valid) {
|
||||
smem_k[SWIZZLE(col, row)] = key_base[global_row * HS + col];
|
||||
smem_v[SWIZZLE(col, row)] = value_base[global_row * HS + col];
|
||||
} else {
|
||||
smem_k[SWIZZLE(col, row)] = 0.0f; // pad with zeros if out of bounds
|
||||
smem_v[SWIZZLE(col, row)] = 0.0f; // pad with zeros if out of bounds
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// step 9 compute one row of S = Qi * Kj^T and apply casual masking
|
||||
float score_reg[BC];
|
||||
float row_max = -FLT_MAX;
|
||||
|
||||
#pragma unroll
|
||||
for (int c = 0; c < BC; c++) {
|
||||
int k_global = j * BC + c;
|
||||
float dot; // compute the dot product of the query and key for this position
|
||||
|
||||
if (!q_valid || k_global > query_global) { // apply casual masking. If the key position is greater than the query position, we mask it out by setting the score to a very large negative number.
|
||||
dot = -FLT_MAX;
|
||||
} else {
|
||||
float acc = 0.0f;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HS; h++) {
|
||||
acc += smem_q[SWIZZLE(h, query_row)] * smem_k[SWIZZLE(h, c)];
|
||||
}
|
||||
dot = acc * scale;
|
||||
}
|
||||
score_reg[c] = dot;
|
||||
row_max = fmaxf(row_max, dot); // local max for this row
|
||||
}
|
||||
|
||||
// step 10 online softmax rescaling
|
||||
float max_new = fmaxf(max_score, row_max);
|
||||
if (max_new == -FLT_MAX) { // if both are -FLT_MAX, we can skip the rescaling because all scores are -FLT_MAX and the output will be zero anyway.
|
||||
continue;
|
||||
}
|
||||
float alpha = (max_score == -FLT_MAX) ? 0.0f : expf(max_score - max_new); // if max_score is -FLT_MAX, we set alpha to 0 to avoid exp overflow. Otherwise, we compute the scaling factor for the old scores.
|
||||
float local_tile = 0.0f; // this will hold the sum of exp for the scores in this tile after rescaling.
|
||||
|
||||
#pragma unroll
|
||||
for (int c = 0; c < BC; c++) {
|
||||
float p = (score_reg[c] == -FLT_MAX) ? 0.0f : expf(score_reg[c] - max_new); // compute the exp for this score after rescaling. If the score is -FLT_MAX, we set it to 0 to avoid exp underflow.
|
||||
local_tile += p;
|
||||
score_reg[c] = p; // store the rescaled exp in score_reg for later use in the weighted sum.
|
||||
}
|
||||
|
||||
// update output acc
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HS; ++h) {
|
||||
float pv = 0.0f;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < BC; ++c) {
|
||||
pv += score_reg[c] * smem_v[SWIZZLE(h, c)];
|
||||
}
|
||||
out_reg[h] = alpha * out_reg[h] + pv;
|
||||
}
|
||||
sum_exp = alpha * sum_exp + local_tile; // update the local tile sum with the contribution from this tile and the previous tiles, scaled by alpha.
|
||||
max_score = max_new; // update the max score for the next iteration
|
||||
|
||||
}
|
||||
if (q_valid) {
|
||||
float inv_l = (sum_exp > 0.0f) ? (1.0f / sum_exp) : 0.0f;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HS; ++h) {
|
||||
out_base[query_global * HS + h] = out_reg[h] * inv_l;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
__global__ void permute_kernel(float* __restrict__ q, float* __restrict__ k, float* __restrict__ v, const float* __restrict__ inp, int B, int N, int NH, int d) {
|
||||
// Implement this
|
||||
|
||||
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
@@ -24,7 +173,7 @@ __global__ void permute_kernel(float* q, float* k, float* v, const float* inp, i
|
||||
}
|
||||
|
||||
|
||||
__global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, int d) {
|
||||
__global__ void unpermute_kernel(const float* __restrict__ inp, float* __restrict__ out, int B, int N, int NH, int d) {
|
||||
// Implement this
|
||||
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
|
||||
@@ -39,39 +188,42 @@ __global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, i
|
||||
|
||||
}
|
||||
|
||||
__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;
|
||||
// __global__ void preattention_scores_kernel(const float* __restrict__ q, const float* __restrict__ k, int B, int T, int NH, int HS, float* __restrict__ 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;
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
// float sum = 0.0f;
|
||||
// const float* __restrict__ q_ptr = q + b * NH * T * HS + nh * T * HS + t1 * HS;
|
||||
// const float* __restrict__ k_ptr = k + b * NH * T * HS + nh * T * HS + t2 * HS;
|
||||
// for (int i = 0; i < HS; i++) {
|
||||
// sum += q_ptr[i] * k_ptr[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;
|
||||
// __global__ void attention_v_kernel(float* __restrict__ vaccum, const float* __restrict__ att, const float* __restrict__ 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];
|
||||
}
|
||||
// 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;
|
||||
// const float* __restrict__ att_ptr = att + b * NH * T * T + nh * T * T + t * T;
|
||||
// for (int t2 = 0; t2 < T; t2++) {
|
||||
// sum += att_ptr[t2] * v[b*NH*T*HS + nh*T*HS + t2*HS + hs];
|
||||
// }
|
||||
|
||||
vaccum[b * NH * T * HS + nh * T * HS + t * HS + hs] = sum;
|
||||
}
|
||||
}
|
||||
// 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) {
|
||||
@@ -83,21 +235,23 @@ void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, i
|
||||
k = qkvr + 1 * B * T * C;
|
||||
v = qkvr + 2 * B * T * C;
|
||||
|
||||
dim3 block(32,1,1);
|
||||
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 = 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);
|
||||
//fa2 — op_7 sweep winner: BR=32, BC=32 was 26% faster than BR=64,BC=64
|
||||
// at T=64 (smaller tiles fit two blocks per SM, doubling occupancy).
|
||||
constexpr int BR = 32; // tile size for queries
|
||||
constexpr int BC = 32; // tile size for keys/values
|
||||
float scale = 1.0f / sqrtf((float)HS);
|
||||
|
||||
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);
|
||||
int number_of_q_tiles = (T + BR - 1) / BR;
|
||||
dim3 fa_grid(B, NH, number_of_q_tiles); // one block per (batch, head, Q tile)
|
||||
dim3 fa_block(BR, 1, 1); // one warp per query in the tile
|
||||
|
||||
float* vaccum = inp;
|
||||
attention_v_kernel<<<grid_preatt,block>>>(vaccum, att, v, B, T, NH, HS);
|
||||
flash_attention_forward_kernel<BR, BC, 64><<<fa_grid, fa_block>>>(q, k, v, vaccum, B, T, NH, scale);
|
||||
|
||||
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user