Lei Mao's Log Book
·
PyTorch Triton内核的透明追踪与编译
💡
原文英文，约2600词，阅读约需10分钟。
📝
内容提要
本文讨论了如何在PyTorch中使用Triton内核进行透明追踪和编译。用户可以通过@triton.jit创建Triton内核，并利用torch.compile和torch.export进行优化。注册的自定义操作可以在追踪和编译中被识别，而未注册的操作需使用TorchDynamo进行追踪。最终，Triton内核可通过AOTInductor进行预编译，以提升性能。
🎯
关键要点
- 用户可以使用@triton.jit创建Triton内核，并在模型中进行训练和推理。
- Triton内核可以被视为不透明的自定义操作，使用@torch.library.register_fake进行注册。
- 要让Triton内核在编译时可见，必须将其实现暴露给编译器。
- 使用torch.compile和torch.export可以优化Triton内核。
- 未注册的操作需要使用TorchDynamo进行追踪。
- 通过AOTInductor可以对Triton内核进行预编译以提升性能。
- 注册的Triton操作可以在torch.export中无论strict参数如何都能成功导出。
- 未注册的Triton操作在strict=True时可以通过TorchDynamo进行追踪并成功导出。
❓
延伸问答
如何在PyTorch中创建Triton内核？
用户可以通过@triton.jit装饰器创建Triton内核，并在模型中进行训练和推理。
Triton内核如何进行优化？
可以使用torch.compile和torch.export对Triton内核进行优化，确保内核实现对编译器可见。
未注册的Triton操作如何进行追踪？
未注册的Triton操作需要使用TorchDynamo进行追踪，以便在编译时被识别。
如何使用AOTInductor进行Triton内核的预编译？
通过AOTInductor可以对Triton内核进行预编译，以提升性能，确保内核实现对编译器可见。
Triton内核在torch.export中的导出条件是什么？
Triton内核可以在torch.export中导出，条件是内核必须注册为自定义操作，并且可以在strict参数为True或False时成功导出。
使用torch.compile时，Triton内核的可见性有什么要求？
Triton内核的实现必须对编译器可见，以便进行优化和编译。
➡️
