拆解大模型几项核心操作背后的数学与 Infra 优化逻辑

📌 一句话摘要

本文从数学原理与 Infra 优化视角，系统拆解大模型核心操作（RMSNorm、Softmax、Causal Mask、Sampling）的设计动机与高性能实现逻辑。

📝 详细摘要

文章以「数学等价变换与精度妥协换取硬件效率」为主线，深入剖析大模型推理与训练中的关键算子。首先从数值稳定性问题出发，解释归一化（LayerNorm → RMSNorm）的演进逻辑，包括方差、标准差、Z-score 等基础统计概念的推导，以及 Pre-Norm 与 Post-Norm 的梯度特性对比。随后聚焦 Softmax，分析 Safe Softmax、缩放因子、方差膨胀与梯度消失的数学成因，并延伸至 Causal Mask 的块稀疏实现。核心篇幅用于讲解 Online Softmax 及其在 FlashAttention 系列中的应用，包括 LSE 压缩、Split-K、反向传播重计算与序列并行。最后介绍 Sampling 阶段的 Gumbel-Max Trick 及其在 vLLM 中的高效实现。全文贯穿数学推导与工程实现之间的权衡讨论，适合具备一定深度学习基础的读者。

💡 主要观点

- Infra 优化的本质是用数学等价变换或精度妥协换取硬件效率。 从 RMSNorm 砍掉均值计算到 Online Softmax 的增量修正，再到 Gumbel-Max 将采样转为 argmax，核心思路始终是减少 HBM 访存、消除数据依赖、适配 GPU 并行架构。

RMSNorm 相比 LayerNorm 的核心优势在于打破数据依赖、减少访存与计算量。 RMSNorm 无需计算均值，只需一次单向 Reduction，省去减法操作与偏置项，在 fused kernel 中寄存器压力更低，成为现代 LLM 的标准选择。

Softmax 的工程实现必须解决数值溢出与梯度消失问题。 Safe Softmax 通过减最大值防止上溢；缩放因子 将点积方差拉回 1，避免 Softmax 退化为 Hard Max 导致梯度消失。

Online Softmax 与 FlashAttention 通过增量修正实现单次 HBM 遍历。 维护 running max 与 running sum，在新块到来时用缩放因子修正历史状态，使 Attention 计算无需保存中间矩阵，将空间复杂度从 O(N²) 降至 O(N)。

Gumbel-Max Trick 将随机采样统一为 argmax，提升 GPU 并行效率。 通过将概率除以指数分布噪声再取 argmax，vLLM 实现了贪心与随机采样的代码路径统一，消除了分支发散与前缀和开销。

💬 文章金句

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

• 神经网络真正关心的并非数值的绝对大小，而是特征之间的相对差异。
• Pre-Norm 治好了训练崩溃，让大模型的规模化（Scaling up）成为可能。然而，Pre-Norm 并非完美，它是用前向传播的妥协换取了反向传播的稳定。
• 在 GPU 底层，除法指令比较昂贵。但由于我们存的是 LSE 标量，重算公式变成了 e^{z_i - LSE}。LSE 利用对数法则化除法为减法。

📊 文章信息

AI 初评：89

来源：腾讯技术工程

作者：腾讯技术工程

分类：人工智能

语言：中文

阅读时间：92 分钟

字数：22971

标签： LLM, AI Infra, 模型训练与推理, FlashAttention, Softmax

阅读完整文章