第 10 课：无状态随机数与 Dropout

本课目标与完成标准

学完后你应能：实现可复现 inverted dropout，解释 seed+offset 的确定性和并行安全。

通过必须同时满足：

• 独立补齐本课唯一的代码填空题，并通过给定检查；
• 三个问答题均说明因果链，而不是只报术语；
• 能指出至少一个正确性边界和一个性能取舍；
• 总分不低于 8/10，且没有一票否决级概念错误。

前置关系

• 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
• 本课在路线中的作用：Triton 常用 counter-based RNG：随机值由 seed 与元素 offset 决定，不维护跨线程共享状态。

核心心智模型

1. 它是什么，解决什么问题

在深度神经网络的正则化技术中，Dropout 以概率 $p$ 随机将神经元激活值置零（Drop），以抑制神经元之间的共适应性（Co-adaptation）。 为了避免在测试推断阶段引入额外的缩放操作，现代深度学习框架普遍采用反向丢弃（Inverted Dropout）： $$y_i = \begin{cases} \frac{x_i}{1 - p}, & \text{with probability } 1 - p \\ 0, & \text{with probability } p \end{cases}$$ 通过在前向训练时直接除以保留概率 $(1 - p)$，保证了输出激活值的数学期望恒等于输入（$\mathbb{E}[y_i] = x_i$）。 在传统并行程序中，伪随机数生成器（PRNG，如线性同余 LCG 或梅森旋转 Mersenne Twister）依赖不断更新的全局内部状态（Stateful RNG）。在海量 GPU 线程并发时，维护共享状态会导致灾难性的线程锁争用或显存瓶颈。 本课解决的核心工程问题是：理解基于计数器的无状态随机数生成算法（Counter-based RNG，如 Philox），在 Triton 中利用 tl.rand(seed, offsets) 实现零共享状态、完全无锁、可任意并发且 100% 确定性可复现的 Dropout 掩码生成，以及如何通过反向重算（Recomputation）免除物化 Mask 的显存开销。

2. 它如何工作

Philox 无状态随机数与 Inverted Dropout 数据通路如图 10-1 与图 10-2 所示：

Philox 无状态随机数生成器

随机种子 seed

绝对元素偏移 offset

tl.rand(seed, offsets)

输出 [0, 1) 均匀分布随机浮点数

图 10-1：Philox 无状态随机数根据种子与偏移即时生成机制

通过 Philox 生成随机数 r

判定 r > p 生成保留掩码 keep_mask

有效值乘以缩放因子 1 / (1 - p)

丢弃项置 0 并写回输出张量

图 10-2：Inverted Dropout 掩码生成与前向缩放执行流

1. 基于计数器的纯函数式映射（Counter-based Philox Engine）： Philox 是借鉴分组密码结构设计的 counter-based 伪随机数生成器（多轮乘法与异或变换，不具备密码学安全性）： $$R_i = \text{Philox}(K = \text{seed}, C = \text{offset}_i)$$ 对于固定的密钥 $K$（即全局种子 seed）和唯一的计数器 $C$（即当前张量元素的全局一维物理偏移 offset），Philox 在固定轮数（Triton 默认 10 轮）内输出一个通过常用统计测试（如 BigCrush）的均匀伪随机整数，tl.rand 再将其映射为区间 $[0, 1)$ 的浮点数。

2. 纯寄存器层伯努利采样（In-Register Bernoulli Sampling）： 内核无需任何跨线程协调，直接并发求值： keep = tl.rand(seed, offsets) > p 生成布尔保留掩码 keep。此过程不读取任何显存，纯粹在 SM 算术逻辑单元（ALU）中完成。

3. 条件缩放与写回： 使用 tl.where(keep, xv / (1.0 - p), 0.0) 对有效元素施加 $\frac{1}{1-p}$ 的期望保持缩放，并置零被丢弃项。最后通过 tl.store 配合尾块掩码 mask 写回全局显存。

4. 反向重算契约（Zero Memory Overhead Recompute）： 在反向传播（Backward）中，原本需要从显存加载保存的 bool 掩码。而在 Triton 中，反向 Kernel 仅需传入完全相同的 seed 和 offsets，直接在寄存器中重新生成一模一样的 keep 掩码，从而节省了保存 Mask 所需的巨额显存带宽与容量（省去 $N \times 1 \text{ Byte}$ 的 HBM 显存）。

3. 正确性条件与常见误区

• 概率定义域严格约束（$0 \le p < 1$）： 若传入 $p = 1.0$，tl.rand 的取值在 $[0, 1)$ 内，keep 恒为 False，本 kernel 的输出全为 0；但 xv / (1.0 - p) 的中间值是 Inf 或 NaN（$0/0$），如果写成 xv * keep / (1 - p) 这类乘法形式，结果会出现 NaN。若传入 $p < 0$，采样没有概率意义。Wrapper 应断言 $0 \le p < 1$，$p = 1$ 时直接返回全零张量。

• 种子-偏移域重叠（Domain Overlap / Seed Collision）： 若模型中有多层连续的 Dropout 算子（如 Attention Dropout 与 MLP Dropout），或者相邻训练 Step 之间复用了相同的 (seed, offsets) 空间，会导致不同层或不同时间步生成的随机掩码完全相同（高度相关）。必须为不同层分配不同的 seed，或者在偏移量上加上全局递增的步长基地址（Base Offset）。

• 非连续张量与网格打散引发的随机数错位： offsets 必须是针对张量逻辑连续排布的一维绝对索引。若网格分块策略在 Forward 和 Backward 中不一致（例如 Forward 按一维展开、Backward 按二维分块），导致同一个逻辑元素两次计算出的 offsets 不相等，反向传播将使用完全错误的 Mask 进行梯度反传，导致模型训练发散。

4. 性能与工程取舍

• 算力换带宽（Compute vs Memory Tradeoff）： 计算 Philox 随机数需要执行几十条整型乘法与异或指令（每轮 2 次 32 位乘法）。Dropout 是 Memory-bound 算子，这部分算术通常能被访存延迟掩盖。重算省掉的是前向写出、反向读回各 $N$ 字节的 bool mask；以 FP32 为例，每个方向的流量从 $9N$ 字节降到 $8N$ 字节，约省 11%，同时省下 $N$ 字节的显存占用。

• 算子融合的极致收益： Dropout 通常紧随激活函数（如 GELU、SiLU）或 Bias 相加之后。在 Triton 中直接将“加偏置 $\to$ 激活 $\to$ Philox 采样 $\to$ 缩放”融合成单个 Kernel（如图 10-2 所示），整个流水线除最初读取和最终写回外，中间所有数据始终停留在寄存器中，消除了所有中间特征图的显存物化。

具体演示

以 $p = 0.2$、缩放因子 $1 / (1 - 0.2) = 1.25$、输入全为 1.0 的张量（$N = 10003$）为例：

• 所有有效位置在经过 Philox 判定后，数值非 0 即 1.25。
• 当传入相同种子 seed=123 运行两次时，torch.equal(y1, y2) 返回严格的 True（确定性位级复现）。
• 统计 10003 个样本的均值： $$\mathbb{E}[y] = (1 - 0.2) \times 1.25 + 0.2 \times 0.0 = 1.0$$ 单个元素的标准差为 $1.25 \times \sqrt{0.8 \times 0.2} = 0.5$，10003 个样本均值的标准差约为 $0.5 / \sqrt{10003} \approx 0.005$，所以断言阈值 0.05 留有约 10 倍标准差的余量。这只验证了期望，不能证明随机数的独立性。

请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。

实践任务：唯一代码填空题

补齐 keep mask。

规则：只能修改 TODO/______ 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。

import torch
import triton
import triton.language as tl

@triton.jit
def dropout_kernel(x, out, n, p, seed,
                   BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    xv = tl.load(x + offsets, mask=mask)
    keep = ______  # TODO: 每元素无状态 Bernoulli
    y = tl.where(keep, xv / (1.0 - p), 0.0)
    tl.store(out + offsets, y, mask=mask)

def dropout(x, p, seed):
    assert x.is_contiguous() and 0.0 <= p < 1.0
    out = torch.empty_like(x)
    n = x.numel()
    dropout_kernel[(triton.cdiv(n, 256),)](x, out, n, p, seed, BLOCK=256)
    return out

x = torch.ones(10003, device="cuda")
y1, y2 = dropout(x, .2, 123), dropout(x, .2, 123)
assert torch.equal(y1, y2)
assert set(torch.unique(y1).tolist()).issubset({0.0, 1.25})
assert abs(y1.mean().item() - 1.0) < .05

检查方法

在 CUDA/Triton 环境运行本单元格；断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。

提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“无状态随机数与 Dropout”的工作机制。

你的答案：

Q2

为什么相同 seed 不同 offsets 能并行无锁地生成随机数？

你的答案：

Q3

融合后 offset 映射改变，如何保持 forward/backward mask 一致？

你的答案：

评分与通过规则

• 代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分；
• Q1～Q3 各 2 分；
• 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。

import torch
import triton
import triton.language as tl

@triton.jit
def dropout_kernel(x, out, n, p, seed,
                   BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    xv = tl.load(x + offsets, mask=mask)
    keep = tl.rand(seed, offsets) > p
    y = tl.where(keep, xv / (1.0 - p), 0.0)
    tl.store(out + offsets, y, mask=mask)

def dropout(x, p, seed):
    assert x.is_contiguous() and 0.0 <= p < 1.0
    out = torch.empty_like(x)
    n = x.numel()
    dropout_kernel[(triton.cdiv(n, 256),)](x, out, n, p, seed, BLOCK=256)
    return out

x = torch.ones(10003, device="cuda")
y1, y2 = dropout(x, .2, 123), dropout(x, .2, 123)
assert torch.equal(y1, y2)
assert set(torch.unique(y1).tolist()).issubset({0.0, 1.25})
assert abs(y1.mean().item() - 1.0) < .05

Q1 参考答案

从三个阶段拆解其运行机制：

1. 输入阶段（种子传递与绝对偏移量生成）：Python 调用层将只读种子 seed 与丢弃概率 $p$ 传入 Kernel。每个 Program 依据 program_id 与 BLOCK 步长生成唯一的全局物理偏移向量 offsets，并在 mask = offsets < n 保护下载入特征数据 xv。

2. 状态变化阶段（无状态 Philox 采样与无分支条件缩放）：

   • 采样：调用 tl.rand(seed, offsets)。底层启动无状态 Philox 密码学置换引擎，仅依据 (seed, offsets) 纯函数式并发计算出一组 $[0, 1)$ 伪随机数。
   • 比较与缩放：执行元素级裁决 keep = rand > p。利用 tl.where 原语，对被保留的元素乘以缩放因子 $\frac{1}{1-p}$，对丢弃项直接置为 0.0。整个过程完全在片上寄存器完成，不依赖任何共享内存或全局锁。

3. 输出阶段（带掩码写回与反向重算支撑）：通过 tl.store 配合 mask 写回输出张量。在反向传播求导时，反向 Kernel 无需读取任何前向保存的掩码张量，仅凭相同的 (seed, offsets) 即可在片上 100% 重新生成完全一致的 Mask，实现零显存额外占用的梯度反传。

Q2 参考答案

传统伪随机数发生器（如 C 语言的 rand()）基于递推公式 $S_{n+1} = f(S_n)$，前一个状态是后一个状态的输入，强行并行必须对全局状态加锁或导致严重的争用。 而以 Philox 为代表的基于计数器的随机数发生器（Counter-based PRNG）从数学原理上打破了这种状态依赖链：

Philox 把 seed 当作密钥 $K$、把元素的全局 offset 当作计数器 $C$，计算 $R = \text{Philox}(K, C)$。固定 $K$ 时它是计数器上的双射，不同 offset 得到不同的输出（详见第 2 节第 1 条）。

第 $i$ 个元素的随机数只由 $(K, i)$ 决定，元素之间没有共享状态和先后依赖。因此各线程可以按任意顺序、在任意 SM 上各自计算，不需要锁或同步；同样的 $(K, i)$ 总是得到同样的结果，这也是反向重算可行的前提。

Q3 参考答案

在大型深度学习框架中，当算子发生深度融合时（例如前向将 Matmul、Bias、GELU 与 Dropout 融合成一个二维网格 Kernel，而反向可能以不同粒度调度）： 若由于分块结构改变导致输入元素对应的 offsets 编号方式发生偏移，Forward 和 Backward 就会产生不一致的随机掩码，破坏梯度反传。保持其严格一致的工业级工程方案为：

1. 基于全局线性绝对坐标（Logical Linear Indexing）： 严禁使用局部 Tile 内部相对偏移。无论网格如何划分（1D 还是 2D），所有线程在调用 tl.rand 前，必须将自身当前处理的数据坐标严格折算为针对原始张量的全局一维平铺绝对线性索引： $\text{global\_offset} = \text{global\_row} \times N + \text{global\_col}$。

2. 种子域隔离与 Base Offset 传递： 若同一个张量需要多次采样，保持全局线性索引不变，通过在 Python 调度层统一分发唯一的全局 seed 和 base_offset（例如 offsets = base_offset + global_idx），保证前向与反向在同一个张量元素上绑定的输入元组 (seed, global_offset) 具有严格且唯一的双射对应关系。

参考资料

• Triton Tutorials
• Triton language API

资料用于建立事实基线；面试回答仍需用自己的语言组织。

← 上一课
第 9 课：Blocked Matmul 与… 下一课 →
第 1 课：GPU 层级、Tensor Co…