Dropout 算子优化与前沿论文阅读

1. Dropout算子背景

Dropout 是深度学习中用于防止过拟合的正则化算子，通过在训练过程中随机将部分神经元输出置零，并对保留元素做缩放，使输出期望不变（Inverted Dropout）。
设丢弃概率为 \(p\)，保留概率为

\[q = 1 - p\]

，对输入元素 \(x_i\)，Dropout 采样

\[m_i \sim \mathrm{Bernoulli}(q)\]

并输出

\[y_i = \frac{m_i}{q}x_i\]

因此

\[\mathbb{E}[y_i] = x_i\]

在训练过程中，前向过程需要生成随机mask,这个mask可以保留，在反向过程中重复使用，也可以使用确定性的随机算法，在反向过程中再次生成。推理过程中不需要使用dropout。

drop算子步骤主要包括三个部分：

生成可复现的伪随机数；
将随机数变成 Bernoulli mask；
应用 mask 并执行 \(1/q\) 缩放。

其中第一步包含了较多的计算。是主要优化与关注的点。

2. Philox：Dropout 中常用的并行随机数算法

Dropout 随机数同时满足三点：统计性质合理、跨运行可复现、适合大规模并行。
串行 PRNG 不适合这个场景，因为线程之间会共享状态并引入同步或顺序依赖。

在GPU上比较常用的随机数算法是Philox, Philox 属于 Counter-Based RNG。它不从上一个随机状态递推下一个随机状态，而是把 counter 和 key 映射成随机输出：

\[R_i=f_{key}(counter_i)\]

每个线程只要知道 seed、offset 和元素索引，就能独立生成随机数。

常见实现是 Philox4x32-N。它输入 4 个 32-bit counter 和 2 个 32-bit key，一次输出 4 个 32-bit 随机字：

\[(c_0,c_1,c_2,c_3),(k_0,k_1)\rightarrow(r_0,r_1,r_2,r_3)\]

单轮 Philox 的核心是整数乘法、拆高低位、异或和 key bump。以 Philox4x32 为例：

\[P_0=M_0c_0,\qquad P_1=M_1c_2\]

\[(H_0,L_0)=\mathrm{split}_{64}(P_0),\qquad (H_1,L_1)=\mathrm{split}_{64}(P_1)\]

更新为

\[\begin{aligned} c_0' &= H_1 \oplus c_1 \oplus k_0 \\ c_1' &= L_1 \\ c_2' &= H_0 \oplus c_3 \oplus k_1 \\ c_3' &= L_0 \end{aligned}\]

Dropout mask 可通过整数阈值生成：

\[T=\lfloor q\cdot 2^{32}\rfloor,\qquad m_i=\mathbf{1}[u_i<T]\]

这里 \(u_i\) 是 Philox 输出的 32-bit 随机字。整数阈值比较避免了逐元素浮点归一化。

3.

Attention 的前向计算可以拆成四步：打分、归一化、随机丢弃、加权求和。设单个 batch、单个 head 内
\(@@MATHJAX_DISPLAY_11@@\)，则标准 Attention 为

\[S=\frac{QK^T}{\sqrt{dH}}\]

对 causal mask 或 padding mask 处理后，执行 Softmax：

\[P=\mathrm{softmax}(S)\]

Dropout作用于 Softmax 之后的概率矩阵 \(P\)，训练阶段采样 mask \(M\)，并做 inverted dropout：

\[\widetilde{P}=\frac{M\odot P}{q},\qquad M_{ij}\sim\mathrm{Bernoulli}(q)\]

随后再与 \(V\) 相乘：

\[O=\widetilde{P}V\]

因此 Attention Dropout 位于 Softmax 与 PV 矩阵乘之间。这个位置很关键：Dropout mask 的形状与注意力概率矩阵一致，而不是与 hidden states 一致。

若 batch size 为 \(B\)，head 数为 \(nH\)，序列长度为 \(SQ\)，则需要随机 mask 的元素数量为

\[N_{mask}=B\cdot nH\cdot SQ^2\]

这就是 Attention Dropout 与普通 activation dropout 的根本区别。普通 Dropout 的随机数规模通常随激活元素数线性增长；Attention Dropout 的随机数规模随序列长度二次增长。

4. Attention Dropout 的瓶颈在哪里

Attention Dropout 的瓶颈不是 \(P\) 上乘 mask 的浮点乘法，而是生成 \(M\) 所需的随机数以及随机结果的存取方式。

如果使用独立 RNG kernel，流程为：

\[RNG \rightarrow M \rightarrow Attention\]

此时至少有两类开销：

Philox 生成随机数的整数计算开销；
mask 写回 HBM、Attention 再读回 mask 的带宽开销。

若 mask 以 1 byte 存储，Attention Dropout 的 mask 容量为

\[C_{mask}=B\cdot nH\cdot SQ^2\ \mathrm{bytes}\]

若 bit-pack 为 1 bit，则容量下降为

\[C_{mask}^{bit}=\frac{B\cdot nH\cdot SQ^2}{8}\ \mathrm{bytes}\]

但 bit-pack 需要额外位操作和解包逻辑，不能只按容量降低比例估算收益。

更重要的是 Philox 本身的计算结构。Philox 依赖整数乘法、异或、counter 更新和多轮混合，主要消耗整数 ALU 与 issue bandwidth。它不是 Tensor Core / MMA 计算，也不是单纯 HBM 带宽问题。

对 Transformer block 中主要算子，论文给出的复杂度趋势为：

\[T_{GEMM}=O(B\cdot SQ\cdot dH^2\cdot nH^2)\]

\[T_{Attention}=O(B\cdot nH\cdot dH\cdot SQ^2)\]

\[T_{RNG}=O(B\cdot nH\cdot SQ^2)\]

这组关系说明：模型维度越大，GEMM 越重；序列越长，Attention 与 RNG 越重。长序列训练会使 RNG 从附属成本变成 Attention 关键路径上的显著成本。

还需要从微架构资源看这个问题。GEMM 主要消耗 Tensor Core / MMA 与 L2 带宽；Philox RNG 主要消耗整数 ALU 与 issue bandwidth；FlashAttention 除 MMA 外，还包含 Softmax、归一化、寄存器读写、特殊函数与控制逻辑，会消耗 register file bandwidth、MUFU/FMA 和 issue bandwidth。

因此单个 kernel 的时间更适合写成多资源瓶颈模型：

\[T_{kernel}=\max(T_{MMA},T_{HBM},T_{L2},T_{RF},T_{issue},T_{ALU},T_{MUFU},T_{FMA},\ldots)\]

这个模型直接决定后续三种方案的优劣：baseline 暴露 RNG；fusion 减少 mask 访存但与 Attention 竞争低层资源；overlap 则试图把 RNG 放到与其资源互补的 GEMM 旁边。

5. Dropout 优化的主要方向

围绕 Attention Dropout，优化方向可以分为三类。

第一类是 RNG 本身优化。重点包括充分利用 Philox4x32 一次输出的 4 个随机字、使用整数阈值比较、优化 counter 到元素的映射、减少不必要的数据格式转换。这类优化直接降低 \(T_{RNG}\)。

第二类是 mask 表示与访存优化。包括 bit-pack、tile 对齐、按块生成和消费 mask。它减少容量和 HBM 读写，但可能增加位操作和解包成本。

第三类是执行调度优化。核心问题是：RNG 应该作为独立 kernel、融入 Attention，还是与其他算子并发。这个选择不能只看是否减少 HBM 访问，而要看算子之间是否竞争同一关键资源。

论文 Reducing the Cost of Dropout in Flash-Attention by Hiding RNG with GEMM 主要讨论第三类问题。

6. FlashAttention 场景：为什么 RNG 成本被暴露

FlashAttention 通过分块计算和在线 Softmax 避免显式写出完整注意力矩阵，显著降低 HBM IO。其 Transformer block 中主要计算包括 QKV GEMM、Attention、Projection GEMM 和 MLP GEMM。

当 Attention 的 IO 被优化后，原先可能被 HBM 开销掩盖的 RNG 成本会更加突出。尤其在训练阶段，Softmax 后的 \(P\) 需要执行 Dropout，随机数规模仍然是 \(B\cdot nH\cdot SQ^2\)。因此，FlashAttention Dropout 的优化目标不是改变 Dropout 语义，而是降低 RNG 对关键路径的影响。

7. Baseline：独立 RNG 顺序执行

Baseline 方案将 RNG 作为独立阶段放在 Attention 前：

\[GEMM\rightarrow RNG\rightarrow Attention\]

对应时间为

\[T_{baseline}=T_{GEMM}+T_{RNG}+T_{Attention\_drop}\]

这种方案的优点是边界清晰：RNG kernel 负责生成 mask，Attention kernel 负责读取并消费 mask。随机数序列和 Attention 计算可以分别验证。

问题也很直接。第一，\(T_{RNG}\) 完全暴露在关键路径上；第二，mask 需要写回 HBM，再由 Attention 读回；第三，随机数规模随 \(SQ^2\) 增长。因此 baseline 在长序列下会快速变差。

8. SOTA：RNG-Attention Fusion

已有 SOTA 方案将 RNG 融入 FlashAttention kernel：

\[GEMM\rightarrow Attention_{RNG\ fused}\]

它的逻辑是：既然 Dropout mask 只被 Attention 消费，就在 Attention 内部生成并立即使用，避免完整 mask 写回 HBM，同时减少 kernel launch 和同步。

这个方案在带宽层面成立，但隐藏 RNG 的能力有限。论文实测显示，RNG-Attention fusion 只能隐藏约 10% 到 20% 的 RNG runtime。原因是 FlashAttention 本身并非纯 Tensor Core-bound。Softmax、归一化、寄存器读写、特殊函数和控制逻辑已经消耗 issue bandwidth、register file bandwidth、MUFU/FMA 等资源；Philox 又消耗整数 ALU 和 issue bandwidth。

所以 fusion 的问题不是方向错误，而是放置位置不理想：它减少了 mask 中间存储，却把 RNG 放进了一个低层资源同样紧张的 Attention kernel。

9. 论文方案：RNG-GEMM Overlap

论文提出将 RNG 从 Attention 中拆出，并与前序 GEMM 并发执行：

执行结构为

\[\begin{aligned} & GEMM \longrightarrow Attention \\ & RNG \ \ \ \nearrow \end{aligned}\]

时间模型为

\[T_{overlap}=\max(T_{GEMM}^{interfere},T_{RNG}^{interfere})+T_{Attention\_drop\_only}\]

该方案成立的关键是资源互补。GEMM 主要占用 Tensor Core / MMA 和 L2；RNG 主要占用整数 ALU 和 issue stage。相比 RNG-Attention fusion，RNG 与 GEMM 的低层资源冲突更弱，因此更适合 overlap。

论文将 workload 划分为三个区域：

\(T_{GEMM}\gg T_{RNG}\)：RNG 可隐藏，但总体收益有限；
\(T_{GEMM}\gtrsim T_{RNG}\)：RNG 足够大，且能被 GEMM 大部分覆盖，是最理想区域；
\(T_{RNG}>T_{GEMM}\)：GEMM 结束后仍有 RNG 尾部暴露。

典型 LLM 训练配置多处在第 2 类区域，所以 RNG-GEMM overlap 能有效缩短关键路径。

代价是需要保存随机结果。若每个元素存 1 bit，则容量为

\[C_{mask}=B\cdot nH\cdot SQ^2\ \mathrm{bits}\]

论文进一步采用 chunk pipeline，只保存当前块或下一块 mask，避免完整 \(SQ^2\) mask 常驻。

10. 实验结论

论文在 H100 / GH100 HBM3 80GB 上进行建模和 CUDA 实测，典型配置包括 head dimension 128、sequence length 2K 到 64K、embedding dimension 4K 到 16K。

主要结果是：

Llama3-like workload：overlap 相比 baseline 获得 1.26x speedup，相比 fusion 获得 1.22x speedup；
GPT-4-like workload：overlap 相比 baseline 约 1.25x，相比 fusion 约 1.2x。

模型与 silicon measurement 的平均误差约为 2%。

Philox 轮数实验表明，Philox 5 runtime 约为 Philox 7 的 81%，Philox 3 runtime 约为 Philox 7 的 67%。但 fusion 场景收益不随轮数线性变化，因为瓶颈更多受 issue 等共享资源限制。

未来 GPU 若继续提升矩阵吞吐，而非矩阵资源提升较慢，RNG 的相对开销会更突出。

11. 对 NPU 的启示

在 NPU 上不能直接套用 H100 结论，但分析方法可以迁移。首先需要确认 Dropout 是否使用 Philox 或类似 CBRNG，Philox4x32 的 4 个输出是否充分利用，mask 是否 bit-pack，以及 RNG 主要受 DDR、UB/L2、Vector/Scalar 还是调度资源限制。

如果 MatMul 主要使用 Cube/Matrix 单元，而 RNG 主要使用 Vector/Scalar 或整数流水，那么 RNG 与 MatMul 也可能具备资源互补性。此时可以考虑类似 RNG-GEMM overlap 的调度：RNG 与前序 MatMul 并发，Attention 等二者完成后消费 mask。

NPU 上更稳妥的做法是建立平台相关模型：

\[T_{kernel}=\max(T_{cube},T_{vector},T_{scalar},T_{UB},T_{L2},T_{DDR},T_{issue},\ldots)\]

再根据 shape 和 profiler 结果在 fusion、overlap、chunk pipeline 之间选择。