AI Infra学习之旅-Transformer知识地图

AI Infra学习路线

写给刚入门AI Infrastructure的同学：你真的需要搞懂Transformer的每个细节吗？

写在前面

最近跟几个刚开始学AI Infra的朋友聊天，大家都有个共同的困惑：明明是搞基础设施的，为什么老有人让我们先把Transformer学透？

说实话，我刚开始也纠结过。看了一堆Transformer的论文和教程，什么Attention机制、位置编码、残差连接…感觉自己要变成算法工程师了。后来跑了几个月的vLLM、搞了一些推理优化，才慢慢想明白：

AI Infra工程师确实要懂Transformer，但不是全懂，而是要懂对地方。

今天这篇文章，我想跟你聊聊：作为AI Infra的学习者，到底应该重点关注Transformer的哪些部分，哪些可以暂时放一放，以及不同方向（推理优化、模型部署、系统架构）的学习侧重点。

一、为什么AI Infra一定要学Transformer？

在回答"学什么"之前，先说说"为什么要学"。

1.1 现代AI系统的核心就是Transformer

看看你做AI Infra会接触的东西：

• 推理框架：vLLM、TensorRT-LLM、TGI、…
• 优化技术：Flash Attention、PagedAttention、KV量化、…
• 部署工具：Triton Inference Server、Ray Serve、…

这些东西，100%都是围绕Transformer模型设计的。

如果你不了解Transformer的工作原理，就像是开飞机但不知道飞机怎么飞的——能起飞，但出了问题根本不知道从哪查。

1.2 "搞硬件"也要懂模型

有人会说：“我是做基础设施的，GPU、内存、网络这些才是我该关心的吧？”

没错，但问题是：

• 内存优化：为什么要优化？因为KV Cache太大了。KV Cache是啥？Transformer推理过程中产生的。
• 计算优化：为什么要Flash Attention？因为原生Attention太慢了。为什么慢？要看Transformer的计算流程。
• 并行策略：为什么要张量并行？因为单卡放不下大模型。为什么放不下？要看Transformer的参数分布。

你优化的对象就是Transformer模型，不懂它怎么优化？

1.3 一个真实的例子

给你讲个我自己的例子。

前阵子我在用vLLM跑一个模型，发现GPU内存占用特别高，batch size稍微大一点就OOM了。当时我只会调参数，试了半天没用。

后来静下心来看了Transformer的KV Cache机制，才明白：

• 每个token都要生成K和V向量
• 这些向量要一直存在GPU显存里
• 序列越长，batch越大，KV Cache占用呈指数级增长

理解了这个，我才知道为什么PagedAttention能省内存（块状分配减少碎片），为什么要用Multi-Query Attention（减少KV头数）。

不懂原理，就只能瞎调参。

二、AI Infra视角下的Transformer知识地图

好了，既然要学，那到底学什么？

我把Transformer的知识点分成三个层级，告诉你哪些必须学、哪些重要、哪些了解就行。

核心必学

这部分不学，你做AI Infra基本寸步难行。

2.1 KV Cache机制

为什么必学？

KV Cache是AI Infra优化的头号目标。不夸张地说，现代LLM推理优化的一半技术都在优化KV Cache。

必须理解的点：

1. KV Cache是怎么产生的？

在Transformer的自回归生成过程中（比如GPT生成文本），每生成一个新token，都需要用到之前所有token的Key和Value向量。

生成第1个词: 需要prompt的KV
生成第2个词: 需要prompt + 第1个词的KV
生成第3个词: 需要prompt + 第1个词 + 第2个词的KV
...

为了避免重复计算，这些KV向量会被缓存起来——这就是KV Cache。

1. KV Cache有多占内存？

假设一个7B参数的模型，hidden size是4096，有32层：

• 每个token的KV向量：2 × 4096 × 32 × 2字节(FP16) ≈ 512KB
• 生成1000个token：512KB × 1000 ≈ 512MB
• Batch size = 32：512MB × 32 = 16GB

16GB！这还只是KV Cache，不包括模型权重和其他激活值。所以你会看到大batch推理很容易OOM。

1. 怎么优化KV Cache？

• PagedAttention：像操作系统的虚拟内存一样，分块管理KV Cache，减少内存碎片（这就是vLLM的核心技术）
• Multi-Query Attention (MQA)：让所有Query头共享一组KV，大幅减少KV Cache大小
• KV量化：把KV向量从FP16量化到INT8甚至INT4
• Prefix Caching：多个请求共享相同的系统prompt对应的KV

实践建议：

去看vLLM的源码，重点看CacheEngine和BlockManager这两个模块。跑一个模型，用nvidia-smi监控显存占用，你会对KV Cache有直观的感受。

2.2 Attention计算流程

为什么必学？

Attention计算是Transformer的性能瓶颈。你做推理优化，90%的时间都在优化Attention。

必须理解的点：

1. Attention的四步计算

# 伪代码
scores = Q @ K.T  / sqrt(d_k)     # 1. 点积 + 缩放
attention_weights = softmax(scores)  # 2. Softmax归一化  
output = attention_weights @ V       # 3. 加权求和

1. 为什么Attention慢？

• 内存瓶颈：Q、K、V都是大矩阵，频繁在HBM（显存）和SRAM（缓存）之间搬运数据
• 计算复杂度：O(n²)，序列越长越慢

1. 优化方向

• Flash Attention：重新组织计算顺序，减少HBM访问次数（vLLM已经集成）
• 稀疏Attention：不是所有token都需要关注所有其他token
• Kernel融合：把多个小算子融合成一个大kernel，减少数据搬运

实践建议：

对比一下开启和关闭Flash Attention的性能差异。在vLLM里很简单：

# 开启Flash Attention（默认）
llm = LLM(model="facebook/opt-125m")

# 关闭Flash Attention
llm = LLM(model="facebook/opt-125m", disable_flash_attn=True)

跑个benchmark，你会发现差距很明显。

2.3 Batch处理与序列并行

为什么必学？

提升推理系统吞吐量的核心手段。

必须理解的点：

1. 静态Batching vs 动态Batching

传统方法：等凑够32个请求再一起推理。问题是：第1个请求可能要等很久。

Continuous Batching（vLLM用的）：来一个处理一个，动态调整batch。

1. 为什么Transformer适合Batching？

因为Attention是独立的！处理batch中的第1个序列和第2个序列，计算是并行的。

实践建议：

用vLLM跑两个实验：

• 单请求推理：测latency（延迟）
• 批量请求推理：测throughput（吞吐量）

你会发现throughput可以提升10倍+，但latency只增加了一点点。

---

重要理解

这部分不是立刻就用，但理解了能帮你少走弯路。

2.4 位置编码（Positional Encoding）

为什么要关注？

影响模型处理长上下文的能力。

Transformer的Self-Attention本身不感知位置，"我爱你"和"你爱我"对它来说是一样的。所以需要位置编码来告诉模型词的顺序。

AI Infra关心什么？

• 原始位置编码（Sinusoidal）：有长度限制，不适合超长上下文
• RoPE（Rotary Position Embedding）：很多新模型（LLaMA等）用这个，支持外推到更长序列
• ALiBi：另一种方案，不需要额外参数

如果你要部署一个支持32K上下文的模型，位置编码方案会直接影响性能和效果。

2.5 Multi-Head Attention

为什么要关注？

影响模型并行策略。

多头注意力意味着计算可以并行。如果一个模型有32个头，你可以把它们分到多个GPU上（张量并行）。

AI Infra关心什么？

• 头数是2的幂次（8、16、32）方便切分
• MQA/GQA（Multi-Query/Grouped-Query Attention）用更少的头，更省内存

2.6 Encoder vs Decoder 架构

为什么要关注？

不同架构的推理特性完全不同。

• Encoder-only (BERT类)：

  • 输入一段文本，输出一段表示
  • 没有自回归生成，不需要KV Cache
  • 推理很快，适合分类、检索任务

• Decoder-only (GPT类)：

  • 自回归生成，每个词依赖前面所有词
  • 需要大量KV Cache
  • 推理慢，但生成能力强

AI Infra关心什么？

你部署BERT和GPT，优化策略完全不同：

• BERT：重点优化计算（算子融合、量化）
• GPT：重点优化内存（KV Cache管理、Paged Attention）

---

了解即可

这部分主要是训练相关，推理时影响不大。

2.7 Masked Attention

训练时用来防止模型"偷看"未来的词。推理时生成本来就是一个词一个词来的，不存在这个问题。

了解就行，不用深究。

2.8 残差连接 & Layer Norm

这些是为了稳定训练过程。推理时它们就是网络的一部分，你优化时treat it as a black box就行。

如果做算子优化，可能会关注Layer Norm的实现。

三、不同AI Infra方向的学习侧重

根据你的具体方向，学习重点不太一样。

推理优化工程师

你的日常工作：

让模型跑得更快、用更少资源、支持更大batch。

Transformer知识优先级：

1. KV Cache管理：必须吃透

   • KV的产生、存储、读取全流程
   • PagedAttention原理和实现
   • Prefix Caching、KV量化

2. Attention计算优化：核心战场

   • Flash Attention的原理（至少懂为什么快）
   • Multi-Query Attention的效果
   • Sparse Attention的适用场景

3. 量化对Transformer的影响：必备技能

   • 哪些部分可以量化（权重、激活、KV Cache）
   • 量化后的精度损失
   • 量化算子的实现

学习路线：

第1周：搞透KV Cache

• 看vLLM的CacheEngine源码
• 画出KV Cache的生命周期图
• 跑实验：监控不同batch下的KV占用

第2-3周：实践优化技术

• 开启/关闭Flash Attention对比
• 尝试不同的block size（PagedAttention参数）
• 实现一个简单的KV Cache管理器（Python即可）

第4周+：深入一个方向

• 如果做kernel优化：学CUDA，看Flash Attention源码
• 如果做系统优化：学调度策略，看vLLM的scheduler
• 如果做量化：学INT8/INT4 kernel实现

模型部署工程师

你的日常工作：

让模型在生产环境稳定运行，支持高并发请求。

Transformer知识优先级：

1. 模型架构与框架对应：必须清楚

   • Encoder-only vs Decoder-only的部署差异
   • 不同框架对Transformer的支持（TensorRT-LLM、vLLM、TGI）
   • 模型格式转换（PyTorch → ONNX → TensorRT）

2. Batch处理机制：影响吞吐量

   • Dynamic Batching的实现
   • Continuous Batching（vLLM的核心）
   • Batch size与latency的trade-off

3. 长上下文支持：客户常问的问题

   • 位置编码方案（RoPE、ALiBi）
   • 长上下文下的内存管理
   • Sliding Window Attention

学习路线：

第1周：熟悉部署流程

• 在vLLM上部署一个模型
• 尝试不同的推理框架（TGI、TensorRT-LLM）
• 对比它们的API和性能

第2周：理解架构差异

• 部署一个BERT模型（如果做任务推理）
• 部署一个GPT模型
• 总结两者的部署要点差异

第3周：优化并发性能

• 测试不同max_batch_size的影响
• 配置动态batching参数
• 监控GPU利用率

第4周+：生产化

• 加入监控（Prometheus + Grafana）
• 实现自动扩缩容
• 做故障演练

AI系统架构师

你的日常工作：

设计支持千亿参数模型的分布式推理系统。

Transformer知识优先级：

1. 模型并行策略：核心技能

   • 张量并行（Tensor Parallelism）在Transformer中的应用
   • 流水线并行（Pipeline Parallelism）的切分点
   • 序列并行（Sequence Parallelism）的适用场景

2. 内存墙与计算墙：系统瓶颈

   • Transformer哪些操作是compute-bound，哪些是memory-bound
   • 不同硬件（A100 vs H100）的特性
   • 通信开销的计算

3. 端到端推理流程：全局视角

   • 从接收请求到返回结果的完整路径
   • 各阶段的瓶颈和优化点
   • SLA保证策略

学习路线：

第1-2周：理解并行策略

• 阅读Megatron-LM、DeepSpeed论文
• 理解Transformer各层如何切分
• 画出并行策略的拓扑图

第3周：分析瓶颈

• 用Nsight分析Transformer的kernel性能
• 区分compute-bound和memory-bound操作
• 计算不同配置下的理论性能上限

第4周+：设计系统

• 设计一个支持千亿参数模型的推理集群
• 考虑成本、延迟、吞吐量的trade-off
• 写技术方案文档

四、学习建议与避坑指南

推荐做法

1. 从实践入手

不要上来就啃论文。先跑起来一个vLLM程序，看看KV Cache在哪，占了多少内存。有了直观感受，再看原理会轻松很多。

2. 关注"为什么影响性能"

学每个概念，都问自己：这个东西为什么会影响推理性能？是影响内存、计算还是通信？

比如学位置编码，AI Infra关心的不是数学公式，而是：

• RoPE为什么比Sinusoidal支持更长上下文？
• 这对内存有什么影响？

3. 动手实验

理论只是理论，跑个实验印象深刻一百倍。

推荐实验：

• 监控KV Cache占用随序列长度的变化
• 对比Flash Attention开启前后的速度
• 测试不同batch size对吞吐量的影响

4. 看开源代码

vLLM、TensorRT-LLM、TGI的代码都开源的。看不懂C++/CUDA没关系，先看Python部分的逻辑。

重点看：

• KV Cache的管理逻辑
• Batch调度策略
• 内存分配器

写在最后

回到开头的问题：AI Infra工程师要不要学Transformer？

要学，但要聪明地学。

不要试图把Transformer的每个细节都搞懂，那是算法工程师的事。你要做的是：

1. 抓住核心：KV Cache、Attention计算、Batch处理
2. 理解影响：每个设计为什么影响性能
3. 动手实践：跑起来，监控起来，优化起来

记住，AI Infra的价值不是"我懂Transformer"，而是"我能让Transformer跑得又快又稳"。

---

参考资料

1. 我的其他文章：

   • 《Transformer详解》 - Transformer基础原理
   • 《PagedAttention解析》 - KV Cache优化的核心技术
   • 《第一个vLLM程序》 - 实践入门

2. 推荐阅读：

   • Attention Is All You Need - Transformer原论文
   • The Illustrated Transformer - 最好的可视化教程
   • vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention - vLLM原论文
   • Flash Attention - Attention优化技术

3. 开源项目：

   • vLLM
   • TensorRT-LLM
   • Text Generation Inference (TGI)