Transformer Architecture Deconstructed: Attention, Geometric Scaling, and Memory IO

The physical barrier modern large language models confront when scaling context length is rarely raw floating-point throughput (FLOPs). It is the memory-access wall between high-bandwidth memory (HBM) and on-chip static random-access memory (SRAM). The standard Scaled Dot-Product Attention operator established by Vaswani et al. (2017) incurs an $O(N^2)$ memory-traffic overhead as sequence length $N$ grows, reducing self-attention primarily to a memory-bound workload.

Understanding the Transformer requires stripping away anthropomorphic cognitive metaphors and grounding the discussion in the geometry of matrix projections alongside the hardware constraints of modern GPUs.

The Physical Reality: Compute-Bound or Memory-Bound?

#

In GPU architecture, operator throughput is dictated by arithmetic intensity: the number of floating-point operations performed per byte of memory transferred (FLOPs/Byte).

A modern accelerator such as the NVIDIA A100 SXM4 40GB delivers 312 TFLOPS of FP16 Tensor Core compute, paired with 1,555 GB/s of HBM2 bandwidth. The hardware balance point (critical intensity) is:

When an operator’s arithmetic intensity falls below this threshold, compute units starve while waiting on memory transfers, placing the workload in a memory-bound regime.

Standard self-attention is constructed from three distinct tensor operations:

1. Attention score computation: $S = Q K^T$
2. Scaling and normalization: $P = \text{softmax}(S / \sqrt{d_k})$
3. Context aggregation: $O = P V$

Let batch size be $B$, head count be $h$, sequence length be $N$, and head dimension be $d_k$. In Step 2, the Softmax operator reads an intermediate matrix $S$ of dimensions $B \times h \times N \times N$ from HBM into SRAM, evaluates exponentiation and division, and writes back matrix $P$ of equal size.

This step performs roughly $O(B \cdot h \cdot N^2)$ floating-point operations while moving $O(B \cdot h \cdot N^2)$ floats across the memory bus. Its arithmetic intensity is merely a few FLOPs per byte—well below the hardware threshold of 200 FLOPs/Byte.

In long-sequence regimes, the vast majority of execution latency is spent round-tripping this quadratic matrix across HBM, while Tensor Cores idle.

The Geometry of Q/K/V Projections and Dot-Product Addressing

#

Self-attention projects input token embeddings into distinct representational subspaces via three learnable linear transformation matrices.

Given an input sequence matrix $X \in \mathbb{R}^{N \times d_{\text{model}}}$, the three projections are defined as:

where $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$.

From Additive Attention to General Matrix Multiplies

#

Prior to Vaswani et al. (2017), sequence-to-sequence models predominantly relied on additive attention (Bahdanau et al., 2014):

Additive attention deploys a single-hidden-layer feedforward network. While dot-product attention shares similar theoretical asymptotic complexity, its engineering performance diverges drastically on accelerator hardware: dot products map directly to General Matrix Multiplies (GEMM), executing natively on Tensor Cores or systolic arrays at peak pipeline utilization, bypassing the nonlinear activation branches required by additive mechanisms.

Dynamic Addressing in Inner-Product Spaces

#

The inner product $q_i \cdot k_j^T = |q_i| |k_j| \cos \theta_{ij}$ evaluates directional alignment and magnitude projection across high-dimensional geometry:

• $Q$ (Query): A semantic probe encoding what context a given token requires.
• $K$ (Key): A structural index defining what features a token offers.
• $V$ (Value): The content payload incorporated into the target token once an alignment match occurs.

This mechanism implements differentiable soft addressing: rather than pulling from memory via discrete pointers, the model extracts a convex combination of Value vectors weighted by probability distributions derived from inner-product affinities.

Softmax Scaling Factor $\sqrt{d_k}$: Derivation and Numerical Stability

#

Vaswani et al. (2017) introduced the scaling factor $\frac{1}{\sqrt{d_k}}$ into the attention formulation:

This term counteracts variance inflation in high-dimensional dot products and preserves gradients across the Softmax operator.

Derivation of Variance Growth with Dimension $d_k$

#

Assume the components $q_m, k_m$ ($m = 1, \dots, d_k$) of Query and Key vectors are mutually independent random variables with zero mean and unit variance:

Consider their dot product $z = q \cdot k = \sum_{m=1}^{d_k} q_m k_m$.

By independence:

The variance of each component product is:

Summing over all $d_k$ independent dimensions:

The standard deviation of $z$ scales as $\sigma = \sqrt{d_k}$. For standard dimensions such as $d_k = 64$ or $128$, raw dot-product magnitudes disperse widely, generating extreme positive and negative values.

Gradient Vanishing in the Softmax Saturation Zone

#

The partial derivative of the Softmax function with respect to input logit $z_i$ is:

When an entry $z_i$ substantially exceeds other logits, $S_i \to 1$ while $S_j \to 0$ for all $j \neq i$.

Substituting these values into the Jacobian:

• For index $i$: $S_i (1 - S_i) \approx 1 \times (1 - 1) = 0$
• For other indices $j$: $S_j (0 - S_j) \approx 0$

The operator enters a saturated regime where the Jacobian elements vanish, collapsing backpropagation gradients to zero.

Dividing by $\sqrt{d_k}$ normalizes the variance back to 1:

This keeps logits inside the high-gradient region of the Softmax curve, maintaining stable gradient propagation during optimization.

Multi-Head Attention: Orthogonal Subspaces and Rank Collapse

#

Convex Hull Constraints in Single-Head Attention

#

Single-head self-attention computes an output that is strictly a convex combination:

Every aggregated representation $y_i$ is confined within the convex hull of the input set ${v_1, \dots, v_N}$.

When a sequence exhibits multiple orthogonal dependencies simultaneously (such as coreference resolution, subject-verb agreement, and temporal clauses), a single Softmax distribution is dominated by the single largest logit, suppressing subtle contextual features.

Rank Collapse in Deep Transformer Stacks

#

Dong et al. (2021) demonstrated in Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth that without residual connections and multi-layer perceptrons, stacked self-attention layers suffer from rapid rank collapse:

Output rows converge toward an identical vector, collapsing the matrix rank to 1. Tokens lose representational variance, destroying the model’s capacity to represent complex sequential structures.

Multi-Head Projections as Subspace Decomposition

#

Vaswani et al. (2017) addressed this geometric limitation via Multi-Head Attention (MHA):

The model splits $d_{\text{model}}$ across $h$ subspaces, each of dimension $d_k = d_{\text{model}} / h$.

Each head projects onto an independent geometric hyperplane:

1. Head 1 captures local morphological collocations.
2. Head 2 tracks long-range syntactic dependencies.
3. The linear projection $W^O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ synthesizes representations from these orthogonal subspaces.

This decomposition expands the expressive capacity of the network without inflating overall compute requirements.

Breaking the Memory Wall: FlashAttention and SRAM Tiling

#

Despite its mathematical elegance, standard self-attention encounters an inescapable bottleneck in practice: the $O(N^2)$ physical memory footprint.

HBM Bottlenecks in Standard Attention

#

GPU memory is organized hierarchically:

• SRAM: On-chip cache (192KB–228KB per SM on A100), delivering over 19 TB/s aggregate bandwidth, but strictly limited in capacity.
• HBM: Off-chip memory (40GB–80GB), providing 1.5–3.3 TB/s bandwidth with substantially higher latency.

Standard execution repeatedly reads and writes large matrices to HBM:

1
2
3

[HBM] Q, K  ---> [SRAM] Compute QK^T      ---> [HBM] Write S (O(N^2) memory)
[HBM] S     ---> [SRAM] Compute Softmax   ---> [HBM] Write P (O(N^2) memory)
[HBM] P, V  ---> [SRAM] Compute PV        ---> [HBM] Write O

For context length $N = 32,768$, batch size $B = 2$, and head count $h = 32$, a single intermediate matrix $S$ in FP16 consumes:

This intermediate allocation exceeds the capacity of an 80GB GPU.

Online Softmax and Blocked Incremental Computation

#

Tri Dao et al. (2022) introduced FlashAttention, restructuring the operator around two core principles: never materialize the $N \times N$ attention matrix to HBM, and maintain execution entirely within SRAM via tiling and online Softmax formulation.

Standard Softmax computes a row-wide maximum to prevent arithmetic overflow:

This introduces an all-to-all dependency across the entire row. Milakov & Gimelshteyn (2018) and Tri Dao et al. (2022) resolved this dependency by formulating an incremental recurrence.

Split a row into two blocks $x^{(1)}$ and $x^{(2)}$. Given the local statistics of the first block:

When the second block $x^{(2)}$ arrives, update the global maximum and denominator:

For the accumulated output vector $O$:

Through this transformation, the GPU loads $Q$ blocks into SRAM and streams $K, V$ blocks through the cache, performing local dot products, online denominator updates, and Value accumulations entirely within SRAM.

No intermediate $N \times N$ matrix ever touches HBM.

Recomputation Strategy in Backward Pass

#

Traditional backpropagation requires caching intermediate activations from the forward pass. Retaining $P \in \mathbb{R}^{N \times N}$ would preserve the $O(N^2)$ footprint.

FlashAttention makes a deliberate engineering trade-off: the forward pass stores only the per-row scalar statistics ($m$ and $d$, occupying $O(N)$ memory) and discards $P$ entirely.

During the backward pass, tiled blocks of attention scores are recomputed on the fly within SRAM using $Q, K, V$ and the cached scalar normalizers.

While this adds roughly 15% more FLOPs, eliminating HBM round-trips yields a 2–4x end-to-end speedup and shrinks memory scaling from $O(N^2)$ to $O(N)$.

Implementation Reference: Python Verification

#

The following script illustrates standard dot-product attention alongside the recurrence logic behind tiled online Softmax:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

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

def scaled_dot_product_attention_reference(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
    """
    Standard scaled dot-product attention reference (O(N^2) memory footprint).
    Q, K, V shapes: [Batch, Heads, SeqLen, HeadDim]
    """
    d_k = Q.size(-1)
    
    # 1. Scaled dot product
    # scores: [B, H, N, N]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. Softmax normalization across last dimension
    attn_weights = F.softmax(scores, dim=-1)
    
    # 3. Value aggregation
    # output: [B, H, N, HeadDim]
    output = torch.matmul(attn_weights, V)
    return output

def online_softmax_tiling_step(q_block: torch.Tensor, k_block: torch.Tensor, v_block: torch.Tensor,
                               prev_m: torch.Tensor, prev_d: torch.Tensor, prev_acc: torch.Tensor):
    """
    Single-tile online recurrence step (core mathematical basis of FlashAttention).
    """
    d_k = q_block.size(-1)
    local_scores = torch.matmul(q_block, k_block.transpose(-2, -1)) / math.sqrt(d_k)
    
    curr_max = torch.max(local_scores, dim=-1, keepdim=True).values
    new_m = torch.maximum(prev_m, curr_max)
    
    alpha = torch.exp(prev_m - new_m)
    beta = torch.exp(curr_max - new_m)
    
    exp_scores = torch.exp(local_scores - new_m)
    new_d = prev_d * alpha + torch.sum(exp_scores, dim=-1, keepdim=True)
    
    new_acc = prev_acc * alpha + torch.matmul(exp_scores, v_block)
    
    return new_m, new_d, new_acc

Production Engineering Trade-offs and Architectural Variants

#

FlashAttention resolves memory bottlenecks during training and prefill phases. During autoregressive decoding, however, the workload shifts into another operational regime.

The KV Cache Memory Wall in Inference

#

During streaming generation, the model consumes a single new token per step, producing a query vector $Q \in \mathbb{R}^{1 \times d_k}$. The workload transitions from compute-dense GEMM to memory-latency-bound GEMV (General Matrix-Vector Multiplication).

To prevent redundant computation across preceding context, keys and values from earlier tokens are retained in memory: the KV Cache.

The memory footprint per request is:

For LLaMA-3-70B (80 layers, 64 heads, $d_k = 128$) at a sequence length of 8,192 tokens:

At this scale, a single concurrent sequence absorbs over a quarter of an 80GB GPU’s memory capacity.

Structural Trade-offs: MQA and GQA

#

To maintain serving concurrency in production, modern architectures adopt architectural compromises:

References and Related Articles

#

• For GPU memory hierarchy architecture and CUDA access optimization, see GPU Acceleration and CUDA Programming Guide.
• For long-context state management and KV Cache mitigation in production agents, see Claude Code Context Compression and State Management.
• For vector embedding spaces and high-dimensional token representations, see AI LLM Foundations: Tokens and Vectors.