ECE 408 sem project

kernels_base/attention_kv_cache_template.cuh

@@ -0,0 +1,86 @@
+#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__