Demystifying Transformers: Scaled Dot-Product, RoPE & FlashAttention-2

1. Scaled Dot-Product Mechanics

At the core of every modern Large Language Model (Llama 3, Mistral, Gemma, GPT-4) lies the Scaled Dot-Product Attention mechanism. Given an input sequence $X \in \mathbb{R}^{N \times d_{\text{model}}}$, we project into three learned subspaces using weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d_{\text{model}} \times d_k}$:

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$

$$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} + M \right) V$$

Where:

• $Q K^T \in \mathbb{R}^{N \times N}$ represents pairwise token affinity scores. Row $i$ encodes how much token $i$ attends to all past and present tokens.
• $M \in \mathbb{R}^{N \times N}$ is the causal attention mask: $M_{ij} = 0$ for $j \le i$, and $M_{ij} = -\infty$ for $j > i$, strictly enforcing autoregressive non-anticipatory decoding.
• Multiplying by $V \in \mathbb{R}^{N \times d_v}$ computes a convex combination of value representations weighted by attention probabilities.

2. Mathematical Proof: Why the $\sqrt{d_k}$ Scaling Factor is Mandatory

Why do we divide by $\sqrt{d_k}$ rather than $d_k$ or $1$? Assume components of $q \in \mathbb{R}^{d_k}$ and $k \in \mathbb{R}^{d_k}$ are independent random variables with zero mean ($\mathbb{E}[q_i] = \mathbb{E}[k_i] = 0$) and unit variance ($\text{Var}(q_i) = \text{Var}(k_i) = 1$):

Consider their dot product $z = q \cdot k = \sum_{i=1}^{d_k} q_i k_i$. By independence:

$$\mathbb{E}[z] = \sum_{i=1}^{d_k} \mathbb{E}[q_i k_i] = \sum_{i=1}^{d_k} \mathbb{E}[q_i]\mathbb{E}[k_i] = 0$$

$$\text{Var}(z) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = \sum_{i=1}^{d_k} \big( \mathbb{E}[q_i^2 k_i^2] - (\mathbb{E}[q_i k_i])^2 \big)$$
$$\text{Var}(z) = \sum_{i=1}^{d_k} \mathbb{E}[q_i^2] \mathbb{E}[k_i^2] = \sum_{i=1}^{d_k} 1 \cdot 1 = d_k$$

The standard deviation of the dot product is $\sigma = \sqrt{d_k}$. For large head dimensions (e.g., $d_k = 128$), $z$ can take values in magnitude $|z| > 30$. Passing large inputs into the $\text{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$ pushes the exponential function into extreme saturation regions where the derivative $\frac{\partial \text{softmax}}{\partial z} \to 0$.

3. Rotary Position Embeddings (RoPE)

Early transformers used absolute sinusoidal embeddings ($x + p_i$) or learned positional vectors. These fail to extrapolate to sequences longer than training window lengths. Modern architectures utilize Rotary Position Embeddings (RoPE), which encode relative position directly through complex rotation:

For a 2D vector $(x_1, x_2)$, rotation by angle $m\theta$ at token position $m$ is given by:

$$R_{\Theta, m}^{2D} = \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{pmatrix}$$

$$\langle R_{\Theta, m} q, \ R_{\Theta, n} k \rangle = q^T R_{\Theta, m}^T R_{\Theta, n} k = q^T R_{\Theta, n - m} k$$

Notice the key property: the inner product between rotated query $m$ and key $n$ depends strictly on the relative distance $(n - m)$ rather than absolute positions. For $d_k$-dimensional vectors, RoPE decomposes the space into $d_k / 2$ orthogonal 2D subspaces with decaying frequency bands:

$$\theta_i = 10000^{-2(i-1)/d_k}, \quad i \in \left\{1, 2, \dots, \frac{d_k}{2}\right\}$$

4. Attention Evolution: MHA vs MQA vs Grouped-Query (GQA)

During autoregressive generation, past tokens must be preserved in GPU memory (the KV-Cache) so we do not recompute $K$ and $V$ at every step. In Multi-Head Attention (MHA), every query head has its own private key and value heads, consuming massive VRAM:

5. FlashAttention-2: GPU Memory Hierarchy & SRAM Tiling

Standard attention implementations are memory-bandwidth bound, not compute bound. An NVIDIA H100 GPU can perform 989 TFLOPS of FP16 tensor core arithmetic, but its High Bandwidth Memory (HBM) bandwidth is limited to 3.35 TB/s.

In naive attention, the $N \times N$ attention matrix $S = Q K^T$ is written to slow HBM, read back to compute $\text{softmax}(S)$, written again to HBM, and read a third time to multiply with $V$. This requires $O(N^2)$ expensive memory roundtrips.

6. Production PyTorch Implementation: GQA with Causal Mask

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

def apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    # x shape: [batch, seq_len, num_heads, head_dim]
    d = x.shape[-1]
    x1 = x[..., :d // 2]
    x2 = x[..., d // 2:]
    rotate_half = torch.cat((-x2, x1), dim=-1)
    return (x * cos) + (rotate_half * sin)

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model: int = 4096, num_q_heads: int = 32, num_kv_heads: int = 8, d_head: int = 128):
        super().__init__()
        self.d_model = d_model
        self.num_q_heads = num_q_heads
        self.num_kv_heads = num_kv_heads
        self.d_head = d_head
        self.num_queries_per_kv = num_q_heads // num_kv_heads
        self.scale = 1.0 / math.sqrt(d_head)

        self.q_proj = nn.Linear(d_model, num_q_heads * d_head, bias=False)
        self.k_proj = nn.Linear(d_model, num_kv_heads * d_head, bias=False)
        self.v_proj = nn.Linear(d_model, num_kv_heads * d_head, bias=False)
        self.out_proj = nn.Linear(num_q_heads * d_head, d_model, bias=False)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        B, S, _ = x.shape

        # 1. Project to Q, K, V
        q = self.q_proj(x).view(B, S, self.num_q_heads, self.d_head)
        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.d_head)
        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.d_head)

        # 2. Inject Rotary Position Embeddings
        q = apply_rotary_pos_emb(q, cos, sin)
        k = apply_rotary_pos_emb(k, cos, sin)

        # 3. Transpose to [B, num_heads, S, d_head]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # 4. Expand KV heads to match Q heads (Grouped-Query Repeat)
        k = k.repeat_interleave(self.num_queries_per_kv, dim=1)
        v = v.repeat_interleave(self.num_queries_per_kv, dim=1)

        # 5. Scaled Dot-Product with Causal Mask
        # Uses FlashAttention underlying C++ kernel if available in PyTorch 2.0+
        out = F.scaled_dot_product_attention(
            q, k, v, 
            attn_mask=None, 
            dropout_p=0.0, 
            is_causal=True
        )

        # 6. Recombine heads and final linear projection
        out = out.transpose(1, 2).contiguous().view(B, S, self.num_q_heads * self.d_head)
        return self.out_proj(out)

7. Memory Footprint & IO Throughput Benchmarks

We benchmarked attention computational throughput (TFLOPS) and KV cache VRAM on an NVIDIA A100-SXM4-80GB across sequence lengths from 4k to 64k tokens:

8. Numerical Hazards & The Attention Sink Phenomenon

Deploying long-context transformers in mission-critical environments reveals distinct numerical instability failure modes: