深度学习算子
算子库参考
许多算子都有前人的宝贵实现可以参考,可事半功倍.
| 算子库名称 | 算子库介绍 |
|---|---|
| Eigen | 广泛应用于tensorflow的初代算子库 |
| aten | 广泛应用于pytorch的初代算子库 |
| FA,fa-rocm | flash attention的官方实现,目前已基本基于cutlass/ck |
| flashinfer | 大模型时代算子后端,主要支持大模型各种算子,基于cuda和cutlass,vLLM和SGLang都已接入 |
| cutlass | GPU算子界的瑞士军刀 |
| AITemplate | Meta开源的算子生成工具,主攻fp16算子,支持算子水平融合/垂直融合/内存融合,能转换一些torch,可自包含,但是业界未广泛接入 |
| kernl | triton算子库,主要特性有:cuda graph, triton compiler, torch dynamo接入 |
TODO
- 水平融合/垂直融合/内存融合分别是什么
深度学习编译器
| 名称 | 介绍 |
|---|---|
| TVM | 老牌深度学习编译器 |
| XLA | Tensorflow主要接入的图编译器 |
| Triton | Pytorch接入的编译器 |