facebookincubator/AITemplate

这是什么？

核心功能

• 01将 fp16 神经网络编译为面向 TensorCore 和 MatrixCore 的 CUDA/HIP C++ 代码。
• 02通过水平、垂直和内存算子融合减少 kernel 启动次数。
• 03生成的模型为自包含二进制，不依赖 cuBLAS、cuDNN 或 TensorRT。
• 04Python 运行时可直接接收 PyTorch 张量，无需额外拷贝。
• 05FX2AIT 可转换 PyTorch 模型，并对不受支持的算子进行部分加速。

适用场景

• 在 A100 等 NVIDIA GPU 上以 fp16 部署 ResNet、BERT 或 Vision Transformer 模型。
• 在 AMD MI 系列 GPU 上运行 Stable Diffusion 推理。
• 加速包含少量 AIT 不支持算子的 PyTorch 模型。

快速上手

1. 1

   克隆仓库及其子模块

   git clone --recursive https://github.com/facebookincubator/AITemplate
2. 2

   进入 python 目录

   cd python
3. 3

   构建 Python wheel 包

   python setup.py bdist_wheel
4. 4

   安装构建好的 wheel 包

   pip install dist/*.whl --force-reinstall

使用前须知

• !NVIDIA 支持仅在 SM80 及以上 GPU（如 Ampere）上测试，较旧的 T4 或 V100 可能无法运行部分内核。
• !AMD 支持仅在 CDNA2 GPU（MI-210/250）上测试，较旧的 CDNA1 GPU（MI-100）可能存在编译器问题。
• !从源码构建需要正确的 CUDA 或 ROCm 编译器，否则会导致性能下降。

收录于专题

标签

常见问题

AITemplate 是什么？

一个将深度神经网络编译为 CUDA/HIP C++ 代码、用于 NVIDIA 和 AMD GPU 高性能推理的 Python 框架。 AITemplate（AIT）是一个 Python 框架，可将深度神经网络转换为 CUDA（NVIDIA）或 HIP（AMD）C++ 代码，用于推理服务。它面向 fp16 计算，充分利用 TensorCore 和 MatrixCore 硬件，支持的模型包括 ResNet、MaskRCNN、BERT、Vision Transformer 和 Stable Diffusion 等。

AITemplate 可以免费使用吗？用的什么许可证？

AITemplate 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。

AITemplate 还在维护吗？

半年内有更新。AITemplate 最近一次提交是在 2026年10月7日。

AITemplate 适合用来做什么？

在 A100 等 NVIDIA GPU 上以 fp16 部署 ResNet、BERT 或 Vision Transformer 模型。 在 AMD MI 系列 GPU 上运行 Stable Diffusion 推理。 加速包含少量 AIT 不支持算子的 PyTorch 模型。

AITemplate 怎么安装和运行？

克隆仓库及其子模块: git clone --recursive https://github.com/facebookincubator/AITemplate 进入 python 目录: cd python 构建 Python wheel 包: python setup.py bdist_wheel 安装构建好的 wheel 包: pip install dist/*.whl --force-reinstall

在哪里获取 AITemplate？

源代码和安装说明在 GitHub 上：https://github.com/facebookincubator/AITemplate