第 6 课：均值方差与 Mask 二次传播

参考答案版：包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。

本课目标与完成标准

学完后你应能：实现 population variance，并解释 padding 在中心化后为何会重新变成非零。

通过必须同时满足：

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

前置关系

课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
本课在路线中的作用：均值和方差是两次 reduction；masked load 的 0 对 sum 安全，但减 mean 后 padding lane 变成 -mean。

核心心智模型

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

在 LayerNorm、RMSNorm 以及批归一化（BatchNorm）中，计算特征的均值 $\mu$ 与方差 $\sigma^2$ 是必经步骤。总体方差（Population Variance）的经典定义为： $\mu = \frac{1}{N} \sum_{i=0}^{N-1} x_i, \quad \sigma^2 = \frac{1}{N} \sum_{i=0}^{N-1} (x_i - \mu)^2$ 在 Triton 向量化内核中，为了满足 SIMD 折半二叉树规约，张量宽度必须扩展至 $2$ 的幂次 $\text{BLOCK}$ 。在上节课中，我们在加载时为 Padding lanes 填充了加法单位元 $0.0$ 。但本课揭示了一个极其隐蔽的致命陷阱：“单位元的代数污染”。当执行中心化操作 $(x - \mu)$ 时，原先填充了 $0.0$ 的 Padding lanes 会变为 $(0.0 - \mu) = -\mu$ ；其平方则变为 $(-\mu)^2 = \mu^2 \neq 0$ 。如果直接对平方项求和，这些原本虚构的 Padding lanes 将把 $(\text{BLOCK} - N) \times \mu^2$ 的巨大虚假数值强行注入方差累加器，造成彻底的数值崩溃。本课解决的核心工程问题是：理解掩码在复合算子数据流中的“二次传播（Mask Re-propagation）”机制，掌握中心化后的掩码重施加（Re-masking）技术，杜绝静默数值污染。

2. 它如何工作

两阶段掩码统计量与二次传播时序如图 6-1 与图 6-2 所示：

flowchart LR
    subgraph Stats["统计量计算两阶段"]
        Sum["第一次归约: sum(x) / valid_count ➔ 得到均值 mu"]
        Var["第二次归约: sum((x - mu)^2) / valid_count ➔ 得到方差 var"]
    end

图 6-1：均值与方差的两阶段掩码统计量结构

flowchart TD
    In["输入块"] --> Pass1["Pass 1: 计算均值并保留 valid_count"]
    Pass1 --> Pass2["Pass 2: 重新遍历差值平方并排除越界元素"]
    Pass2 --> Out["同时输出均值与无偏方差"]

图 6-2：两阶段掩码二次传播计算时序流

Pass 1（带掩码加载与均值规约）：
- 使用 cols = tl.arange(0, BLOCK) 与 mask = cols < N。
- values = tl.load(..., mask=mask, other=0.0)。
- 在片上调用 tl.sum(values, axis=0) 并除以真实逻辑长度 $N$ （绝对不能除以硬件向量长度 $\text{BLOCK}$ ），得到严格无偏的均值标量 $\mu$ 。
中心化与掩码重施加（Re-masking）：
- 寄存器中所有 lane 同时广播减去标量 $\mu$ ：diff = values - mean。
- 此时，前 $N$ 个有效 lane 为 $x_i - \mu$ ；后 $\text{BLOCK} - N$ 个 Padding lane 变为 $-\mu$ 。
- 计算元素平方项 diff_sq = diff * diff。此时 Padding lane 上的数值为 $\mu^2$ 。
- 核心二次传播防护：通过 tl.where(mask, diff_sq, 0.0) 强制对平方项重新施加掩码。将后 $\text{BLOCK} - N$ 个 lane 携带的 $\mu^2$ 重新强行抹平为加法单位元 $0.0$ 。
Pass 2（方差规约与双标量写回）：
- 对清洗后的 centered_sq 向量执行二次 tl.sum，除以真实长度 $N$ ，得到位级精准的总体方差 $\sigma^2$ 。
- 将均值与方差标量分别通过 tl.store 写回显存。

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

归一化分母混淆（除 $N$ vs 除 $\text{BLOCK}$ ）：若在计算均值或方差时误用除以 BLOCK，由于 $\text{BLOCK} \ge N$ ，算出的统计量将被无辜缩小 $N / \text{BLOCK}$ 倍。必须除以真实有效元素数 $N$ 。
样本方差 vs 总体方差（Bessel 校正）： PyTorch 默认的 torch.var 使用无偏样本方差（Unbiased Sample Variance，分母为 $N - 1$ ）。而神经网络算子（如 LayerNorm）规范定义使用总体方差（Population Variance，分母为 $N$ ，对应 PyTorch 的 unbiased=False）。当 $N=1$ 时，样本方差的分母为 0 导致除零错误，必须在系统边界做契约防御。
遗漏 tl.where 的二次掩码：这是新手最容易犯的隐蔽 Bug。只要 $\mu \neq 0$ 且存在 Padding（ $N < \text{BLOCK}$ ），不重施掩码就会导致方差偏大 $(\text{BLOCK}-N)\mu^2 / N$ 。当输入均值很大时，方差误差可达数个数量级。

4. 性能与工程取舍

片上两遍计算（Two-pass）的极致性能：在 Triton 中，数据加载到寄存器后，均值与方差的两轮规约完全在片上寄存器文件内部展开，中间数据 values 和 centered_sq 根本不需要写入全局显存（HBM）。访存量仅为读取矩阵一次（ $M \times N$ 元素），相比于在 PyTorch 中分别调用 x.mean() 和 x.var() 导致的 2 次全局显存扫描，显存带宽节约了整整 50%。
数值稳定性与 Welford 算法：两阶段算法（Two-pass）需要先求均值再求方差。如果遇到“大均值、小方差”极端分布（例如所有元素在 $10^6$ 附近，波动仅 $10^{-4}$ ），两次减法可能在浮点有限精度下产生轻微的灾难性抵消。此时若采用 Welford 增量流式算法，数值条件数最好，但 Welford 的 Warp 间归约公式较为复杂（需要维护 $(n, \mu, M_2)$ 三元组）。

具体演示

以 $N = 3, \text{BLOCK} = 4$ ，输入行向量为 $[1.0, 2.0, 3.0]$ 为例：

Pass 1 均值：
- 加载向量：values = [1.0, 2.0, 3.0, 0.0]（Padding 位被赋予 0.0）。
- 求和： $1 + 2 + 3 + 0 = 6.0$ 。
- 均值： $\mu = 6.0 / 3 = 2.0$ 。正确。
中心化陷阱（若无二次掩码）：
- 差值向量：values - mean = [1-2, 2-2, 3-2, 0-2] = [-1.0, 0.0, 1.0, -2.0]。
- 平方向量：[1.0, 0.0, 1.0, 4.0]。注意第 3 位变成了 $4.0$ ！
- 若直接求和： $1 + 0 + 1 + 4 = 6.0 \to \sigma^2 = 6.0 / 3 = 2.0$ 。而真实方差为 $(1+0+1)/3 = 2/3 \approx 0.667$ 。误差高达 $200\%$ ！
二次掩码保护（tl.where）：
- 施加掩码：tl.where(mask, diff_sq, 0.0) = [1.0, 0.0, 1.0, 0.0]。第 3 位被精准重置为 0.0。
- 求和： $1 + 0 + 1 + 0 = 2.0$ 。
- 方差： $\sigma^2 = 2.0 / 3 \approx 0.667$ 。与理论数学解位级完全契合。

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

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

补齐二次 mask。

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

Code

import torch
import triton
import triton.language as tl

@triton.jit
def mean_var_kernel(x, means, variances, N: tl.constexpr,
                    stride_m: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N
    values = tl.load(x + row * stride_m + cols, mask=mask, other=0.0)
    mean = tl.sum(values, axis=0) / N
    centered_sq = tl.where(______, (values - mean) * (values - mean), 0.0)  # TODO
    var = tl.sum(centered_sq, axis=0) / N
    tl.store(means + row, mean)
    tl.store(variances + row, var)

def mean_var(x):
    assert x.ndim == 2 and x.stride(1) == 1
    M, N = x.shape
    means = torch.empty(M, device=x.device, dtype=x.dtype)
    variances = torch.empty_like(means)
    mean_var_kernel[(M,)](x, means, variances, N, x.stride(0),
                          BLOCK=triton.next_power_of_2(N))
    return means, variances

for shape in ((3, 17), (4, 128)):
    x = torch.randn(shape, device="cuda")
    m, v = mean_var(x)
    torch.testing.assert_close(m, x.mean(1), atol=1e-4, rtol=1e-4)
    torch.testing.assert_close(v, x.var(1, unbiased=False), atol=2e-4, rtol=2e-4)

检查方法

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

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

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“均值方差与 Mask 二次传播”的工作机制。

你的答案：

Q2

只在 load 时 mask 为什么仍会污染方差？

你的答案：

Q3

大均值小方差时如何改进数值稳定性？

你的答案：

评分与通过规则

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

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

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

参考答案（仅 answer 分支）

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

Code

import torch
import triton
import triton.language as tl

@triton.jit
def mean_var_kernel(x, means, variances, N: tl.constexpr,
                    stride_m: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N
    values = tl.load(x + row * stride_m + cols, mask=mask, other=0.0)
    mean = tl.sum(values, axis=0) / N
    centered_sq = tl.where(mask, (values - mean) * (values - mean), 0.0)
    var = tl.sum(centered_sq, axis=0) / N
    tl.store(means + row, mean)
    tl.store(variances + row, var)

def mean_var(x):
    assert x.ndim == 2 and x.stride(1) == 1
    M, N = x.shape
    means = torch.empty(M, device=x.device, dtype=x.dtype)
    variances = torch.empty_like(means)
    mean_var_kernel[(M,)](x, means, variances, N, x.stride(0),
                          BLOCK=triton.next_power_of_2(N))
    return means, variances

for shape in ((3, 17), (4, 128)):
    x = torch.randn(shape, device="cuda")
    m, v = mean_var(x)
    torch.testing.assert_close(m, x.mean(1), atol=1e-4, rtol=1e-4)
    torch.testing.assert_close(v, x.var(1, unbiased=False), atol=2e-4, rtol=2e-4)

Q1 参考答案

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

输入阶段（网格绑定与加法零元注入）：按行发射 Program，列维度以 2 的幂次构造 cols = tl.arange(0, BLOCK)。在掩码 mask = cols < N 保护下，通过带 other=0.0 的 tl.load 一次性将该行数据拉入寄存器向量 values，超出 $N$ 的 Padding lane 注入加法单位元 0.0。
状态变化阶段（两阶段归约与二次掩码阻断）：
- 阶段 A：调用 tl.sum 计算有效总和，除以真实列宽 $N$ 得到行均值标量 $\mu$ 。
- 阶段 B：执行广播减法 values - mean 并求平方。由于 Padding lane 上的 0.0 变为 $-\mu$ 后平方变成 $\mu^2$ ，立即通过 tl.where(mask, diff_sq, 0.0) 强制将 Padding lane 的虚假能量清零，使加法单位元在数据流中重新闭环。
- 阶段 C：调用二次 tl.sum 计算清洗后的方差平方和，除以真实列宽 $N$ 得到总体方差 $\sigma^2$ 。
输出阶段（双标量写回）：将算出的均值标量与方差标量，通过各自独立的 tl.store 写回到目标全局内存地址 means + row 与 variances + row。

Q2 参考答案

在初始加载阶段（tl.load）通过 other=0.0 确实成功阻断了未初始化内存的污染，使加法归约满足 $X_i + 0.0 = X_i$ 。但是，计算方差的前提是执行中心化偏移运算 $(x_i - \mu)$ ：对于有效的 $N$ 个数据，计算的是真实的差值 $(x_i - \mu)$ ；但对于多余的 $\text{BLOCK} - N$ 个 Padding lanes，其寄存器数值为 0.0，参与运算后变成了 $(0.0 - \mu) = -\mu$ 。随后的平方运算使其变成了 $(-\mu)^2 = \mu^2$ 。只要均值 $\mu \neq 0$ ， $\mu^2$ 就必然是一个大于 0 的正数！这意味着，多余的每个 Padding lane 都会虚构出一个数值为 $\mu^2$ 的非零偏差。在后续执行 tl.sum 时，这些本该被完全忽略的无效 lane 会向累加器注入总量为 $(\text{BLOCK} - N) \times \mu^2$ 的巨大伪方差能量，导致计算出的方差虚高，产生严重的数学错误。

Q3 参考答案

当输入数据呈现“大均值、小方差”特征时（例如数据集中在 $10^6 \pm 0.01$ ），直接使用单精度 FP32 的两阶段朴素公式会导致严重的灾难性数值抵消（Catastrophic Cancellation）。工程改进策略包括：

寄存器累加器提升为 FP64 计算：在计算均值与平方和时，将寄存器内的 values 显式转换为高精度的双精度浮点数 tl.float64。求出 mean 与二次中心化平方和后再转换回 FP32 写回，彻底抹平浮点尾数截断误差。
采用流式 Welford 增量递推算法：改用单遍（One-pass）Welford 在线均值方差算法： $\mu_k = \mu_{k-1} + \frac{x_k - \mu_{k-1}}{k}, \quad M_{2, k} = M_{2, k-1} + (x_k - \mu_{k-1})(x_k - \mu_k)$ 在 Warp Shuffle 归约时采用 Welford Combine 算子合成两个子块的统计量。该算法从代数结构上保证了差值始终在局部小量之间累积，条件数最优，理论上杜绝了浮点相消问题。

参考资料

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

← 上一课
第 5 课：行归约与 Padding 单位元 下一课 →
第 7 课：稳定 Softmax 与融合

--- title: "第 6 课：均值方差与 Mask 二次传播" jupyter: python3 --- ::: {.callout-note appearance="simple"} **参考答案版**：包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。 ::: ## 本课目标与完成标准学完后你应能：实现 population variance，并解释 padding 在中心化后为何会重新变成非零。通过必须同时满足： - 独立补齐本课唯一的代码填空题，并通过给定检查； - 三个问答题均说明因果链，而不是只报术语； - 能指出至少一个正确性边界和一个性能取舍； - 总分不低于 8/10，且没有一票否决级概念错误。 ## 前置关系 - 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念 - 本课在路线中的作用：均值和方差是两次 reduction；masked load 的 0 对 sum 安全，但减 mean 后 padding lane 变成 `-mean`。 ## 核心心智模型 ### 1. 它是什么，解决什么问题在 LayerNorm、RMSNorm 以及批归一化（BatchNorm）中，计算特征的均值 $\mu$ 与方差 $\sigma^2$ 是必经步骤。总体方差（Population Variance）的经典定义为： $$ \mu = \frac{1}{N} \sum_{i=0}^{N-1} x_i, \quad \sigma^2 = \frac{1}{N} \sum_{i=0}^{N-1} (x_i - \mu)^2 $$ 在 Triton 向量化内核中，为了满足 SIMD 折半二叉树规约，张量宽度必须扩展至 $2$ 的幂次 $\text{BLOCK}$。在上节课中，我们在加载时为 Padding lanes 填充了加法单位元 $0.0$。但本课揭示了一个极其隐蔽的致命陷阱：**“单位元的代数污染”**。当执行中心化操作 $(x - \mu)$ 时，原先填充了 $0.0$ 的 Padding lanes 会变为 $(0.0 - \mu) = -\mu$；其平方则变为 $(-\mu)^2 = \mu^2 \neq 0$。如果直接对平方项求和，这些原本虚构的 Padding lanes 将把 $(\text{BLOCK} - N) \times \mu^2$ 的巨大虚假数值强行注入方差累加器，造成彻底的数值崩溃。本课解决的核心工程问题是：**理解掩码在复合算子数据流中的“二次传播（Mask Re-propagation）”机制，掌握中心化后的掩码重施加（Re-masking）技术，杜绝静默数值污染**。 ### 2. 它如何工作两阶段掩码统计量与二次传播时序如图 6-1 与图 6-2 所示： ```{mermaid} flowchart LR subgraph Stats["统计量计算两阶段"] Sum["第一次归约: sum(x) / valid_count ➔ 得到均值 mu"] Var["第二次归约: sum((x - mu)^2) / valid_count ➔ 得到方差 var"] end ``` 图 6-1：均值与方差的两阶段掩码统计量结构 ```{mermaid} flowchart TD In["输入块"] --> Pass1["Pass 1: 计算均值并保留 valid_count"] Pass1 --> Pass2["Pass 2: 重新遍历差值平方并排除越界元素"] Pass2 --> Out["同时输出均值与无偏方差"] ``` 图 6-2：两阶段掩码二次传播计算时序流 1. **Pass 1（带掩码加载与均值规约）**： - 使用 `cols = tl.arange(0, BLOCK)` 与 `mask = cols < N`。 - `values = tl.load(..., mask=mask, other=0.0)`。 - 在片上调用 `tl.sum(values, axis=0)` 并除以**真实逻辑长度 $N$**（绝对不能除以硬件向量长度 $\text{BLOCK}$），得到严格无偏的均值标量 $\mu$。 2. **中心化与掩码重施加（Re-masking）**： - 寄存器中所有 lane 同时广播减去标量 $\mu$：`diff = values - mean`。 - 此时，前 $N$ 个有效 lane 为 $x_i - \mu$；后 $\text{BLOCK} - N$ 个 Padding lane 变为 $-\mu$。 - 计算元素平方项 `diff_sq = diff * diff`。此时 Padding lane 上的数值为 $\mu^2$。 - **核心二次传播防护**：通过 `tl.where(mask, diff_sq, 0.0)` 强制对平方项重新施加掩码。将后 $\text{BLOCK} - N$ 个 lane 携带的 $\mu^2$ 重新强行抹平为加法单位元 $0.0$。 3. **Pass 2（方差规约与双标量写回）**： - 对清洗后的 `centered_sq` 向量执行二次 `tl.sum`，除以真实长度 $N$，得到位级精准的总体方差 $\sigma^2$。 - 将均值与方差标量分别通过 `tl.store` 写回显存。 ### 3. 正确性条件与常见误区 - **归一化分母混淆（除 $N$ vs 除 $\text{BLOCK}$）**：若在计算均值或方差时误用除以 `BLOCK`，由于 $\text{BLOCK} \ge N$，算出的统计量将被无辜缩小 $N / \text{BLOCK}$ 倍。必须除以真实有效元素数 $N$。 - **样本方差 vs 总体方差（Bessel 校正）**： PyTorch 默认的 `torch.var` 使用无偏样本方差（Unbiased Sample Variance，分母为 $N - 1$）。而神经网络算子（如 LayerNorm）规范定义使用总体方差（Population Variance，分母为 $N$，对应 PyTorch 的 `unbiased=False`）。当 $N=1$ 时，样本方差的分母为 0 导致除零错误，必须在系统边界做契约防御。 - **遗漏 `tl.where` 的二次掩码**：这是新手最容易犯的隐蔽 Bug。只要 $\mu \neq 0$ 且存在 Padding（$N < \text{BLOCK}$），不重施掩码就会导致方差偏大 $(\text{BLOCK}-N)\mu^2 / N$。当输入均值很大时，方差误差可达数个数量级。 ### 4. 性能与工程取舍 - **片上两遍计算（Two-pass）的极致性能**：在 Triton 中，数据加载到寄存器后，均值与方差的两轮规约完全在片上寄存器文件内部展开，中间数据 `values` 和 `centered_sq` 根本不需要写入全局显存（HBM）。访存量仅为读取矩阵一次（$M \times N$ 元素），相比于在 PyTorch 中分别调用 `x.mean()` 和 `x.var()` 导致的 2 次全局显存扫描，显存带宽节约了整整 50%。 - **数值稳定性与 Welford 算法**：两阶段算法（Two-pass）需要先求均值再求方差。如果遇到“大均值、小方差”极端分布（例如所有元素在 $10^6$ 附近，波动仅 $10^{-4}$），两次减法可能在浮点有限精度下产生轻微的灾难性抵消。此时若采用 Welford 增量流式算法，数值条件数最好，但 Welford 的 Warp 间归约公式较为复杂（需要维护 $(n, \mu, M_2)$ 三元组）。 ## 具体演示以 $N = 3, \text{BLOCK} = 4$，输入行向量为 $[1.0, 2.0, 3.0]$ 为例： 1. **Pass 1 均值**： - 加载向量：`values = [1.0, 2.0, 3.0, 0.0]`（Padding 位被赋予 0.0）。 - 求和：$1 + 2 + 3 + 0 = 6.0$。 - 均值：$\mu = 6.0 / 3 = 2.0$。正确。 2. **中心化陷阱（若无二次掩码）**： - 差值向量：`values - mean = [1-2, 2-2, 3-2, 0-2] = [-1.0, 0.0, 1.0, -2.0]`。 - 平方向量：`[1.0, 0.0, 1.0, 4.0]`。注意第 3 位变成了 $4.0$！ - 若直接求和：$1 + 0 + 1 + 4 = 6.0 \to \sigma^2 = 6.0 / 3 = 2.0$。而真实方差为 $(1+0+1)/3 = 2/3 \approx 0.667$。误差高达 $200\%$！ 3. **二次掩码保护（`tl.where`）**： - 施加掩码：`tl.where(mask, diff_sq, 0.0) = [1.0, 0.0, 1.0, 0.0]`。第 3 位被精准重置为 0.0。 - 求和：$1 + 0 + 1 + 0 = 2.0$。 - 方差：$\sigma^2 = 2.0 / 3 \approx 0.667$。与理论数学解位级完全契合。请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。 ## 实践任务：唯一代码填空题补齐二次 mask。规则：只能修改 `TODO`/`______` 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。 ```{python} #| course_role: exercise import torch import triton import triton.language as tl @triton.jit def mean_var_kernel(x, means, variances, N: tl.constexpr, stride_m: tl.constexpr, BLOCK: tl.constexpr): row = tl.program_id(0) cols = tl.arange(0, BLOCK) mask = cols < N values = tl.load(x + row * stride_m + cols, mask=mask, other=0.0) mean = tl.sum(values, axis=0) / N centered_sq = tl.where(______, (values - mean) * (values - mean), 0.0) # TODO var = tl.sum(centered_sq, axis=0) / N tl.store(means + row, mean) tl.store(variances + row, var) def mean_var(x): assert x.ndim == 2 and x.stride(1) == 1 M, N = x.shape means = torch.empty(M, device=x.device, dtype=x.dtype) variances = torch.empty_like(means) mean_var_kernel[(M,)](x, means, variances, N, x.stride(0), BLOCK=triton.next_power_of_2(N)) return means, variances for shape in ((3, 17), (4, 128)): x = torch.randn(shape, device="cuda") m, v = mean_var(x) torch.testing.assert_close(m, x.mean(1), atol=1e-4, rtol=1e-4) torch.testing.assert_close(v, x.var(1, unbiased=False), atol=2e-4, rtol=2e-4) ``` ### 检查方法在 CUDA/Triton 环境运行本单元格；断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。 ### Q1 不要背定义：请从输入、状态变化和输出三个阶段解释“均值方差与 Mask 二次传播”的工作机制。 **你的答案：** ### Q2 只在 load 时 mask 为什么仍会污染方差？ **你的答案：** ### Q3 大均值小方差时如何改进数值稳定性？ **你的答案：** ## 评分与通过规则 - 代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分； - Q1～Q3 各 2 分； - 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。 ```{python} #| course_role: solution import torch import triton import triton.language as tl @triton.jit def mean_var_kernel(x, means, variances, N: tl.constexpr, stride_m: tl.constexpr, BLOCK: tl.constexpr): row = tl.program_id(0) cols = tl.arange(0, BLOCK) mask = cols < N values = tl.load(x + row * stride_m + cols, mask=mask, other=0.0) mean = tl.sum(values, axis=0) / N centered_sq = tl.where(mask, (values - mean) * (values - mean), 0.0) var = tl.sum(centered_sq, axis=0) / N tl.store(means + row, mean) tl.store(variances + row, var) def mean_var(x): assert x.ndim == 2 and x.stride(1) == 1 M, N = x.shape means = torch.empty(M, device=x.device, dtype=x.dtype) variances = torch.empty_like(means) mean_var_kernel[(M,)](x, means, variances, N, x.stride(0), BLOCK=triton.next_power_of_2(N)) return means, variances for shape in ((3, 17), (4, 128)): x = torch.randn(shape, device="cuda") m, v = mean_var(x) torch.testing.assert_close(m, x.mean(1), atol=1e-4, rtol=1e-4) torch.testing.assert_close(v, x.var(1, unbiased=False), atol=2e-4, rtol=2e-4) ``` ### Q1 参考答案从三个阶段拆解其运行机制： 1. **输入阶段（网格绑定与加法零元注入）**：按行发射 Program，列维度以 2 的幂次构造 `cols = tl.arange(0, BLOCK)`。在掩码 `mask = cols < N` 保护下，通过带 `other=0.0` 的 `tl.load` 一次性将该行数据拉入寄存器向量 `values`，超出 $N$ 的 Padding lane 注入加法单位元 0.0。 2. **状态变化阶段（两阶段归约与二次掩码阻断）**： - 阶段 A：调用 `tl.sum` 计算有效总和，除以真实列宽 $N$ 得到行均值标量 $\mu$。 - 阶段 B：执行广播减法 `values - mean` 并求平方。由于 Padding lane 上的 0.0 变为 $-\mu$ 后平方变成 $\mu^2$，立即通过 `tl.where(mask, diff_sq, 0.0)` 强制将 Padding lane 的虚假能量清零，使加法单位元在数据流中重新闭环。 - 阶段 C：调用二次 `tl.sum` 计算清洗后的方差平方和，除以真实列宽 $N$ 得到总体方差 $\sigma^2$。 3. **输出阶段（双标量写回）**：将算出的均值标量与方差标量，通过各自独立的 `tl.store` 写回到目标全局内存地址 `means + row` 与 `variances + row`。 ### Q2 参考答案在初始加载阶段（`tl.load`）通过 `other=0.0` 确实成功阻断了未初始化内存的污染，使加法归约满足 $X_i + 0.0 = X_i$。但是，计算方差的前提是执行**中心化偏移运算** $(x_i - \mu)$：对于有效的 $N$ 个数据，计算的是真实的差值 $(x_i - \mu)$；但对于多余的 $\text{BLOCK} - N$ 个 Padding lanes，其寄存器数值为 0.0，参与运算后变成了 $(0.0 - \mu) = -\mu$。随后的平方运算使其变成了 $(-\mu)^2 = \mu^2$。只要均值 $\mu \neq 0$，$\mu^2$ 就必然是一个大于 0 的正数！这意味着，多余的每个 Padding lane 都会虚构出一个数值为 $\mu^2$ 的非零偏差。在后续执行 `tl.sum` 时，这些本该被完全忽略的无效 lane 会向累加器注入总量为 $(\text{BLOCK} - N) \times \mu^2$ 的巨大伪方差能量，导致计算出的方差虚高，产生严重的数学错误。 ### Q3 参考答案当输入数据呈现“大均值、小方差”特征时（例如数据集中在 $10^6 \pm 0.01$），直接使用单精度 FP32 的两阶段朴素公式会导致严重的**灾难性数值抵消（Catastrophic Cancellation）**。工程改进策略包括： 1. **寄存器累加器提升为 FP64 计算**：在计算均值与平方和时，将寄存器内的 `values` 显式转换为高精度的双精度浮点数 `tl.float64`。求出 `mean` 与二次中心化平方和后再转换回 FP32 写回，彻底抹平浮点尾数截断误差。 2. **采用流式 Welford 增量递推算法**：改用单遍（One-pass）Welford 在线均值方差算法： $$ \mu_k = \mu_{k-1} + \frac{x_k - \mu_{k-1}}{k}, \quad M_{2, k} = M_{2, k-1} + (x_k - \mu_{k-1})(x_k - \mu_k) $$ 在 Warp Shuffle 归约时采用 Welford Combine 算子合成两个子块的统计量。该算法从代数结构上保证了差值始终在局部小量之间累积，条件数最优，理论上杜绝了浮点相消问题。 ## 参考资料 - [Triton Tutorials](https://triton-lang.org/main/getting-started/tutorials/) - [Triton language API](https://triton-lang.org/main/python-api/triton.language.html) 资料用于建立事实基线；面试回答仍需用自己的语言组织。 ::: ::: {.course-pagination} [**← 上一课** 第 5 课：行归约与 Padding 单位元](lesson05_row_sum.qmd){.pagination-link .pagination-prev} [**下一课 →** 第 7 课：稳定 Softmax 与融合](lesson07_softmax.qmd){.pagination-link .pagination-next} :::