AITemplate

| |

AITemplate (AIT) 是一个 Python 框架，可将深度神经网络转换为 CUDA（NVIDIA GPU）/ HIP（AMD GPU）C++ 代码，以实现闪电般的推理服务。AITemplate 的亮点包括：

• 高性能：在主要模型上接近 NVIDIA GPU 的 fp16 TensorCore 或 AMD GPU 的 MatrixCore 性能上限，包括 ResNet、MaskRCNN、BERT、VisionTransformer、Stable Diffusion 等。
• 统一、开放且灵活。无缝支持适用于 NVIDIA GPU 或 AMD GPU 的 fp16 深度神经网络模型。完全开源，乐高式易扩展的高性能原语可用于新模型的支持。相比现有方案，它对两种 GPU 平台都支持更为全面的融合操作。

关于 AITemplate 的更多信息

出色的向后兼容性

AITemplate 不依赖于第三方库或运行时环境，例如 cuBLAS、cuDNN、rocBLAS、MIOpen、TensorRT、MIGraphX 等。每个模型都会被编译成一个自包含的可移植二进制文件，可在任何具有相同硬件的软件环境中使用。

水平融合

AITemplate 提供独特的高级水平融合功能。它可以将并行的 GEMM、LayerNorm 以及其他输入形状不同的算子融合为单个 GPU 内核。

垂直融合

AITemplate 具有强大的垂直融合能力。它可以将大量操作融合为 TensorCore/MatrixCore 操作，例如逐元素运算、归约和布局变换等。AITemplate 还支持连续风格的 TensorCore / MatrixCore 操作融合。

内存融合

AITemplate 提供创新的内存融合技术。它可以将 GEMM、LayerNorm 等算子与后续的拼接、拆分、切片等内存操作融合为单个算子。

无需 PyTorch 即可工作

AITemplate 生成的 Python 运行时可以接受 PyTorch 张量作为输入和输出，而无需额外的拷贝。对于没有 PyTorch 的环境，AITemplate 的 Python/C++ 运行时是自包含的。

扩展无忧

AITemplate 提供了一种简单的代码生成扩展方法。要向 AITemplate 中添加新的算子或新的融合内核，大多数情况下只需添加两个 Python 文件：一个用于图节点定义，另一个用于后端代码生成。文本头文件中的 CUDA/HIP 内核可以直接用于代码生成。

FX2AIT

FX2AIT 是一款基于 Python 的工具，可将 PyTorch 模型转换为 AITemplate（AIT）引擎，以实现闪电般的推理服务。借助 FX2AIT 内置的 AITLowerer，即使模型中包含 AITemplate 不支持的算子，也可以实现部分 AIT 加速。

FX2AIT 的主要特性包括：

• 易于转换：FX2AIT 只需一个 PyTorch 模型和输入即可完成转换，生成用于推理服务的“AITModule”输出。
• 更广泛的兼容性：AITemplate 并非支持所有 PyTorch 算子。FX2AIT 的 AITLowerer 为包含不支持算子的模型提供了部分 AIT 转换的解决方案。更多信息请参阅 fx2ait/fx2ait/example/03_lowering_split。

更多详细信息请访问：https://github.com/facebookincubator/AITemplate/tree/main/fx2ait。

安装

硬件要求：

• NVIDIA：AIT 仅在 SM80+ GPU（Ampere 等）上进行了测试。并非所有内核都适用于旧的 SM75/SM70（T4/V100）GPU。
• AMD：AIT 仅在 CDNA2（MI-210/250）GPU 上进行了测试。对于旧的 CDNA1（MI-100）GPU，可能存在编译器问题。

克隆代码

克隆代码时，请使用以下命令同时克隆子模块：

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

Docker 镜像

我们强烈建议使用 Docker 来运行 AITemplate，以避免意外使用错误版本的 NVCC 或 HIPCC。

• CUDA：./docker/build.sh cuda
• ROCM：DOCKER_BUILDKIT=1 ./docker/build.sh rocm

这将构建一个标签为 ait:latest 的 Docker 镜像。

从源码安装

以下命令将为 AITemplate 创建一个 Python wheel 包。请确保已安装正确的 CUDA/ROCm 编译器。

• CUDA：CUDA 11.6
• ROCm：我们已在 ROCm 5.2.3 上进行了测试，并使用 docker/Dockerfile.rocm#L87-L96 中的命令定制了 HIPCC。

使用错误的编译器会导致性能下降。

请在继续下一步之前，确认所有子模块均已正确克隆。

cd python
python setup.py bdist_wheel
pip install dist/*.whl --force-reinstall

开始使用

请参阅 AITemplate 文档，以获取 API 参考。

以下是几篇入门教程：

• 01：如何使用 AIT 推理 PyTorch 模型
• 02：如何向 AIT 代码生成中添加算子
• 03：如何可视化 AIT 的优化过程

示例与性能

AITemplate 提供以下模型模板及在 A100/MI-250 上的参考性能数据：

• 01_ResNet-50 使用 PyTorch Image Models (TIMM)
• 02_MaskRCNN-FPN 使用 Detectron2
• 03_BERT 使用 Hugging Face Transformer
• 04_Vision Transformer 使用 PyTorch Image Models (TIMM)
• 05_Stable Diffusion 使用 Hugging Face Diffusers

发布

当前的所有开发更新都可以在 AITemplate 仓库中查看。发布并不遵循固定时间表，只有在重大功能发布时才会打上标签。

中期计划：

• 更好的动态形状支持：重点解决 Transformer 中的动态序列问题，并增加符号化形状支持。
• 更多自动图遍历：减少手动重写模型以获得最佳性能的工作量。
• 量化：fp8/int8/int4。
• Gemm 的稀疏化剪枝。
• PT2 集成：Aten2AIT 正在积极开发中。

长期计划：

• 自动 ONNX、Open-XLA 等格式的模型转换。
• 在 AVX2/AVX-512 上为 AMD Epyc CPU 提供可组合的内核扩展。

贡献

请查阅我们的 贡献指南，了解如何为该项目做出贡献。

团队

AITemplate 当前由 Meta 的工程师维护：Ying Zhang、Yang Chen、Terry Chen、Mu-Chu Lee、Max Podkorytov、Adnan Akhundov。

AITemplate 由 Meta 的工程师共同创建：Bing Xu、Ying Zhang、Hao Lu、Yang Chen 和 Terry Chen，同時也得到了其他多位優秀工程師的重要貢獻。在此特別提及的部分人員包括：Mike Iovine、Mu-Chu Lee、Scott Wolchok、Oleg Khabinov、Shirong Wu、Huamin Li、Hui Guo、Zhijing Li、Max Podkorytov。我們亦感謝 Andrew Tulloch、Yinghai Lu 和 Lu Fang 提供的寶貴討論。

FX2AIT 和 Aten2AIT 由 Meta 的工程師共同創建並維護：Wei Wei、Shirong Wu 和 Zhijing Li。

致謝

AITemplate 團隊與 NVIDIA 的 CUTLASS 團隊（由 Andrew Kerr 和 Haicheng Wu 領導）以及 AMD 的 Composable Kernel 團隊（由 Chao Liu 和 Jing Zhang 領導）保持緊密合作。我們共同設計了許多針對各個平台的先進 GPU 最佳化技術，而這些成果的實現離不開彼此的密切協作。

授權條款

AITemplate 依據 Apache 2.0 許可證 授權。

AITemplate 快速上手指南

AITemplate (AIT) 是一个由 Meta 开源的 Python 框架，旨在将深度神经网络转换为高性能的 CUDA (NVIDIA) 或 HIP (AMD) C++ 代码，以实现极速推理。它不依赖 cuBLAS、cuDNN 等第三方运行时库，生成的模型为自包含的二进制文件，支持高度融合优化。

环境准备

硬件要求

• NVIDIA GPU: 仅测试支持 SM80+ 架构（如 Ampere 系列的 A100, RTX 3090/4090 等）。旧版 SM75/SM70 (如 T4, V100) 可能无法运行部分内核。
• AMD GPU: 仅测试支持 CDNA2 架构（如 MI-210, MI-250）。CDNA1 (MI-100) 可能存在编译器问题。

软件依赖

• 操作系统: Linux
• 编译器:
  • NVIDIA: CUDA 11.6+ (需安装 nvcc)
  • AMD: ROCm 5.2.3+ (需安装定制版 hipcc)
• Python: 建议 Python 3.8+
• Git: 用于克隆代码及子模块
• Docker (强烈推荐): 为避免编译器版本冲突，建议使用官方提供的 Docker 镜像进行构建和运行。

安装步骤

方法一：使用 Docker（推荐）

此方法可自动配置正确的 NVCC/HIPCC 环境，避免手动安装编译器的繁琐步骤。

1. 克隆项目代码（务必包含子模块）：

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

2. 构建 Docker 镜像：

   • NVIDIA (CUDA):

     ./docker/build.sh cuda
     

   • AMD (ROCm):

     DOCKER_BUILDKIT=1 ./docker/build.sh rocm
     

   构建完成后，镜像标签为 ait:latest。

方法二：源码安装

如果您已在宿主机正确配置了编译器环境，可直接构建 Python Wheel 包。

1. 克隆代码并进入 Python 目录：

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

2. 构建并安装 Wheel 包：

   python setup.py bdist_wheel
   pip install dist/*.whl --force-reinstall
   

   注意：请确保使用的 CUDA/ROCm 编译器版本符合要求，否则会导致性能下降或编译失败。

基本使用

AITemplate 通常配合 FX2AIT 工具使用，将 PyTorch 模型直接转换为 AIT 引擎。以下是最简化的转换与推理流程：

1. 准备 PyTorch 模型

假设你有一个标准的 PyTorch 模型 model 和对应的输入张量 inputs。

2. 使用 FX2AIT 进行转换

利用 fx2ait 将模型编译为 .ait 格式或直接生成可调用的模块。

import torch
from fx2ait import AITLowerer, AITModule
from fx2ait.tools.common_fx2ait import change_batch_size

# 1. 定义你的 PyTorch 模型 (示例为 ResNet-50)
# model = torchvision.models.resnet50(pretrained=True)
# model.eval()

# 2. 定义输入示例 (Batch Size, Channels, Height, Width)
# inputs = (torch.randn(1, 3, 224, 224).cuda(),)

# 3. 创建 AITLowerer 实例
lowerer = AITLowerer(
    model=model,
    inputs=inputs,
    # 可选配置：指定动态 Shape 或其他优化参数
    # dynamic_shapes=[{0: "batch"}], 
)

# 4. 执行转换，生成 AITModule
ait_module = lowerer.lower()

# 5. 推理执行 (输入输出均为 PyTorch Tensor，无额外拷贝)
output = ait_module(*inputs)

3. 独立运行 (无 PyTorch 环境)

如果部署环境没有安装 PyTorch，AITemplate 生成的二进制文件也可通过其自带的 C++/Python 运行时独立加载运行，具体可参考生成的 .so 或 .dll 文件接口。

更多资源

• 官方文档: AITemplate Documentation
• 示例代码: 查看 examples/ 目录下的 ResNet-50, BERT, Stable Diffusion 等完整案例。
• 教程: 官方提供了如何添加算子、可视化优化过程等详细教程。