先记住一句话
先证明瓶颈，再按“成熟库 → 编译融合 → Triton → CUDA extension → CUTLASS/CuTe”逐级增加实现成本；每一级都要守住 reference、梯度、动态 shape 和真实 workload benchmark。
1. 什么时候值得写
- profile 显示许多小 kernel/中间 HBM round trips；
- 现有库不支持特殊 layout、稀疏结构或 fused epilogue；
- shape distribution 稳定，专用 kernel 可以利用先验；
- 量化/解量化、sampling、ragged sequence 有明显 framework overhead。
如果标准库已经 compute-bound 且接近硬件上限，重写通常只增加维护负担。
2. 工具选择梯子
| 层级 | 适合 | 代价 |
|---|---|---|
| cuBLAS/cuDNN/CUB 等 | 标准 dense/primitives | 定制边界较少 |
| torch.compile | 可捕获的 PyTorch 组合 | graph break/shape guards |
| Triton | 规则 dense、fusion、快速 autotune | 仍需理解 mapping/memory |
| CUDA C++ extension | 不规则控制、底层原语、完整控制 | 代码、build、ABI 成本最高 |
| CUTLASS/CuTe | GEMM/conv 类高性能模板 | 模板与架构专门化复杂 |
3. Triton 的编程心智
一个 Triton program instance 处理一个数据 block；program_id 选 tile，向量化 pointer arithmetic 生成 addresses，用 mask 处理边界。编译器把 block program 映射到 GPU threads。它省去大量 thread-level boilerplate，但 tile、warps、stages、coalescing 和 occupancy 仍需调优。
pid → compute block offsets
load pointers with mask
compute vector/tile expression
store result with mask4. CUDA extension 的接口契约
host wrapper 要检查 device、dtype、shape、stride、alignment 与 contiguous 假设；使用当前 CUDA device/stream，而不是私建默认 stream；kernel launch 后保留可诊断的 error path。kernel 则明确 grid/block、shared memory、vector width 和边界条件。
5. PyTorch dispatcher 与 autograd
算子先定义 schema，再注册 CPU/CUDA implementations。若需要训练，应注册 backward/autograd formula 并用 gradcheck；若要进入 torch.compile，还需 FakeTensor/meta implementation 描述 output 的 shape、dtype、device，而不访问真实数据。opcheck 检查注册契约，不代替数值测试。
6. Forward 快不够
保存给 backward 的 tensor 会占显存；recompute 能省显存但多算。in-place 写入可能破坏 autograd versioning；非连续 stride、channels-last、空 tensor、极小/极大 shape、NaN/inf 都要定义行为。训练算子的目标应是 forward + backward + optimizer step 的总时间。
7. Autotune 与 shape 分布
block sizes、num warps、pipeline stages 对 shape/hardware 敏感。autotune 应在有限 config space 中按 shape key 缓存；不要让首次编译时间混入 steady-state latency，也不要只优化一个漂亮的 square shape。真实请求中常见的是一组长尾形状。
8. CUTLASS/CuTe 的位置
当核心是 GEMM-like mainloop + custom epilogue，CUTLASS 提供 architecture-specific MMA、async copy、scheduler 与 layouts，通常比从裸 CUDA 造矩阵乘更合理。CuTe 用 layout algebra 描述 tensor 与 thread/value mapping，能力强但学习曲线陡。
9. 正确性与发布清单
- 与高精度 reference 比 forward；
- finite difference/gradcheck 比 backward；
- 覆盖 dtype/device/stride/boundary/dynamic shapes；
- compute-sanitizer 查越界与 race；
- benchmark warmup、同步、编译缓存和端到端；
- 为不同 CUDA/PyTorch/GPU 架构构建并测试 wheel，或采用受控 JIT cache。
10. 四个 custom-op 决策计算
例 1：fusion 收益上限
三个 40 μs elementwise kernels 各有 5 μs launch，总 135 μs。理想融合成一个 70 μs kernel+5 μs launch 为 75 μs，speedup=135/75=1.8×。
例 2：Triton program count
N=1,000,000、BLOCK_SIZE=1024，programs=ceil(N/1024)=977；最后 program 仅处理 576 elements，其余 lanes 由 mask 禁用。
例 3：Amdahl 检查
Hot op 占 step 25%，即使 custom kernel 让它 4×，总 speedup=1/(0.75+0.25/4)=1.23×。若维护成本很高，可能不值得。
例 4：autotune 成本摊销
12 configs 每个 benchmark 20 次、单次 0.2 ms，纯运行约 48 ms（编译另算）。若之后执行 1M 次，每次省 5 μs，总省 5 s，调优成本可摊销；只运行 100 次则不行。
Triton 不是“不用懂 CUDA”。它提升表达层级，但 warp 数、寄存器、shared memory、访存合并、tile reuse 与 numerical semantics 仍决定结果。
自测
1. 为什么 fused kernel 常比逐算子快？
减少 launch 和中间 tensor 的 HBM 读写，并可能在 register/shared memory 中直接传递中间值。
2. FakeTensor/meta implementation 做什么？
不计算真实数据，只告诉编译器输出 metadata，使 graph capture、shape propagation 和编译成为可能。
3. 什么时候首选 CUTLASS 而非裸 CUDA？
问题是 GEMM/conv 类且需要自定义 layout、scheduler 或 epilogue，同时希望复用成熟的架构级矩阵乘 pipeline。
官方资料
PyTorch Custom C++ and CUDA Operators覆盖 dispatcher、autograd、torch.compile 与测试；Triton 官方教程展示 vector add、fused softmax 与 matmul；NVIDIA CUTLASS Overview说明 CuTe 与 GEMM building blocks。