LLM 的算法实现
一个 decoder-only 大模型,拆开看其实只有五个模块:注意力、归一化、位置编码、前馈网络,再加上把它们串成两阶段推理的 prefill / decode 调度。这套结构从 2017 年的 Transformer [1] 到今天几乎没变,所有「新架构」都是在这五个模块里的某一个做局部替换。本文一个模块一章,把每个模块的核心公式写清楚,最后用最新的 Qwen3.5 [24] 和 DeepSeek-V4 [25] 两个真实模型,看它们分别在这些模块上做了什么选型。
全篇沿用一套符号:
| 符号 | 含义 |
|---|---|
| batch size | |
| 序列(prompt)长度 | |
| 层数 | |
| hidden dim(模型主维度) | |
| 词表大小 | |
| Q 头数 | |
| KV 头数 | |
| 每头维度 | |
| FFN 中间维度 |
张量 shape 按 PyTorch 习惯写成 [B, ..., H];权重矩阵按「输入维 × 输出维」写成 。
Attention — Q/K/V · Scaled Dot-Product · MHA/MQA/GQA/MLA
注意力是 Transformer 里唯一让不同 token 交换信息的模块,其余模块都逐 token 独立。它先把每个 token 的隐状态 投影成查询、键、值三份:
- — 归一化后的输入
- — Q 投影; — K、V 投影
- ,,再 reshape 出 head 维
Scaled Dot-Product Attention
单头内,注意力就是「用 Q 和所有 K 的相似度做加权,去平均 V」:
- — 每对 的相似度矩阵
- — 缩放因子,防止内积过大把 softmax 推进梯度近乎为零的饱和区
- — causal mask,上三角为 (位置 只能看 的位置)
- softmax 沿 K 维归一化成注意力权重,再与 加权求和回到
多头(Multi-Head)把这个操作复制 份,每份在自己的 维子空间里做,最后拼接投回 :
不同头学到不同的关注模式(近邻、句法、指代……),这是注意力表达力的来源。
四种变体:只动 K、V 一侧
推理时, 用完即弃,但每个历史位置的 、 都要留着给后续 token 用 —— 这就是 KV Cache。它随序列线性增长,是长上下文显存和 decode 带宽的主要压力。四种注意力变体的差别,全部集中在「有多少套独立的 K、V,以及 K、V 是否被低秩压缩」,Q 侧始终是 个独立头。
MHA(Multi-Head Attention,原版):每个 Q 头配一套独立 K、V,共 套。表达力最强,cache 也最大。
MQA(Multi-Query Attention)[2]:所有 个 Q 头共用同一套 K、V,只剩 1 套:
cache 收缩到 ,比 MHA 小 倍,但表达力受限,大模型直接用容易掉点。
GQA(Grouped-Query Attention)[3]:把 个 Q 头切成 组,组内共享 K、V —— 是 MHA 与 MQA 之间的连续插值:
退化为 MHA, 退化为 MQA。kernel 里只算 套 K、V,算分数时沿 group 维 broadcast,不真复制张量。这是目前 dense 与多数 MoE 模型(Llama 3、Qwen3、Mistral)的默认选择。
MLA(Multi-head Latent Attention)[4]:GQA 只按头数线性缩 cache;MLA 直接把 K、V 共同压成一个低秩潜向量 ,再为 RoPE 单开一条共享小分支。分四步:
内容分支 —— K、V 共享一份 down-projection,入 cache 的只有潜向量 :
Q 端同样走低秩(训练省显存):
RoPE 解耦分支 —— K 侧只算一份共享的 维向量,所有 head 共用:
拼接后算注意力(内容部分 + RoPE 部分沿 head 维拼起来):
关键在于:MLA 实际缓存的只有 和 ,共 个数。推理时用矩阵结合律把上投影 折进 :
这样 K 内容部分根本不用重建,注意力直接在缓存的潜向量上算。但 RoPE 的旋转角依赖绝对位置,无法折进固定权重 —— 一旦把它加在重建后的 K 上折叠就失效,所以位置信号必须拆成独立的 维浅分支。
KV 维度对比
| 变体 | 单 token cache(per layer,fp16) | 代表模型 | |
|---|---|---|---|
| MHA | GPT-2/3, Llama 1/2 7B | ||
| MQA | PaLM, Falcon | ||
| GQA | 分组 | Llama 3, Qwen3, Mistral | |
| MLA | 低秩压缩 | DeepSeek V2/V3 |
一句话记:MHA → MQA/GQA 是同时省 params、FLOPs、cache;MLA 是加了几段投影、用 params/FLOPs 换 cache —— 把单 token cache 从 KB 量级压到约 1 KB。
线性注意力与混合堆叠
上面四种变体都还是 softmax 注意力,对序列长度是 ,cache 随长度线性膨胀。线性注意力换掉 softmax,用一个可结合的特征映射 把注意力改写成对「状态矩阵」的递推,复杂度降到 ,KV Cache 也压成一个大小固定、与序列长度无关的状态:
- — 到第 步为止累积的状态,等价于把全部历史 KV「压」进一个矩阵
- decode 每步只更新并读一次 ,不再随 cache 变长而变慢
朴素线性注意力记不住长程细节(状态只增不减、无法精确改写)。Gated DeltaNet [26] 给递推补了两样东西:一个门控衰减 (借鉴 Mamba2,让状态可快速遗忘),和一条 delta 规则( 控制对特定 key 的精确改写,而非只做累加):
实践中没人拿纯线性注意力单打:它擅长扛长序列、却弱于精确检索。主流做法是混合堆叠 —— 大多数层用线性注意力压成本,每隔几层插一层全注意力补回检索能力。下面 Qwen3.5 一章就是这种「3 层线性 + 1 层全」的配比。
稀疏化与压缩:让每个 query 只看一部分 KV
full softmax 让每个 query 看全部历史 KV;要在超长上下文里省下 cache 与 FLOPs,有两条正交的路线:
压缩(compression)—— 把连续 个 token 的 KV 事先聚合成 1 个「压缩条目」,注意力在压缩条目上算,序列的有效长度缩到 :
稀疏(sparse selection)—— 不压缩,但为每个 query 只挑最相关的 top- 个历史条目参与注意力,其余跳过。DeepSeek 的做法是一个轻量 lightning indexer 先为每个 query 算一组廉价的 index 分数 ,再取 top-:
两条路线可叠加:先压缩、再在压缩条目上做 top- 稀疏选择,就是下面 DeepSeek-V4 的 CSA;把压缩率拉满、放弃稀疏选择、对压缩后的条目保留稠密注意力,就是它的 HCA。再配一条 sliding-window 分支补局部细节、一组 attention sink 让每个头的注意力权重和可不为 1,就构成了 V4 的混合注意力。
Normalization — LayerNorm vs RMSNorm · pre-norm
归一化稳住每一层输入的数值尺度,让深层网络能训得动。标准 LayerNorm 在特征维上做零均值、单位方差,再仿射:
- — 可学习的 scale / shift
- — 防除零小常数(通常 )
- — 逐元素乘
RMSNorm [5](Llama / Qwen / DeepSeek 的主流选择)去掉均值中心化和偏置,只用均方根缩放:
分母是 的均方根(root mean square),故得名。省了均值、也省了 ,计算量和参数都约减半,实证对质量几乎无损。
位置:pre-norm 而非 post-norm。原始 Transformer 是 post-norm(norm 在残差相加之后),深层容易梯度不稳;现代模型一律用 pre-norm —— norm 放在残差分支内部,主干直接跨过:
pre-norm 让残差主干成为一条不被 norm 干扰的「高速公路」,深到上百层也能稳定训练。
残差之外:hyper-connections。pre-norm 把每层输出加回单一残差主干。Hyper-Connections [27] 把这条主干拓宽成 路并行残差流,用可学习(且依赖输入)的系数矩阵 在层间做混合:
它放宽了「固定深度顺序 + 单一残差」的约束,缓解梯度消失与表示塌缩之间的拉锯。DeepSeek-V4 的 mHC 进一步把残差映射 约束到双随机矩阵流形(谱范数 、映射非扩张),以换取深层堆叠时的数值稳定 —— 见下面 DeepSeek 一章。
位置编码 — 绝对位置 → RoPE → 长度外推
注意力本身对顺序不敏感(打乱 token 结果不变),位置信息必须显式注入。原始 Transformer 用绝对位置编码:给每个位置 造一个 sinusoidal 向量,加到 embedding 上。缺点是只编码绝对位置,且难以外推到训练没见过的长度。
RoPE(旋转位置编码)[6] 换了思路:不加向量,而是让位置以旋转的形式作用在 Q、K 上,使内积自动只保留相对位置。把每个头的 维切成 个二维子空间,位置 对第 对坐标乘一个 2D 旋转矩阵:
- — token 的绝对位置; — 二维子空间索引
- — 第 个子空间的角速度, 是频率衰减基(Llama 取 ,Qwen3 取 )
为什么旋转能编码相对位置?把整块 维旋转记作分块对角矩阵 ,它正交且满足 ,因此:
注意力算 时,每对 的分数只依赖差 —— 绝对位置在内积里被消掉,留下相对位置。频率谱 让 个子空间分到从快到慢的角速度: 周期约 个 token 管近邻, 周期几万 token 管远距离。RoPE 只作用于 Q、K,不作用于 V。
长度外推。训练只见过 ,一旦推理长度超出,低频子空间进入训练分布外,注意力会退化。常见解法都在改 :
- Position Interpolation:,同比例压缩所有频率,简单但牺牲高频精度。
- NTK-aware:只缩放低频、保留高频,等价于放大 。
- YaRN[7]:分频段处理 —— 高频不变、低频按 PI 缩放、中间平滑过渡,再叠一个温度修正抵消注意力熵漂移。Llama 3.1、Qwen3 的长上下文扩展用的都是它。
激活函数与 FFN — GELU · SwiGLU · GeGLU · MoE
注意力之后是逐 token 的前馈网络(FFN),负责非线性变换与知识存储。经典 FFN 先升维、过激活、再降维:
激活 是逐元素的标量非线性。主流选择:
- 是标准正态 CDF,GELU [9] 用于 GPT-2/3、BERT
- SiLU(又名 Swish)形状接近 GELU 但更简单,是 Llama / PaLM 系的门控激活
门控 FFN:GLU 家族
真正的工程拐点,是把激活从「套在投影上」换成「套在门控上」。GLU 家族 [8] 用两个独立升维投影,一个当门、一个当值,逐元素相乘:
选谁就叫什么变体: 是 GeGLU(T5 v1.1), 就是最常用的 SwiGLU(PaLM、Llama、Qwen、DeepSeek):
SwiGLU 比经典 FFN 多一个投影(三矩阵 vs 两矩阵),为对齐参数预算,实现里一般把中间维设成 。
Mixture-of-Experts(MoE)
dense FFN 每个 token 都过同一对权重,参数和 FLOPs 全付。MoE [10] 把 FFN 复制成 份「专家」,每个 token 只挑 top- 份算 —— 总参数线性放大,单 token 激活的参数和 FLOPs 几乎不变,用容量换知识量,不换 compute:
- — router,把隐状态映到 个专家的 logit
- — 被选中的 top- 专家集合; 是归一化后的 combine 权重
- 每个 通常就是一个 SwiGLU,仅权重不共享
朴素 router 会「塌缩」到少数热门专家,需要负载均衡。GShard/Switch 用辅助损失 ( 是路由占比、 是平均概率)软约束;DeepSeek 则改用无辅助损失方案 [12]:给每个专家维护偏置 ,TopK 基于 ,被选少的调高、被选多的调低,只影响选择、不污染梯度。这两块(细粒度 MoE、无损失均衡)是下面 DeepSeek 一章的核心。
Prefill 与 Decoding — GEMM vs GEMV · 算力 vs 带宽
自回归推理分两个阶段,它们的维度和瓶颈截然不同。
Prefill:一次性吃进 个 prompt token,并行走一遍所有层,产出第一个输出 token 和完整 KV Cache。主线张量 shape 是 ,所有矩阵乘都是矩阵 × 矩阵(GEMM),算术强度高,算力受限(compute-bound),很像训练的一次 forward。
Decoding:每步只输入上一步生成的 1 个 token,从 cache 读全部历史 K、V,算出下一个 token。主线 shape 是 ,矩阵乘退化成矩阵 × 向量(GEMV),绝大部分时间花在把权重从显存搬进计算单元,显存带宽受限(memory-bound)。
| 位置 | Prefill | Decode 每步 |
|---|---|---|
| input_ids | ||
| Q | ||
| K/V(从 cache) | ||
| attention scores | ||
| 运算性质 | GEMM(矩 × 矩) | GEMV(矩 × 量) |
| 瓶颈 | 算力 | 内存带宽 |
KV Cache 的形状与增长
每层缓存一对张量,随生成不断变长:
单 token、单层的 cache 大小(fp16,GQA)为 。以 为例即 4 KB/层;32 层、4096 token 的请求约 512 MB —— 长上下文下这就是显存的主要占用,也是 MLA、KV 量化、PagedAttention 等一切 cache 优化的战场。
吞吐的物理上限
decode 每步都要把整个模型权重扫一遍。设参数量 、fp16,权重字节约 ,则单请求()的 decode 吞吐上限:
一个 8B 模型 fp16 权重约 16 GB,H100 @ 3 TB/s 下单 token 光搬权重就要约 ms —— 这是单请求 decode 的物理下限,与算力几乎无关。突破口是 continuous batching:把 个请求拼成一个大 GEMV,同一批权重被 个请求共享,算术强度 ×,吞吐近线性上涨,直到算力或 attention 带宽先撞墙。
计算复杂度上,有没有 KV Cache 差三个量级。带 cache 时,生成 个 token 的总量:
第一项是逐 token 的线性层(decode 每步恒定),第二项是注意力(随 cache 线性增长)。无 cache 则要 ,每步都重算全部历史。「越生成越慢」的本质就是第二项里 越来越长。工程上 FlashAttention [13](把 softmax 与 matmul 融成一个 kernel、显存从 降到 )、PagedAttention [14](分页管理 cache、消除碎片)、投机解码 [15](小模型 draft、大模型一次验证 个)都在这套骨架上做局部手术,公式本身不变。
以 Qwen 为例 — Qwen3.5-397B-A17B
Qwen3.5 [24](2026 年发布,旗舰原生多模态)是阿里最新一代开源模型,旗舰 Qwen3.5-397B-A17B(397B 总参数、每 token 激活约 17B)。相比上一代 Qwen3,它在注意力上换了骨架:不再是清一色的全注意力,而是线性注意力(Gated DeltaNet)与全注意力混合堆叠。拆到五个模块(以下只谈语言主干):
| 模块 | Qwen3.5 的选型 |
|---|---|
| 注意力 | 混合:每 4 层 = 3 层 Gated DeltaNet(线性注意力)+ 1 层门控全注意力(GQA) |
| 归一化 | RMSNorm,pre-norm, |
| 位置编码 | RoPE(),原生上下文 |
| 激活 / FFN | SwiGLU |
| 稀疏化 | 细粒度 MoE,512 路由专家、top-10 + 1 共享专家 |
关键架构数字(取自官方 config):60 层,hidden ;全注意力层 、、head_dim (与 hidden 解耦);线性注意力层 16 个 key 头、64 个 value 头、头维 、卷积核宽 ;专家中间维 ,词表 (byte-level BPE,较 Qwen3 的约 15 万扩到约 25 万)。
三处模块级看点:
- 混合注意力:每 4 层里 3 层用 Gated DeltaNet —— 状态递推、cache 是固定大小的状态矩阵,扛住长序列成本;每 4 层插 1 层全注意力补精确检索。全序列的注意力成本因此随长度近线性、而非平方增长。公式见前面「线性注意力与混合堆叠」小节。
- 门控注意力(gated attention):全注意力层的输出再过一个逐通道 sigmoid 门 ,抑制 attention sink 与大幅激活,稳住大模型训练。
- 细粒度 MoE + 共享专家:512 个路由专家里每 token 选 10 个,外加 1 个所有 token 必过的共享专家 —— 相比 Qwen3 的「无共享专家」,Qwen3.5 转向了 DeepSeek 式的共享专家配置。
一句话:Qwen3.5 = 「线性/全注意力混合 + 门控注意力 + RoPE + SwiGLU + 细粒度 MoE(含共享)」,把长上下文的成本从注意力这一侧压下去。
以 DeepSeek 为例 — DeepSeek-V4-Pro / Flash
DeepSeek-V4 [25](2026 年 4 月,预览版)给出两个 MoE 模型:V4-Pro(1.6T 总 / 49B 激活)与 V4-Flash(284B / 13B 激活),都原生支持 100 万 token 上下文。相比 V3 [18] / V3.2 [19],它保留 DeepSeekMoE 与 MTP,但在注意力和残差两处换了新设计。
混合注意力:CSA + HCA(取代 MLA)。V3 的招牌是 MLA(把 K、V 压成低秩潜向量);V4 更进一步,不再用 MLA,而是把「压缩」和「稀疏」两条路线(前面「稀疏化与压缩」小节)拧在一起,做成两种可交错堆叠的注意力层:
- CSA(Compressed Sparse Attention):先把每 个 token 的 KV 压成 1 个条目,再用 lightning indexer 为每个 query 只选 top- 个压缩条目做注意力(Pro 、Flash )—— 压缩 × 稀疏叠加。
- HCA(Heavily Compressed Attention):把压缩率拉到 (远大于 CSA),但不再稀疏选择,对压缩后的条目保留稠密注意力。
核心注意力用 shared-KV MQA(压缩条目同时当 K 和 V),Q 端走低秩压缩(:Pro 、Flash ),另配一条 的 sliding-window 分支补局部依赖、一组可学习 attention sink 让每个头的注意力权重和可不为 1。RoPE 只作用在每个向量的最后 64 维。KV 存储用混精度:RoPE 那几维用 BF16、其余用 FP8,cache 再砍近半。效果:1M 上下文下 V4-Pro 只需 V3.2 的约 27% 单 token FLOPs、10% KV cache,Flash 更降到约 10% FLOPs、7% KV cache;相对 BF16 GQA(head_dim 128)基线,KV cache 约压到 2%。
DeepSeekMoE(细粒度 + 共享专家)[11] 沿用 V3 的思路,配置略调:MoE 铺满所有 Transformer 层(不再留 dense FFN),前 3 个 MoE 层改用 hash 路由(按 token id 的固定哈希指派专家);路由打分函数从 V3 的 sigmoid 换成 ;负载均衡仍走无辅助损失偏置法 [12],再加一个弱的序列级均衡损失。Pro 每层 1 共享 + 384 路由专家(中间维 ),Flash 1 共享 + 256 路由专家(中间维 ),都是每 token 激活 6 个路由专家。
另外两处模块:
- mHC(Manifold-Constrained Hyper-Connections):把单一残差主干拓成 路并行残差流,并把残差混合矩阵约束到双随机矩阵流形(谱范数 、非扩张),换深层堆叠的数值稳定。原理见前面「残差之外」一段。
- MTP(多 token 预测):深度 1 的浅层模块预测「下下个」token,训练时给主干更密的监督信号,推理时可直接当投机解码的 draft 提吞吐 —— 与 V3 相同。
架构数字汇总:V4-Pro 61 层、hidden ;V4-Flash 43 层、hidden ;词表 ;均用 FP8 训练,后训练对 MoE 专家权重与 indexer QK 做 FP4 量化。把三个模型并排:Qwen3.5 走「线性/全注意力混合 + 细粒度 MoE」,DeepSeek-V4 走「压缩+稀疏混合注意力(CSA/HCA)+ 细粒度 MoE + mHC + MTP」—— 同一张骨架,不同的取舍。
总结 — 五个模块,一张骨架
decoder-only LLM 的推理,说到底就是五个模块的组合:注意力(谁看谁)、归一化(稳数值)、位置编码(编顺序)、FFN/MoE(存知识)、prefill/decode(怎么调度)。每个模块都有一条主公式,变体只是在其中一处做替换 —— GQA/MLA 换的是注意力的 KV 结构,RMSNorm 换的是归一化,RoPE/YaRN 换的是位置编码,SwiGLU/MoE 换的是 FFN。理解了这五条公式,再看 Qwen3.5 与 DeepSeek-V4 的差别,不过是同一张骨架上不同的选型表;后面读任何推理优化论文,也都能定位到它动的是哪个模块。
参考资料 — 论文 · 技术报告
- Vaswani et al. Attention Is All You Need (NeurIPS 2017). arxiv.org/abs/1706.03762
- Shazeer. Fast Transformer Decoding: One Write-Head is All You Need (2019) — MQA. arxiv.org/abs/1911.02150
- Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (EMNLP 2023). arxiv.org/abs/2305.13245
- DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (2024) — MLA 首次提出. arxiv.org/abs/2405.04434
- Zhang & Sennrich. Root Mean Square Layer Normalization (NeurIPS 2019) — RMSNorm. arxiv.org/abs/1910.07467
- Su et al. RoFormer: Enhanced Transformer with Rotary Position Embedding (2021) — RoPE. arxiv.org/abs/2104.09864
- Peng et al. YaRN: Efficient Context Window Extension of Large Language Models (2023). arxiv.org/abs/2309.00071
- Shazeer. GLU Variants Improve Transformer (2020) — SwiGLU / GeGLU. arxiv.org/abs/2002.05202
- Hendrycks & Gimpel. Gaussian Error Linear Units (GELUs) (2016). arxiv.org/abs/1606.08415
- Shazeer et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (ICLR 2017). arxiv.org/abs/1701.06538
- DeepSeek-AI. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models (2024) — 细粒度 + 共享专家. arxiv.org/abs/2401.06066
- Wang et al. Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts (2024). arxiv.org/abs/2408.15664
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (NeurIPS 2022). arxiv.org/abs/2205.14135
- Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023) — vLLM. arxiv.org/abs/2309.06180
- Leviathan et al. Fast Inference from Transformers via Speculative Decoding (ICML 2023). arxiv.org/abs/2211.17192
- Qwen Team. Qwen3 Technical Report (2025). arxiv.org/abs/2505.09388
- Qwen Team. Qwen2.5 Technical Report (2024). arxiv.org/abs/2412.15115
- DeepSeek-AI. DeepSeek-V3 Technical Report (2024) — MLA + MoE + MTP + FP8. arxiv.org/abs/2412.19437
- DeepSeek-AI. DeepSeek-V3.2-Exp: Boosting Long-Context Efficiency with DeepSeek Sparse Attention (2025). github.com/deepseek-ai/DeepSeek-V3.2-Exp
- Meta AI. The Llama 3 Herd of Models (2024). arxiv.org/abs/2407.21783
- Jiang et al. Mixtral of Experts (Mistral AI, 2024). arxiv.org/abs/2401.04088
- Fedus et al. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (JMLR 2022). arxiv.org/abs/2101.03961
- Adam Casson. Transformer Inference Arithmetic — 推理 FLOPs / KV cache 心算. kipp.ly
- Qwen Team. Qwen3.5-Omni Technical Report (2026) — Gated DeltaNet + 门控全注意力混合. arxiv.org/abs/2604.15804
- DeepSeek-AI. DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence (2026) — CSA/HCA + mHC. arxiv.org/abs/2606.19348
- Yang et al. Gated Delta Networks: Improving Mamba2 with Delta Rule (ICLR 2025). arxiv.org/abs/2412.06464
- Zhu et al. Hyper-Connections (ICLR 2025, ByteDance). arxiv.org/abs/2409.19606