Self-attention, CUDA kernels and FlashAttention

Modern Transformers rely on Self-Attention matrix math, high-level GPU compilers like Triton, and memory-aware CUDA kernels like FlashAttention to process massive token sequences at peak hardware speed.

1. How Self-Attention Works Mathematically in a Vision Transformer

In a Vision Transformer (ViT), an image (e.g., 224×224 pixels) is sliced into a grid of 16×16 patches, producing N = 196 flattened patches. Each patch is linearly projected into a vector of dimension D (e.g., 768) and combined with a positional embedding, forming an input matrix X ∈ R^(N × D).

• Creating Queries, Keys, and Values. The model multiplies the patch matrix X by three learned weight matrices (W_Q, W_K, W_V ∈ R^(D × d_k)) to create three distinct representations for every image patch:

Q=XW_(Q), K=XW_(K), V=XW_(V)

• Query (Q): Represents what a specific image patch is looking for (e.g., "I am a fur texture, looking for eyes or ears nearby").
• Key (K): Advertises what information a patch contains (e.g., "I contain a sharp curved edge").
• Value (V): Holds the actual visual feature payload that will be passed forward if a match occurs.

• Calculating Attention Scores. To measure how strongly every patch relates to every other patch, the model computes the dot product between all Queries and Keys:

Scores=(QK^(T))/(sqrt(d_(k)))

Dividing by sqrt(d_(k)) prevents large dot products from pushing the upcoming softmax function into regions with tiny gradients. This produces an N × N attention matrix where entry (i, j) scores how much patch i should attend to patch j.

• Softmax Weighting and Aggregation. Applying row-wise Softmax converts raw scores into probabilities that sum to 1. Multiplying these probabilities by the Value matrix V produces a blended representation for each patch:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

In Multi-Head Self-Attention, this exact process runs in parallel across multiple heads (e.g., 12 heads), allowing different heads to focus on different spatial distances, colors, or object parts simultaneously.

2. Hands-On Code Comparison: Raw CUDA C++ vs. OpenAI Triton

To see how GPU programming has evolved, consider a simple element-wise vector addition (C = A + B) for N elements.

Raw CUDA C++ (Thread-Level Management)

In CUDA C++, you program from the perspective of a single scalar thread. You must manually calculate thread IDs, check array bounds, and write boilerplate host-side launch configurations.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (Block-Level Python Programming)

In Triton, you program from the perspective of a block (tile) of data directly in Python. Triton’s compiler automatically handles thread synchronization, memory coalescing, and vectorization under the hood.

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: Solving the Transformer Memory Bottleneck

Standard Self-Attention hits a severe wall when processing high-resolution images or long text sequences because the attention matrix S = QK^T grows quadratically (N × N) with sequence length N.

• The Memory Bandwidth Bottleneck (HBM vs. SRAM). Modern NVIDIA GPUs have two main memory tiers:

• HBM (High Bandwidth Memory / VRAM): Large (e.g., 80 GB) but relatively slow to read and write.
• SRAM (On-Chip Shared Memory): Tiny (~20 MB total across all Streaming Multiprocessors) but roughly 10× to 20× faster than HBM.

In standard attention, the GPU computes S = QK^T, writes the massive N × N matrix out to slow HBM, reads it back to apply Softmax, writes the probabilities back to HBM, and reads them again to multiply by V. The Tensor Cores spend most of their time waiting idle for memory transfers.

• Tiling (Never Materializing the N × N Matrix). Introduced by Tri Dao et al. (2022), FlashAttention uses kernel fusion to keep the entire attention calculation inside fast on-chip SRAM. It slices Q, K, and V into small blocks that fit inside SRAM, computes attention tile-by-tile, and writes only the final output matrix back to HBM. The giant N × N intermediate matrix is never saved to VRAM.

• Online Softmax Trick. Normally, Softmax requires knowing the maximum value and sum of exponentials across an entire row before normalizing any single element. FlashAttention uses an online Softmax algorithm that tracks running maximums and normalization factors block-by-block, rescaling partial outputs mathematically as each new tile is processed.

• Recomputation During Backpropagation. Because FlashAttention does not save the N × N attention matrix during the forward pass, it uses a clever trade-off during training: it saves only the small normalization statistics in HBM and recomputes the attention tiles on the fly inside fast SRAM during backpropagation. Doing extra math in SRAM proves much faster than reading large matrices from HBM, yielding 2× to 4× wall-clock speedups while shrinking memory usage from quadratic O(N²) to linear O(N).