深度学习算子

主页 >> 深度学习算子 | 大模型预训练 | 常用网址 | 实用工具

算子库参考

许多算子都有前人的宝贵实现可以参考,可事半功倍.

算子库名称 算子库介绍
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接入的编译器