万字入门 AI Infra：深入理解大模型中的数学与 Infra 优化

📌 一句话摘要

本文从数学原理和工程优化两个维度，深入拆解了大模型推理中 RMSNorm、Softmax、Causal Mask 等核心操作的底层逻辑，揭示了 AI Infra 优化的本质是用数学等价变换和精度妥协换取硬件效率。

📝 详细摘要

本文是一篇面向 AI Infra 入门与进阶的深度技术文章，旨在帮助读者融会贯通理解大模型从数学、算法到硬件系统工程的完整链路。文章以「AI 不需要脑子」的犀利观点开篇，引出大模型推理中数学逻辑并不复杂，但推理系统（如 vLLM）却异常复杂的核心矛盾。作者随后聚焦于大模型中的几个核心操作：RMSNorm、Softmax、Causal Mask 和 Sampling，逐一拆解其背后的数学原理与 Infra 优化逻辑。文章从方差、标准差等基础统计学概念讲起，逐步推导到 LayerNorm 和 RMSNorm 的数学公式与工程实现差异，强调了 RMSNorm 在减少计算量和访存开销上的优势。接着深入分析了 Softmax 的概率归一化机制、Safe Softmax 的平移技巧以及 Attention 中除以 √dₖ 的方差缩放原理。文章还详细解释了 Causal Mask 的实现与优化，以及 Online Softmax 和 FlashAttention 如何通过 Kernel 融合和分块计算突破内存墙。全文贯穿一个核心观点：AI Infra 优化本质上是在用数学上的等价变换，或对精度的适度妥协，去换取更高的硬件利用率和极致的推理速度。文章最后预告了后续关于矩阵乘法和 FlashAttention 的内容。

💡 主要观点

- AI Infra 优化的本质是用数学等价变换和精度妥协换取硬件效率。 文章通过 RMSNorm、Online Softmax 等案例，展示了如何通过数学上的等价变换（如砍掉均值计算）或对精度的适度妥协，来减少计算量和访存开销，从而提升硬件利用率和推理速度。

RMSNorm 通过砍掉均值计算，在保持效果的同时显著提升了 Infra 效率。 相比 LayerNorm，RMSNorm 省去了计算均值和减均值的步骤，打破了数据依赖，减少了寄存器/SRAM 占用和 ALU 指令，使得 Kernel 融合更丝滑，是现代大模型的标准选择。

Softmax 的工程实现（Safe Softmax 和 Online Softmax）是解决数值稳定性和内存墙的关键。 Safe Softmax 通过减去最大值防止溢出；Online Softmax 通过维护局部最大值和指数和，将 Softmax 的多次 HBM 访存减少为一次，为 FlashAttention 的 Kernel 融合奠定了基础。

Attention 中的 √dₖ 缩放是防止梯度消失的数学保障，其有效性更多来自实验验证而非严格假设。 虽然推导基于 q 和 k 独立同分布的理想假设，但 √dₖ 在模型初始化时稳住了 Logits 量级，防止了 Softmax 退化到 Hard Max 导致的梯度消失，是大模型实验科学的典型体现。

Causal Mask 在 FlashAttention 中通过块稀疏调度实现零 HBM 开销。 FlashAttention 不再生成全局 L×L 掩码矩阵，而是在 Tile 调度层面根据行列索引跳过对角线以上的分块，仅在边界块内处理掩码，将计算量和访存量降低近 50%。

💬 文章金句

- Infra 优化，本质上就是在用数学上的等价变换，或者对精度的适度妥协，去换取更高的硬件利用率和极致的推理速度。

• RMSNorm 不需要算 μ，直接计算每个元素的平方和即可。这意味着只需要一次单向的 Reduction 操作。
• Softmax 的核心特性是：它只关心数值之间的'相对差值'，根本不关心数值的绝对大小。
• Pre-Norm 治好了训练崩溃，让大模型的规模化（Scaling up）成为可能。
• 不管黑猫白猫，能捉老鼠的就是好猫。

📊 文章信息

AI 初评：92

来源：腾讯云开发者

作者：腾讯云开发者

分类：人工智能

语言：中文

阅读时间：88 分钟

字数：21885

标签： AI Infra, 大模型推理, RMSNorm, Softmax, Causal Mask

阅读完整文章