flowchart LR
PtrA["指针 a_ptr + offsets"] -->|tl.load| RegA["向量寄存器块 a"]
PtrB["指针 b_ptr + offsets"] -->|tl.load| RegB["向量寄存器块 b"]
RegA & RegB -->|元素级 + 运算| RegC["结果向量块 c"]
RegC -->|tl.store| PtrC["写回 c_ptr + offsets"]
第 2 课:Load/Store 与向量加法
参考答案版:包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。
本课目标与完成标准
学完后你应能:实现双输入逐元素 kernel,理解 pointer arithmetic、mask 与 wrapper 契约。
通过必须同时满足:
- 独立补齐本课唯一的代码填空题,并通过给定检查;
- 三个问答题均说明因果链,而不是只报术语;
- 能指出至少一个正确性边界和一个性能取舍;
- 总分不低于 8/10,且没有一票否决级概念错误。
前置关系
- 课程前置:Python、PyTorch 张量、CUDA 基本线程/内存概念
- 本课在路线中的作用:Triton pointer tensor 与 value tensor 逐 lane 对应;load/store 接受同 shape mask。
核心心智模型
1. 它是什么,解决什么问题
Triton 与传统 CUDA 在内存模型上的本质区别在于:CUDA 编写的是单个线程(Thread-level)的标量指针寻址与操作,而 Triton 编写的是以 Program 块为单位、针对向量/张量块(Block of Tensors)的指针与数据操作。 在 Triton 的向量加法中,指针不再是一个标量地址,而是一个指针张量(Pointer Tensor):ptr + offsets。它与值张量(Value Tensor)在维度和形状上一一对应。解决的核心工程问题是如何将高维张量在显存中的线性地址映射为硬件向量加载指令,并在非块大小整数倍的数据尾部通过谓词掩码(Predicate Mask)安全切断非法的内存读写,防止越界访存(Out-of-Bounds Memory Fault)。
2. 它如何工作
整个向量加法的数据流与控制流遵循紧密的协同契约(如图 2-1 与图 2-2 所示):
图 2-1:Triton 向量块级加载、计算与写回数据通路
flowchart TD
Host["Python Host: launch vector_add[(n_blocks,)](a, b, c, n)"] --> Grid["GPU JIT 编译并根据 BLOCK_SIZE 分发网格"]
Grid --> Compute["各 Program 块并行读取与执行 SIMD 加法"]
Compute --> Sync["无全局同步,独立完成并退出"]
图 2-2:Triton 向量加法端到端发射与并行执行流程
主机端发射与网格映射:Python Wrapper 校验张量形状、数据类型与内存连续性后,以
grid = (cdiv(n, BLOCK),)发射 Kernel。每个 Program 分配唯一的pid = tl.program_id(0)。块级基准指针与偏移量生成: 生成的
offsets是一维 1D 张量,与基地址张量x_ptr相加形成 128/256 个连续内存地址构成的指针张量。谓词生成与带掩码的合并加载(Masked Coalesced Load): 执行
xv = tl.load(x + offsets, mask=mask)和yv = tl.load(y + offsets, mask=mask)。在硬件底层,Triton 编译器(基于 MLIR Triton Dialect)将连续指针降低为宽位宽的全局内存向量合并读取(如 128-bitLDG.E.128),对mask == False的 lane 则自动抑制内存事务或填入默认填充值(默认为 0.0)。寄存器级 SIMD 加法与写回: 向量寄存器块直接执行元素加法
zv = xv + yv,并在写回时再次施加mask:tl.store(out + offsets, zv, mask=mask)。非法 lane 被硬件写屏蔽抑制,杜绝踩踏显存。
3. 正确性条件与常见误区
- 指针与掩码的形状对齐:
tl.load和tl.store的mask形状必须与指针张量完全一致。如果加载时使用了掩码,但写回tl.store时遗漏了mask=mask,尾部 Program 会将无效寄存器值写入相邻显存地址,引发严重的显存越界污染或未定义行为。 - 连续性(Contiguity)契约:上述一维线性偏移公式
offsets假定张量在物理显存中是密集且连续分布的(stride = 1)。若传入切片或转置后的非连续张量(如x[:, 0]),x + offsets将直接跳跃跨越错误的步长,读取完全风马牛不相及的内存数据。 - 类型与广播安全:两输入张量的 dtype 必须严格匹配;若混合
fp16与fp32,必须显式调用x.to(tl.float32),否则在 JIT 编译期会报类型不匹配错误。
4. 性能与工程取舍
访存带宽受限(Memory Bound)与算子融合: 向量加法的算术强度(Arithmetic Intensity)极低:每个元素完成 1 次浮点加法(1 FLOP),需要从显存读取 2 个数据并写回 1 个数据(对于 FP32 为 字节)。算术强度仅为 。在 A100/H100 等现代加速器上,算子性能完全卡死在 HBM 显存带宽上限。
块大小(
BLOCK)与硬件向量化对齐:BLOCK应选择为 128、256、512 或 1024。必须保证为 32(Warp 大小)的倍数以维持合并访存(Memory Coalescing)。过小的BLOCK(如 16)无法打满 GPU Warp 与内存总线带宽;过大的BLOCK(如 4096)会增加寄存器压力,降低 Occupancy。融合的倍增收益:若后续紧接 ReLU、GELU 或 Scale 等逐元素算子,在 Triton 中直接在寄存器层完成
tl.maximum(xv + yv, 0.0),可彻底消除一整次显存写出和重新读取,直接减少 50% 的全局显存带宽开销。
具体演示
以 、 为例:
- 网格大小: 个 Program 实例。
- Program 0 (
pid=0):处理索引 , 全为 True。发射全宽合并访存。 - Program 1 (
pid=1):处理索引 , 全为 True。发射全宽合并访存。 - Program 2 (
pid=2):处理索引 。仅第 0 个 lane(索引 512)的 为 True,其余 255 个 lane 的 全部为 False。硬件安全读取单个标量,计算后仅写回第 512 位,剩余 255 个 lane 被硬件抑制,保证绝对内存安全。
请在阅读后先合上这一节,用自己的语言复述“输入状态 → 中间状态 → 输出状态”,再做练习。
实践任务:唯一代码填空题
补齐写回值。
规则:只能修改 TODO/______ 所在位置;不要删除断言或放宽误差。代码注释说明了每个边界条件。
Code
import torch
import triton
import triton.language as tl
@triton.jit
def add_kernel(x, y, out, n: tl.constexpr, BLOCK: tl.constexpr):
offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
mask = offsets < n
xv = tl.load(x + offsets, mask=mask)
yv = tl.load(y + offsets, mask=mask)
tl.store(out + offsets, ______, mask=mask) # TODO: 逐元素结果
def add(x, y):
assert x.is_contiguous() and y.is_contiguous() and x.shape == y.shape
out = torch.empty_like(x)
n = x.numel()
add_kernel[(triton.cdiv(n, 256),)](x, y, out, n, BLOCK=256)
return out
for n in (1, 513, 8193):
x, y = torch.randn(n, device="cuda"), torch.randn(n, device="cuda")
torch.testing.assert_close(add(x, y), x + y)检查方法
在 CUDA/Triton 环境运行本单元格;断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。
提交时请给出:补齐后的代码、实际运行输出(环境不可用时注明“仅静态审查”)以及对失败用例的解释。
Q1
不要背定义:请从输入、状态变化和输出三个阶段解释“Load/Store 与向量加法”的工作机制。
你的答案:
Q2
wrapper 不检查 contiguous,二维 view 会有什么地址解释错误?
你的答案:
Q3
如何把 add+relu 融合,能省多少中间读写?
你的答案:
评分与通过规则
- 代码 4 分:正常输入 2 分,边界输入 1 分,解释实现 1 分;
- Q1~Q3 各 2 分;
- 一票否决:结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。
需要提示时按四级机制请求:概念区域 → 具体方向 → 关键局部 → 完整答案。