How Did I Write My First FlashAttention Kernel in CUDA ?

A developer implemented a FlashAttention kernel in CUDA on an NVIDIA RTX 4050 GPU, achieving a 41.6x speedup over standard attention by keeping computation on-chip in SRAM instead of materializing the N×N attention matrix in global memory. The implementation follows the 2022 FlashAttention algorithm from Tri Dao et al., which uses the online softmax trick from Milakov & Gimelshein's 2018 work to compute exact attention in a single pass. For a sequence length of N = 16,384, the standard N×N matrix holds over 268 million floating-point numbers, exceeding 1 GB of memory for a single attention head.

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.