飞书练习 3 · 算子主线入口

torch.compile 与 Triton

两条路要一起走：框架怎么自动生成 kernel（compile），以及你怎么手写 tile（Triton）。后者才是求职主战场。

代码：04 / 04b / 04c · 需要 CUDA

0. 先建立地图

手写 Triton                    torch.compile
─────────────                  ──────────────
你写 @triton.jit               你写 PyTorch 模型
你算 offset + mask             FX 追踪成图
你选 BLOCK_SIZE / grid         Inductor 融合、选模板
compiler 把 tile 降到 GPU      生成 Triton 或 CUDA
                               第二次运行走缓存

Triton 的卖点不是「用 Python 代替 CUDA 语法」，而是：你按 tile 写访存与计算，compiler 负责把一个 tile 拆到 warp / 做 coalescing。CUDA 里那些 thread idx、shared memory 排布，很多被藏进 compiler。

飞书强调：tl.program_id 要理解成 tile id，不要当成 CUDA 的 thread id。 一个 program 一次 tl.load 一整块；越界用 mask 挡住。

1. 手写：vector add（必跑 04b）

每个 GPU program 处理 BLOCK_SIZE 个元素。第 pid 个 program 覆盖区间：

offsets = pid * BLOCK_SIZE + [0, 1, …, BLOCK_SIZE-1]
mask    = offsets < n          # 尾块：n 往往不能整除 BLOCK
tile 切向量 pid=0 pid=1 pid=2 尾块+mask 一条长向量，按 BLOCK_SIZE 切成 tiles 每个 pid：tl.load → 寄存器里相加 → tl.store 橙色：offsets 会超出 n，mask=False 的 lane 不写
图 1 · 04b 的网格：program_id 是第几块砖，不是第几根线程
python 04b_triton_vector_add.py
# 期望：n=100000 不能整除 1024，allclose True

对照代码里三件事说清就算过关：

  1. grid = cdiv(n, BLOCK_SIZE) 开多少个 program
  2. offsets 怎么从 pid 算出来
  3. 为什么必须有 mask（否则尾块越界）

2. 手写：融合 add+relu（04c）

Eager：t = x+y 写回 HBM，再 relu(t) 再读一遍。融合：max(x+y, 0) 在同一 kernel 里做完，中间数不用进全局内存。

这就是「算子融合」的最小例子，也是 Inductor 想自动干的事。

python 04c_triton_fused_add_relu.py
# 打印 fused vs eager 是否 allclose，以及粗计时

数字随卡而变，重点是趋势：少一次中间张量的 store/load。大模型里 Softmax、RMSNorm、残差加法都是同一类动机。

3. 自动生成：torch.compile（04）

PyTorch 模型 / 函数
  → Dynamo/FX 抓图（看见 add、relu、matmul…）
  → Inductor 决定融合、调度、选 Triton 模板
  → 生成 @triton.jit 或 CUDA
  → 第一次慢（编译），第二次走 cache
python 04_torch_compile_observe.py
# 日志很吵时：
TORCH_LOGS=output_code python 04_torch_compile_observe.py 2>&1 | tee dumps/compile.log
grep -n triton dumps/compile.log

脚本会：

生成代码通常比手写啰嗦（一堆 tl.static_range、对齐、自动调 BLOCK）。你要找的仍是那三件套：pid、load/store、mask。

CPU 上看不到 Triton。无 GPU 时 04/04b/04c 会跳过，这是预期，不是脚本坏了。

4. 和 CUDA / 编译器的边界（面试口述）

你写什么compiler 写什么
CUDAthread、warp、shared、sync很少
Tritontile 的 offset、mask、算法lane 映射、coalesce、流水
torch.compilePyTorch 图连 tile 策略也代劳（质量不一定够）

飞书「设计思想」里的原话精神：Triton compiler 假定数据已经在 SRAM 语义上按 tile 处理；DRAM 进出由你（或 Inductor）用 load/store 写明。优化两大绞肉机仍是 fusion 和 tiling。

5. 建议动手顺序

  1. 04b 跑通，改 BLOCK_SIZE=256 看 grid 变大
  2. 04c 看融合；把 tl.maximum 改成只 add，对比计时（中间仍可能被写回）
  3. 04 对照 Inductor 生成文本和你的 04c 像不像
  4. 回飞书摘录 12/13：JIT compile & run，不必先读完整份 Triton 源码

6. 通关（对着代码说 2 分钟）

  1. program_id 是 tile，不是 thread
  2. mask 解决什么
  3. 融合省的是哪一次 HBM 往返
  4. compile 的链路：FX → Inductor → Triton

7. 可观看精读（上站）

B 站合集、GPU MODE Lecture 14/29、Reduce / Matmul 短片等，已压成图文：Triton 精读导览。音视频只在本地 gpu_triton/。

← 张量并行 · 下一篇：Triton 精读 →