1. Dropout 算子：从正则化语义到算子实现

Dropout 是神经网络训练中广泛使用的正则化方法。其基本思想是在训练阶段以概率 p 随机丢弃部分元素，使模型不能过度依赖某些局部激活模式，从而降低过拟合风险。推理阶段通常不再执行随机丢弃。

工程实现中更常见的是 Inverted Dropout。设保留概率为：

1
keep_prob = 1 - p

对输入元素 x_i，Dropout 首先生成一个 Bernoulli 随机变量：

1
keep_i = random_i < keep_prob

随后执行 mask 与缩放：

1
y_i = keep_i ? x_i / keep_prob : 0

这样可以保证训练阶段输出期望与推理阶段保持一致：

1
E[y_i] = keep_prob * x_i / keep_prob + (1 - keep_prob) * 0 = x_i

因此，Dropout 算子可以分解为两个部分：

  1. 随机 mask 生成：为每个元素生成随机数，并根据 keep_prob 判定是否保留；
  2. mask 应用与缩放：对保留元素乘以 1 / keep_prob，对丢弃元素置零。

对于普通 activation dropout，元素规模通常与激活张量大小线性相关。然而在 Attention Dropout 中，Dropout 作用于 Softmax 之后的注意力概率矩阵，其规模随序列长度平方增长。

若 batch size 为 B，head 数为 nH，序列长度为 SQ，则 Attention Dropout 所需 mask 的元素数近似为：

1
N_mask = B * nH * SQ^2

这意味着当序列长度从 8K 增长到 16K 时，mask 规模会增长 4 倍。对于长上下文训练，Dropout 不再只是一个边缘的 element-wise 操作，其随机数生成成本可能成为 Attention 层的重要组成部分。

2. Philox 随机数生成算法

Dropout 的随机性不要求物理意义上的真随机，但对伪随机数生成器有较强工程约束：

  1. 统计性质可接受：随机数应能支持近似 Bernoulli 采样；
  2. 结果可复现：相同 seed、相同 offset 和相同元素位置应得到相同结果；
  3. 并行友好：不同线程不应依赖全局串行 RNG 状态；
  4. 支持跳转访问：线程能够根据元素索引直接计算对应随机数。

Philox 属于 Counter-Based Random Number Generator（CBRNG）。与线性递推式 PRNG 不同，Philox 不通过前一随机状态生成后一随机状态，而是将 key 与 counter 输入一个确定性混合函数：

1
R_i = f_key(counter_i)

因此，第 i 个随机数可以由 counter_i 直接计算得到。这一性质非常适合 GPU/NPU 上的并行 Dropout：每个线程或线程组只需要知道 seed、offset 与元素索引，就可以独立生成对应随机数，无需串行维护共享状态。

2.1 Philox4x32 的输入与输出

常见实现为 Philox4x32-N，其中：

  • 4x32 表示 counter 由 4 个 32-bit 字组成；
  • key 通常由 2 个 32-bit 字组成；
  • 一次 Philox 调用输出 4 个 32-bit 随机字；
  • N 表示混合轮数，例如 Philox4x32-10、Philox4x32-7。

可抽象表示为：

1
(c0, c1, c2, c3), (k0, k1)  ->  (r0, r1, r2, r3)

这意味着一次 Philox 调用天然能够服务 4 个 dropout 元素。如果实现每次只消费其中 1 个或 2 个随机字，就会降低 RNG 生成效率。此前给出的算子优化笔记中也强调过这一点：Philox4x32 的 4 个输出应尽可能充分利用，否则每个元素承担的 Philox 轮函数成本会被放大。

2.2 Philox 单轮计算过程

以 Philox4x32 为例，每一轮的核心操作包括整数乘法、高低位拆分、异或和 key bump。给定两个常数：

1
2
M0 = 0xD2511F53
M1 = 0xCD9E8D57

首先执行两个 32x32 到 64-bit 的整数乘法：

1
2
P0 = M0 * c0
P1 = M1 * c2

将乘积拆分为高 32-bit 与低 32-bit：

1
2
(H0, L0) = split64(P0)
(H1, L1) = split64(P1)

随后构造下一轮 counter：

1
2
3
4
c0' = H1 ^ c1 ^ k0
c1' = L1
c2' = H0 ^ c3 ^ k1
c3' = L0

每轮结束后，key 通过 Weyl 常数递增：

1
2
k0 += 0x9E3779B9
k1 += 0xBB67AE85

完整 Philox 会重复上述过程若干轮。轮数越高，随机混合越充分，但整数指令数量也越多。论文主要分析 Philox 7，并额外比较 Philox 5 与 Philox 3，以评估 RNG 成本下降后 overlap 方案的收益变化。

2.3 从 Philox 输出到 Dropout Mask

Philox 输出为 32-bit 无符号整数。生成 dropout mask 时，可以将随机整数映射到 [0, 1) 区间，再与 keep_prob 比较：

1
2
3
u = philox(counter, key)
x = u / 2^32
keep = x < keep_prob

工程中也可将阈值转为整数，避免浮点转换：

1
2
threshold = round(keep_prob * 2^32)
keep = u < threshold

对于 Dropout，整数阈值方法的离散误差约为 2^-32，通常可以忽略。由此可见，Dropout 的随机 mask 生成主要由 Philox 的整数混合过程决定，而非由后续 mask 应用决定。

3. 复杂度与带宽分析：为什么 Attention Dropout 会变贵

在分析论文提出的优化之前，需要先明确三个不同层次的成本：随机数数量、存储带宽和计算资源。

3.1 随机数规模复杂度

Attention Dropout 的随机数规模为：

1
N_rng = B * nH * SQ^2

因此，其随机数数量随序列长度平方增长。相比之下，Transformer Block 中 GEMM 的复杂度与 embedding dimension、head 数和序列长度相关。论文中给出的趋势可概括为：

1
2
3
Runtime(GEMM)      = O(B * SQ * dH^2 * nH^2)
Runtime(Attention) = O(B * nH * dH * SQ^2)
Runtime(RNG)       = O(B * nH * SQ^2)

其中 dH 是 head dimension。可以看到：

  • 增大 embedding dimension 时，GEMM 增长更显著；
  • 增大 sequence length 时，Attention 与 RNG 增长更显著；
  • RNG 与 Attention 都具有 SQ^2 项，但 RNG 不含 dH，其相对占比取决于序列长度与模型宽度的组合。

这解释了为什么长序列训练中 RNG 可能从辅助成本变为显著成本。

3.2 带宽成本：mask 写回与读取

如果采用独立 RNG kernel，则随机结果通常需要写入显存，随后由 Attention kernel 读取。若每个元素只存 1 bit mask，理论容量为：

1
B * nH * SQ^2 bits

如果以 byte 或更宽格式存储，容量还会进一步放大。对于长序列单卡训练，这一存储需求可能不可忽略。

不过，带宽并不是唯一瓶颈。即使 mask 被 bit-pack，Philox 本身仍包含大量整数乘法与混合指令。因此，RNG kernel 很可能同时具有计算瓶颈与 issue 瓶颈，而不只是 HBM 带宽瓶颈。

3.3 计算与微架构资源成本

传统 roofline 往往关注 FLOPS 与 HBM 带宽，但本文讨论的问题需要更细粒度的资源视角。论文将以下资源纳入分析：

资源 / Limiter 含义 相关算子
MMA Math Tensor Core / 矩阵计算吞吐 GEMM、Attention MatMul
HBM Read BW 主存读带宽 GEMM、Attention、mask 读取
L2 Read BW L2 到 SM 的读带宽 tiled GEMM、Attention
RF Read BW Register File 读带宽 Softmax、Attention 非矩阵部分
Issue Stage 指令发射带宽 RNG、Softmax、控制逻辑
ALU Pipe 整数或普通算术流水 Philox RNG
MUFU Pipe 特殊函数单元 Softmax 中指数相关操作
FMA Pipe 普通浮点算术流水 Attention 非 Tensor Core 部分

对单个 kernel，可写成：

1
T_kernel = max(T_MMA, T_HBM, T_L2, T_RF, T_issue, T_ALU, T_MUFU, T_FMA, ...)

这一模型的关键含义是：两个算子能否有效融合或并发，不取决于它们名称是否不同，也不只取决于是否减少 HBM 读写，而取决于它们是否竞争同一组关键硬件资源。

4. Baseline 方案：顺序执行 RNG 与 Attention

论文中的 baseline 是最直接的执行策略：先运行前序 GEMM，再运行独立 RNG kernel，最后运行 Attention kernel 消费随机结果并应用 dropout。

其执行顺序可表示为：

1
GEMM -> RNG -> Attention

对应运行时间近似为：

1
T_baseline = T_GEMM + T_RNG + T_Attention_with_drop

4.1 为什么这样做

该方案具有较强的工程直接性：

  1. RNG 与 Attention 分离，kernel 边界清晰；
  2. Attention 只需要读取已经生成的 mask 或随机 bit；
  3. 正确性容易验证，随机数序列和 dropout mask 可以独立测试；
  4. 对现有算子库的改动较少。

4.2 问题在哪里

Baseline 的主要问题在于 RNG 位于关键路径上。Attention 必须等待随机结果生成完成，因此 RNG 延迟无法被隐藏。此外，独立 RNG 方案还需要写回随机结果，并在 Attention 中再次读取，带来额外存储访问。

当 SQ 较大时，T_RNG 随 SQ^2 增长，baseline 总时间中暴露的 RNG 成本会迅速增加。因此，baseline 是正确但低效的参考实现。

5. SOTA 方案：RNG-Attention Fusion

现有 SOTA 优化通常采用 RNG-Attention Fusion：不再预先生成完整 mask，而是在 FlashAttention kernel 内部生成随机数，并立即用于 Softmax 后的 dropout。

执行顺序可表示为：

1
GEMM -> Attention_with_RNG_fused

5.1 为什么这种方案看起来合理

Fusion 的动机来自典型 kernel fusion 原则：

  1. 减少随机 mask 写回 HBM；
  2. 避免 Attention 再从 HBM 读取完整 mask；
  3. 减少 kernel launch 和跨 kernel 同步；
  4. 缩短中间数据生命周期。

从带宽角度看，该方案确实比 baseline 更自然，因为 dropout mask 的生产者和消费者位于同一 kernel 内。

5.2 为什么 fusion 仍然有限

论文的核心发现之一是：RNG-Attention Fusion 只能隐藏约 10% 到 20% 的 RNG runtime。原因在于 Attention kernel 本身并非纯 Tensor Core-bound。除 QK^T 和 PV 之外，FlashAttention 还包含 Softmax、reduce、归一化、指数函数近似、寄存器访问、同步和控制逻辑。

这些非矩阵部分会消耗：

  • Issue Stage；
  • Register File bandwidth；
  • MUFU / FMA / ALU pipeline；
  • warp 调度与指令混合资源。

Philox RNG 同样依赖大量整数 ALU 指令和 issue bandwidth。因此，将 RNG 融入 Attention 并不等于将 RNG 放入一个“空闲”的执行空间；相反，它可能被放入了一个已经在低层资源上较为拥挤的 kernel。

换言之，fusion 减少了中间存储，但没有消除 RNG 与 Attention 在关键微架构资源上的竞争。

6. 论文方案：RNG-GEMM Overlap

论文提出的方案是 RNG-GEMM Overlap。其主要思想是将 RNG 从 Attention 中剥离，并提前与前序 GEMM 并发执行。Attention 在 GEMM 与 RNG 都完成后启动，只负责消费随机结果并应用 dropout。

执行结构可表示为：

1
2
GEMM  ---------> Attention
RNG   ---并行---^

对应运行时间近似为：

1
2
T_overlap = max(T_GEMM_with_interference, T_RNG_with_interference)
          + T_Attention_drop_only

6.1 为什么 RNG 可以和 GEMM 重叠

该方案成立的关键在于 GEMM 与 RNG 的资源瓶颈更具互补性。

论文对 H100 上主要 kernel 的资源限制进行了分析：

  • GEMM 主要受 MMA math / Tensor Core 和 L2 带宽限制；
  • RNG 主要受 Issue Stage 与 ALU pipeline 限制；
  • Attention 除 MMA 外，还显著受 RF bandwidth 与 Issue Stage 限制。

因此，RNG 与 Attention 的冲突更强，而 RNG 与 GEMM 的冲突相对较弱。论文实测显示，为 RNG 预留少量 register 和 shared memory 后，GEMM 性能变化很小。并发执行时，虽然 RNG 本身可能变慢，但只要 T_RNG 小于或接近 T_GEMM，其大部分延迟仍可被 GEMM 覆盖。

6.2 哪些 workload 最适合 overlap

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

  1. GEMM 远长于 RNG：RNG 可被隐藏，但原本占比不高，因此总体收益有限；
  2. RNG 小于但接近 GEMM：RNG 足够大，值得隐藏，同时又能被 GEMM 大部分覆盖，是最理想区域；
  3. RNG 长于 GEMM：GEMM 结束后仍有 RNG 尾部暴露，overlap 收益下降。

典型 LLM 训练配置常落在第二类区域，因此 RNG-GEMM overlap 能获得较明显收益。

6.3 额外存储与流水化优化

将 RNG 从 Attention 中拆出后，需要存储随机结果供 Attention 消费。这会引入额外容量需求。论文进一步提出按 sequence chunk 进行 pipeline，避免一次性保存完整 SQ x SQ 随机结果。

可抽象为：

1
RNG(chunk i+1)  ||  Attention consume(chunk i)

通过分块与 double buffering，随机结果存储可从全量矩阵降低为局部缓冲，从而缓解长序列下的容量压力。

7. 实验设置与结果

论文在 H100 / GH100 HBM3 80GB 上进行建模与 CUDA 实测验证，使用 FP8 precision。典型参数包括：

  • head dimension：128；
  • batch size：1；
  • sequence length：2K 到 64K；
  • embedding dimension：4K 到 16K；
  • GEMM tile：128x128x128；
  • RNG / Attention tile：64x128x128；
  • RNG：主要采用 Philox 7，并对比 Philox 5 与 Philox 3。

7.1 加速比结果

论文报告的代表性结果为：

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

该结果说明，fusion 虽然能够减少部分存储访问，但在典型训练 workload 中，利用 GEMM 隐藏 RNG 延迟可以取得更高收益。

7.2 模型与 silicon measurement 的一致性

论文实现了三种 CUDA 版本：

  1. baseline：QKV_GEMM -> RNG -> Attention；
  2. fusion：QKV_GEMM -> fused Attention + RNG；
  3. overlap：使用两个 CUDA stream 并发执行 GEMM 与 RNG，Attention 等二者完成后启动。

模型与硅上测量结果的平均误差约为 2%，标准差约为 0.04。这说明细粒度资源模型能够较好解释不同方案之间的性能差异。

7.3 Philox 轮数变化的影响

论文比较了 Philox 7、Philox 5 与 Philox 3：

  • Philox 5 runtime 约为 Philox 7 的 81%；
  • Philox 3 runtime 约为 Philox 7 的 67%。

减少 Philox 轮数会降低 RNG 绝对开销，因此 overlap 可隐藏的绝对时间也会减少。但 fusion 的收益并不会严格按轮数线性变化，因为 fusion 场景更强地受 issue stage 等共享资源影响。

7.4 面向未来 GPU 的趋势

论文还使用 Blackwell B100 GPU 模型分析未来硬件趋势。若矩阵计算吞吐继续快速增长，而非矩阵资源、指令发射与随机数生成相关资源没有同比例提升，则 GEMM 时间可能缩短，RNG 相对占比上升。此时，RNG 相关优化的重要性会进一步提高。

8. 对 NPU 场景的启示

以华为 NPU 为例，是否存在完全相同的问题需要依赖具体芯片、CANN 编译器、算子实现和 profiler 数据确认。不过，从算子结构上看，Attention Dropout 中的 RNG 成本同样值得重点分析。

8.1 需要首先确认的事实

在 NPU 上，应首先确认以下问题：

  1. Dropout 是否使用 Philox 或类似 CBRNG；
  2. Philox4x32 一次输出的 4 个随机字是否被充分利用；
  3. mask 是 bit-pack、int8 还是更宽格式；
  4. Dropout/RNG kernel 的瓶颈是 DDR、L2/UB 搬运，还是 Vector/Scalar 执行；
  5. fused attention 中 Softmax、mask、格式转换是否已使 Vector/Scalar 或调度资源接近饱和。

若当前实现未充分复用 Philox4x32 的 4 个输出，则应优先优化 RNG 生成粒度。该优化通常不改变 Bernoulli 采样的统计正确性，但会改变随机流排列，因此需要同步更新确定性测试与 golden reference。

8.2 是否可以借鉴 RNG-GEMM Overlap

如果 NPU MatMul 主要由 Cube/Matrix 单元限制，而 RNG 主要消耗 Vector/Scalar 或通用整数流水，则两者可能具有资源互补性。在这种情况下，可以考虑将 RNG 与前序 MatMul 并发执行，并让 Attention 依赖两者完成。

但该方案需要满足三类条件：

  1. 编译器或 runtime 支持 task/stream 级并发与依赖表达；
  2. 随机结果布局与后续 Attention tile 对齐，避免额外搬运；
  3. 长序列场景采用 chunk pipeline，避免全量 SQ^2 mask 写出。

因此，NPU 上更合理的策略不是直接照搬 GPU 方案，而是建立类似的资源模型：

1
T_kernel = max(T_cube, T_vector, T_scalar, T_UB, T_L2, T_DDR, T_issue, ...)

然后根据 shape、资源占用与 pipeline 能力，在 fusion、overlap 和 chunk pipeline 之间自动选择。

9. 总结

这篇论文围绕 FlashAttention 开启 dropout 后 RNG 成本上升的问题，提出了基于资源互补性的系统优化方法。其逻辑链条可以概括为：

  1. Attention Dropout 的随机数规模为 B * nH * SQ^2，长序列下成本显著；
  2. Philox 适合并行可复现 RNG，但其整数混合过程带来明显 ALU 与 issue 开销；
  3. Baseline 顺序执行 RNG，会使随机数生成延迟完全暴露在关键路径上；
  4. SOTA 的 RNG-Attention Fusion 能减少 mask 中间存储，但无法充分隐藏 RNG，因为 Attention 与 RNG 共享 issue、register file 等低层资源；
  5. 论文提出的 RNG-GEMM Overlap 利用 GEMM 与 RNG 的资源互补性，将 RNG 延迟隐藏在前序 GEMM 中；
  6. 在典型 LLM workload 上，overlap 相比 baseline 和 fusion 均取得更高加速比；
  7. 对 GPU 与 NPU 等加速器而言，未来矩阵计算吞吐提升后，随机数生成、指令发射、寄存器访问等非矩阵资源可能成为更重要的优化对象。

这篇论文的更一般启示是：算子优化不能只依据“是否可以 fusion”进行判断。正确的问题应当是：融合后是否会竞争同一关键资源，拆分后是否能利用互补资源，以及在具体 shape 下关键路径如何变化。

参考资料

  • Haiyue Ma, Jian Liu, Ronny Krashinsky. Reducing the Cost of Dropout in Flash-Attention by Hiding RNG with GEMM. arXiv:2410.07531v2.

分享文章

快来参与讨论吧~

本站由 Zane Jiang 使用 Stellar 1.33.1 主题创建，一款很棒的 Hexo 主题！

总访问 次 || 本页访问 次
总访客 人 || 本页访客 人