# How Did I Write My First FlashAttention Kernel in CUDA ? > Source: > Published: 2026-10-12 04:19:13+00:00 From demystifying the online softmax trick to hitting a **41.6x** speedup on an NVIDIA RTX 4050 GPU — a beginner’s hands-on journey into hardware-aware deep learning. Transformers have completely reshaped the landscape of artificial intelligence. At the beating heart of every Transformer model lies the self-attention mechanism. It empowers language models to capture intricate contextual relationships across sequences. Yet, self-attention is simultaneously a double-edged sword: its algorithmic complexity is **𝒪(N²)**, meaning memory consumption and computation explode quadratically as sequence lengths grow. Over the years, numerous researchers attempted to alleviate this burden using algorithmic approximation ([Sparse Attention](https://arxiv.org/abs/1904.10509), [Linformer](https://arxiv.org/abs/2006.04768), [Performer](https://arxiv.org/abs/2009.14794), etc.). While mathematically clever, many failed to dominate in practice because the quality trade-offs were unacceptable or they were not hardware-friendly. When purely algorithmic solutions hit a ceiling, researchers looked at the hardware. In 2022, [Tri Dao et al.](https://arxiv.org/abs/2205.14135) introduced FlashAttention — an exact attention algorithm that redesigns attention around GPU memory hierarchy and IO-awareness. │ Note: Throughout my learning and debugging journey, I paired with Google Antigravity as my AI pair programmer, which helped me unravel the math and pinpoint subtle GPU memory pitfalls. To understand why FlashAttention is necessary, we must look at how modern GPUs operate. A GPU is not just a massive compute engine, it is a hierarchy of memory spaces with vastly different bandwidths: Standard attention computes: In standard frameworks (naive PyTorch or CUDA implementations): For a sequence length of N = 16, 384, an N × N matrix contains over 268 million floating-point number (over 1 GB of memory for just a single attention head). The GPU spends nearly all its time waiting to transfer bytes across the memory bus rather than computing floating-point operations. It is overwhelmingly memory bound. FlashAttention solves this via Kernel Fusion: it never materializes the N × N matrix in Global Memory. Everything stays on-chip in SRAM. Why couldn’t we simply tile attention like regular matrix multiplication?Because Softmax is not naturally associative. To compute the denominator, you traditionally need to see all N elements first. This normally requires 3 passes over the data: In 2018, [Milakov & Gimelshein](https://arxiv.org/abs/1805.02867) proved that you can maintain a running maximum and running denominator dynamically. If you have processed tokens up to index i — 1 with running max mᵢ₋₁ and denominator d′ᵢ₋₁, and you encounter a new element xᵢ by following those steps: Tri Dao noticed something brilliant: we do not actually need the attention matrix A! We only care about the final output vector O = AV. By defining a surrogate running output vector o′ᵢ: As we sweep through the sequence, we continuously rescale our running output vector and blend in the new token’s contribution. When we reach the final token o′ₙ , is mathematically identical to standard attention — computed entirely in a single pass without ever writing N × N scores to VRAM! Before writing a single line of CUDA code, let’s look at the high-level algorithm. The genius of FlashAttention lies in transforming a global operation over huge N × N matrices into iterative tile-by-tile computations. Here is the pseudo-algorithm from the paper: Divide the large input matrices **Q**, **K**, and **V** from HBM into smaller blocks fitting directly into SRAM. **Q** and **O** are split into row blocks of size **Bᵣ × d**, while **K** and **V** are split into blocks of size **B꜀ × d**. Statistics vectors **m** (row maximums) and **ℓ** (row sum of exponentials) are initialized to **-∞** and **0**. 2. Loading blocks into SRAM: Execute a double loop over the tiles. The outer loop loads a column block **Kⱼ**, **Vⱼ** into fast on-chip SRAM once. The inner loop iterates over row blocks, streaming **Qᵢ**, the current accumulator **Oᵢ**, and local statistics **mᵢ**, **ℓᵢ** into SRAM to process tile by tile without writing intermediate full attention matrices to VRAM. 3. Computing local attention logits: Multiply the current query block with the key block directly in SRAM: ***Sᵢⱼ = QᵢKⱼᵀ*** This produces a tiny **Bᵣ × B꜀** tile of raw attention scores rather than the massive **N × N** matrix. 4. Applying the online softmax update: Find the local row maximum **m̃ᵢⱼ = rowmax(Sᵢⱼ)** and compute unnormalized exponentials **P̃ᵢⱼ = exp(Sᵢⱼ — m̃ᵢⱼ)** along with their row sums **ℓ̃ᵢⱼ**. Update the global row statistics on-chip using online softmax: mᵢⁿᵉʷ = max(mᵢ, m̃ᵢⱼ) ℓᵢⁿᵉʷ = eᵐⁱ⁻ᵐⁱⁿᵉʷ ℓᵢ + eᵐ̃ⁱʲ⁻ᵐⁱⁿᵉʷ ℓ̃ᵢⱼ 5. Updating the output accumulator (**Oᵢ**): Rescale the previous partial output to align with the new max, add the contribution of the current tile **P̃ᵢⱼVⱼ,** and normalize by **ℓᵢⁿᵉʷ**: Oᵢ ← diag(ℓᵢⁿᵉʷ)⁻¹ (diag(ℓᵢ)eᵐⁱ⁻ᵐⁱⁿᵉʷ Oᵢ + eᵐ̃ⁱʲ⁻ᵐⁱⁿᵉʷ P̃ᵢⱼVⱼ) Write updated **Oᵢ, mᵢ,** and **ℓᵢ** back to HBM for the next iteration. Translating the above theoretical algorithm to physical GPU hardware requires a rigorous understanding of the CUDA Memory Hierarchy and Thread Hierarchy: 2. Threads within a block: Collaborative parallelism 3. Memory hierachy: Where data lives Writing CUDA kernels is humbling. Here are the real bugs I encountered and resolved: Initially, I allocated four tiles directly in Shared Memory: s_q, s_k, s_v, and s_s (each 64 × 64 × 4 bytes = 16 KB), totaling 64 KB. When compiling, nvcc threw a hard fatal error: ``` ptxas error : Entry function uses too much shared data (0x10000 bytes, 0xc000 max) ``` The default hardware limit for statically allocated Shared Memory per thread block is **48 KB (****0xc000)**, but my kernel requested **64 KB (****0x10000)**. **The Fix:** Notice that thread tid only ever reads its own row of attention logits during the online softmax update. There was zero cross-thread communication needed for matrix S: This eliminated 16 KB of SRAM, bringing total shared memory footprint to exactly **48 KB (3 × 16 KB)** and fitting inside the default hardware budget. In my first draft, I attempted to load Key and Value rows into Shared Memory like this: ``` s_k[tid][kv_idx] = K[global_kv_row * d + kv_idx]; // Bug! ``` **The Bug:** Instead of copying the full head dimension vector (**d = 64**), each thread copied only a single float at column kv_idx. This left the remaining 63 elements filled with uninitialized garbage data in Shared Memory, completely corrupting subsequent dot-product results. **The Fix:** I replaced the single assignment with an explicit loop across all **d** channels so each thread populates its complete row: ``` for (int k = 0; k < d; k++) { s_k[tid][k] = K[global_kv_row * d + k]; s_v[tid][k] = V[global_kv_row * d + k];} ``` This guarantees the entire feature vector is loaded before computing attention scores. Here is the clean, working implementation of Tiled FlashAttention: ``` #define D 64 // Head dimension#define Q_SIZE 64 // Block size for Q (Br)#define KV_SIZE 64 // Block size for K, V (Bc)__global__ void flash_attention( const float *Q, const float *K, const float *V, float *O, int N, int d) { // 3 * 64 * 64 * 4 bytes = 48 KB (fits default Shared Memory limit) __shared__ float s_q[Q_SIZE][D]; __shared__ float s_k[KV_SIZE][D]; __shared__ float s_v[KV_SIZE][D]; int tid = threadIdx.x; // Each thread handles 1 row of Q int global_q_row = blockIdx.x * Q_SIZE + tid; int n_kv_blocks = (N + KV_SIZE - 1) / KV_SIZE; float scale = 1.0f / sqrtf((float)d); // Online Softmax statistics in registers float m_global = -1e30f; float l_global = 0.0f; float o[D] = {0.0f}; float s_row[KV_SIZE]; // Local attention logits for thread's row if (global_q_row < N) { // Step 1: Load Q tile into Shared Memory once for (int k = 0; k < d; k++) { s_q[tid][k] = Q[global_q_row * d + k]; } __syncthreads(); // Step 2: Loop over all KV blocks for (int kv_idx = 0; kv_idx < n_kv_blocks; kv_idx++) { int global_kv_row = kv_idx * KV_SIZE + tid; // Load K and V tiles into Shared Memory (with bounds check) if (global_kv_row < N) { for (int k = 0; k < d; k++) { s_k[tid][k] = K[global_kv_row * d + k]; s_v[tid][k] = V[global_kv_row * d + k]; } } else { for (int k = 0; k < d; k++) { s_k[tid][k] = 0.0f; s_v[tid][k] = 0.0f; } } __syncthreads(); // Step 3: Compute Q * K^T for thread's row for (int i = 0; i < KV_SIZE; i++) { float sum = 0.0f; for (int j = 0; j < d; j++) { sum += s_q[tid][j] * s_k[i][j]; } s_row[i] = sum * scale; } // Step 4: Find local maximum in current tile float m_local = -1e30f; for (int i = 0; i < KV_SIZE; i++) { if (s_row[i] > m_local) { m_local = s_row[i]; } } float m_new = fmaxf(m_local, m_global); // Step 5: Safe exponentiation and tile sum float l_local = 0.0f; for (int i = 0; i < KV_SIZE; i++) { float p = expf(s_row[i] - m_new); s_row[i] = p; l_local += p; } // Step 6: Online Softmax updates float alpha = (m_global <= -1e20f) ? 0.0f : expf(m_global - m_new); float l_new = l_global * alpha + l_local; // Step 7: Rescale previous output vector float rescale_old = (l_global * alpha) / l_new; for (int i = 0; i < d; i++) { o[i] *= rescale_old; } // Step 8: Accumulate (P * V) / l_new float rescale_new = 1.0f / l_new; for (int i = 0; i < KV_SIZE; i++) { float p_val = s_row[i] * rescale_new; for (int j = 0; j < d; j++) { o[j] += p_val * s_v[i][j]; } } m_global = m_new; l_global = l_new; __syncthreads(); } // Step 9: Write final output vector to Global Memory (HBM) for (int i = 0; i < d; i++) { O[global_q_row * d + i] = o[i]; } }} ``` I evaluated the implementation on an NVIDIA GeForce RTX 4050 Laptop GPU with a sequence length of N = 1024 and head dimension d = 64: ``` Implementation | Execution Time (… | Speedup vs. CPU | Numerical Difference--------------------------|-------------------|-----------------|----------------------- CPU Standard Attention | 143.81 ms | 1.0x (Baseline) | 0.0 GPU Naive FlashAttention | 7.63 ms | 18.8x | 3.02 × 10⁻⁹ GPU Tiled FlashAttention | 3.45 ms | 41.6x | 6.40 × 10⁻⁹ ``` Writing this kernel taught me three core lessons: If you are diving into CUDA or exploring how high-throughput inference engines like vLLM, TensorRT-LLM, and FlashAttention operate under the hood, there is no substitute for writing a kernel from scratch and profiling the execution yourself. Happy coding! [How Did I Write My First FlashAttention Kernel in CUDA ?](https://pub.towardsai.net/how-did-i-write-my-first-flashattention-kernel-in-cuda-bb5727e66ce5) was originally published in [Towards AI](https://pub.towardsai.net) on Medium, where people are continuing the conversation by highlighting and responding to this story.