第 1 课：Program、Block 与 Mask

本课目标与完成标准

学完后你应能：建立 Triton program instance 的 SPMD 心智模型并实现任意长度 fill。

通过必须同时满足：

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

前置关系

• 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
• 本课在路线中的作用：Triton 的 program 类似一个向量化 CTA：一个 program 同时处理 BLOCK 个 lane，而不是显式写单线程代码。

核心心智模型

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

Triton 编程范式（Program、Block 与 Mask） 把 CUDA C++ 以单线程（SIMT Thread-level）为单位的编程，换成以数据块为单位的 分块级单程序多数据（Block-level SPMD） 编程。

CUDA SIMT 痛点 vs. Triton 块级抽象（结合图 1-1 与图 1-2）： - 传统 CUDA 的心智负担：开发者必须在脑海中手动折算“Warp 内 32 个线程的交错排布、寄存器 Bank Conflict、Shared Memory Padding”，极易出错； - Triton 的分块心智模型（图 1-1）：

1D Program 网格 (tl.program_id(0))

Program 0: 处理 [0, BLOCK_SIZE)

Program 1: 处理 [BLOCK_SIZE, 2*BLOCK_SIZE)

Program k: 处理 [k*BLOCK, (k+1)*BLOCK)

偏移计算: offsets = k * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

图 1-1：Triton SPMD Block 划分与块偏移指针代数

• 开发者编写的代码不是针对“单个线程”，而是针对一个 Program 实例（相当于 CUDA 的一个 CTA / ThreadBlock）；
• 变量不是单个标量，而是多维向量/块张量（Block of Data），例如 tl.arange(0, BLOCK) 一次性声明了一个长度为 BLOCK 的向量切片；
• 编译器（Triton-MLIR 编译器）自动负责将其映射到底层的物理 Warp、分配通用寄存器与 Tensor Core 指令；
• 边界防御武器：Mask（图 1-2）：

计算逻辑元素索引 offsets

生成布尔掩码: mask = offsets < N

安全加载/存储: tl.store(ptr + offsets, val, mask=mask)

边界外非法越界被自动抑制 (无任何非法访存)

图 1-2：Triton tl.store 边界掩码 (Mask) 保护执行流程图

真实世界的数据规模 $N$ 极少能被 2 的幂次分块大小（如 128、256）整除。Triton 不会自动做边界检查，需要开发者显式传入布尔掩码（Mask），由编译器生成带谓词的访存指令，屏蔽尾部越界访问。

2. 它如何工作（结合图 1-1 与图 1-2）

以一维数组填充操作（Fill Kernel）为例，执行流程如下：

1. 全局网格调度（Grid Launch）：
   • 算法通过天花板除法确定启动的 Program 实例总数：$\text{grid} = \lceil N / \text{BLOCK} \rceil$；
   • 每个 Program 拥有全局唯一的身份标识 pid = tl.program_id(0)；
2. 块级指针向量代数（图 1-1）：
   • 基础偏移量：当前 Program 的起始全局索引为 $\text{start} = \text{pid} \times \text{BLOCK}$；
   • 向量下标生成：offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)；
   • 此时 offsets 是一个在各 Lane 上均匀递增的 1D 整数向量：$[\text{start}, \text{start}+1, \dots, \text{start}+\text{BLOCK}-1]$；
3. 谓词掩码解算与边界防御（图 1-2）：
   • 生成布尔向量掩码：mask = offsets < N；
   • 对前部的完整块，mask 全为 True；而在最后一个尾块中，超出真实长度 $N$ 的尾部槽位自动被置为 False；
4. 安全向量写入（tl.store）：
   • 执行 tl.store(out + offsets, value, mask=mask)；
   • 编译出的底层汇编指令带有 @p 谓词保护，谓词为假的 lane 不发出写操作。

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

1. 尾部 Program 越界写入引发静默内存撕裂：
   • 致命误区：在分配了 cdiv(N, BLOCK) 个 Program 的情况下，若省略 mask=mask；
   • 因果必然：最后一个 Program 必定会向 $[N, \text{grid} \times \text{BLOCK})$ 的未知非法区域强行写入数据，直接踩踏 PyTorch 显存池中属于其他张量的活跃内存，造成不可逆的脏数据污染；
2. constexpr 误用引发编译风暴（Cache Thrashing）：
   • 很多初学者误把运行时动态传入的浮点标量（如 value）标记为 tl.constexpr；
   • 严重后果：Triton 会把所有 constexpr 的具体取值作为 JIT 编译特化的缓存 Key。若传入不同的数值，Triton 每次都会被迫重新触发完整的 Python-to-PTX 编译耗时（数十毫秒到数秒），导致严重的运行时卡顿。

4. 性能与工程取舍（Block 尺寸与 JIT 调优）

1. BLOCK 尺寸的权衡艺术：
   • 过大 BLOCK（如 1024/2048）：单个 Program 所需寄存器文件急剧膨胀，容易引发寄存器溢出（Spill），且在数据量小的时候无法激活足够的 SM 核心；
   • 过小 BLOCK（如 16/32）：无法打满 GPU 的 128 字节合并访存事务，指令发射开销过高；
   • 工业标准：通常将 BLOCK 设置为 2 的整数次幂（128、256 或 512），并通过 @triton.autotune 自动化搜索最优组合。

具体演示：尾块物理槽位与掩码过滤实证

考虑数据规模 $N = 10000$，分块尺寸 $\text{BLOCK} = 256$：

• Program 划分： $$\text{grid} = \lceil 10000 / 256 \rceil = 40 \text{ 个 Program (ID 为 } 0 \sim 39\text{)}$$

• 完全满载 Program（ID = $0 \sim 38$）：

  • 每个 Program 覆盖 256 个元素；
  • 累计已处理有效元素：$39 \times 256 = 9984$ 个元素；
  • 对应的 mask 均为 256 个全 True。

• 边界尾块 Program（最后一个 ID = 39）：

  • 起始偏移：$\text{start} = 39 \times 256 = 9984$；
  • 向量下标范围：$[9984, 9985, \dots, 10239]$；
  • 该块内有效数据数量：$10000 - 9984 = 16$ 个元素；
  • 生成的掩码形态：前 16 个槽位为 True，后 $256 - 16 = 240$ 个槽位全部被精确置为 False。
  • 硬件据此精准写入仅存的 16 个有效数字，其余 240 个非法显存地址被 100% 抑制，捍卫了系统边界完整性。

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

补齐 program 的向量下标。

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

import torch
import triton
import triton.language as tl

@triton.jit
def fill_kernel(out, n, value, BLOCK: tl.constexpr):
    offsets = ______  # TODO: program 起点 + lane
    mask = offsets < n
    tl.store(out + offsets, value, mask=mask)

def fill(n, value):
    out = torch.empty((n,), device="cuda", dtype=torch.float32)
    fill_kernel[(triton.cdiv(n, 256),)](out, n, value, BLOCK=256)
    return out

for n in (1, 257, 10000):
    got = fill(n, 3.25)
    torch.testing.assert_close(got, torch.full_like(got, 3.25))

检查方法

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

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

Q1

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

你的答案：

Q2

删除 store mask 会在哪个 program 越界？

你的答案：

Q3

value 改成 tl.constexpr 会怎样影响编译缓存？

你的答案：

评分与通过规则

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

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

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

参考答案（仅 answer 分支）

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

import torch
import triton
import triton.language as tl

@triton.jit
def fill_kernel(out, n, value, BLOCK: tl.constexpr):
    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    tl.store(out + offsets, value, mask=mask)

def fill(n, value):
    out = torch.empty((n,), device="cuda", dtype=torch.float32)
    fill_kernel[(triton.cdiv(n, 256),)](out, n, value, BLOCK=256)
    return out

for n in (1, 257, 10000):
    got = fill(n, 3.25)
    torch.testing.assert_close(got, torch.full_like(got, 3.25))

Q1 参考答案

1. 输入阶段（Grid Partitioning & Block Offset Ingestion）：
   • 调度系统根据张量总长 $N$ 和分块大小 $\text{BLOCK}$，通过天花板除法 $\lceil N / \text{BLOCK} \rceil$ 启动对应数量的 Program 实例；
   • 每个 Program 根据自己的唯一标识 tl.program_id(0) 计算所属瓦片的起始基准，并结合向量生成原语 tl.arange(0, BLOCK) 构造出覆盖整个 Block 的一维整数偏移张量 offsets = pid * BLOCK + tl.arange(0, BLOCK)；
2. 状态变化阶段（Vectorized Predication Evaluation）：
   • Program 对整个偏移向量执行逐元素逻辑比较 mask = offsets < n，在寄存器中生成布尔条件张量；
   • 内部满载的 Program 获得全为 True 的掩码；最后一个负责处理边界余数的 Program，超出 $n$ 的尾部槽位被精准置为 False；
3. 输出阶段（Predicated Hardware Vector Store）：
   • 执行 tl.store(out + offsets, value, mask=mask)；
   • Triton 编译器将其编译为带 @p 谓词保护的写指令（地址对齐且连续时可向量化为 128 位写）。只有掩码为 True 的位置被写入，越界地址不会被访问。

Q2 参考答案

若在 tl.store 中删除了 mask=mask，越界访存会精确发生在最后一个 Program 实例中（即 $\text{Program ID} = \lceil N / \text{BLOCK} \rceil - 1$），引发严重的内存破坏：

1. 越界物理区间的数学定位：
   • 设总长度 $N = 10000$，$\text{BLOCK} = 256$。总计启动 $P = \lceil 10000 / 256 \rceil = 40$ 个 Program，编号为 $0 \sim 39$；
   • 前 39 个 Program（$pid = 0 \dots 38$）处理的物理范围是 $[0, 39 \times 256) = [0, 9984)$，均严格落在有效区间 $[0, 10000)$ 之内，绝对不会越界；
   • 唯一发生越界的实体是最后一个 Program（$pid = 39$）： 该 Program 计算出的偏移量为 $[39 \times 256, 40 \times 256) = [9984, 10240)$；
   • 其中合法索引仅有 $9984 \sim 9999$（共 16 个元素）；从索引 $10000$ 到 $10239$（240 个浮点数、即 960 字节）超出了 out 的范围；
2. 毁灭性后果：
   • 若无 Mask 保护，这 240 个浮点数会强行写入显存，直接覆写 PyTorch 内存池中紧邻该张量的其他活跃权重、激活值或梯度数据，引发极难排查的静默数值错乱（Silent Data Corruption）或触发 cudaErrorIllegalAddress 崩溃。

Q3 参考答案

tl.constexpr 形参的具体取值会作为常量折叠进生成的代码，同时是 JIT 缓存 key 的一部分（详见第 3 节第 2 条）。因此：

1. 每个新取值都触发一次完整编译：value 依次传入 3.25、3.26、0.0 时，每次都要重新走 Triton IR → LLVM → PTX/cubin 的编译，单次耗时通常为几十毫秒到秒级，而 kernel 本身只需微秒级；
2. 缓存膨胀：内存缓存和 ~/.triton/cache 中会积累大量只差一个常量的 kernel 二进制。

准则：tl.constexpr 只用于决定代码形状的元参数（如 BLOCK_SIZE、dtype 选择、开关）；随数据变化的标量（如 value、n）作为普通运行时参数传递。

参考资料

• Triton Tutorials
• Triton language API

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

← 上一课
第 14 课：DeepEP 通信算子 下一课 →
第 2 课：Load/Store 与向量加法