不能直接混用。torch.compile默认后端不理解手写Triton kernel，需封装为torch.autograd.Function或通过torch._custom_ops注册；否则会报错、fallback或性能下降。
PyTorch 2.0 的 torch.compile 和 Triton 算子能直接混用吗？
不能直接混用。PyTorch 2.0 的 torch.compile 默认后端（如 aot_eager、inductor）会尝试将整个图转成可优化的中间表示，而手写的 Triton kernel 是独立编译的 CUDA 函数，torch.compile 不会“理解”或“内联”它——除非你显式把它封装成一个 torch.autograd.Function 子类，并确保 forward/backward 中调用的是标准 tensor 操作或已注册的自定义算子。
常见错误现象：RuntimeError: Unsupported op TritonKernel 或编译后性能反而下降，因为 Inductor 绕过了你的 kernel，改用默认实现。
- 必须把 Triton kernel 封装进
torch.autograd.Function，且forward返回普通Tensor，不能返回 raw pointer 或triton.runtime.jit.CompiledKernel - 若想让
torch.compile“看到”并可能融合你的 kernel，需通过torch._custom_ops注册为可导出算子（PyTorch 2.3+ 更成熟，2.0 支持有限） - 更稳妥的做法：关闭对 Triton 算子部分的编译，用
torch.compiler.disable装饰器隔离
如何封装 Triton kernel 为可训练的 torch.autograd.Function
这是最常用、兼容性最好的方式。核心是让 forward 调用 kernel[(grid,)](…)，backward 手动实现梯度逻辑（或用 @triton.jit 写反向 kernel）。
使用场景：你需要精细控制内存布局、避免冗余 transpose，或已有高性能 Triton 实现（如 custom softmax、block-wise matmul）。
立即学习“Python免费学习笔记（深入）”；
- forward 中所有输入 tensor 必须在同一个 device 上，且 dtype 明确；Triton kernel 参数里不要传 Python int/float，统一用
tensor.numel()或tensor.shape[0]等 tensor 属性 - backward 若复用同一 kernel，注意梯度 tensor 的 layout：比如 forward 输入是
[M, K]，backward 对权重的梯度常是[K, N]，需显式.t()或.permute()对齐 - 别忘了在
ctx.save_for_backward()中保存需要的中间结果（如 forward 的 softmax 输出），否则 backward 拿不到
示例片段：
class TritonLinearFn(torch.autograd.Function):
@staticmethod
def forward(ctx, x, w, b):
out = torch.empty(x.shape[0], w.shape[1], device=x.device, dtype=x.dtype)
grid = lambda META: (triton.cdiv(x.shape[0], META['BLOCK_SIZE_M']),)
_linear_kernel[grid](x, w, b, out, x.shape[0], w.shape[1], x.shape[1],
x.stride(0), x.stride(1),
w.stride(0), w.stride(1),
BLOCK_SIZE_M=32, BLOCK_SIZE_N=32, BLOCK_SIZE_K=32)
ctx.save_for_backward(x, w, b)
return out
为什么 torch.compile 有时会“绕过”你的 Triton 算子？
Inductor 在图捕获阶段会对每个 node 做 pattern matching，只对它认识的 ATEN ops（如 aten.mm、aten.add）做 fusion 和 lowering。你封装的 torch.autograd.Function 在 IR 中是一个 opaque node，Inductor 默认不展开——它既不优化，也不替换，只是原样保留调度。
性能影响：如果你的 Triton kernel 本身比 Inductor 生成的 fused kernel 慢（比如没做 bank conflict 优化、block size 不匹配硬件），整体性能就会变差；反之，若 kernel 高度定制（如稀疏 attention），则收益明显。
- 可通过
torch._dynamo.explain(model, *args)查看是否被 capture，以及哪些 subgraph 被 fallback（含 “backend fallback” 字样即未被编译） - 强制启用 Inductor 并打印 debug 信息：
TORCHINDUCTOR_DEBUG=1 python script.py - 若发现 Triton 算子频繁 fallback，检查是否用了 unsupported dtype（如
torch.bfloat16在旧版 Triton 中需手动 cast）或动态 shape（Triton kernel 需静态grid，但 PyTorch 2.0 支持 symbolic shape +triton.language.static_assert）
PyTorch 2.0 下 Triton 算子部署到 TorchScript 或 LibTorch 是否可行？
不可行。Triton kernel 编译依赖 Python runtime 和 CUDA driver API，生成的是 Python-callable 对象，不是可序列化的 C++ kernel。TorchScript 无法 trace 进 triton.kernel 调用，LibTorch 更无 Triton 运行时支持。
容易踩的坑：有人尝试用 torch.jit.script 包裹封装类，结果报错 Unsupported node type: triton.language.semantic 或 No attribute 'kernel'。
- 生产部署 Triton 算子，正确路径是：用 Triton 自带的
triton.compile+kernel.save导出为.so，再用 C++ 加载并绑定到 LibTorch tensor —— 但这已脱离 PyTorch 生态，需自己管理 memory 和 stream - 若必须走 TorchScript 流程，唯一办法是把 Triton kernel 功能降级为等效的 PyTorch ops（如用
torch.einsum+torch.softmax），再交给torch.jit.script - PyTorch 2.0 的
torch.export（beta）同样不支持 Triton，export 后的 GraphModule 里该 node 会变成 placeholder，运行时报错
真正要落地，得接受一个事实：Triton 和 PyTorch 的部署栈目前是两条线。写 Triton 是为了极致性能，而放弃一部分可移植性。
