3. LLM基础回顾：Transformer架构、自注意力机制、KV Cache原理
好，咱们直接进入正题。这一章我打算把Transformer的老底儿翻一翻。你可能会说：“这玩意儿我都看烂了。” 但别急，我重点讲的是跟多GPU推理强相关的三个点：架构怎么拆、自注意力为什么是瓶颈、KV Cache到底省了什么。
我个人习惯是，讲原理之前先画个图。下面这张SVG图，是我自己梳理的Transformer核心逻辑，你看一眼，心里就有谱了。
3.1 Transformer 架构：解码器才是主角
现在大语言模型，清一色都是 Decoder-only 架构。为什么？说白了，自回归生成天然适合这种单向结构。你想想看，GPT系列、LLaMA、ChatGLM，哪个不是这样？
一个标准的Transformer Decoder层，核心就两块：
- 自注意力层 (Self-Attention)：让每个token看到它之前的所有token。
- 前馈网络 (FFN)：对每个token做独立的非线性变换。
嗯，这里要注意，LayerNorm 和 残差连接 是标配。没有它们，网络深了根本训不动。我在项目中遇到过，有人把Pre-Norm和Post-Norm搞混了，结果推理时精度对不上，排查了半天。
核心要点： 推理时，我们只关心 前向计算。反向传播？那是训练的事。所以，注意力机制的计算量，直接决定了推理延迟。
3.2 自注意力机制：O(n²) 的甜蜜与痛苦
自注意力，公式很简单：
Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V但它的计算量，跟序列长度 n 是 平方关系。为什么？
- Q * K^T：得到一个 n x n 的注意力矩阵。这一步是 O(n²)。
- softmax + * V：同样是 O(n²)。
你想想看，当 n=4096 时，这个矩阵有 1600 万个元素。当 n=8192 时，直接翻四倍。我刚开始做推理优化时，总觉得瓶颈在FFN，后来一profiling才发现，注意力那块才是吃显存和算力的大户。
我的经验： 在多GPU场景下，自注意力的 QKV 投影矩阵，通常用 张量并行 切分。把 Q、K、V 的权重按列切到不同GPU上，各自算一部分，最后再all-reduce。这样能有效降低单卡显存压力。
3.3 KV Cache：推理加速的基石
终于到重点了。KV Cache，说白了就是 用空间换时间。
自回归生成时，每生成一个新token，都要重新计算所有历史token的注意力。但仔细想想，前面token的 K 和 V 矩阵，每次计算都是一样的！
所以，我们把这些 K 和 V 缓存起来。每次只算新token的 Q、K、V，然后跟缓存里的 K、V 做注意力计算。这样，计算量就从 O(n²) 降到了 O(n)。
具体流程是这样的：
- Prefill 阶段：处理整个prompt，计算所有token的 K、V，并存入缓存。
- Decode 阶段：每步只算一个新token，从缓存中取出历史 K、V，做注意力。
我曾经踩过的坑： KV Cache 的大小跟 batch size × 序列长度 × 层数 × 头数 × 维度 成正比。在 70B 模型上，如果 batch size 开到 64，序列长度 4096，单是 KV Cache 就能吃掉 40GB 显存。所以，多GPU推理时，KV Cache 的显存分配一定要提前规划好，不然很容易 OOM。
3.4 多GPU视角下的KV Cache管理
在多GPU推理中，KV Cache 的管理变得复杂了。我总结了几种常见策略：
| 策略 | 说明 | 适用场景 |
|---|---|---|
| 张量并行 (TP) | 每个GPU只存一部分注意力头对应的 KV Cache | 单机多卡，通信带宽高 |
| 流水线并行 (PP) | 每个GPU存自己负责的那几层 Transformer 的 KV Cache | 跨机多卡，通信延迟高 |
| 分布式 KV Cache | 用专门的缓存服务管理，GPU只负责计算 | 大规模部署，需要动态调度 |
我个人习惯是，在单机8卡场景下，优先用 张量并行。因为 NVLink 带宽足够，KV Cache 的 all-reduce 开销可以接受。但如果是跨机，那就得用流水线并行，减少通信次数。
避坑指南： 我曾经在跨机推理时，用了张量并行，结果发现通信成了瓶颈。后来改成流水线并行，把 KV Cache 按层切分，每台机器只负责几层，延迟直接降了 30%。所以，选哪种策略，一定要看你的硬件拓扑。
3.5 小结：理解这些，你才能玩转多GPU
好了，这一章的内容就这些。总结一下：
- Transformer 的 Decoder-only 架构，决定了推理是自回归的。
- 自注意力的 O(n²) 复杂度，是推理延迟的瓶颈。
- KV Cache 通过缓存历史 K、V，把复杂度降到 O(n)。
- 多GPU下，KV Cache 的切分策略直接影响性能。
嗯，这些概念听起来简单，但实际落地时，细节特别多。下一章，我会带你看看 TensorRT-LLM 是怎么把这些原理变成代码的。到时候，咱们手把手调一个多GPU推理的例子。