{"slug": "how-did-i-write-my-first-flashattention-kernel-in-cuda", "title": "How Did I Write My First FlashAttention Kernel in CUDA ?", "summary": "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.", "body_md": "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.\n\nTransformers 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.\n\nOver 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.\n\nWhen 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.\n\n│ 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.\n\nTo 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:\n\nStandard attention computes:\n\nIn standard frameworks (naive PyTorch or CUDA implementations):\n\nFor 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.\n\nFlashAttention solves this via Kernel Fusion: it never materializes the N × N matrix in Global Memory. Everything stays on-chip in SRAM.\n\nWhy couldn’t we simply tile attention like regular matrix multiplication?Because Softmax is not naturally associative.\n\nTo compute the denominator, you traditionally need to see all N elements first. This normally requires 3 passes over the data:\n\nIn 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:\n\nTri 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′ᵢ:\n\nAs 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!\n\nBefore 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.\n\nHere is the pseudo-algorithm from the paper:\n\nDivide 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**.\n\n2. Loading blocks into SRAM:\n\nExecute 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.\n\n3. Computing local attention logits:\n\nMultiply the current query block with the key block directly in SRAM:\n\n***Sᵢⱼ = QᵢKⱼᵀ***\n\nThis produces a tiny **Bᵣ × B꜀** tile of raw attention scores rather than the massive **N × N** matrix.\n\n4. Applying the online softmax update:\n\nFind 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:\n\nmᵢⁿᵉʷ = max(mᵢ, m̃ᵢⱼ) ℓᵢⁿᵉʷ = eᵐⁱ⁻ᵐⁱⁿᵉʷ ℓᵢ + eᵐ̃ⁱʲ⁻ᵐⁱⁿᵉʷ ℓ̃ᵢⱼ\n\n5. Updating the output accumulator (**Oᵢ**):\n\nRescale the previous partial output to align with the new max, add the contribution of the current tile **P̃ᵢⱼVⱼ,** and normalize by **ℓᵢⁿᵉʷ**:\n\nOᵢ ← diag(ℓᵢⁿᵉʷ)⁻¹ (diag(ℓᵢ)eᵐⁱ⁻ᵐⁱⁿᵉʷ Oᵢ + eᵐ̃ⁱʲ⁻ᵐⁱⁿᵉʷ P̃ᵢⱼVⱼ)\n\nWrite updated **Oᵢ, mᵢ,** and **ℓᵢ** back to HBM for the next iteration.\n\nTranslating the above theoretical algorithm to physical GPU hardware requires a rigorous understanding of the CUDA Memory Hierarchy and Thread Hierarchy:\n\n2. Threads within a block: Collaborative parallelism\n\n3. Memory hierachy: Where data lives\n\nWriting CUDA kernels is humbling. Here are the real bugs I encountered and resolved:\n\nInitially, 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.\n\nWhen compiling, nvcc threw a hard fatal error:\n\n```\nptxas error : Entry function uses too much shared data (0x10000 bytes, 0xc000 max)\n```\n\nThe default hardware limit for statically allocated Shared Memory per thread block is **48 KB (****0xc000)**, but my kernel requested **64 KB (****0x10000)**.\n\n**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:\n\nThis eliminated 16 KB of SRAM, bringing total shared memory footprint to exactly **48 KB (3 × 16 KB)** and fitting inside the default hardware budget.\n\nIn my first draft, I attempted to load Key and Value rows into Shared Memory like this:\n\n```\ns_k[tid][kv_idx] = K[global_kv_row * d + kv_idx]; // Bug!\n```\n\n**The Bug:**\n\nInstead 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.\n\n**The Fix:**\n\nI replaced the single assignment with an explicit loop across all **d** channels so each thread populates its complete row:\n\n```\nfor (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];}\n```\n\nThis guarantees the entire feature vector is loaded before computing attention scores.\n\nHere is the clean, working implementation of Tiled FlashAttention:\n\n```\n#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];        }    }}\n```\n\nI evaluated the implementation on an NVIDIA GeForce RTX 4050 Laptop GPU with a sequence length of N = 1024 and head dimension d = 64:\n\n```\n 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⁻⁹\n```\n\nWriting this kernel taught me three core lessons:\n\nIf 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.\n\nHappy coding!\n\n[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.", "url": "https://wpnews.pro/news/how-did-i-write-my-first-flashattention-kernel-in-cuda", "canonical_source": "https://pub.towardsai.net/how-did-i-write-my-first-flashattention-kernel-in-cuda-bb5727e66ce5?source=rss----98111c9905da---4", "published_at": "2026-10-12 04:19:13+00:00", "updated_at": "2026-10-12 04:38:58.870345+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-research"], "entities": ["NVIDIA RTX 4050", "CUDA", "FlashAttention", "Tri Dao", "Milakov", "Gimelshein", "Google Antigravity", "PyTorch"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/how-did-i-write-my-first-flashattention-kernel-in-cuda", "markdown": "https://wpnews.pro/news/how-did-i-write-my-first-flashattention-kernel-in-cuda.md", "text": "https://wpnews.pro/news/how-did-i-write-my-first-flashattention-kernel-in-cuda.txt", "jsonld": "https://wpnews.pro/news/how-did-i-write-my-first-flashattention-kernel-in-cuda.jsonld"}}