Transformer Architecture Deep Dive

In late 2022, a Bengaluru legal-tech team fine-tuning a 7-billion-parameter open model to review long commercial contracts (8,000–32,000 tokens per document) kept hitting an out-of-memory crash during training. Their first instinct was to blame the GPU's compute rating — they upgraded to a card with a higher advertised FLOPS number. The crash didn't go away, and training got barely faster. The problem was never a shortage of arithmetic throughput. It was that standard self-attention, implemented the textbook way, forces the GPU to write and re-read an entire N×N matrix of attention scores to slow off-chip memory for every layer, every head, every batch — and at N = 32,000 that matrix has over a billion entries. The fix that solved this problem, FlashAttention, changed nothing about what attention computes. It changed only where the arithmetic happens and what gets written to slow memory versus fast memory. That distinction — computation versus memory movement — is the subject of this chapter's central mechanism. The second half looks at a different production lever entirely: instead of making one feed-forward block per layer, modern serving-scale models often route each token to a handful of specialist feed-forward blocks out of a much larger pool, activating only a fraction of total parameters per token.

Why a GPU can be starving for data while its cores sit idle

A GPU has two very different kinds of memory. HBM (high-bandwidth memory) is the large pool — tens of gigabytes — that holds your model's weights and activations between operations; on an NVIDIA A100 this is roughly 40 GB with a transfer rate around 1.5 TB/s. SRAM is the tiny, on-chip memory distributed across the GPU's streaming multiprocessors — on an A100, roughly 20 MB in total — but it moves data at roughly 19 TB/s, an order of magnitude faster than HBM. Every arithmetic operation a GPU core performs has to pull its operands from somewhere; if that somewhere is HBM, the core often finishes the arithmetic before the next operand has even arrived, and it sits idle waiting. This is called being memory-bandwidth-bound rather than compute-bound, and it is the normal state for attention, whose core operations (matrix multiply, elementwise softmax, another matrix multiply) each move roughly as much data as they compute arithmetic on.

Standard, unoptimized self-attention makes this worse than it needs to be. Given Q, K, V matrices each of shape N×d already sitting in HBM, the naive algorithm computes S = QKᵀ (an N×N matrix), writes all of S back to HBM, reads S back from HBM to apply softmax row-wise, writes the resulting probability matrix P back to HBM, reads P and V back from HBM to compute O = PV, and finally writes O back to HBM. Every one of those N×N-sized intermediate tensors — S and P — makes a full round trip through the slow memory pool, even though neither S nor P is ever needed again after the next step consumes it. For N = 32,000 and float32 storage, S alone is roughly 32,000² × 4 bytes ≈ 4 GB, written once and read once, on top of the same again for P — and that is for a single attention head, in a single layer, for a single sequence in the batch. This N² memory traffic, not the N² multiply-adds (which GPUs chew through easily), is what stalls training on long contexts.

FlashAttention: computing the exact same softmax without ever storing the full matrix

FlashAttention (Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré, "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS 2022) restructures the computation so the N×N score matrix is never materialized in HBM at all — not even once. The trick is to process K and V in small tiles that fit entirely inside SRAM, compute a partial attention output for each tile, and combine the partial results into a running total using an identity called online softmax. The output is mathematically identical to standard attention — this is an exact reformulation, not an approximation — because softmax and the weighted sum over values can both be expressed as a running accumulation that gets rescaled every time a new, possibly larger, maximum score is seen.

Concretely, for a query vector q attending over key/value pairs streamed in blocks, the algorithm keeps three running quantities: m (the largest score seen so far, for numerical stability), l (the running sum of exponentiated, max-shifted scores — the softmax denominator), and O (the running, unnormalized weighted sum of values — the softmax numerator). When a new tile arrives, the algorithm computes that tile's local max and local sum, figures out the new global max, rescales the *old* running l and O by a correction factor that accounts for the shift in max, adds in the new tile's contribution (itself shifted to the new max), and moves on. After the last tile, dividing O by l gives the exact softmax-weighted output — no full row of scores was ever kept around simultaneously, and nothing beyond one tile's worth of numbers ever left SRAM until the very end. The backward pass recomputes the tiles it needs rather than storing them, trading a modest amount of extra arithmetic (cheap) for a large reduction in HBM traffic (expensive) — precisely the IO-aware tradeoff the paper's title names.

Worked example: tracing online softmax by hand and confirming it matches the exact answer

Take a single query q = [1, 0] attending over four keys and values, split into two tiles of two (Block A = positions 1–2, Block B = positions 3–4), matching how a real FlashAttention kernel would tile a longer sequence:

k1 = [1, 0]   v1 = [1, 0]
k2 = [0, 1]   v2 = [0, 1]
k3 = [1, 1]   v3 = [1, 1]
k4 = [0, 0]   v4 = [2, 2]

First, the ground truth: compute standard attention directly. The raw scores are s_i = q·k_i (the 1/√d scaling is omitted here purely to keep the arithmetic legible; it rescales every score by a constant and does not change which mechanism is being demonstrated). That gives s = [1, 0, 1, 0]. Exponentiating: e¹ ≈ 2.718282 for positions 1 and 3, e⁰ = 1 for positions 2 and 4. The softmax denominator is 2(2.718282) + 2(1) = 7.436564, giving weights w1 = w3 ≈ 0.365529 and w2 = w4 ≈ 0.134471 (each pair sums to 0.5, and all four sum to 1, as required). The output is the weighted sum of the values:

Output_x = w1(1) + w2(0) + w3(1) + w4(2) = 2w1 + 2w4 = 2(0.365529 + 0.134471) = 1.000000
Output_y = w1(0) + w2(1) + w3(1) + w4(2) = w1 + 3w2 = 0.365529 + 3(0.134471) = 0.768942

So the exact answer is [1.000000, 0.768942]. Now trace the tiled, online-softmax version block by block. Initialize m = −∞, l = 0, O = [0, 0].

Block A (scores 1, 0): local max m_A = 1. Local exponentials: e^(1−1) = 1, e^(0−1) ≈ 0.367879. Local sum l_A ≈ 1.367879. Local weighted output O_A = 1·[1,0] + 0.367879·[0,1] = [1, 0.367879]. Since the running max was −∞, the update simply initializes: m = 1, l ≈ 1.367879, O ≈ [1, 0.367879].

Block B (scores 1, 0): local max m_B = 1, matching the running max, so no rescale of the old accumulator is needed this round (the correction factor e^(m_old−m_new) = e⁰ = 1). Local sum l_B ≈ 1.367879. Local weighted output O_B = 1·[1,1] + 0.367879·[2,2] ≈ [1.735759, 1.735759]. Combine: l = 1.367879 + 1.367879 ≈ 2.735759; O ≈ [1 + 1.735759, 0.367879 + 1.735759] = [2.735759, 2.103638].

Final normalization: Output = O / l = [2.735759/2.735759, 2.103638/2.735759] = [1.000000, 0.768940]. Matching the direct computation to five decimal places (the tiny residual is float rounding from carrying six digits by hand) — confirming the tiled, block-streamed computation is exact, not approximate, exactly as the mechanism above claims.

The same computation as runnable code, using the identical numbers, with every variable defined before use:

import numpy as np

def flash_attention_tile(q, K_full, V_full, block_size):
    d = q.shape[0]
    n = K_full.shape[0]
    m_running = -np.inf
    l_running = 0.0
    O_running = np.zeros(d)

    for start in range(0, n, block_size):
        K_blk = K_full[start:start + block_size]
        V_blk = V_full[start:start + block_size]

        scores = K_blk @ q                     # tile scores, shape (block_size,)
        m_blk = np.max(scores)
        p_blk = np.exp(scores - m_blk)          # local, max-shifted exponentials
        l_blk = np.sum(p_blk)
        O_blk = p_blk @ V_blk                   # local unnormalized weighted sum

        m_new = max(m_running, m_blk)
        correction_old = np.exp(m_running - m_new)   # 0.0 when m_running is -inf
        correction_new = np.exp(m_blk - m_new)
        l_running = l_running * correction_old + l_blk * correction_new
        O_running = O_running * correction_old + O_blk * correction_new
        m_running = m_new

    return O_running / l_running

q = np.array([1.0, 0.0])
K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.0, 0.0]])
V = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, 2.0]])

print(flash_attention_tile(q, K, V, block_size=2))

Since m_running starts at −∞, correction_old evaluates to np.exp(-inf) = 0.0 on the first iteration (not NaN — the difference −∞ − 1 stays −∞, and exp of −∞ is a clean 0.0 in NumPy), so the first block simply initializes the accumulators, exactly as traced by hand above. Running this prints [1. 0.76894142] — matching the hand-traced result to four decimal places.

Diagram: the tiling loop inside the GPU memory hierarchy

The misconception worth killing

The single most common misconception about FlashAttention is that it approximates attention to save memory — that it's in the same family as linear attention, Performer's kernel approximations, or sparse/local attention patterns that genuinely change which tokens attend to which. It is not. Every one of those methods changes the function being computed: linear attention replaces the softmax kernel with a different, cheaper one; sparse attention drops some token pairs from the computation entirely. FlashAttention computes precisely the same softmax-weighted sum as standard attention — the worked trace above shows the two agreeing to five decimal places, because the online-softmax rescaling is an algebraic identity, not a truncation. What changes is purely the *order and locality* of the arithmetic: tiles small enough to fit in SRAM, processed one at a time, with the mathematically necessary rescaling folded in as each new tile arrives. A student who believes FlashAttention "loses some accuracy for speed" will incorrectly hedge against using it in any setting where exactness matters (evaluation benchmarks, deterministic replay, alignment-sensitive fine-tuning) — when in fact it is a drop-in replacement for standard attention with an identical forward-pass output up to floating-point rounding, which is why it shipped inside PyTorch's scaled_dot_product_attention and inside virtually every serious training and inference stack within a year of publication.

Sparse activation: Mixture-of-Experts routing

FlashAttention makes the attention sub-layer cheaper to move through memory. A separate, independent lever attacks the feed-forward (FFN) sub-layer instead, by making most of the model's parameters sit idle for any given token. In a standard dense transformer, every token passes through the *same* FFN weights at every layer. A Mixture-of-Experts (MoE) transformer replaces that single FFN with E independently-weighted FFN "experts" plus a small router network, and routes each token to only a handful of them — commonly the top 2 — while attention sub-layers stay dense and shared as before. The idea traces to Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean's "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer" (ICLR 2017), and was adapted into the transformer FFN block by William Fedus, Barret Zoph, and Noam Shazeer's Switch Transformer (JMLR 2022).

Mechanically: the router computes gate logits g = softmax(W_r x) over the E experts for token x, keeps only the top-k logits, renormalizes those k weights to sum to 1, and the FFN output is the weighted sum of just those k experts' outputs — the other E−k experts do no work at all for that token. Because the router's top-k choice is a hard, discontinuous decision, naive training tends to collapse: whichever experts get picked early receive more gradient signal, get better, and get picked even more, until a handful of experts absorb almost all traffic while the rest sit unused. Switch Transformer's fix is an auxiliary load-balancing loss added to the training objective, of the form L_aux = α·E·Σᵢ fᵢ·Pᵢ, where fᵢ is the fraction of tokens in a batch actually routed to expert i and Pᵢ is the router's average softmax probability assigned to expert i across the batch — this loss is minimized when routing is close to uniform across experts, and α is a small weighting coefficient tuned so the balancing term doesn't dominate the main language-modeling loss.

The parameter arithmetic is the whole point, and it is worth deriving directly rather than quoting it. Take a modest FFN with model dimension d_model = 512 and an inner expansion dimension of 2048 (the standard 4× ratio), giving two weight matrices per expert of shape 512×2048 and 2048×512 — 512×2048×2 = 2,097,152 parameters per expert, ignoring biases. With E = 8 experts, the layer's total FFN parameter count is 8 × 2,097,152 ≈ 16.78 million. But with top-2 routing, only 2 × 2,097,152 ≈ 4.19 million of those parameters do any arithmetic for a given token. Compare that to a dense model built from a single FFN of the same per-expert size: it has 2.1 million total FFN parameters, and all 2.1 million are active every token. The MoE layer therefore holds 8× the parameter capacity of that single dense FFN while spending only 2× the compute per token — a 4× improvement in "learned capacity per unit of inference FLOPs," which is the entire economic argument for building MoE models at serving scale: you pay compute for k active experts, but you pay VRAM for all E of them, because routing is data-dependent per token and any expert might be needed on the next one.

The production-scale version of this is Mixtral 8x7B (Albert Q. Jiang et al., "Mixtral of Experts," Mistral AI, 2024), which applies top-2-of-8 routing to only the FFN blocks of each transformer layer while sharing attention and embedding parameters across all experts. The "8x7B" name suggests 56 billion parameters, but because attention and embeddings are shared rather than duplicated per expert, the actual total is about 46.7 billion parameters, of which only about 12.9 billion are active for any given token. That gap between the marketing-friendly name and the real parameter count is itself worth noting precisely because it is such a common point of confusion: multiplying "8 experts" by "7B per expert" overcounts every non-FFN parameter eight times over.

Active recall

Attempt all six before reading the worked answers.

1. In the online-softmax update, why must the algorithm track a running maximum m and rescale the previous accumulator every time a new tile's local max exceeds it, rather than simply summing exp(score) directly as each tile arrives?
2. If SRAM can only hold a tile of size 2 (as in the worked example) and the sequence length grows from N = 4 to N = 8 tokens, how many tiles does the inner loop process, and does the final output differ from computing the full 8×8 attention matrix in one shot?
3. In the worked example, suppose key vector k4 changes from [0, 0] to [1, 1] (identical to k3), while v4 stays [2, 2]. Recompute the exact attention output directly, then re-trace the two-tile online-softmax computation and confirm the two still agree. Which of the four running quantities from Block A's pass (m, l, O after Block A) need to be recomputed, and which stay untouched?
4. A Mixture-of-Experts FFN layer has 8 experts, each with 2.1 million parameters, plus 2 million parameters in shared attention/embedding components. With top-2 routing, what is the layer's total parameter count and its active-per-token parameter count? Does inference latency track total or active parameters — and does GPU memory footprint track the same one?
5. FlashAttention's backward pass recomputes attention tiles instead of storing them from the forward pass — meaning it does strictly more floating-point operations than standard attention, not fewer. Why does it still train faster in wall-clock time?
6. True or false: FlashAttention is an approximation of full attention, similar in spirit to linear attention or Performer's kernel methods, trading a small amount of accuracy for memory savings. Justify your answer.

Worked answers

1. Exponentiating a raw attention score without shifting it first risks numerical overflow: scores can be large (unbounded dot products, especially before the 1/√d scaling or with poorly normalized activations), and e^(large number) can exceed the range a float32 can represent, turning into inf and then NaN once divided. Subtracting the running max before exponentiating keeps every exponentiated term in (0, 1], which is numerically safe. But because that shift value keeps changing as new, larger scores appear in later tiles, everything computed under the *old* shift has to be corrected by a factor of e^(old_max − new_max) before being combined with the new tile — that correction is exactly what steps 3 and 4 in the diagram perform. Skipping the rescale would silently produce a softmax normalized against the wrong denominator.
2. With tile size 2 and N = 8, the inner loop processes 4 tiles instead of 2 (positions 1–2, 3–4, 5–6, 7–8), applying the same running-max/running-sum/running-output update after each. The final output is identical to computing the full 8×8 matrix directly — online softmax is an exact algebraic reassociation of the same sum, and this holds for any number of tiles, not just two. What changes with more tiles is only the memory-traffic profile: the full method would have written and read a 64-entry score matrix in one shot, while the tiled method never holds more than one tile (here, 2 scores) in memory at once, at the cost of one extra loop iteration.
3. New score s4 = q·k4 = [1,0]·[1,1] = 1 (previously 0). Direct computation: scores become [1, 0, 1, 1], so three of the four positions now tie for the maximum. Exponentials are [e, 1, e, e] with e ≈ 2.718282; the denominator is 3e + 1 ≈ 9.154846, giving weights w1=w3=w4 ≈ 0.296923 and w2 ≈ 0.109232, and since w1+w2+w3+w4 = 1 with three of the four weights equal, Output_y = w2(1)+w3(1)+w4(2) = w2 + 3w1 = 1 exactly (the y-numerator collapses to the full weight sum). Output_x = w1(1)+w3(1)+w4(2) = w1+w3+2w4 = 4w1 ≈ 1.18769. So the new exact output is [1.18769, 1.00000]. Re-tracing the tiles: Block A is untouched (k1, k2, v1, v2 didn't change), so m after Block A stays 1, l stays ≈1.367879, and O stays ≈[1, 0.367879] exactly as before — that pass never needs to be recomputed. Only Block B changes: its scores are now [1, 1] (both keys tie), giving local max m_B = 1, local sum l_B = 1+1 = 2, and local output O_B = 1·[1,1] + 1·[2,2] = [3, 3]. Combining with Block A's unchanged accumulator (no rescale needed since both maxes are 1): l = 1.367879 + 2 = 3.367879, O = [1+3, 0.367879+3] = [4, 3.367879]. Final output = [4/3.367879, 3.367879/3.367879] = [1.18769, 1.00000], matching the direct computation exactly. The ripple is real but contained: changing one key/value pair only forces recomputation of the tile that pair lives in, plus the final combine step — every earlier tile's cached (m, l, O) is reused untouched, which is exactly the property that makes tiling worthwhile at scale.
4. Total FFN parameters = 8 × 2.1M = 16.8M, plus 2M shared = 18.8M total. Active-per-token = 2 × 2.1M (top-2) + 2M shared = 6.2M active. Inference latency (compute time) tracks active parameters, since idle experts perform no multiplications for that token. GPU memory footprint tracks *total* parameters, not active ones, because which experts get used is decided per token at runtime — the router might send the very next token to any of the 8 experts, so all of them must be resident in VRAM even though only 2 fire per forward pass. This is why MoE models are cheap to run but not cheap to host.
5. Wall-clock time on a GPU is usually set by whichever resource is the bottleneck, and for attention that resource is HBM bandwidth, not FLOPs — GPUs have far more arithmetic throughput available than they have memory bandwidth to feed it. FlashAttention's backward pass spends extra compute cycles (cheap, plentiful) recomputing tiles instead of spending HBM bandwidth (expensive, scarce) reading back a stored N×N matrix from the forward pass. Since the bottleneck resource shrinks even though the non-bottleneck resource grows slightly, total wall-clock time goes down — a textbook case of optimizing the actual constraint rather than the one that's easiest to count.
6. False. FlashAttention computes the mathematically exact same softmax-weighted output as standard attention — the worked trace in this chapter confirms the tiled and direct computations agree to five decimal places, and this agreement is not a coincidence of the specific numbers chosen but a consequence of online softmax being an algebraic identity (rescale-and-combine), not a truncation or kernel substitution. Linear attention and Performer genuinely change the function being computed by replacing the softmax kernel; FlashAttention only changes the order in which arithmetic happens and what gets written to slow memory. The "accuracy for speed" framing describes an entirely different category of technique.

Think About It

Think about this: How would you explain transformer architecture deep dive to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where transformer architecture deep dive is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting transformer architecture deep dive to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind transformer architecture deep dive, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.