GPU系统拆解-14-从 PyTorch 到 CUDA Extension：把算子真正接入工程

本文是「GPU系统拆解」系列第 14 篇。
系列导读：GPU系统拆解-00-导读：从架构认知到推理系统的学习路线
上一篇：GPU系统拆解-13-多 GPU 与通信：并行策略、拓扑与扩展代价

这一篇的目标不是再讲 CUDA 语法，而是把前面学过的 GPU、kernel、profiling 和推理系统知识，真正落到“我怎么把一个想法做成可运行、可 benchmark、可 profile、可面试讲清楚的算子工程”上。学完这一篇之后，你应该知道什么样的算子值得自己写、最小工程链路是什么、为什么很多手写 kernel 反而更慢，以及怎样把一个自定义算子讲成完整工程闭环。

1. 先给结论

学 PyTorch -> CUDA extension 的意义，不是学会写一个 demo，而是学会把性能想法真正接进框架和系统。
自定义算子的价值通常不在“比 PyTorch 更低级”，而在“更贴近你的 shape、数据布局和系统关键路径”。
一个最小 extension 工程，至少要有：Python 调用层、C++ binding 层、launcher 层、CUDA kernel 层。
很多手写 kernel 比 PyTorch 慢，不是因为想法错了，而是因为你在和高度优化库对比，或者 benchmark / 访存 / 瓶颈判断出了问题。
真正值得自己写的，通常是高频、小而热、可融合、shape 较固定、默认实现没有照顾好的算子。
一个成熟的自定义算子工程，不只是 kernel 本身，还包括正确性验证、稳定 benchmark、profile 和收益判断。

2. 为什么要学从 PyTorch 到 CUDA extension

前面学的很多东西都偏“硬件视角”：

warp 和 block 怎么执行
shared memory 为什么有用
coalescing 和 occupancy 怎么看
Tensor Core 为什么适合矩阵乘
decode 为什么更像 memory problem

但在真实 AI infra / 推理 岗位里，常见问题不是：

你会不会解释一个 kernel。

而是：

你能不能把一个优化想法真的落到框架里，并证明它值得做。

这就是 extension 的意义。它是：

框架层和 CUDA 层之间的桥。

3. 先建立一个最小调用链脑图

可以先把整个路径记成这样：

Python / PyTorch
  -> Python wrapper / torch op
  -> C++ binding
  -> launcher
  -> CUDA kernel
  -> GPU 执行

真正重要的是理解每一层在干什么。

3.1 Python 层

这一层通常负责：

提供用户接口
做 shape / dtype 检查
组织 benchmark 和测试
方便和模型代码集成

3.2 C++ binding 层

这一层负责把 Python 调用接到 C++ / CUDA。

它通常会做：

接收 torch::Tensor
检查设备、dtype、contiguous
准备输出 tensor
调用 launcher

3.3 Launcher 层

这一层通常负责：

计算 grid / block
取指针
绑定 stream
调 CUDA kernel

3.4 Kernel 层

这一层才是你真正优化的地方：

并行映射
访存模式
shared memory
warp primitive
向量化
Tensor Core 路径

所以一个很实用的理解是：

Python / C++ 层决定“这个算子怎么接进系统”
kernel 层决定“它到底快不快”

4. 为什么很多时候要自己写算子

4.1 通用算子不一定适合你的工作负载

PyTorch 内置算子很强，但它必须兼容很多不同情况：

各种 shape
各种 dtype
各种边界条件
各种数据布局

所以它未必对你的固定场景最优。

特别在推理里，很多常见需求都很像：

fused bias + activation
fused add + norm
特定 KV cache gather / scatter
特定 layout 的 attention 前后处理

这些场景的共同点是：

单个操作不大，但调用频繁；如果不融合，就会被 launch 和中间读写拖慢。

4.2 推理系统追求的是整体路径，而不是单点漂亮

推理里经常关心的是：

TTFT
ITL
token/s
P99 延迟
显存占用

所以写自定义算子往往不是为了炫技，而是为了：

少一次 kernel launch
少一次中间 tensor 写回
少一次无意义的数据搬运
让系统关键路径更短

5. 一个最小工程结构应该长什么样

最常见的结构大致是：

my_extension/
  setup.py
  bindings.cpp
  kernel.cu
  wrapper.py
  test_op.py

这些文件的职责很清楚：

setup.py：构建 extension
bindings.cpp：暴露给 Python 调用的接口
kernel.cu：launcher 和 CUDA kernel
wrapper.py：Python 端封装
test_op.py：正确性和 benchmark

这里最重要的工程习惯是：

test 和 benchmark 不是可选项，而是算子工程的一部分。

6. 最小示例真正应该学什么

很多教程会用向量加法作为最小示例。

这个例子本身不重要，重要的是你要看懂它背后的三个点。

6.1 你写的不是普通 C++ 函数，而是框架可调用算子

这和“写一个 CUDA 作业”完全不同。

作业里你只关心 kernel；
工程里你要关心：

外部怎么调用
输入怎么检查
输出怎么组织
怎么方便测试和集成

6.2 算子设计一定要带 shape 思维

很多 GPU 优化不是“算法更复杂”，而是“shape 特征很清楚”。

例如：

hidden size 固定
head dim 固定
top-k 固定
page size 固定

当 shape 很固定时，自定义 kernel 往往更有价值。

6.3 要习惯 launcher + kernel 的二层分工

不要把所有逻辑全堆进 kernel。更好的分工是：

launcher 负责参数准备和 launch 配置
kernel 只专注执行

这样更清晰，也更容易调优。

7. 为什么很多手写 kernel 反而比 PyTorch 慢

这是面试很爱问的点。

7.1 你在和高度优化库比

很多时候你以为在和“一个 Python 算子”比，实际上你是在和下面这些比：

cuBLAS
cuDNN
vendor-tuned fused kernel
更成熟的 Triton / CUTLASS / TensorRT 路径

所以如果你的 kernel 只是“正确但朴素”，跑不过它们很正常。

7.2 你的访存模式可能很差

高频问题通常不是数学逻辑，而是：

没有 coalescing
shared memory 用得不对
寄存器压力太高
block 配置不合理
launch 太碎

7.3 你可能 benchmark 错了

常见错误包括：

没 warmup
没 synchronize
把 host 开销混进来
shape 太小，结果主要测到 launch overhead

7.4 你优化了局部，却没优化系统路径

例如某个小操作快了 2 倍，但它只占总时间 3%，那系统整体几乎不会变。

这说明你没有先做 profile，没有先找真正热点。

8. 什么样的算子值得自己写

这题非常重要，因为它体现的是工程判断力。

8.1 高频、小而热、容易融合的算子

例如：

add + norm
bias + activation
某些小型后处理 kernel

这类算子的问题常常不是 FLOPs 不够，而是：

中间结果来回写
launch 次数太多

8.2 shape 很固定的算子

如果你的部署场景非常固定，例如：

hidden size 固定
head dim 固定
top-k 固定

那通用实现为了兼容各种情况付出的代价，你也许可以绕开。

8.3 默认实现没有照顾你的数据布局

尤其在 attention、KV cache、paged memory 场景，layout 对性能影响非常大。

8.4 系统关键路径上的热点算子

例如：

attention 相关
norm 相关
sampling / top-k
KV cache gather / scatter

9. 一个很典型的推理算子思路：fused add + norm

这类算子很有代表性，因为它体现了推理优化的核心逻辑：

不是做更复杂的数学，而是减少中间结果回写和重复读写。

如果原始路径是：

1 2	x = x + residual x = layer_norm(x)

拆成两个 kernel 时，可能会发生：

add 写回一次
norm 再读一次
norm 输出再写一次

而 fused kernel 的思路是：

一次读入需要的数据
在寄存器或更近层级中继续处理
最后尽量只写一次结果

这类融合的收益通常不在“数学更少”，而在：

launch 更少
global round-trip 更少
locality 更好

10. 写自定义算子时，kernel 设计该怎么想

这里给一个非常实用的五步法。

10.1 第一步：先判断瓶颈

先问：

是 compute-bound
还是 memory-bound
还是 launch-bound

不同瓶颈，对应完全不同的优化方向。

10.2 第二步：确定并行映射

例如：

一个 thread 一个元素
一个 warp 一行
一个 block 一个 token / 一个 row / 一个 tile

并行映射决定：

coalescing 是否自然
reduction 是否方便
shared memory 值不值得用
occupancy 是否容易受限

10.3 第三步：确定数据搬运路径

明确：

哪些数据进寄存器
哪些进 shared memory
哪些必须从 global memory 读
有没有重复读机会可以缓存

10.4 第四步：确定同步范围

thread 内：寄存器
warp 内：优先 warp primitive
block 内：shared memory + __syncthreads()
block 间：通常拆成多 kernel 或借助更高层机制

10.5 第五步：先 correctness，再 benchmark，再 profile

顺序不要反。

正确顺序通常是：

正确性
稳定 benchmark
profile 定位瓶颈
再优化

11. benchmark 和正确性验证为什么必须独立看

11.1 benchmark 至少要做到什么

至少要有：

warmup
torch.cuda.synchronize()
多次迭代平均

否则测出来的往往只是异步提交时间或初始化噪声。

11.2 正确性验证至少要覆盖什么

至少要做：

和 baseline 对比
多 shape 测试
多 dtype 测试
边界条件测试

真正能上线的算子，正确性和性能同等重要。

12. 什么时候不该自己写

这同样重要。

12.1 已有成熟库已经非常好

如果热点就是标准 GEMM、标准卷积，而且成熟库已经很强，自己写通常性价比不高。

12.2 热点根本不在这里

如果系统瓶颈不在这个算子，写了也很难有整体收益。

12.3 shape 不固定，维护成本太高

如果输入变化太大，你可能要为很多边界情况付出大量维护成本。

12.4 正确性和可维护性风险大于收益

如果收益很小，但验证和维护代价很大，工程上就不值得。

所以一个成熟判断是：

不是所有能写 custom op 的地方，都值得写 custom op。

13. 和 TensorRT、Triton、CUTLASS 的关系怎么理解

这一题经常被追问。

13.1 TensorRT

更偏推理引擎和图级优化，它不只是写 kernel，而是在做更完整的推理执行优化。

13.2 Triton

更像一种更高层的 GPU kernel 编写方式，适合很多算子级优化场景。

13.3 CUTLASS

更偏高性能矩阵和 Tensor Core 路径的模板化组件库。

13.4 手写 PyTorch CUDA extension

更适合做：

特定场景定制
框架集成
小而热的定制 kernel

这几者不是简单替代关系，更像不同抽象层。

14. 这一篇必须记住的几句话

自定义算子的价值，通常不在“更低级”，而在“更贴近你的 shape、layout 和系统关键路径”。
一个最小 extension 工程，至少包括 Python 层、C++ binding 层、launcher 和 kernel 层。
很多手写 kernel 比 PyTorch 慢，不是因为 custom op 没价值，而是因为你在和高度优化库比，或者 benchmark / 访存 / 瓶颈判断出了问题。
值得自己写的，通常是高频、小而热、可融合、shape 较固定的算子。
一个成熟的算子工程，不只是 kernel 本身，还包括 correctness、benchmark、profile 和收益判断。
不是所有能写 custom op 的地方，都值得写。

15. 精简版面试表达

如果面试官问你怎么理解 PyTorch 到 CUDA extension 的工程落地，可以这样答：

我会先看这个算子是不是系统关键路径上的热点，shape 是否比较固定，默认实现是否没有照顾我的数据布局或融合需求。如果值得做，我会把它拆成 Python 接口、C++ binding、launcher 和 CUDA kernel 四层，先做正确性，再做稳定 benchmark，再用 profile 看它到底是 compute-bound、memory-bound 还是 launch-bound。很多时候自定义算子的真正收益不在单个 kernel 的绝对速度，而在于减少 launch、减少中间结果回写、让整个推理路径更贴近 GPU 的执行方式。