Ml-Infra

上一篇从高斯消元推到了低秩分解，最后用 LoRA 举了个例子——\(\Delta W = BA\)。当时想清楚的是数学层面为什么两个小矩阵可以代替一个大矩阵，但 LoRA 本身没有展开。

PagedAttention 那篇写完之后，我觉得 KV Cache 的内存管理问题算是理清了——分页、按需分配、Block Table 映射，思路很干净。末尾我提到下一步要看 Continuous Batching 和 Chunked Prefill，也就是在多个并发请求之间怎么调度 GPU 的计算资源。

上一篇理清了 Prefill 和 Decode 的因果链之后，我开始翻 Decode 阶段的加速方案。结论已经很清楚了：在单请求、小 Batch 的典型场景下，Decode 阶段是 memory-bound 的——每生成一个 token 就要把整个模型权重从显存搬一遍，计算单元大部分时间在等数据（Batch 较大或模型结构不同时，瓶颈也可能转向 compute-bound，但单次请求的延迟优化场景下 memory-bound 是常态）。

前一篇写完之后，KV Cache 在我脑子里的形象已经变了——它不再是一次 forward 里的临时中间量，而是一个跨多轮存活、随 Decode 持续增长、每个请求各自独立的状态。 但"知道它是有状态的"和"知道怎么管它"是两件事。当我开始想具体的内存管理问题的时候，发现事情比想象的复杂。

最近在从传统模型 Serving 转向 LLM Inference，一个让我停下来想了很久的问题是：为什么 LLM 的推理系统看起来和我之前熟悉的那套完全不一样？ 我之前做推荐模型和 CV 模型的部署，排序模型的 Serving 架构虽然在工程上也有不少复杂性（远程 embedding 读取、稀疏参数缓存等），但核心计算模式不难概括：一个请求进来，跑一次 forward，返回结果。LLM Serving 里冒出来一大堆新概念——推理要分成 Prefill 和 Decode 两个阶段，中间要维护 KV Cache，请求调度要做 Continuous Batching，显存管理要用 PagedAttention——我花了些时间才把它们背后的因果链串起来。

写完 Multi-Head Attention 和 Causal Attention 之后，我一直觉得 Attention 这条线已经差不多了——公式拆过了，routing space 的 factorization 也想明白了，Causal Mask 的因果约束也理清了。 直到最近看 Llama 2 的模型配置，发现一行让我停下来的参数：

最近学习 MIT 6.S191 Lecture 9，陆续把 DDP、ZeRO、FSDP2、Tensor Parallelism、Pipeline Parallelism、Sequence Parallelism、Context Parallelism，以及 Expert Parallelism 过了一遍。

上一篇把单个 MoE Layer 的执行流走了一遍：Router 选 Expert、Token Dispatch 按 Expert 重排、Grouped GEMM 统一调度、Combine 加权合回原始 Token 顺序。整条链路假设所有 Expert 都在同一张 GPU 上。

上一篇 MoE 笔记把架构家族理了一遍——稀疏、稠密、Soft MoE，gate 在前还是在后，决定了算力形态和适用场景。 但那篇有一段伪代码一直让我不太舒服： 1for prob, idx in zip(topk_probs.unbind(1), topk_idx.unbind(1)): 2 out += prob.unsqueeze(1) * self.experts[idx](x) 逻辑上没问题。但想象一下真实场景：一个 batch 有几千个 token，每个 token 选 2 个 expert，就是上万个 token-expert 计算任务。用 Python for loop 逐个调 expert 的 forward，GPU 利用率会低得离谱。

上一篇 TP 写完之后，我对 Megatron 的 Column Parallel + Row Parallel 已经比较清楚了：前一层按输出特征切分，后一层按输入特征切分，两层之间不需要 AllGather，只在最后做一次 AllReduce 把部分和加起来。

起因 # 上一篇搞清楚了 Pipeline Parallelism 的调度机制：GPipe 先 Forward 后 Backward，1F1B 通过交替调度降低 Activation 暂存压力，但理想条件下 Bubble Time 相同。

起因 # 上一篇把 ZeRO 的三个 Stage 走了一遍之后，我以为关于参数分片这件事已经想得差不多了——ZeRO-3 / FSDP 把参数、梯度和 Optimizer States 全部分片存储到多张 GPU 上，单卡 Model States 从 112 GB 降到约 28 GB，问题解决。

起因 # 上一篇 Tensor Parallelism 写完之后，我对 TP 的基本原理有了底：把单个算子内部的矩阵乘法拆到多张 GPU 上执行，每张卡只算部分输出，最后通过 AllReduce 聚合。TP 分片的是计算本身。

起因 # 上一篇搞清楚了 AllReduce、ReduceScatter 和 AllGather 三种通信原语，末尾提到 DDP 和 ZeRO Stage 2 选择了不同的通信原语组合，根本原因是参数更新的责任分配不同。

起因 # 之前几篇文章从模型显存组成一路算到 7B 模型的训练账本，结论很直接：BF16 mixed precision + AdamW 的 Model States 就有 112 GB，一张 80 GB A100 连放都放不下。Checkpointing + Offloading 两个一起用可以勉强塞进去，但训练速度要被 PCIe 传输和额外重计算拖累。

起因 # 上一篇讨论完 Checkpointing 和 Offloading 之后，结尾留了一句：单卡优化的代价已经不可接受时，下一步自然是考虑让多张 GPU 分担训练状态。

起因 # 上一篇算完了 7B 模型的训练显存账本。结论很直接：BF16 mixed precision + AdamW 的 Model States 就有 112 GB，加上 30 GB 的 Activations，单卡至少需要 145 GB——一张 80 GB 的 A100 连 Model States 都放不下。

起因 # 之前分别写了推理和训练的显存心智模型和混合精度训练中的显存组成，反复出现一个粗估：BF16 mixed precision + AdamW 大约 16 bytes/parameter。但这个数字每次都只是"1B × 16 = 16 GB"的单位换算，没有真正对着一个具体的 Transformer 架构把账从头到尾算一遍。

最近在看 MIT 6.S191 Lecture 9 的时候，反复听到一句话：Data Movement Matters。 之前写过 Arithmetic Intensity 和 Compute/Memory Bound，也刚写完 FlashAttention 的分块和 Online Softmax。这些文章里"HBM"和"片上存储"出现了很多次，但我发现自己其实一直没有认真想过一个问题：HBM 到底是什么？它全称叫 High Bandwidth Memory——高带宽内存——既然带宽已经很高了，为什么还要费这么大劲减少对它的访问？

之前的几篇把 Self-Attention 的公式拆得比较清楚了，从 QKV 的计算含义到 Multi-Head 的分工。另一条线是 GPU 性能——Arithmetic Intensity、memory bound、HBM 带宽这些概念也已经聊过。关于 HBM 和片上存储的物理区别，可以参考GPU 存储层级那篇。