第 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 内核中，tl.arange 要求块宽度为 2 的幂，因此张量宽度需扩展至 $\text{BLOCK} = \text{next\_power\_of\_2}(N)$ 。在上节课中，我们在加载时为 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 所示：

统计量计算两阶段

第一次归约: sum(x) / valid_count ➔ 得到均值 mu

第二次归约: sum((x - mu)^2) / valid_count ➔ 得到方差 var

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

输入块

Pass 1: 计算均值并保留 valid_count

Pass 2: 重新遍历差值平方并排除越界元素

同时输出均值与总体方差（分母 N）

图 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，PyTorch 返回 NaN（并给出警告），调用方需要在边界处理这种情况。
遗漏 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 算法：本课的两遍算法（先求 $\mu$ ，再求 $\sum (x_i-\mu)^2$ ）数值上是稳定的；容易出问题的是单遍公式 $E[x^2] - E[x]^2$ ，在“大均值、小方差”（例如元素在 $10^4$ 附近、波动约 $10^{-2}$ ）时两个大数相减会发生灾难性抵消。整行放不进一个 Program、必须流式单遍处理时，应使用 Welford 算法，它的跨块合并需要维护 $(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 参考答案

other=0.0 只保证 padding lane 在第一次求和时是单位元。中心化之后 padding lane 变为 $0 - \mu = -\mu$ ，平方后为 $\mu^2$ ，只要 $\mu \neq 0$ 就不再是 0。第二次 tl.sum 会多加 $(\text{BLOCK} - N)\mu^2$ ，方差偏大 $(\text{BLOCK} - N)\mu^2 / N$ （数值例子详见“具体演示”第 2 步）。所以每次逐元素变换可能破坏单位元时，都要在归约前用 tl.where(mask, ..., 0.0) 重新施加掩码。

Q3 参考答案

先区分公式：单遍公式 $E[x^2] - E[x]^2$ 在“大均值、小方差”时（例如数据在 $10^4 \pm 0.01$ ）会发生灾难性数值抵消（Catastrophic Cancellation）；本课的两遍公式先减均值再平方，已经避开了这一问题，剩余误差主要来自 $\mu$ 本身的舍入。若数据绝对值大到 FP32 无法表示其波动（例如 $10^6$ 附近的 ulp 为 0.0625， $\pm 0.01$ 的波动在输入阶段就已丢失），任何算法都无法恢复，只能提高输入精度。工程改进策略包括：

寄存器累加器提升为 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)$ 跨块合并时使用 Welford/Chan 的并行合并公式合成两个子块的 $(n, \mu, M_2)$ 。它的优势是单遍、可流式，并且避免了 $E[x^2] - E[x]^2$ 的抵消；精度与两遍算法相当。

参考资料

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

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

--- title: "第 6 课：均值方差与 Mask 二次传播" jupyter: python3 --- ## 本课目标与完成标准学完后你应能：实现 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 内核中，`tl.arange` 要求块宽度为 2 的幂，因此张量宽度需扩展至 $\text{BLOCK} = \text{next\_power\_of\_2}(N)$。在上节课中，我们在加载时为 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 所示： ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-lr"><div class="fd-sg" data-nid="Stats" data-depth="0"><div class="fd-sg-title">统计量计算两阶段</div><div class="fd-stack"><div class="fd-comp fd-lr"><div class="fd-rank fd-few"><div class="fd-node fd-box" data-nid="Sum">第一次归约: sum(x) / valid_count ➔ 得到均值 mu</div><div class="fd-node fd-box" data-nid="Var">第二次归约: sum((x - mu)^2) / valid_count ➔ 得到方差 var</div></div></div></div></div></div></div></div> ``` <p class="caption" align="center"><em>图 6-1：均值与方差的两阶段掩码统计量结构</em></p> ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank"><div class="fd-node fd-box" data-nid="In">输入块</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Pass1">Pass 1: 计算均值并保留 valid_count</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Pass2">Pass 2: 重新遍历差值平方并排除越界元素</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Out">同时输出均值与总体方差（分母 N）</div></div></div></div><i class="fd-e" data-s="In" data-t="Pass1" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Pass1" data-t="Pass2" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Pass2" data-t="Out" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 6-2：两阶段掩码二次传播计算时序流</em></p> 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，PyTorch 返回 NaN（并给出警告），调用方需要在边界处理这种情况。 - **遗漏 `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 算法**：本课的两遍算法（先求 $\mu$，再求 $\sum (x_i-\mu)^2$）数值上是稳定的；容易出问题的是单遍公式 $E[x^2] - E[x]^2$，在“大均值、小方差”（例如元素在 $10^4$ 附近、波动约 $10^{-2}$）时两个大数相减会发生灾难性抵消。整行放不进一个 Program、必须流式单遍处理时，应使用 Welford 算法，它的跨块合并需要维护 $(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 参考答案 `other=0.0` 只保证 padding lane 在第一次求和时是单位元。中心化之后 padding lane 变为 $0 - \mu = -\mu$，平方后为 $\mu^2$，只要 $\mu \neq 0$ 就不再是 0。第二次 `tl.sum` 会多加 $(\text{BLOCK} - N)\mu^2$，方差偏大 $(\text{BLOCK} - N)\mu^2 / N$（数值例子详见“具体演示”第 2 步）。所以每次逐元素变换可能破坏单位元时，都要在归约前用 `tl.where(mask, ..., 0.0)` 重新施加掩码。 ### Q3 参考答案先区分公式：单遍公式 $E[x^2] - E[x]^2$ 在“大均值、小方差”时（例如数据在 $10^4 \pm 0.01$）会发生**灾难性数值抵消（Catastrophic Cancellation）**；本课的两遍公式先减均值再平方，已经避开了这一问题，剩余误差主要来自 $\mu$ 本身的舍入。若数据绝对值大到 FP32 无法表示其波动（例如 $10^6$ 附近的 ulp 为 0.0625，$\pm 0.01$ 的波动在输入阶段就已丢失），任何算法都无法恢复，只能提高输入精度。工程改进策略包括： 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) $$ 跨块合并时使用 Welford/Chan 的并行合并公式合成两个子块的 $(n, \mu, M_2)$。它的优势是单遍、可流式，并且避免了 $E[x^2] - E[x]^2$ 的抵消；精度与两遍算法相当。 ## 参考资料 - [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} [**← 上一课**<br>第 5 课：行归约与 Padding 单位元](lesson05_row_sum.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 7 课：稳定 Softmax 与融合](lesson07_softmax.qmd){.pagination-link .pagination-next} :::