87 lines
3.5 KiB
Plaintext
87 lines
3.5 KiB
Plaintext
#ifndef __ATTENTION_KV_CACHE_TEMPLATE_CUH__
|
|
#define __ATTENTION_KV_CACHE_TEMPLATE_CUH__
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
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
|
|
|
|
// --- Suggested Kernel Placeholders (TODO: Add your regular attention kernel here, and any other kernels you may need) ---
|
|
|
|
/**
|
|
* 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(/* args */) {
|
|
// Implement logic to copy K and V from the input QKV buffer into g_kv_cache
|
|
}
|
|
|
|
/**
|
|
* 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(/* args */) {
|
|
// Implement T=1 attention logic using the KV cache
|
|
}
|
|
|
|
void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, int T, int C, int NH) {
|
|
// 1. Check if KV-caching is enabled, used by the verification script
|
|
if (!g_enable_kv_cache) {
|
|
// TODO: Call your standard Baseline Attention here (Milestone 1/2)
|
|
return;
|
|
}
|
|
|
|
// 2. Initialize KV Cache if necessary
|
|
if (g_kv_cache == nullptr) {
|
|
// TODO: Allocate the KV cache, based accordingly to the model parameters
|
|
// Option 1: Allocate a very big kv cache to store all the tokens
|
|
// Option 2: Dynamically resize the kv cache repeatedly for each new token
|
|
}
|
|
|
|
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.
|
|
g_layer_idx++;
|
|
if (g_layer_idx == 12) {
|
|
g_layer_idx = 0;
|
|
g_is_prefill = false; // Transition to Decode phase for the next token
|
|
}
|
|
} 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.
|
|
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__
|