Building Self-Attention and Transformers from Scratch: Tensor Math, FlashAttention, and CUDA Kernels

1. The Conceptual Shift: From RNNs to Parallel Self-Attention

Before Transformers, natural language processing relied on Recurrent Neural Networks (RNNs), LSTMs, and GRUs. An RNN processes input tokens sequentially, updating a hidden state vector $h_t$ at each step $t$:

RNNs suffered from two critical flaws: sequential compute dependency (preventing GPU parallelization across sequence length $N$) and vanishing/exploding gradients across long contexts. Transformers solved both by replacing recurrence with pairwise matrix operations.

2. Mathematical Formulation of Scaled Dot-Product Attention

Given an input token sequence matrix $\mathbf{X} \in \mathbb{R}^{N imes d_{ ext{model}}}$, where $N$ is the sequence length and $d_{ ext{model}}$ is the hidden dimension, we project $\mathbf{X}$ into three separate latent spaces using learned weight matrices $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{d_{ ext{model}} imes d_k}$:

• Query ($\mathbf{Q}$): What information the current token is searching for.
• Key ($\mathbf{K}$): The indexing label that each token advertises to queries.
• Value ($\mathbf{V}$): The actual information content conveyed by each token.

The Scaled Dot-Product Attention formula calculates the affinity matrix between Queries and Keys, scales the dot products, applies a softmax normalization, and computes a weighted average of Values:

3. Step-by-Step Numerical Walkthrough: 2-Token Self-Attention Trace

Let's trace a worked numerical example with sequence length $N=2$ tokens ("AI", "Tools") and dimension $d_k=2$:

Input Embeddings X:
  Token 1 ("AI"):    [1.0, 2.0]
  Token 2 ("Tools"): [3.0, 0.0]

Projection Weights (W_q, W_k, W_v = Identity for simplicity):
  Q = X * W_q = [[1.0, 2.0], [3.0, 0.0]]
  K = X * W_k = [[1.0, 2.0], [3.0, 0.0]]
  V = X * W_v = [[0.5, 1.5], [2.0, 1.0]]

Step 1: Compute Raw Dot Product Scores S = Q * K^T
  S_11 = (1*1 + 2*2) = 5.0    S_12 = (1*3 + 2*0) = 3.0
  S_21 = (3*1 + 0*2) = 3.0    S_22 = (3*3 + 0*0) = 9.0
  S = [[5.0, 3.0], [3.0, 9.0]]

Step 2: Scale by 1 / sqrt(d_k) = 1 / sqrt(2) ≈ 0.7071
  S_scaled = [[3.5355, 2.1213], [2.1213, 6.3639]]

Step 3: Softmax across rows
  Row 1: exp(3.5355)=34.31, exp(2.1213)=8.34 -> Sum = 42.65
  Softmax Row 1 = [0.8045, 0.1955]

  Row 2: exp(2.1213)=8.34, exp(6.3639)=580.50 -> Sum = 588.84
  Softmax Row 2 = [0.0142, 0.9858]

Step 4: Multiply by Values V -> Output Matrix O = Softmax * V
  O_1 = 0.8045 * [0.5, 1.5] + 0.1955 * [2.0, 1.0] = [0.793, 1.402]
  O_2 = 0.0142 * [0.5, 1.5] + 0.9858 * [2.0, 1.0] = [1.979, 1.007]

4. Multi-Head Attention (MHA)

Single-head attention averages relationships across all tokens into a single subspace. Multi-Head Attention (MHA) splits the model dimension $d_{ ext{model}}$ into $h$ parallel heads, each with dimension $d_k = d_{ ext{model}} / h$:

5. Causal Masking (Decoder Self-Attention)

In autoregressive language modeling (GPT-style decoders), a model generating token at position $t$ must NOT be allowed to attend to future tokens $t+1, t+2, \dots, N$. To enforce causality, we add an Upper Triangular Causal Mask $\mathbf{M}$ filled with $-\infty$ above the diagonal before taking the softmax:

Since $e^{-\infty} = 0$, the softmax assigns exactly zero weight to future tokens!

6. Positional Encodings: Sinusoidal vs. RoPE (Rotary Position Embedding)

Self-attention is completely permutation-invariant—swapping token positions produces identical outputs. To inject sequence order, positional information must be added.

6.1 Vaswani Sinusoidal Positional Encoding

The original Transformer used fixed sine and cosine functions of varying frequencies:

6.2 Rotary Position Embedding (RoPE)

Modern LLMs (Llama 3, Mistral) use RoPE (Rotary Position Embedding). Instead of adding a vector to token embeddings, RoPE rotates the 2D Query and Key vector pairs in the complex plane by an angle proportional to position $m$:

RoPE ensures that the inner product $\langle \mathbf{R}_m \mathbf{q}_m, \mathbf{R}_n \mathbf{k}_n angle$ depends solely on the relative distance $m-n$, enabling robust context extrapolation!

7. Complete PyTorch Implementation from Scratch

Below is a clean, production-grade PyTorch implementation of a complete Causal Multi-Head Attention block:

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, max_seq_len: int = 2048):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Key, Query, Value projections combined into one linear layer
        self.c_attn = nn.Linear(d_model, 3 * d_model, bias=False)
        # Output projection
        self.c_proj = nn.Linear(d_model, d_model, bias=False)

        # Causal mask buffer
        mask = torch.full((max_seq_len, max_seq_len), float("-inf"))
        mask = torch.triu(mask, diagonal=1)
        self.register_buffer("causal_mask", mask)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.size() # Batch, Time-steps (seq_len), Channels (d_model)

        # Calculate Q, K, V
        qkv = self.c_attn(x) # (B, T, 3 * d_model)
        q, k, v = qkv.split(self.d_model, dim=2)

        # Reshape for multi-head attention: (B, n_heads, T, d_k)
        q = q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # Raw attention scores Q * K^T / sqrt(d_k)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.d_k)) # (B, n_heads, T, T)

        # Apply causal mask
        att = att + self.causal_mask[:T, :T]

        # Softmax & dropout
        att = F.softmax(att, dim=-1)

        # Multiply by V
        y = att @ v # (B, n_heads, T, d_k)

        # Re-assemble heads into (B, T, d_model)
        y = y.transpose(1, 2).contiguous().view(B, T, C)

        # Output projection
        return self.c_proj(y)

8. Layer Normalization & Residual Connections: Pre-LN vs. Post-LN

Transformer blocks rely on residual connections ($\mathbf{x} + ext{SubLayer}(\mathbf{x})$) and Layer Normalization. The placement of LayerNorm defines two distinct architectures:

• Post-LN (Original 2017 Paper): Applies LayerNorm AFTER the residual addition: $\mathbf{x}_{l+1} = ext{LN}(\mathbf{x}_l + ext{Attention}(\mathbf{x}_l))$. Requires warmup learning rates; unstable for deep networks ($>30$ layers).
• Pre-LN (Modern Standard - GPT-3, Llama): Applies LayerNorm BEFORE the sub-layer: $\mathbf{x}_{l+1} = \mathbf{x}_l + ext{Attention}( ext{LN}(\mathbf{x}_l))$. Allows unimpeded gradient flow through the residual identity stream, enabling training of 100B+ parameter models!

9. Hardware Memory Bottlenecks: HBM vs. SRAM in Attention

Why is standard self-attention slow on GPUs? The $O(N^2)$ attention matrix $\mathbf{S} = \mathbf{Q} \mathbf{K}^T$ must be materialized in High Bandwidth Memory (HBM). For sequence length $N=8192$, storing $\mathbf{S}$ for 32 heads in FP16 consumes:

GPU compute units (Tensor Cores) run thousands of times faster than HBM read/write bandwidth. Standard attention is **memory-bandwidth bound**, spending 90% of execution time copying $N imes N$ matrices to and from HBM.

10. FlashAttention: Tiling and Online Softmax Kernel

Designed by Tri Dao et al., FlashAttention restructures attention to be memory-aware. It avoids writing the $N imes N$ intermediate matrix to HBM by computing attention in small blocks that fit entirely inside high-speed GPU **SRAM** (L1 Cache):

10.1 The Online Softmax Algorithm

To compute softmax incrementally over tiled blocks without seeing all $N$ tokens at once, FlashAttention tracks running row maximums $m(x)$ and normalizers $l(x)$:

FlashAttention reduces memory I/O from $O(N^2)$ to $O(N)$, achieving 2x to 4x speedups in training and inference!

11. Inference Optimizations: KV-Caching

During autoregressive inference, generating one token at a time requires computing Key and Value projections for all previous tokens. Without caching, step $T$ recomputes keys and values for tokens $1 \dots T-1$, leading to $O(T^2)$ redundant operations.

KV-Caching stores previously computed $K$ and $V$ tensors in GPU RAM. At step $T$, the model projects $Q, K, V$ only for the single new incoming token, appends the new $K, V$ to the cache, and computes attention over the cached history in $O(T)$ time!

12. Multi-Query Attention (MQA) & Grouped-Query Attention (GQA)

While KV-caching speeds up compute, huge KV-caches exhaust GPU memory when serving thousands of concurrent users. Two variants solve this memory pressure:

• Multi-Query Attention (MQA): All Query heads share a single Key and Value head. Reduces KV-cache size by $h imes$ (e.g. 32x), but can slightly degrade model quality.
• Grouped-Query Attention (GQA - Llama 3): Query heads are divided into $G$ groups (e.g. 8 groups of 4 Q-heads). Each group shares one K/V head. Offers 8x KV-cache compression with zero loss in accuracy!

13. SwiGLU Feed-Forward Networks (FFN)

In addition to self-attention, every Transformer block contains a Position-Wise Feed-Forward Network (FFN). Modern LLMs use SwiGLU (Swish Gated Linear Unit) activation instead of traditional ReLU:

SwiGLU provides smooth gradient flow and allows the model to act as a dynamic key-value memory store for factual knowledge learned during pre-training.

14. Softmax Derivation & Gradient Computation in Backpropagation

Understanding gradient propagation through the attention layer is essential for custom CUDA kernel design. Given row vector $\mathbf{s}$, the softmax function is defined as:

The Jacobian derivative of $P_i$ with respect to logit $S_k$ is given by:

where $\delta_{ik}$ is the Kronecker delta. When backpropagating upstream loss gradients $\mathbf{g} = rac{\partial L}{\partial \mathbf{P}}$, the gradient with respect to raw score vector $\mathbf{S}$ simplifies to:

This formulation allows FlashAttention backward passes to recompute gradients in SRAM without storing intermediate activations in HBM!

15. Cross-Attention in Encoder-Decoder Architectures (T5, Whisper)

While Decoder-only LLMs use causal self-attention, Encoder-Decoder models (such as T5 or Whisper for speech recognition) utilize Cross-Attention to bridge the prompt encoder and response decoder:

In cross-attention, Queries originate from the generated decoder states, while Keys and Values are extracted from the fixed encoder output representations $\mathbf{H}_{ ext{encoder}}$. This allows the decoder to query the full source sequence at every generation step.

16. Low-Rank Adaptation (LoRA / QLoRA) for Self-Attention

Fine-tuning all parameters of a 70B Transformer is computationally prohibitive. LoRA (Low-Rank Adaptation) freezes the pre-trained weight matrices $\mathbf{W}_0 \in \mathbb{R}^{d imes k}$ and decomposes weight updates into two low-rank matrices $\mathbf{A} \in \mathbb{R}^{r imes k}$ and $\mathbf{B} \in \mathbb{R}^{d imes r}$ with rank $r \ll d$:

/* PyTorch LoRA Injection for Query Projection Linear Layer */
class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, r=8, alpha=16):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=False)
        self.lora_A = nn.Parameter(torch.zeros(r, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
        self.scaling = alpha / r
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    def forward(self, x):
        result = self.linear(x)
        lora_update = (x @ self.lora_A.T) @ self.lora_B.T
        return result + lora_update * self.scaling

LoRA reduces trainable parameter counts by over 99%, allowing full fine-tuning of Llama 3 on consumer GPUs!

17. PagedAttention and vLLM Memory Architecture

In high-throughput LLM serving systems like vLLM, traditional KV-cache memory allocation suffers from severe fragmentation (up to 80% wasted RAM). Inspired by operating system virtual memory paging, PagedAttention breaks continuous KV-caches into non-contiguous fixed-size physical memory pages (e.g. 16 tokens per block):

18. Mixture of Experts (MoE) Architecture (Mixtral 8x7B)

To scale parameter count without scaling per-token compute costs, modern architectures use Mixture of Experts (MoE). Instead of passing every token through a dense Feed-Forward Network (FFN), a gating router network computes top-$k$ expert routing probabilities:

In Mixtral 8x7B, each token is routed to 2 out of 8 experts per layer, giving the capacity of 47 Billion parameters while using the compute cost of a 13 Billion parameter model!

19. Speculative Decoding & Fast Token Generation

LLM generation is heavily memory-bandwidth bound during autoregressive decoding. Speculative Decoding uses a fast, small "Draft Model" (e.g. 1B parameter) to generate a candidate sequence of $K$ tokens quickly. The large target model (e.g. 70B parameter) runs a single parallel forward pass over all $K$ tokens in SRAM, accepting or rejecting tokens via target probability verification. This speeds up generation by 2x to 3x without changing output distribution!

20. Tensor Parallelism (Megatron-LM) vs. Pipeline Parallelism

When a model exceeds the memory of a single GPU (e.g. 70B parameter model in FP16 requires 140GB RAM), model parallelism splits weights across GPUs:

• Tensor Parallelism (Column & Row Parallel): Splits linear layer matrices across GPUs within the same Transformer block. Query/Key/Value weight matrices are split column-wise, while the output projection matrix is split row-wise, requiring `All-Reduce` communication over NVLink after every layer.
• Pipeline Parallelism (DeepSpeed): Assigns consecutive groups of Transformer layers to different GPUs in a pipeline chain. Inter-GPU communication is limited to passing activation vectors between layer boundaries.

21. Attention Sinks & StreamingLLM: Infinite Context Window Generation

When generating ultra-long text streams beyond trained context limits, standard sliding window KV-caching fails, causing perplexity spikes. Researchers discovered the Attention Sink phenomenon: self-attention softmax allocates massive attention scores ($pprox 40\%$) to the first 4 initial tokens of a sequence regardless of their semantic content, simply to act as a numerical normalization sink. StreamingLLM preserves the first 4 initial tokens alongside a sliding window KV-cache, enabling models to generate millions of tokens continuously with stable perplexity!

22. Compute and Memory Estimation: The FLOPs Formula

For a Transformer with $P$ parameters trained on $D$ tokens, the total floating-point operations (FLOPs) required during forward and backward passes is governed by Kaplan's rule of thumb:

For inference forward pass per token, the FLOPs required is approximately $2 \cdot P$ FLOPs per token. Knowing these metrics allows systems engineers to calculate exact GPU cluster requirements and compute-to-memory throughput ratios!

23. Visualizing Attention Weights in Python with Matplotlib

To inspect what token relationships an attention head has learned, developers visualize the attention matrix $\mathbf{A} \in \mathbb{R}^{T imes T}$ as a heatmap:

import matplotlib.pyplot as plt
import seaborn as sns

def plot_attention_matrix(tokens, attention_weights_2d):
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        attention_weights_2d.detach().cpu().numpy(),
        xticklabels=tokens,
        yticklabels=tokens,
        cmap="Blues",
        annot=True,
        fmt=".2f"
    )
    plt.xlabel("Key Tokens")
    plt.ylabel("Query Tokens")
    plt.title("Self-Attention Weight Distribution Heatmap")
    plt.tight_layout()
    plt.savefig("attention_heatmap.png")

24. Vision Transformers (ViT): Adapting Self-Attention for Computer Vision

The self-attention mechanism is not limited to text processing. Dosovitskiy et al. introduced Vision Transformers (ViT) by reshaping a 2D image $\mathbf{I} \in \mathbb{R}^{H imes W imes C}$ into a sequence of flattened 2D patches $\mathbf{x}_p \in \mathbb{R}^{N imes (P^2 \cdot C)}$, where $(P, P)$ is the spatial patch resolution (e.g., $16 imes 16$ pixels):

Each patch is linearly projected into dimension $d_{ ext{model}}$, prepended with a learnable `[CLS]` classification token, and passed through standard Transformer encoder blocks. ViT achieves state-of-the-art accuracy on ImageNet without using convolutional layers!

25. Mathematical Proof of RoPE Relative Position Invariance

Let $\mathbf{R}_{\Theta, m}$ be the 2D complex rotation matrix for position $m$. The inner product between a Query at position $m$ and a Key at position $n$ is calculated as:

Because rotation matrices are orthogonal ($\mathbf{R}^T( heta) = \mathbf{R}(- heta)$) and additive under composition ($\mathbf{R}(lpha) \mathbf{R}(eta) = \mathbf{R}(lpha + eta)$), we have:

Thus, the dot product depends ONLY on the relative distance $(n-m)$, proving RoPE's exact mathematical relative positional invariance!

26. Complete Autoregressive Generation Loop with KV-Cache in PyTorch

Below is a working implementation of an efficient autoregressive text generation loop utilizing KV-caching:

@torch.no_grad()
def generate_autoregressive(model, prompt_ids: torch.Tensor, max_new_tokens: int = 50):
    B, T = prompt_ids.size()
    kv_cache = None
    curr_input = prompt_ids

    for step in range(max_new_tokens):
        # Forward pass returning logits and updated KV cache
        logits, kv_cache = model(curr_input, kv_cache=kv_cache)
        
        # Select last token logits
        next_token_logits = logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        
        # Append to generated sequence
        prompt_ids = torch.cat([prompt_ids, next_token], dim=1)
        # Next input is ONLY the newly generated single token!
        curr_input = next_token
        
    return prompt_ids

27. Complete Decoder-Only Transformer Architecture Definition in PyTorch

Below is a production-grade, modular implementation of a complete Llama-style Decoder-Only Transformer Block:

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        norm_x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm_x * self.weight

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, ffn_dim: int):
        super().__init__()
        self.attn_norm = RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ffn_norm = RMSNorm(d_model)
        
        # SwiGLU Feed-Forward Network
        self.w1 = nn.Linear(d_model, ffn_dim, bias=False)
        self.w2 = nn.Linear(ffn_dim, d_model, bias=False)
        self.w3 = nn.Linear(d_model, ffn_dim, bias=False)

    def forward(self, x):
        # Pre-LN Self Attention Residual Connection
        h = x + self.attn(self.attn_norm(x))
        
        # SwiGLU FFN Residual Connection
        ffn_in = self.ffn_norm(h)
        ffn_out = self.w2(F.silu(self.w1(ffn_in)) * self.w3(ffn_in))
        return h + ffn_out

28. Complete PyTorch Training Loop with Mixed Precision (AMP)

Below is a full PyTorch training loop incorporating `torch.cuda.amp.autocast` for FP16 training of Transformer blocks:

# PyTorch Mixed Precision Training Step
model = CausalSelfAttention(d_model=768, n_heads=12).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
scaler = torch.cuda.amp.GradScaler()

x = torch.randint(0, 50257, (8, 512)).cuda() # Synthetic token batch
embedding = nn.Embedding(50257, 768).cuda()

for step in range(100):
    optimizer.zero_grad()
    
    with torch.cuda.amp.autocast(dtype=torch.float16):
        hidden = embedding(x)
        output = model(hidden)
        loss = output.sum() # Dummy loss objective
        
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

29. Architectural Comparison Matrix of Attention Mechanisms

30. Production Engineering Summary Checklist

• Use Pre-LN: Always place LayerNorm before self-attention blocks to ensure stable gradient flow during deep network training.
• Leverage GQA: Replace standard MHA with Grouped-Query Attention to shrink KV-cache footprint by 75% during inference serving.
• Enable FlashAttention: Use F.scaled_dot_product_attention in PyTorch 2.0+ to automatically invoke FlashAttention-2 kernels.
• Implement RoPE: Use Rotary Position Embeddings to support seamless context window extensions (e.g. 32K to 128K tokens).

31. Developer FAQ

Written by Professor Pixel · CodingPancake AI Engineering Series