第 5 课：行归约与 Padding 单位元

本课目标与完成标准

学完后你应能：实现非 2 次幂长度 row sum，解释 reduction 的 padding 单位元。

通过必须同时满足：

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

前置关系

• 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
• 本课在路线中的作用：Triton tl.sum 在 program 内归约 tensor；常把 BLOCK 取为 N 的下一 2 次幂。

核心心智模型

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

在科学计算与深度学习中，跨维度规约（Reduction，如求和、均值、方差、Softmax 分母）是高频核心模式。对于矩阵 $X \in \mathbb{R}^{M \times N}$，行求和操作（Row Sum）输出一个 $M$ 维列向量： $$y_i = \sum_{j=0}^{N-1} X_{i, j}$$ 传统 CUDA 实现中，行归约需要由一个 Thread Block 协作完成，开发者必须精细管理共享内存、多轮 __syncthreads() 同步壁垒以及 Warp Shuffle 洗牌指令。 而在 Triton 中，规约通过内置原语 tl.sum 在 Program 块内对张量直接完成。 本课解决的核心工程问题是：当矩阵列数 $N$ 不是 2 的幂次时，如何通过填充“代数单位元（Algebraic Identity Element）”，让 SIMD 折半树状规约硬件与任意物理形状完美解耦，保证计算的绝对正确性。

2. 它如何工作

行归约的代数保护与片上树状折半数据通路如图 5-1 与图 5-2 所示：

矩阵第 row 行数据

超出真实列数 N 的元素填充单位元 0.0 (加法零元)

tl.sum(values, axis=0) 对一行求和

输出当前行总和标量

图 5-1：Triton 行归约结构与 Padding 加法单位元保护图

Program 绑定矩阵单行

分块加载行内数据并应用掩码

Triton 片上折半树状规约

单值写回输出向量 out[row]

图 5-2：行归约内部树状累加与写回流程

1. 行绑定与 2 的幂次对齐： 网格以行维度划分 grid = (M,)，Program $i$ 绑定矩阵第 $i$ 行。Triton 要求 tl.arange(0, BLOCK) 的长度（以及块张量的各维）是 2 的幂，因此取 BLOCK = triton.next_power_of_2(N)，多出的 lane 用 mask 处理。

2. 带单位元的安全合并加载（Identity-padded Masked Load）： 构造谓词 mask = cols < N。在执行 tl.load 时，必须显式指定 other=0.0： values = tl.load(x + row * stride_m + cols, mask=mask, other=0.0) 对于有效索引 $j < N$，加载真实的显存数值；对于超出 $N$ 的 Padding lanes（数量为 $\text{BLOCK} - N$），硬件直接将其在寄存器中初始化为加法单位元 $0.0$。

3. 片上 SIMD 树状归约（Reduction Tree）： 调用 tl.sum(values, axis=0) 时，编译器先在每个线程内部累加它持有的元素，再用 Warp 内 Shuffle 折半相加（如 $\text{lane}[i] += \text{lane}[i+16]$、$\text{lane}[i+8] \dots$），多个 warp 之间再经共享内存合并。因为 Padding 处的数值恒为加法零元： $$X_{i,j} + 0.0 = X_{i,j}$$ 因此这部分虚构的 lane 参与任何层级的两两加法，都不会污染真实的数值总和。

4. 标量结果写回： 归约完成后产生一个标量寄存器值，通过 tl.store(out + row, ...) 单播写回全局内存目标向量。

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

• 单位元与代数算子的严格对齐（Identity Mismatch）： 单位元的选择必须严格满足代数群论中该算子的零元性质：
  • 求和规约（Sum）：单位元为加法零元 $0.0$；
  • 乘积规约（Prod）：单位元为乘法单位元 $1.0$；
  • 最大值规约（Max）：单位元为负无穷大 $-\infty$（float('-inf')）；
  • 最小值规约（Min）：单位元为正无穷大 $+\infty$（float('inf')）。 若在求和时误设 other=1.0，最终结果将必然产生偏移误差 $+ (\text{BLOCK} - N)$；若在求 Max 时误设 other=0.0，如果该行数据全为负数（如 $[-5, -10, -3]$），Padding 的 0 将反客为主被错误判定为最大值。
• Hardware Block 尺寸上限： Triton 块张量的元素数上限是 TRITON_MAX_TENSOR_NUMEL（$2^{20}$），但实际可用的 BLOCK 远小于此，受寄存器文件限制：fp32 下几千到一两万个元素以上就会出现严重的寄存器溢出（Spill）或编译变慢。若列数 $N = 65536$，应改为单 Program 内沿列循环分块累加，或用多 Program 两阶段规约（Two-pass Reduction）。

4. 性能与工程取舍

• 一 Program 一行 vs 循环分块规约： 当 $N \le 4096$ 时，“一行一 Program”模型能够将整行数据装入寄存器，归约只需一次全局读，不需要全局内存中继（跨 warp 的部分由编译器经少量共享内存完成）。

• Padding 算力浪费与对齐开销： 若 $N = 1025$，其下一 2 的幂为 $2048$。此时有 $2048 - 1025 = 1023$ 个无效 lane（近 50% 的算力与寄存器被单位元填充）。这是 2 的幂形状约束的代价；它们不访问内存，所以带宽不浪费。浪费过大时可改为用较小的 BLOCK 沿列循环。

具体演示

以 $N = 100$ 为例：

• 计算 BLOCK = triton.next_power_of_2(100) = 128。
• 构建 cols = [0, 1, ..., 99, 100, ..., 127]。
• mask = [True x 100, False x 28]。
• 执行 tl.load(..., mask=mask, other=0.0) 后，寄存器向量前 100 个 lane 为矩阵真实值，后 28 个 lane 严格置为 0.0。
• 执行 7 级折半相加网络（$128 \to 64 \to 32 \to 16 \to 8 \to 4 \to 2 \to 1$），最终得到 $\sum_{j=0}^{99} X_{i,j}$，多余 28 个 lane 贡献为 0。由于求和顺序与 PyTorch 不同，结果与 x.sum(1) 在舍入误差范围内一致，而不是逐位相同，所以断言使用了容差。

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

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

补齐 masked load 的 other。

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

import torch
import triton
import triton.language as tl

@triton.jit
def row_sum_kernel(x, out, 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=______)  # TODO: 加法单位元
    tl.store(out + row, tl.sum(values, axis=0))

def row_sum(x):
    assert x.ndim == 2 and x.stride(1) == 1
    M, N = x.shape
    out = torch.empty((M,), device=x.device, dtype=x.dtype)
    row_sum_kernel[(M,)](x, out, N, x.stride(0), BLOCK=triton.next_power_of_2(N))
    return out

for shape in ((1, 1), (4, 100), (7, 257)):
    x = torch.randn(shape, device="cuda")
    torch.testing.assert_close(row_sum(x), x.sum(1), atol=1e-4, rtol=1e-4)

检查方法

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

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

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“行归约与 Padding 单位元”的工作机制。

你的答案：

Q2

为什么 padding 有 28 lanes 而不是“只处理 100 lanes”？

你的答案：

Q3

若归约是 max，other 应改成什么？

你的答案：

评分与通过规则

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

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

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

参考答案（仅 answer 分支）

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

import torch
import triton
import triton.language as tl

@triton.jit
def row_sum_kernel(x, out, 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)
    tl.store(out + row, tl.sum(values, axis=0))

def row_sum(x):
    assert x.ndim == 2 and x.stride(1) == 1
    M, N = x.shape
    out = torch.empty((M,), device=x.device, dtype=x.dtype)
    row_sum_kernel[(M,)](x, out, N, x.stride(0), BLOCK=triton.next_power_of_2(N))
    return out

for shape in ((1, 1), (4, 100), (7, 257)):
    x = torch.randn(shape, device="cuda")
    torch.testing.assert_close(row_sum(x), x.sum(1), atol=1e-4, rtol=1e-4)

Q1 参考答案

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

1. 输入阶段（网格绑定与幂次对齐生成）：Python 包装层为每个矩阵行发射一个独立的 Program，核函数内部利用 cols = tl.arange(0, BLOCK) 生成向上对齐到 2 的幂次的列索引向量。在 mask = cols < N 保护下，通过 tl.load 结合 other=0.0 加载数据，有效索引从显存读入真实数据，越界索引在寄存器中初始化为加法单位元 0.0。

2. 状态变化阶段（寄存器折半树状归约）：调用 tl.sum 时，底层硬件调度执行 Warp 内 Shuffle 指令，按严格平衡的二叉树网络展开折半累加（例如 128 维经过 7 轮步长递减累加折叠为 1 个标量）。由于后 28 个 lane 携带的是加法单位元 0.0，代数上满足 $x + 0 = x$，折半树中任意节点与 0.0 相加均不改变真实数据累加和，以极高流水线吞吐在寄存器内完成归约。

3. 输出阶段（标量定点写回）：折半树汇聚出的单精度标量值，通过 tl.store(out + row, ...) 直接写入目标一维张量的第 row 个位置。整个过程各行之间完全数据隔离，零线程间争用。

Q2 参考答案

直接原因是 Triton 的形状约束：tl.arange(0, BLOCK) 和块张量的各维必须是 2 的幂，所以无法声明长度为 100 的块，只能取 BLOCK = next_power_of_2(100) = 128。编译器把 128 个元素按固定布局分配给 warp 内各线程，归约的 shuffle 步长也按 2 的幂展开（详见第 2 节第 3 条）。

多出的 28 个 lane 用 mask 屏蔽访存，再用 other=0.0 填入加法单位元，使它们参与归约但不改变结果。代价是约 22% 的寄存器和算术空转，不产生额外的内存流量。

Q3 参考答案

若归约算子从求和（Sum）改为最大值（Max），other 必须从加法单位元 $0.0$ 修改为最大值代数单位元：负无穷大 $-\infty$（在 Triton/Python 中写作 float('-inf')）。

• 数学因果链：最大值代数运算 $(\mathbb{R} \cup \{-\infty\}, \max)$ 的单位元是 $-\infty$，对于任意实数 $x$，恒有 $\max(x, -\infty) = x$。
• 错误后果：若错误保留 other=0.0，一旦输入行中的所有有效元素全部为负数（例如一组对数概率值 $[-3.5, -7.2, -1.8]$），无效的 Padding lanes 携带的 0.0 会在折半比较中强行胜出（$\max(-1.8, 0.0) = 0.0$），导致整个内核输出完全荒谬的错误最大值 0.0。

参考资料

• Triton Tutorials
• Triton language API

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

← 上一课
第 4 课：二维 Grid、Stride 与… 下一课 →
第 6 课：均值方差与 Mask 二次传播