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,309 @@
|
||||
#ifndef __ATTENTION_KV_CACHE_TEMPLATE_CUH__
|
||||
#define __ATTENTION_KV_CACHE_TEMPLATE_CUH__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
#include "softmax.cuh"
|
||||
#include <float.h>
|
||||
#include <math.h>
|
||||
|
||||
extern bool g_enable_kv_cache; // Toggles KV-caching logic
|
||||
extern float* g_kv_cache; // Pointer to the GPU-allocated KV cache buffer
|
||||
extern int g_layer_idx; // Tracks the current transformer layer, useful for knowing when to switch to decode
|
||||
extern bool g_is_prefill; // True during the initial prompt, False during token generation
|
||||
extern int g_current_pos; // Tracks the absolute sequence position during Decode, helps index into the KV cache
|
||||
constexpr int MAX_SEQ_LEN = 1024; // gpt 2 default
|
||||
// --- Suggested Kernel Placeholders (TODO: Add your regular attention kernel here, and any other kernels you may need) ---
|
||||
//all the tests are running out of mem to store KV so implementing dynamic alloc
|
||||
static int s_cache_b = 0;
|
||||
static int s_cache_nh = 0;
|
||||
static int s_cache_hs = 0;
|
||||
|
||||
|
||||
__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;
|
||||
|
||||
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(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;
|
||||
|
||||
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_];
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// base line kernels
|
||||
|
||||
__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;
|
||||
|
||||
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* __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;
|
||||
const float* __restrict__ att_ptr = att + b * NH * T * T + nh * T * T + t * T;
|
||||
const float* __restrict__ v_ptr = v + b * NH * T * HS + nh * T * HS;
|
||||
for (int i = 0; i < T; i++) {
|
||||
sum += att_ptr[i] * v_ptr[i * HS + hs];
|
||||
}
|
||||
vaccum[b * NH * T * HS + nh * T * HS + t * HS + hs] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// helper kernel to calculate index into 1D cache buffer
|
||||
__device__ inline int get_cache_index(int layer, int is_v, int b, int nh, int t, int B, int NH, int HS) {
|
||||
// our layout layer -> k/v -> batch -> head -> time -> head_size
|
||||
return ((((layer * 2 + is_v) * B + b) * NH + nh) * MAX_SEQ_LEN + t) * HS;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Kernel to store the newly computed Key and Value into the global KV cache.
|
||||
* CAN BE FUSED WITH QKV MATMUL
|
||||
* In Prefill: Stores T tokens for the current layer.
|
||||
* In Decode: Stores 1 token at g_current_pos for the current layer.
|
||||
*/
|
||||
__global__ void store_kv_kernel(float* cache, const float* k, const float* v,int layer, int B, int NH, int HS, int T, int current_pos, bool is_prefill) {
|
||||
// Implement logic to copy K and V from the input QKV buffer into g_kv_cache
|
||||
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;
|
||||
|
||||
//check where to write in the cache
|
||||
int cache_idx;
|
||||
if(is_prefill){
|
||||
cache_idx = t;
|
||||
} else {
|
||||
cache_idx = current_pos;
|
||||
}
|
||||
|
||||
// Write K and V to the cache
|
||||
int cache_source_idx = b * NH * T * HS + nh * T * HS + t * HS + hs;
|
||||
cache[get_cache_index(layer, 0, b, nh, cache_idx, B, NH, HS) + hs] = k[cache_source_idx]; // is_v = 0 for K
|
||||
cache[get_cache_index(layer, 1, b, nh, cache_idx, B, NH, HS) + hs] = v[cache_source_idx]; // is_v = 1 for V
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Kernel to compute the attention for the Decode phase (T=1).
|
||||
* This involves:
|
||||
* 1. Extracting the current Query (Q)
|
||||
* 2. Gathering all previous Keys (K) and Values (V) from the cache
|
||||
* 3. Computing attention using the new Q and cached KV values.
|
||||
*/
|
||||
__global__ void decode_attention_kernel(float* vaccum, const float* q, const float* cache, int layer, int B, int NH, int HS, int current_pos) {
|
||||
// Implement T=1 attention logic using the KV cache
|
||||
int x = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
|
||||
if(x >= B*NH*HS){
|
||||
return;
|
||||
}
|
||||
|
||||
int b = (x / HS / NH);
|
||||
int nh = (x / HS) % NH;
|
||||
int hs = x % HS;
|
||||
|
||||
float max_val = -FLT_MAX;
|
||||
float scale = rsqrtf((float)HS);
|
||||
const float* __restrict__ q_ptr = q + b * NH * HS + nh * HS;
|
||||
|
||||
// Q * K^T for all previous tokens in the cache
|
||||
for(int t = 0; t <= current_pos; t++){
|
||||
const float* __restrict__ k_ptr = cache + get_cache_index(layer, 0, b, nh, t, B, NH, HS); // is_v = 0 for K
|
||||
float score = 0.0f;
|
||||
for(int i = 0; i < HS; i++){
|
||||
score += q_ptr[i] * k_ptr[i];
|
||||
}
|
||||
score *= scale;
|
||||
if(score > max_val){
|
||||
max_val = score;
|
||||
}
|
||||
}
|
||||
|
||||
// Compute softmax
|
||||
float sum_exp = 0.0f;
|
||||
float output = 0.0f;
|
||||
|
||||
for(int t = 0; t <= current_pos; t++){
|
||||
const float* __restrict__ k_ptr = cache + get_cache_index(layer, 0, b, nh, t, B, NH, HS);
|
||||
float score = 0.0f;
|
||||
for(int i = 0; i < HS; i++){
|
||||
score += q_ptr[i] * k_ptr[i];
|
||||
}
|
||||
score *= scale;
|
||||
|
||||
float exp_score = expf(score - max_val);
|
||||
sum_exp += exp_score;
|
||||
|
||||
const float* __restrict__ v_ptr = cache + get_cache_index(layer, 1, b, nh, t, B, NH, HS);
|
||||
output += exp_score * v_ptr[hs];
|
||||
}
|
||||
vaccum[b * NH * HS + nh * HS + hs] = output / sum_exp;
|
||||
}
|
||||
|
||||
//alloc and realloc kv cache
|
||||
static void allocate_kv_cache(int B, int NH, int HS){
|
||||
if(g_kv_cache != nullptr && s_cache_b == B && s_cache_nh == NH && s_cache_hs == HS){
|
||||
// If the cache exists and the dimensions match, do nothing
|
||||
return;
|
||||
}
|
||||
|
||||
// If the cache exists but the dimensions don't match, free and reallocate
|
||||
if(g_kv_cache != nullptr){
|
||||
cudaFree(g_kv_cache);
|
||||
g_kv_cache = nullptr;
|
||||
}
|
||||
|
||||
size_t total_cache_size = 12 * 2 * B * NH * MAX_SEQ_LEN * HS * sizeof(float); // 12 layers, K and V, max seq len
|
||||
cudaMalloc(&g_kv_cache, total_cache_size);
|
||||
cudaMemset(g_kv_cache, 0, total_cache_size); // 0 the cache
|
||||
|
||||
s_cache_b = B;
|
||||
s_cache_nh = NH;
|
||||
s_cache_hs = HS;
|
||||
}
|
||||
|
||||
|
||||
void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, int T, int C, int NH) {
|
||||
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(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);
|
||||
|
||||
// 1. Check if KV-caching is enabled, used by the verification script
|
||||
if (!g_enable_kv_cache) {
|
||||
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);
|
||||
|
||||
float scale = 1.0f / sqrtf((float)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((B * NH * T * HS + block.x - 1) / block.x); // <<< ADD
|
||||
attention_v_kernel<<<grid_attv, block>>>(vaccum, att, v, B, T, NH, HS); // <<< grid_attv not grid_preatt
|
||||
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
|
||||
return;
|
||||
}
|
||||
|
||||
allocate_kv_cache(B, NH, HS); // Ensure KV cache is allocated with the right dimensions
|
||||
// store K and V in the cache
|
||||
store_kv_kernel<<<grid, block>>>(g_kv_cache, k, v, g_layer_idx, B, NH, HS, T, g_current_pos, g_is_prefill);
|
||||
|
||||
|
||||
if (g_is_prefill) {
|
||||
// --- PREFILL PHASE ---
|
||||
// During prefill, we process the entire input sequence (T > 1).
|
||||
|
||||
// A. Store the K and V results from this layer into the cache.
|
||||
// The projections are available in 'inp' (the QKV matmul output).
|
||||
|
||||
// B. Run standard baseline attention for the initial prompt tokens.
|
||||
|
||||
// C. Update the global layer index to track progress through the model.
|
||||
|
||||
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);
|
||||
|
||||
float scale = 1.0f / sqrtf((float)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((B * NH * T * HS + block.x - 1) / block.x);
|
||||
attention_v_kernel<<<grid_attv, block>>>(vaccum, att, v, B, T, NH, HS);
|
||||
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
|
||||
|
||||
g_layer_idx++;
|
||||
if (g_layer_idx == 12) {
|
||||
g_layer_idx = 0;
|
||||
g_is_prefill = false; // Transition to Decode phase for the next token
|
||||
g_current_pos = T; // Set current position to the end of the prefill sequence
|
||||
}
|
||||
} else {
|
||||
// --- DECODE PHASE ---
|
||||
// During decode, we only process the single newest token (T = 1).
|
||||
|
||||
// A. Store the single new Key and Value for this step into the cache at g_current_pos.
|
||||
|
||||
// B. Run specialized Decode Attention.
|
||||
// This requires reading the history of Keys and Values from the cache
|
||||
// to compute the output for the current token.
|
||||
// TODO: Don't forget to scale your attention scores
|
||||
// TODO: Ensure the output shapes match the next layer
|
||||
|
||||
// C. Update the global state.
|
||||
|
||||
float *vaccum = inp;
|
||||
dim3 grid_dec((B * NH * HS + block.x - 1) / block.x);
|
||||
decode_attention_kernel<<<grid_dec, block>>>(vaccum, q, g_kv_cache, g_layer_idx, B, NH, HS, g_current_pos);
|
||||
unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
|
||||
|
||||
g_layer_idx++;
|
||||
if (g_layer_idx == 12) {
|
||||
g_layer_idx = 0;
|
||||
g_current_pos++; // Advance the sequence position
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#endif // __ATTENTION_KV_CACHE_TEMPLATE_CUH__
|
||||
|
||||
Reference in new issue
Block a user