AI Infra学不完学不完
最近读了腾讯写的文章,简单自己做了下总结和理解,又有了新的输入,不愧是腾讯👍
一、学到的核心知识
1. 大模型推理的两大基石
- Continuous Batching(连续批处理):传统批处理要求一个 batch 内所有请求"同时开始、同时结束",但 LLM 输出长度不可预测,会浪费 GPU。vLLM 把调度粒度从 request level 下沉到 token level——每一步看哪些请求"差几个 token",就调度几个 token,把短请求填进 batch 的空隙里。
- Paged Attention:传统做法一次为请求申请全部显存,碎片严重;Paged Attention 仿照操作系统的页表,KV Cache 按固定
block_size(默认 16)分块,用block_table维护虚拟→物理块映射,按需分配。代价是 block 间访存不再连续会有 L2 Cache miss,但用少量访存带宽换大幅吞吐提升,整体是赚的。
2. Prefill vs Decode 的本质差异
- Prefill:一次性处理完整 prompt,QKV_Proj/O_Proj/MLP 全是稠密 GEMM,算力密集 (Compute-bound)。
- Decode:自回归每次只算 1 个 token,矩阵乘退化成矩阵向量乘 (GEMV),瓶颈在读权重和 KV Cache 的访存带宽 (Memory-bound)。
- Continuous Batching 的真正价值:通过 token 级调度让多请求复用模型权重,把 Decode 阶段的 GEMV 重新变回 GEMM,摊薄权重访存开销。
3. FlashAttention:通过 Online Softmax(维护 Running Max 和 Running Sum + 动态缩放因子)解除全局数据依赖,把 Q·K → Softmax → ·V 融合成一个 CUDA Kernel,在 SRAM 里分块完成计算,不再把 N×N 中间矩阵写回 HBM,是打破内存墙的关键。
4. 一些工程细节也很有趣
- RoPE 不实时算 sin/cos,而是引擎启动时预计算
cos_sin_cache——SFU 算超越函数吞吐只有 FP32 ALU 的 1/4,用几 MB 的小表换掉每次 forward 几百万次调用,空间换时间。 - Fused QKV 和 Fused Gate/Up 都是把多个 Linear 按列拼起来,一次宽 GEMM 替代多次窄 GEMM,减少 Kernel Launch。
- TP 切分时 QKV/Gate-Up 走 Column Parallelism,O_Proj/Down_Proj 走 Row Parallelism(最后 AllReduce 聚合)——理解了 shape 流,TP 一下就懂了。
二、个人理解
读完最大的震动是:大模型推理的数学其实并不深奥,难的全是工程 Trade-off。
Paged Attention 牺牲访存连续性换显存利用率,FlashAttention 用重计算换 IO,最后一层 MLP 明明可以只算最后一个 token 但 vLLM 选择"老老实实全算一遍"——每一个决策都是物理极限下的权衡,没有"最优解"只有"在当前硬件约束下最划算的解"。
另一个收获是看到了 Attention 为什么是大模型的命门:其他模块都能通过 Continuous Batching 摊薄访存,只有 Attention 因为 KV Cache 不能跨请求共享 + 计算随上下文长度线性膨胀,始终无法均摊。这就解释了为什么 DeepSeek MLA、Linear Attention、Flash-Decoding 这些工作都在围绕 KV Cache 做文章——根因都在这里。
"AI Infra 的东西真的多,根本学不完"特别有共鸣。比起 Agent 圈里不断重命名的"上下文工程""harness 工程",AI Infra 这边每个名词背后都有扎实的数学和工程支撑。下一步我打算啃 FlashAttention 论文 + MLA 原理,把这篇当作整个 Infra 知识树的索引。
原文指路:AI Infra入门干货总结:大模型是如何高效推理的 ·