3.1 Why AI Infra engineers must understand Transformer

Transformer is the shared skeleton of every major LLM today, and almost everything an AI Infra engineer does — from low-level kernels to distributed training to serving — sits on that skeleton. This article covers what AI Infra is, how Transformer became the common base, and the precise mapping between the two, so you can answer “what am I actually optimizing?” with a clear picture.

📑 Contents

• 1. What is AI Infra
• 2. How Transformer became the common base
• 3. How AI Infra layers map to Transformer modules
• 4. End-to-end case: one LLM inference request
• 5. Learning advice: building a complete knowledge system
• Self-check list
• References

---

1. What is AI Infra

1.1 One-sentence definition

AI Infra (AI Infrastructure) can be summed up in one sentence: use systems engineering to feed hardware compute to the model efficiently.

If training and running a large model is like throwing a banquet for ten thousand people, algorithm researchers design the menu — which dishes and which recipes. AI Infra engineers build the kitchen, schedule the cooks, and run the supply chain. Half of a good meal is the recipe; the other half is whether the kitchen runs smoothly — enough stoves, fast plating, ingredients not piled up in the warehouse. AI Infra is that logistics system.

1.2 Where AI Infra sits in the stack

To see AI Infra’s job boundary, draw the whole AI stack as layers:

AI Infra sits between the “algorithm and model” layer and the “hardware” layer. Its job is to be the translator for both sides: turn the model’s compute needs into instructions the hardware can run well. Upward, it needs to know what the model looks like and how much compute and memory each module needs. Downward, it needs GPU memory hierarchy, interconnect topology, and parallelism. And that “upper model,” today, is almost always Transformer.

1.3 Day-to-day work for an AI Infra engineer

The exact work varies by team, but the core jobs usually fall into these buckets:

🍎 CUDA operator development and optimization

This is the most “low-level” work and talks straight to the GPU. Typical tasks: write efficient CUDA kernels for Attention (implement or improve FlashAttention), tune GEMM to use Tensor Cores well, fuse many small ops into one kernel to cut HBM traffic. You can only do this if you know the object you are optimizing. FlashAttention, for example, targets the compute and memory access pattern of $QK^T$ and Softmax in Self-Attention. If you do not know how Self-Attention runs, you cannot tell what FlashAttention is doing.

🍎 Distributed training

When the model no longer fits on one GPU, you split it across cards or machines. That means tensor parallelism (split a matmul along a dimension across GPUs), pipeline parallelism (put different layers on different GPUs), data parallelism (different batches per GPU, then sync gradients), and so on. Every split point follows Transformer structure: tensor parallelism often cuts along Attention heads; pipeline parallelism cuts along the stack of Decoder Blocks.

🍎 Inference deployment and optimization

After training, serving user requests efficiently is another large job. Core challenges include KV Cache memory management (each request caches Keys and Values from Attention), Continuous Batching (dynamic batching to raise GPU utilization), quantization (lower-precision weights and cache to save memory and bandwidth), Speculative Decoding (a small model drafts several tokens; the large model verifies them in one pass). Every one of these targets a concrete module inside Transformer.

🍎 Performance analysis and system tuning

Use Nsight Systems, Nsight Compute, torch.profiler, and similar tools to find bottlenecks in training or inference: compute-bound vs bandwidth-bound, and which hot kernels are worth fixing. That again requires knowing which Transformer module each kernel belongs to. Otherwise a long-running kernel name tells you nothing about what it is doing.

---

2. How Transformer became the common base

2.1 Before Transformer: the RNN/LSTM trap

Before Transformer, the mainstream for sequence data (text, speech, time series) was RNN (recurrent neural networks) and its improved cousin LSTM (long short-term memory).

RNN’s idea is simple: walk the sequence one step at a time, combine the current input with the previous “memory” (hidden state), and pass a new memory forward. Like reading a novel word by word and updating your mental model of the plot after every sentence.

That design has two fatal weaknesses:

❌ No parallel training. Each step depends on the previous output, so the whole sequence is serial. For a 2048-token sequence, RNN needs 2048 sequential steps and cannot use the thousands of GPU cores sitting idle. Like a one-lane road: no matter how many cars (GPU cores) you have, only one gets through at a time.

❌ Forgetting long-range dependencies. LSTM’s gates ease vanishing gradients in plain RNNs, but information still decays over long distances. Imagine a long talk: a key definition from the opening is fuzzy by the conclusion an hour later. LSTM faces a similar forgetting problem and drops off badly past a few thousand tokens.

Together those two become a bottleneck: you cannot train efficiently, and you cannot use long context well. That is what Transformer set out to fix.

2.2 Attention Is All You Need: a paradigm shift

In 2017, Vaswani et al. proposed Transformer in “Attention Is All You Need.” The core move was to replace recurrence entirely with Self-Attention, which attacks both problems at the root.

Self-Attention lets every token interact with every other token: compute pairwise “attention weights,” then aggregate by those weights. Information from token 1 to token 2000 does not walk through 1999 steps; it goes in one hop. In the novel metaphor, you stop reading word by word and instead see the whole page and jump to the important paragraphs.

Two fundamental wins follow:

⭕ Natural parallelism. Attention weights for all token pairs can be computed at once — essentially one matmul $QK^T$. Matmul is what GPUs do best; thousands of CUDA Cores and Tensor Cores can join in. Training speed jumps qualitatively versus RNN.

⭕ Global information access. Every token can “see” every other token; there is no long-hop decay. The model learns which positions to attend to, instead of being forced to pass information in a fixed order.

The all-to-all design has a cost: compute is $O(N^2)$ in sequence length $N$. In real training, the speedup from parallelism still outweighs the extra $O(N^2)$ work, especially with strong GPU matmul. And the memory and compute pain of $O(N^2)$ became a core battlefield for AI Infra (FlashAttention, KV Cache management, and so on).

2.3 From the original Transformer to the LLM family

The original Transformer was Encoder-Decoder, aimed at machine translation. Later work branched in several directions:

✳️ BERT (2018): Encoder only, with bidirectional attention (each token sees past and future) and masked language model pretraining, which broke through on NLU. BERT proved the “large-scale pretrain + downstream fine-tune” paradigm.

✳️ GPT series (2018–2023): Decoder only, with unidirectional attention (each token sees only past tokens) and autoregressive pretraining, pushing text generation further. GPT-2 showed large LMs can generate text well; GPT-3 backed scaling law (more parameters + more data = better results); GPT-4 showed large gains on multimodal understanding and harder reasoning.

✳️ LLaMA (2023): Meta’s open Decoder-only model, with practices that stuck: Pre-Norm (RMSNorm), SwiGLU, RoPE. Opening LLaMA gave the community one shared architecture to study and optimize, and pushed the AI Infra toolchain hard — vLLM, SGLang, Megatron-LM and friends defaulted first to LLaMA-like layouts.

✳️ Later evolution (Mistral, Qwen, DeepSeek, and others): On top of LLaMA they added GQA (Grouped-Query Attention), MoE (Mixture of Experts), MLA (Multi-head Latent Attention), and other variants. The core skeleton stayed Transformer Decoder.

Looking across that history, one fact stands out: no matter how names, scale, or applications change, the compute units stay the same Transformer modules — Self-Attention, FFN, LayerNorm, positional encoding. That is why AI Infra engineers must know Transformer: the object you optimize barely changes for years; what changes is scale and detail.

---

3. How AI Infra layers map to Transformer modules

With AI Infra’s scope and Transformer’s place in history clear, the core question is: which Transformer module does each AI Infra optimization hit?

The table below is the full mapping. Each row is one concrete AI Infra technique and the Transformer module it acts on.

Why these correspondences exist, one by one:

3.1 CUDA operator optimization layer

1️⃣ FlashAttention and Self-Attention’s $QK^T$ / $PV$

Self-Attention builds an $N \times N$ attention matrix ($N$ = sequence length). A naive implementation writes the full matrix to HBM, using $O(N^2)$ memory. Growing sequence length from 2K to 128K multiplies that memory by 4096. FlashAttention tiles the matrix so each tile finishes $QK^T$, Softmax, and $PV$ in on-chip SRAM, and never writes the full $N \times N$ matrix to HBM. It hits Self-Attention’s two core matmuls directly, so you need the Attention compute path before FlashAttention makes sense.

2️⃣ Fused Softmax and Attention’s normalization step

Softmax over each row of the score matrix normally needs three passes (max, sum of exp, normalize). Online Softmax merges the first two into two passes; FlashAttention folds that into tiling for a block-level single pass. Softmax sits in the middle of the Attention pipeline; fusing it with the surrounding matmuls cuts a lot of intermediate HBM traffic.

3️⃣ LayerNorm kernel fusion

Each Transformer Decoder Block has two LayerNorm ops (in Pre-Norm, before Attention and before FFN). LayerNorm itself is light on compute but scans the feature vector twice (mean/variance, then normalize), so it is classic Memory Bound work. CUDA practice often fuses Residual Add + LayerNorm into one kernel to drop one full HBM round trip — and that fusion only makes sense if you know where residual and LayerNorm sit in the Block.

4️⃣ Fuse RoPE into the Attention kernel

Rotary positional encoding (RoPE) rotates adjacent dimension pairs of Q and K. The compute is small but the call is frequent. Fusing it into the Attention kernel avoids launching a separate GPU kernel for a tiny op. That only works if you know RoPE sits after the Q/K linear projections and before $QK^T$.

3.2 Distributed training layer

1️⃣ Tensor parallelism (TP) and Attention heads / FFN matrices

TP splits a large matmul across GPUs. Transformer has two natural cut points. First, Multi-Head Attention: 32 heads can go evenly onto 4 GPUs (8 heads each) with little cross-GPU communication until the final output projection. Second, FFN: Megatron-LM column-splits $W_{up}$ and row-splits $W_{down}$ so the intermediate activation does not need cross-GPU communication. Both strategies depend on knowing multi-head Attention and the FFN up-down structure.

2️⃣ Pipeline parallelism (PP) and stacked Decoder Blocks

A large model stacks tens of Decoder Blocks (LLaMA-2-7B has 32). PP puts contiguous segments on different GPUs — e.g. layers 1–8 on GPU 0, 9–16 on GPU 1, and so on — and data flows through like a pipeline. This works because every Decoder Block has the same I/O shape ($(N, d_{model})$); one Block’s output is the next Block’s input, so segmentation is natural.

3️⃣ ZeRO memory optimization and all parameter matrices

ZeRO removes “duplicate storage” in training. Standard data parallel keeps full parameters, gradients, and optimizer state on every GPU; ZeRO shards those across ranks. It touches every parameter matrix — Attention’s $W_Q / W_K / W_V / W_O$ and FFN’s $W_{gate} / W_{up} / W_{down}$. Knowing the parameter split (FFN ~2/3, Attention ~1/3) matters when you estimate ZeRO’s memory wins and communication cost.

3.3 Inference deployment layer

1️⃣ KV Cache and Self-Attention’s K/V

Autoregressive generation emits one token per step, but Attention still needs the new Query against Keys of all past tokens. Without a cache you recompute every historical K and V every step — pure waste. KV Cache stores per-layer K and V in GPU memory for reuse. This is lifecycle management of K and V inside Self-Attention.

2️⃣ PagedAttention and KV Cache fragmentation

Requests have different lengths, so KV Cache sizes differ. Naive preallocation fragments memory: short requests waste space; long ones may be rejected for lack of contiguous room. PagedAttention borrows OS virtual-memory paging, splits KV Cache into fixed-size pages, and allocates on demand. The root issue is dynamic growth: each new token adds one “row” to the cache.

3️⃣ Quantization and all weight matrices / KV Cache

Quantization uses fewer bits per value — e.g. FP16 (16-bit) down to INT4 (4-bit) cuts weight memory to 1/4. It hits every Transformer weight matrix ($W_Q, W_K, W_V, W_O, W_{gate}, W_{up}, W_{down}$) and KV Cache. Knowing each matrix’s value distribution (e.g. whether some FFN layers have large activation outliers) matters for picking a quantization scheme.

4️⃣ Speculative Decoding and the serial bottleneck of autoregressive generation

Autoregressive generation can emit only one token at a time, so much GPU compute sits idle. Speculative Decoding uses a small, fast Draft model to guess several tokens, then the large model verifies them in one pass — keep the hits, resample the misses. It targets the serial nature of Transformer Decoder generation by speeding serial generation with parallel verification.

---

4. End-to-end case: one LLM inference request

To ground the mapping, follow one request from user question to answer: how data flows through each Transformer module, and which AI Infra optimizations sit on each step.

Suppose we serve LLaMA-2-7B (32 Decoder Blocks, $hiddendim = 4096$, 32 Attention heads) with vLLM. The user sends: “Please explain what Transformer is in one sentence.”

4.1 Tokenize and Embedding

User input → Tokenizer → Token IDs: [3529, 671, 19, 13563, ...] (~10 tokens)
Token IDs → Embedding layer → input tensor: (10, 4096)

AI Infra link: Tokenizer runs on CPU. Under extreme concurrency, CPU tokenize can become the bottleneck; use Nsight Systems to check for GPU idle gaps.

4.2 Prefill: process the whole prompt in parallel

All 10 input tokens enter the model together and walk 32 Decoder Blocks. In each layer:

LayerNorm: normalize the input.

(10, 4096) → RMSNorm → (10, 4096)

AI Infra optimization: fuse Residual Add + RMSNorm into one CUDA kernel to cut one HBM round trip.

QKV linear projections: build Q, K, V.

Q = X * W_Q: (10, 4096) x (4096, 4096) = (10, 4096) → reshape (10, 32, 128)
K = X * W_K: same as above
V = X * W_V: same as above

AI Infra optimization: the three GEMMs are a TP split point — with 4-way TP each GPU computes QKV for 8 heads.

RoPE: rotate Q and K.

Q, K → rotate by position → Q', K'  (shape unchanged)

AI Infra optimization: RoPE is cheap but frequent, usually fused into the following Attention kernel.

Self-Attention:

S = Q' * K'^T / sqrt(128): (10, 128) x (128, 10) = (10, 10)  [per head]
A = softmax(mask(S)):       (10, 10)
O = A * V:                  (10, 10) x (10, 128) = (10, 128)  [per head]
concat 32 heads → (10, 4096)
output projection: (10, 4096) x (4096, 4096) = (10, 4096)

AI Infra optimization: FlashAttention fuses $QK^T$, Scale, Mask, Softmax, and $PV$ in one kernel and finishes the work in SRAM via tiling. Prefill’s batch dimension is larger (10 tokens), so this is Compute Bound.

Cache K and V: write this layer’s K and V into KV Cache.

AI Infra optimization: vLLM’s PagedAttention allocates KV Cache in pages. These 10 tokens’ KV land in preallocated physical pages.

Residual + LayerNorm + FFN:

residual add: h = input + attention_output    (10, 4096)
RMSNorm(h):                                (10, 4096)
FFN (SwiGLU):
  gate = h * W_gate: (10, 4096) x (4096, 11008) = (10, 11008)
  up   = h * W_up:   (10, 4096) x (4096, 11008) = (10, 11008)
  mid  = Swish(gate) * up:                         (10, 11008)
  down = mid * W_down: (10, 11008) x (11008, 4096) = (10, 4096)
residual add: output = h + down                 (10, 4096)

AI Infra optimization: FFN’s three large matrices ($W_{gate}$, $W_{up}$, $W_{down}$) dominate parameters (~2/3). Under TP, column-split $W_{gate}$ and $W_{up}$, row-split $W_{down}$. Under quantization, these three are the main INT4/INT8 targets.

After all 32 layers: final output + Final LayerNorm + LM Head → vocab distribution → sample the first output token.

Prefill is done. Time from request to first character is TTFT (Time To First Token).

4.3 Decode: generate token by token

From the second output token on, Decode runs. Each step has only 1 new token:

Each step:
  new token → Embedding → (1, 4096)
  each layer:
    LayerNorm → QKV projection → RoPE
    new Q (1, 32, 128) × cached K (n, 32, 128)^T → (1, n) attention scores [n = current total length]
    Softmax → weighted V → output projection
    append new K, V to KV Cache
    residual + LayerNorm + FFN + residual
  LM Head → sample → emit one token

Key differences from Prefill:

• Matmul collapses to matrix-vector (batch dim = 1); GPU compute is badly underused
• Most time goes to moving KV Cache from HBM: Memory Bound
• After each step, KV Cache grows by one row

This stage is where AI Infra optimizations densest:

4.4 Full-path sketch

The full data flow and its AI Infra mapping, in one figure:

The end-to-end case makes the point: every hop of the request — Embedding, Attention, FFN, autoregressive generation — is where some AI Infra technique lands. Without Transformer’s internals you cannot place the bottleneck, and you cannot judge whether an optimization worked.

---

5. Learning advice: building a complete knowledge system

Once “why Transformer” is clear, the next question is how to learn efficiently. Here is a path that matches this article series.

5.1 See the tree before the forest

Do not start with the FlashAttention or Megatron-LM papers — you will drown in unfamiliar symbols. The right order is:

1. Step 1: know every Transformer module. Read the series piece Transformer architecture: a quick start. Be able to draw a full Decoder Block on a whiteboard, label every shape, and write the Attention formula from memory. That is the base for everything later.

2. Step 2: build the module–optimization map. That is section 3 of this article. For every new optimization, ask first: “Which Transformer module does it hit? What is that module’s bottleneck?”

3. Step 3: go deep by layer. Depending on your job, dig into CUDA ops, distributed training, and/or inference. Going deep in any of those also deepens your sense of Transformer architecture.

5.2 Keep watching architecture evolution

The Transformer skeleton is stable, but details keep moving: GQA changes KV head count, MoE changes how FFN activates, MLA changes how KV Cache is compressed. Each variant shifts AI Infra strategy. Follow tech reports from mainstream models (LLaMA, DeepSeek, Qwen, Mistral), note architecture changes, and think through the training and serving impact.

---

🎯 Self-check list

After this article, check whether you really see the AI Infra–Transformer link:

• Explain in your own words what AI Infra is, where it sits in the AI stack, and what sits above and below it
• State RNN/LSTM’s two core weaknesses (serial, forgetting) and how Transformer fixes them with Self-Attention
• List at least 3 CUDA operator optimizations and the Transformer modules they hit, and explain the mapping
• List at least 3 distributed training techniques and the Transformer modules they hit, and explain the mapping
• List at least 3 inference optimizations and the Transformer modules they hit, and explain the mapping
• Given an LLM inference request, trace data through Transformer modules and name the AI Infra optimizations on each step
• Answer “why AI Infra engineers must understand Transformer” with at least two concrete reasons (not a vague “because it is important”)

📚 References

Papers

• Attention Is All You Need (Vaswani et al., 2017)：https://arxiv.org/abs/1706.03762 — Original Transformer paper
• BERT: Pre-training of Deep Bidirectional Transformers (Devlin et al., 2018)：https://arxiv.org/abs/1810.04805 — Encoder-only pretraining paradigm
• Language Models are Unsupervised Multitask Learners (Radford et al., 2019, GPT-2)：https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf — GPT-2 technical report
• Language Models are Few-Shot Learners (Brown et al., 2020, GPT-3)：https://arxiv.org/abs/2005.14165 — Scaling Law validation for large LMs
• LLaMA: Open and Efficient Foundation Language Models (Touvron et al., 2023)：https://arxiv.org/abs/2302.13971 — Milestone open large model
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)：https://arxiv.org/abs/2205.14135 — Memory-aware Attention optimization
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)：https://arxiv.org/abs/1909.08053 — Tensor and pipeline parallelism
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al., 2019)：https://arxiv.org/abs/1910.02054 — Memory optimization
• Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)：https://arxiv.org/abs/2309.06180 — vLLM / PagedAttention

Tutorials and blogs

• The Illustrated Transformer (Jay Alammar)：https://jalammar.github.io/illustrated-transformer/ — Visual Transformer intro
• Andrej Karpathy: Let’s build GPT from scratch：https://www.youtube.com/watch?v=kCc8FmEb1nY — Write GPT from scratch
• 3Blue1Brown: But what is a GPT： https://www.youtube.com/watch?v=wjZofJX0v4M — Intuitive Transformer
• The Annotated Transformer (Harvard NLP)：https://nlp.seas.harvard.edu/annotated-transformer/ — Paper line-by-line with PyTorch