AITemplate Documentation

Technical Documentation: facebookincubator/AITemplate

ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference

Overview (`README.md`) Copy Section

AITemplate

|
|

AITemplate (AIT) is a Python framework that transforms deep neural networks into CUDA (NVIDIA GPU) / HIP (AMD GPU) C++ code for lightning-fast inference serving. AITemplate highlights include:

• High performance: close to roofline fp16 TensorCore (NVIDIA GPU) / MatrixCore (AMD GPU) performance on major models, including ResNet, MaskRCNN, BERT, VisionTransformer, Stable Diffusion, etc.

• Unified, open, and flexible. Seamless fp16 deep neural network models for NVIDIA GPU or AMD GPU. Fully open source, Lego-style easily extendable high-performance primitives for new model support. Supports a significantly more comprehensive range of fusions than existing solutions for both GPU platforms.

More about AITemplate Copy Chapter

Excellent Backward Capability Copy Section

AITemplate doesn't depend on third-party libraries or runtimes, such as cuBLAS, cuDNN, rocBLAS, MIOpen, TensorRT, MIGraphX, etc. Each model is compiled into a self-contained portable binary, which can be used on any software environment with the same hardware.

Horizontal Fusion Copy Section

AITemplate provides unique advanced horizontal fusion. AITemplate can fuse parallel GEMM, LayerNorm, and other operators with different input shapes into a single GPU kernel.

Vertical Fusion Copy Section

AITemplate provides strong vertical fusion. AITemplate can fuse a large range of operations into TensorCore/MatrixCore operations, such as elementwise operations, reductions, and layout permutations. AITemplate also provides back-to-back style TensorCore / MatrixCore operation fusion.

Memory Fusion Copy Section

AITemplate provides innovative memory fusions. AITemplate can fuse GEMM, LayerNorm, and other operators, followed by memory operations such as concatenation, split, and slice into a single operator.

Working w/wo PyTorch Copy Section

The AITemplate-generated Python runtime can take PyTorch tensors as inputs and outputs without an extra copy. For environments without PyTorch, the AITemplate Python/C++ runtime is self-contained.

Extensions without suffering Copy Section

AITemplate provides a straightforward approach for making an extension in codegen. To add a new operator or a new fused kernel into AITemplate, most of the time one only needs to add two Python files: one for a graph node definition and another for the backend codegen. The CUDA/HIP kernel in a text header file can be directly utilized in the codegen.

FX2AIT Copy Chapter

FX2AIT is a Python-based tool that converts PyTorch models into AITemplate (AIT) engine for lightning-fast inference serving. Using FX2AIT's built-in AITLowerer, partial AIT acceleration can be achieved for models with unsupported operators in AITemplate.

Key features of FX2AIT include:

• Easy Conversion: FX2AIT requires only a PyTorch model and input for conversion, generating an "AITModule" output for inference serving.
• Expanded Support: AITemplate does not support all PyTorch operators. FX2AIT's AITLowerer offers a solution for partial AIT conversion for models with unsupported operators. Check the fx2ait/fx2ait/example/03_lowering_split for more information.

More info can be found from https://github.com/facebookincubator/AITemplate/tree/main/fx2ait.

Installation Copy Chapter

Hardware requirements:

• NVIDIA: AIT is only tested on SM80+ GPUs (Ampere etc). Not all kernels work with old SM75/SM70 (T4/V100) GPUs.
• AMD: AIT is only tested on CDNA2 (MI-210/250) GPUs. There may be compiler issues for old CDNA1 (MI-100) GPUs.

Clone the code Copy Section

When cloning the code, please use the following command to also clone the submodules:

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

Docker Image Copy Section

We highly recommend using AITemplate with Docker to avoid accidentally using a wrong version of NVCC or HIPCC.

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

This will build a docker image with tag ait:latest.

From Source Copy Section

The following command will create a Python wheel for AITemplate. Please ensure you have correct CUDA/ROCm compiler installed.

• CUDA: CUDA 11.6
• ROCm: We tested on ROCm 5.2.3 with a customized build HIPCC with the command in docker/Dockerfile.rocm#L87-L96

Incorrect compiler will lead performance regression.

Please check all submodules are cloned correctly before go to next step.

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

Getting Started Copy Chapter

Check out the AITemplate Documentation for API reference.

There are a few tutorials for onboarding:

• 01: How to inference a PyTorch model with AIT
• 02: How to add an op to AIT codegen
• 03: How to visualize AIT's optimization

Examples & Performance Copy Chapter

AITemplate provides the following model templates & reference performance data on A100/MI-250:

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

Release Copy Chapter

All current development updates can be seen in the AITemplate repository. Releases are not on a set schedule and will only be tagged for significant feature releases.

Mid-term plan:

• Better dynamic shape support: Focus on the dynamic sequence in Transformers. Add symbolic shape support.
• More automatic graph passes: Relief manual rewrite models to obtain the best performance.
• Quantization: fp8/int8/int4.
• Sparsity pruning for Gemm.
• PT2 integration: Aten2AIT is under active development.

Long-term plan:

• Automatic ONNX, Open-XLA and other format model conversion.
• Composable Kernel CPU extension on AVX2/AVX-512 for AMD Epyc CPU.

Contributing Copy Chapter

Check our contributing guide to learn about how to contribute to the project.

The Team Copy Chapter

AITemplate is currently maintained by Meta engineers: Ying Zhang, Yang Chen, Terry Chen, Mu-Chu Lee, Max Podkorytov, Adnan Akhundov.

AITemplate is co-created by Meta engineers: Bing Xu, Ying Zhang, Hao Lu, Yang Chen, and Terry Chen, with major contributions coming from other talented engineers. A non-exhaustive list to mention is Mike Iovine, Mu-Chu Lee, Scott Wolchok, Oleg Khabinov, Shirong Wu, Huamin Li, Hui Guo, Zhijing Li, Max Podkorytov. We also want to thank Andrew Tulloch, Yinghai Lu, Lu Fang for the valuable discussions.

FX2AIT and Aten2AIT are co-created and maintained by Meta engineers: Wei Wei, Shirong Wu and Zhijing Li.

Acknowledgements Copy Chapter

AITemplate team works closely with NVIDIA CUTLASS Team (led by Andrew Kerr, Haicheng Wu) and AMD Composable Kernel Team (led by Chao Liu, Jing Zhang). We co-designed many advanced GPU optimizations specialized for each platform, and nothing is possible without our close collaboration.

License Copy Chapter

AITemplate is licensed under the Apache 2.0 License.

Overview (`docker/README.md`) Copy Section

Docker + AITemplate

AITemplate provides a Docker image with all test, benchmark, and documentation dependencies installed.

Build CUDA Docker Image Copy Chapter

bash docker/build.sh cuda
This will build a CUDA 11 docker image with tag: ait:latest

Build ROCM Docker Image Copy Chapter

DOCKER_BUILDKIT=1 bash docker/build.sh rocm
This will build a RCOM 5 docker image with tag: ait:latest

Running Unit Tests in Docker Copy Chapter

docker run --gpus all ait:latest bash /AITemplate/tests/nightly/unittest.sh

Launching CUDA Docker Copy Chapter

docker run --gpus all -it ait:latest

Launching ROCM Docker Copy Chapter

docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined ait:latest

Common questions: Copy Chapter

• Q: When building ROCm Docker, I hit this error => ERROR [internal] load metadata for docker.io/library/ubuntu:20.04, what shall I do?

  A: Run docker pull docker.io/library/ubuntu:20.04 to pull base image manually, then re-run ./docker/build.sh rocm

Overview (`docs/README.md`) Copy Section

AITemplate Documentation

Build locally Copy Chapter

1. Install AITemplate

2. Install Sphinx

pip install autodocsumm
pip install sphinx_rtd_theme
pip install sphinx_gallery
pip install sphinxcontrib-inlinesyntaxhighlight
pip install sphinx_toolbox

3. Build HTML

make html

Overview (`fx2ait/README.md`) Copy Section

FX2AIT for AITemplate

FX2AIT is a Python-based tool that converts PyTorch models into AITemplate (AIT) engine for lightning-fast inference serving.
AITLowerer built on top of FX2AIT is able to perform AIT conversion on PyTorch model with AIT unsupported operators. Model can enjoy partial AIT acceleration using AITLowerer.

FX2AIT highlights include:

• Automatic Conversion: FX2AIT only need PyTorch model and input as input for conversion. The output can be used for inference serving directly.
• Expanded Support: AITemplate doesn't cover all operators PyTorch provides. FX2AIT provided AITLowerer as solution to support partial AIT conversion for models with AIT unsupported operators. For more information, please check example/03_lowering_split.

Installation Copy Chapter

Hardware requirement:

• NVIDIA: FX2AIT is based on AIT, thus the hardware requirement is same as AIT. AIT is only tested on SM80+ GPUs (Ampere etc). Not all kernels work with old SM75/SM70 (T4/V100) GPUs.

From Source Copy Section

The following command will create a Python wheel for AITemplate. Please ensure you have correct CUDA compiler installed.

• CUDA: CUDA 11.6
• cuDNN: v8.7.0 for CUDA 11.x
  download source: https://developer.nvidia.com/rdp/cudnn-download
  installation guidance: https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html

Incorrect compiler will lead performance regression.

cd fx2ait
python setup.py install

Docker Image Copy Section

We highly recommend using AITemplate with Docker to avoid accidentally using a wrong version of NVCC or HIPCC.

• CUDA: ./docker/build.sh cuda

This will build a docker image with tag ait:latest.

Examples Copy Chapter

AITemplate provides the following getting started tutorials:

• 01: How to inference a PyTorch Transformer model with FX2AIT
• 02: How to inference a PyTorch vision model with FX2AIT
• 03: How to inference a general PyTorch model with AIT unsupported operator using AIT Lowerer

Run Example and Test Copy Section

Example command:

cd fx2ait
python example/03_lowering_split/test_lower.py
python test/test_ait_lower.py

Overview (`static/README.md`) Copy Section

About

This directory contains all of the C++ sources that are static during AITemplate compilation, including the bulk of the runtime implementation.

C++ Implementation Copy Chapter

`Model` v.s. `ModelContainer` Copy Section

These are the two main classes involved in the C++ runtime implementation.

• The bulk of the runtime implementation is in Model.
• ModelContainer stores a set of shared constants and a collection of Models. Almost all functions in model_interface.h forward to a method on ModelContainer. When Run is invoked, ModelContainer looks for an available Model, or blocks until one is available (see the section on asynchronous predictions). It then forwards the run request to the runtime.

Code Structure Copy Section

Some important files:

• include/model_interface.h: The interface that we expose in the compiled .so
• include/model_container.h: The bulk of the ModelContainer implementation.

Some files are generated at compile time. These include:

• model-generated.h: The implementation for Model.
• model_container_base.cu: A small part of the implementation for ModelContainer needs to be codegened. So ModelContainer inherits from ModelContainerBase, and ModelContainerBase's implementation lives in this file. See model_container.h for more details.

All codegen templates can be found in backend/main_templates.py. The codegen implementation is in backend/codegen.py.

Note that many of the headers in this directory rely on generated code and thus cannot be #included in external projects. The exception is model_interface.h.

Python `Model` Copy Chapter

Model is a collection of Python bindings to the C++ AIT runtime. This section describes the API.

`AITData` Copy Section

This class represents a contiguous blob of memory that AIT will use as a tensor. It is simply a named tuple with these fields:

• data_ptr: int: An unowned pointer to GPU memory. In general, all of the APIs expect that this pointer will be valid for the entire duration of the call.
• shape: List[int]: The shape of the tensor.
• dtype: str: The tensor's dtype; one of "float32", "float16", "int32", "int64". Note that most ops only support float16 at this stage.

If using AITemplate with PyTorch, AITDatas can be constructed with the torch_to_ait_data utility:

x = torch.randn(3, 3, 3).half().cuda()
# Equivalent to AITData(x.data_ptr(), [3, 3, 3], "float16")
x_ait = torch_to_ait_data(x)

If PyTorch is not available, Model provides a set of functions for copying, allocating, and freeing GPU memory. See the docstrings in compiler/model.py for more information.

`run` Copy Section

run takes a set of inputs and outputs as AITDatas. Both arguments can be passed as either an ordered list or a dictionary (mapping name to tensor).

# Arguments as a dictionary
module.run(
  {"input0": in0_ait, "input1": in1_ait},
  {"output0": out0_ait, "output1": out1_ait},
)

# Arguments as an ordered list. Note that you might need to query
# the index mapping.
input_name_to_idx = module.get_input_name_to_index_map()
output_name_to_idx = module.get_output_name_to_index_map()

inputs = [None] * len(input_name_to_idx)
outputs = [None] * len(output_name_to_idx)

for name in input_name_to_idx:
  inputs[input_name_to_idx[name]] = ait_inputs[name]

for name in output_name_to_idx:
  outputs[output_name_to_idx[name]] = ait_outputs[name]

module.run(inputs, outputs)

One important caveat is that the output must be its maximum size. This is because of dynamic shapes - the size of the output may vary, but its shape is not inferred until inference time. The maximum shape can be queried with the get_output_maximum_shape():

# Can use either name or index.
name_to_idx = module.get_output_name_to_idx()
max_shape = module.get_output_maximum_shape(name_to_idx["output"])
max_shape = module.get_output_maximum_shape("output")

Model.run returns a dictionary of output AITDatas with (possibly dynamic) shapes that the runtime inferred.

Nullptr Inputs/Outputs

In general, inputs are allowed to be null if they are size 0 (e.g. at least one dimension is 0). The runtime enforces this with a check before any kernels are launched.

If (input_name == nullptr && dim0 * dim1 * … * dimN != 0) {
  throw std::runtime_error(“input_name cannot be null!”);
}

This is convenient since torch.data_ptr() returns null for size zero tensors. The dynamic shape computation is skipped if the lower bound of the tensor’s size is positive.

Constants

There are two types of constants in AIT; bound and unbound constants. A bound constant is known at compile time and may participate in constant folding. Bound constants are copied into GPU memory at model loading time. Values for bound constants may be provided by passing a dictionary (mapping constant name to AIT tensor) to compile_model.

Unbound constants, on the other hand, do not participate in constant folding and must be provided before running the model. These must be set via Model.set_constant:

module.set_constant("my_constant", AITData(...))
# The pointer in the the tensor must live for the entire duration of run()
module.run(...)

Constants are read-only and shared with all runtimes in the ModelContainer.

`run_with_tensors`

run_with_tensors is a convenience method with the same interface as run, except it can take lists of torch.Tensors:

input0 = torch.randn(input0_shape).cuda().half()
output0 = torch.empty(output0_shape).cuda().half()
# Returns a dictionary of reshaped outputs
result = module.run_with_tensors([input0], [output0])

Streams and Asynchronous Predictions

A pointer to a stream can optionally be passed to run. If none is given, the prediction happens on the default stream 0. If the sync argument is set to True, the stream is synchronized before run() returns. sync is True by default.

Multiple predictions can happen at the same time (on the same or different streams). Under the hood, there is a fixed-size pool of runtime objects. When all the runtimes are used, run() blocks until one is available.
The size of this pool can be configured with the num_runtimes option in Model's constructor.

CUDA Graph

Run also takes a graph_mode option. If set to true, the runtime will try to use CUDA graphs to run the model. graph_mode is not supported on ROCm.

The following is a high level overview of how graph mode works:

1. Each Model has an internal stream used for graph capturing. The model first runs all ops on this stream in capture mode. No kernel launches happen during this stage.
2. If this is the first run, a graph is instantiated via cudaGraphInstantiate.
3. On subsequent runs, we try to avoid the relatively expensive cudaGraphInstantiate call by updating the graph executor (cudaGraphExecUpdate). However, a new graph may still be instantiated if the topology of the graph somehow changed between runs.
4. Once we have the graph executor, we launch a single kernel on the stream that the user provided to run().

Graph mode is mainly beneficial when there are many small kernel launches. A lot of overhead can be avoided since there is only a single kernel launch in graph mode.

Overview (`examples/01_resnet-50/README.md`) Copy Section

ResNet-50

In this example, we will demo how to use AITemplate for inference on the ResNet-50 model from PyTorch Image Models (TIMM).

We will demo two usages:

• Using AIT to accelerate PyTorch inference
• Using AIT standalone without PyTorch

Code structure Copy Chapter

modeling
    resnet.py              # ResNet definition using AIT's frontend API
weight_utils.py            # Utils to convert TIMM R-50 weights to AIT
infer_with_torch.py        # Example to accelerate PyTorch, and seamlessly use with other PyTorch code
infer_with_numpy.py        # Dump TIMM weights to Numpy and use AIT & Numpy without 3rdparties
benchmark_pt.py            # Benchmark code for PyTorch
benchmark_ait.py           # Benchmark code for AIT

Multi-GPU profiling Copy Chapter

AIT requires to do profiling to decide best algorithms for CUTLASS and CK.
To enable multiple GPUs profiling, use the environment variable CUDA_VISIBLE_DEVICES on NVIDIA platform and HIP_VISIBLE_DEVICES on AMD platform.

For example, CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 benchmark_ait.py.

Benchmark is fast once the profilers are built.

Reference Speed vs PyTorch Eager Copy Chapter

A100-40GB / CUDA 11.6.2 Copy Section

PT = PyTorch 1.12 Eager

MI-250 / ROCm 5.2.3 / HIPCC-10736 Copy Section

PT = PyTorch 1.12 Eager

1 GCD

2 GCDs

Note for Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• To benchmark MI-250, the first step is to run python3 benchmark_ait.py to generate all necessary model dynamic library files with single GCD. Then run ./benchmark_mi250.sh {batch_size} to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch.
• To benchmark MI-250 1 GCD, we lock the frequency with command rocm-smi -d x --setperfdeterminism 1700, where x is the GPU id.
• To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command rocm-smi -d x --resetperfdeterminism, where x is the GPU id.
• Performance results are what we can reproduce and for reference only. It should not be used for other purposes.

Overview (`examples/02_detectron2/README.md`) Copy Section

Getting Started with AIT for the Inference of Detectron2 Based Models

This document describes the usage of AIT for detectron2 vision models such as mask RCNN and faster RCNN.

For an end-to-end example with the API, see prepare_and_run_rcnn.sh which covers how to prepare and run inference with mask_rcnn_R_50_FPN.

Create the AIT Model from a Config File Copy Chapter

1. Pick a model and its config file from configs, for example, mask_rcnn_R_50_FPN.yaml.

2. Build the AIT Model by running compile_model.py with the config file, for example,

cfg=examples/02_detectron2/configs/mask_rcnn_R_50_FPN.yaml
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 examples/02_detectron2/compile_model.py   \
  --config $cfg \
  --batch 1

All parameters in the built AIT model are not initialized, and therefore are filled with random values. We will initialize these parameters in the following step (i.e., exporting the weights of the pre-trained model to the AIT model). Check tmp/mask_rcnn_R_50_FPN/params.json for the list of parameters in the AIT model and their shapes.

Download the Detectron2 Pre-trained Model, and Export the Weights to the AIT Model Copy Chapter

1. For example, download Detectron2 mask_rcnn_R_50_FPN pre-trained model and save it to tmp/pt_mask_rcnn_R_50_FPN.pkl:

wget https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl -O tmp/pt_mask_rcnn_R_50_FPN.pkl

2. Export the weights from the pre-trained model to AIT model by running tools/convert_pt2ait.py:

python3 examples/02_detectron2/tools/convert_pt2ait.py  \
  --d2-weight tmp/pt_mask_rcnn_R_50_FPN.pkl \
  --ait-weight tmp/ait_mask_rcnn_R_50_FPN.pt \
  --model-name mask_rcnn_R_50_FPN

The weights are exported to AIT and saved as tmp/ait_mask_rcnn_R_50_FPN.pt for inference run.

Download Inference DataSet and Run AIT Model Copy Chapter

1. For example, download the COCO 2017 Dataset:

mkdir -p ~/.torch/datasets/coco

wget https://dl.fbaipublicfiles.com/detectron2/annotations/coco/val2017_100.tgz -O ~/.torch/datasets/coco/val2017_100.tgz
tar xzf ~/.torch/datasets/coco/val2017_100.tgz -C ~/.torch/datasets/coco && rm -f ~/.torch/datasets/coco/val2017_100.tgz

2. Run inference of the AIT model on the inputs with demo.py:

python3 examples/02_detectron2/demo.py \
  --weight tmp/ait_mask_rcnn_R_50_FPN.pt \
  --config examples/02_detectron2/configs/mask_rcnn_R_50_FPN.yaml \
  --batch 1 --input "~/.torch/datasets/coco/val2017/*.jpg" \
  --confidence-threshold 0.5 \
  --display \
  --cudagraph

Multi-GPU profiling Copy Chapter

AIT requires to do profiling to decide best algorithms for CUTLASS and CK.
To enable multiple GPUs profiling, set the environment variable CUDA_VISIBLE_DEVICES on NVIDIA platform and HIP_VISIBLE_DEVICES on AMD platform with available GPU ids.

Results Copy Chapter

PT = PyTorch 1.12 Eager

A100-40GB / CUDA 11.6 Copy Section

• Input size: 448x608

• Input size: 800x1344

MI-250 / ROCm 5.2.3 / HIPCC-10736 Copy Section

PT = PyTorch 1.12 Eager

1 GCDs

• Input size: 448x608

• Input size: 800x1344

2 GCDs

• Input size: 448x608

• Input size: 800x1344

Sample outputs Copy Section

Note for Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• To benchmark MI-250, the first step is to run python3 benchmark_ait.py to generate all necessary model dynamic library files with single GCD. Then run ./benchmark_mi250.sh {batch_size} to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch.
• To benchmark MI-250 1 GCD, we lock the frequency with command rocm-smi -d x --setperfdeterminism 1700, where x is the GPU id.
• To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command rocm-smi -d x --resetperfdeterminism, where x is the GPU id.
• Performance results are what we can reproduced. It should not be used for other purposes.

Overview (`examples/03_bert/README.md`) Copy Section

BERT

This directory contains an AIT demo for the BERT language representation model.

Only bert-base-uncased is included.

Prerequisites Copy Chapter

Install the dependencies:

python3 -m pip install transformers click torch

Benchmarking Copy Chapter

To run a basic benchmark, use benchmark.py:

python3 examples/03_bert/benchmark_ait.py

There are two options for hidden activations, gelu and fast_gelu (fast_gelu by default).
gelu is not supported on AMD hardware yet.

python3 examples/03_bert/benchmark_ait.py --activation gelu
python3 examples/03_bert/benchmark_ait.py --activation fast_gelu

The batch size and sequence length can also be configured via the command line:

python3 examples/03_bert/benchmark_ait.py --batch-size 1 --seq-length 128

PyTorch eager mode benchmarks can also be run:

python3 examples/03_bert/benchmark_pt.py

To benchmark BERT embeddings, run benchmark with --encoders-only False

Quick Demo Copy Chapter

To run a quick demo with a simple prompt, use demo.py:

python3 examples/03_bert/demo.py --prompt "The quick brown fox jumps over the lazy dog."

The demo prints out the resulting logits. The demo only works with sequence length <= 512.

Multi-GPU profiling Copy Chapter

AIT requires to do profiling to decide best algorithms for CUTLASS and CK.
To enable multiple GPUs profiling, use the environment variable CUDA_VISIBLE_DEVICES on NVIDIA platform and HIP_VISIBLE_DEVICES on AMD platform.

Reference Speed vs PyTorch Eager Copy Chapter

PT = PyTorch 1.12 Eager
OOM = Out of Memory

A100-40GB / CUDA 11.6.2 Copy Section

• Sequence length 64

• Sequence length 128

• Sequence length 384

• Sequence length 1024

• Sequence length 4096

MI-250 / ROCm 5.2.3 / HIPCC-10736 Copy Section

1 GCD

• Sequence length 64

• Sequence length 128

• Sequence length 384

• Sequence length 1024

• Sequence length 4096

2 GCDs

• Sequence length 64

• Sequence length 128

• Sequence length 384

• Sequence length 1024

• Sequence length 4096

Note Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• To benchmark MI-250, the first step is to run python3 benchmark_ait.py to generate all necessary model dynamic library files with single GCD. Then run ./benchmark_mi250.sh {batch_size} to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch.
• To benchmark MI-250 1 GCD, we lock the frequency with command rocm-smi -d x --setperfdeterminism 1700, where x is the GPU id.
• To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command rocm-smi -d x --resetperfdeterminism, where x is the GPU id.
• PyTorch Eager result doesn't reflect BetterTransformer, mainly due to BetterTransformer integration to TIMM/Transformer package is not yet landed.
• Performance results are what we can reproduced. It should not be used for other purposes.

Overview (`examples/04_vit/README.md`) Copy Section

Vision Transformer (VIT)

In this example, we will demo how to lower a pretrained Vision Transformer from TIMM, and run inference in AITemplate. We tested on two variants of Vision Transformer: Base version with 224x224 input / patch 16, and Large version with 384x384 input / patch 16.

Code structure Copy Chapter

modeling
    vision_transformer.py    # VIT definition using AIT's frontend API
weight_utils.py              # Utils to convert TIMM VIT weights to AIT
verification.py              # Numerical verification between TIMM and AIT
benchmark_pt.py              # Benchmark code for PyTorch
benchmark_ait.py             # Benchmark code for AITemplate

Reference Speed vs PyTorch Eager Copy Chapter

A100-40GB / CUDA 11.6.2 Copy Section

PT = PyTorch 1.12 Eager

• vit_base_patch16_224

• vit_large_patch16_384

MI-250 / ROCm 5.2.3 / HIPCC-10736 Copy Section

PT = PyTorch 1.12 Eager

1 GCD

• vit_base_patch16_224

• vit_large_patch16_384

2 GCDs

• vit_base_patch16_224

• vit_large_patch16_384

Note for Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• To benchmark MI-250, the first step is to run python3 benchmark_ait.py to generate all necessary model dynamic library files with single GCD. Then run ./benchmark_mi250.sh {batch_size} to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch.
• To benchmark MI-250 1 GCD, we lock the frequency with command rocm-smi -d x --setperfdeterminism 1700, where x is the GPU id.
• To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command rocm-smi -d x --resetperfdeterminism, where x is the GPU id.
• PyTorch Eager result doesn't reflect BetterTransformer, mainly due to BetterTransformer integration to TIMM/Transformer package is not yet landed.
• Performance results are what we can reproduce. It should not be used for other purposes.

Overview (`examples/05_stable_diffusion/README.md`) Copy Section

Stable Diffusion Example Copy Chapter

In this example, we show how to build fast AIT modules for CLIP, UNet, VAE models, and benchmark/run them.

Build Dependencies Copy Section

First, clone, build, and install AITemplate per the README instructions.

This AIT stable diffusion example depends on diffusers, transformers, torch and click. You could install them using pip.

Verify the library versions. We have tested transformers==4.25, diffusers==0.11[torch] and torch==1.12.

>>> import transformers
>>> transformers.__version__
'4.25.0'
>>> import diffusers
>>> diffusers.__version__
'0.11.0'
>>> import torch
>>> torch.__version__
'1.12.0+cu113'

Download the diffusers pipeline files Copy Section

Optionally, you can use Hugging Face access token. You can register in Hugging Face Hub to obtain an access token for the Stable Diffusion weights. See user access tokens for more info. Your access tokens are listed in your Hugging Face account settings.

stable-diffusion model has two variants - base and regular.
For example:

• stabilityai/stable-diffusion-2-1-base - image resolution 512x512
• stabilityai/stable-diffusion-2-1 - image resolution 768x768

python3 scripts/download_pipeline.py \
--model-name "stabilityai/stable-diffusion-2-1-base"

# Optionally, you can use access token
--token ACCESS_TOKEN

Build AIT modules for CLIP, UNet, VAE Copy Section

Build the AIT modules by running compile.py.

Set correct width and height depending on the model variant

python3 scripts/compile.py --width 512 --height 512

It generates three folders: ./tmp/CLIPTextModel, ./tmp/UNet2DConditionModel, ./tmp/AutoencoderKL. In each folder, there is a test.so file which is the generated AIT module for the model.

Alternative build script

python3 scripts/compile_alt.py --width 64 1536 --height 64 1536 --batch-size 1 4 --clip-chunks 6

This compiles modules with dynamic shape. In the example, modules will work with width in range 64-1536px, batch sizes 1-4. Clip chunks refers to the number of tokens accepted by UNet in multiples of 77, 1 chunk = 77 tokens, 3 chunks = 231 tokens.
By default, compile_alt.py does not include model weights (constants) with the compiled module, to include the model weights in the compiled module use --include-consants True.

Alternative pipeline

The original pipeline requires a diffusers model local dir, and relies directly on StableDiffusionPipeline. This pipeline builds similar functionality without directly using StableDiffusionPipeline, and is capable of loading model weights from either diffusers or CompVis models to compiled aitemplate modules.

• AITemplate modules are created
• Model weights are loaded, converted/mapped, then applied to AITemplate module
• Tokenizer is created from openai/clip-vit-large-patch14.
• Scheduler is created from hf-hub-or-path.
• Loading CLIPTextModel from ckpt requires the appropriate hf-hub-or-path to be specified i.e. runwayml/stable-diffusion-v1-5 for SD1.x checkpoints, stabilityai/stable-diffusion-2-1 for SD2.x checkpoints.

python3 scripts/demo_alt.py \
--hf-hub-or-path runwayml/stable-diffusion-v1-5 \
--ckpt v1-5-pruned-emaonly.ckpt

python3 scripts/demo_alt.py \
--hf-hub-or-path stabilityai/stable-diffusion-2-1 \
--ckpt v2-1_768-ema-pruned.ckpt

--ckpt takes preference over --hf-hub-or-path if both are specified

Multi-GPU profiling

AIT needs to do profiling to select the best algorithms for CUTLASS and CK.
To enable multiple GPUs for profiling, use the environment variable CUDA_VISIBLE_DEVICES on NVIDIA platform and HIP_VISIBLE_DEVICES on AMD platform.

Benchmark Copy Section

This step is optional. You can run benchmark.py to measure throughput for each of the subnets.

Benchmark script supports base model variant only for now - 512x512

python3 src/benchmark.py

Verify Copy Section

This step is optional. You can verify numerical correctness for each of the subnets.

HUGGINGFACE_AUTH_TOKEN=ACCESS_TOKEN python3 -m unittest src/test_correctness.py

Run Models Copy Section

Run AIT models with an example image:

Set correct width and height depending on the model variant

python3 scripts/demo.py --width 512 --height 512

Img2img demo:

Set correct width and height depending on the model variant

python3 scripts/demo_img2img.py --width 512 --height 512

Check the resulted image: example_ait.png

Sample outputs Copy Section

Command: python3 scripts/demo.py --prompt "Mountain Rainier in van Gogh's world"

Command: python3 scripts/demo.py --prompt "Sitting in a tea house in Japan with Mount Fuji in the background, sunset professional portrait, Nikon 85mm f/1.4G"

Command: scripts/demo.py --prompt "A lot of wild flowers with North Cascade Mountain in background, sunset professional photo, Unreal Engine"

Results Copy Chapter

PT = PyTorch 1.12 Eager

OOM = Out of Memory

A100-40GB / CUDA 11.6, 50 steps Copy Section

• PT: 17.50 it/s
• AIT: 42.45 it/s

RTX 3080-10GB / CUDA 11.6, 50 steps Copy Section

• PT: OOM
• AIT: 24.51 it/s

MI-250 1 GCD, 50 steps Copy Section

• PT: 12.43 it/s
• AIT: 15.60 it/s

Batched Version Copy Chapter

A100-40GB / CUDA 11.6 Copy Section

• Stable Diffusion with AIT batch inference, 50 steps

• AIT Faster rendering, 25 steps

IMG2IMG Copy Chapter

A100-40GB / CUDA 11.6, 40 steps Copy Section

Note for Performance Results Copy Section

• For all benchmarks we render the images of size 512x512
• For img2img model we only support fix input 512x768 by default, stay tuned for dynamic shape support
• For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• To benchmark MI-250 1 GCD, we lock the frequency with command rocm-smi -d x --setperfdeterminism 1700, where x is the GPU id.
• Performance results are what we can reproduced & take reference. It should not be used for other purposes.

Overview (`examples/08_esrgan/README.md`) Copy Section

ESRGAN Example Copy Chapter

In this example, we show how to build fast AIT modules for ESRGAN models, and benchmark/run them.

Build Dependencies Copy Section

First, clone, build, and install AITemplate per the README instructions.

This AIT ESRGAN example depends on torch, click and optionally safetensors. You could install them using pip.

Download the ESRGAN model Copy Section

We have tested the following ESRGAN models.

RealESRGAN_x4plus

Model architecture:

num_in_ch: 3,
num_out_ch: 3,
num_feat: 64,
num_block: 23,
num_grow_ch: 32,
scale: 4,

RealESRGAN_x4plus_anime_6B

Model architecture:

num_in_ch: 3,
num_out_ch: 3,
num_feat: 64,
num_block: 6,
num_grow_ch: 32,
scale: 4,

RealESRGAN_x2plus

Model architecture:

num_in_ch: 3,
num_out_ch: 3,
num_feat: 64,
num_block: 23,
num_grow_ch: 32,
scale: 2,

A database of ESRGAN models can be found here.

Safetensors versions are supported.

Build AIT modules for ESRGAN Copy Section

Build the AIT modules by running compile.py.

Usage: compile.py [OPTIONS]

Options:
  --model-path TEXT               model path. supports torch or safetensors
  --width <INTEGER INTEGER>...    Minimum and maximum width
  --height <INTEGER INTEGER>...   Minimum and maximum height
  --batch-size <INTEGER INTEGER>...
                                  Minimum and maximum batch size
  --include-constants BOOLEAN     include constants (model weights) with
                                  compiled model
  --num-in-ch INTEGER             Number of in channels
  --num-out-ch INTEGER            Number of out channels
  --num-feat INTEGER              Number of intermediate features
  --num-block INTEGER             Number of RRDB layers
  --num-grow-ch INTEGER           Number of channels for each growth
  --scale INTEGER                 Scale
  --use-fp16-acc BOOLEAN          use fp16 accumulation
  --convert-conv-to-gemm BOOLEAN  convert 1x1 conv to gemm
  --work-dir TEXT                 Work directory
  --model-name TEXT               Model name
  --help                          Show this message and exit.

Use --num-in-ch, --num-out-ch, --num-feat, --num-block, --num-grow-ch, --scale options to set the ESRGAN model architecture. The default values are for RealESRGAN_x4plus architecture.

--width and --height require a minimum and maximum value, the compiled module supports the range of resolutions. However, with 2x model architecture, only static shape is supported, the maximum value for each dimension is used. Defaults are 64 and 1024.

--batch-size is supported for 4x model architecture, provide minimum and maximum values. Default is 1 1.

Use --include-constants False to compile the module without model weights.

AIT modules are compatible with all ESRGAN models with the same model architecture. This can simplify deployment by compiling a module without model weights then applying weights at runtime by using AIT mapped weights (see map_rrdb in ./modeling/rrdbnet.py) with the module's set_many_constants_with_tensors.

In our tests an ESRGAN module compiled with weights is approximately ~38MB, and ~6.5MB without.

Examples:

python compile.py --model-path "RealESRGAN_x4plus.safetensors"

python compile.py --model-path "RealESRGAN_x4plus_anime_6B.pth" --num-block 6 --model-name RealESRGAN_x4plus_anime_6B

python compile.py --model-path "RealESRGAN_x2plus.pth" --scale 2 --model-name RealESRGAN_x2plus --width 512 512 --height 512 512

Multi-GPU profiling

AIT needs to do profiling to select the best algorithms for CUTLASS and CK.
To enable multiple GPUs for profiling, use the environment variable CUDA_VISIBLE_DEVICES on NVIDIA platform and HIP_VISIBLE_DEVICES on AMD platform.

Run Models Copy Section

demo.py provides example usage of ESRGAN modules.

Usage: demo.py [OPTIONS]

Options:
  --module-path TEXT        the AIT module path
  --input-image-path TEXT   path to input image
  --output-image-path TEXT  path to output image
  --scale INTEGER           Scale of ESRGAN model
  --help                    Show this message and exit.

--scale must match the scale of the model architecture.

Limitations:

• Demo does not support multiple images/batch size.
• Demo does not support images with alpha channel.

Overview (`tests/ci_profile_cache/README.md`) Copy Section

Profiling Database for CI (Deprecated)

Profile Cache DB for CI is deprecated. Now CI will select the algorithm with the smallest tiling size and smallest alignments for CI.

The selection function is defined at: backend/target.py: Target:select_minimal_algo and specialized in each backend target implementation.

Overview (`fx2ait/fx2ait/example/01_transformer_model/README.md`) Copy Section

Transfomer encoder

In this example, we will demo how to use FX2AIT for inference on the transformer encoder block from pytorch.

Code structure Copy Chapter

test_transformer_encoder.py     # Transformer encoder block definition using torch API
../benchmark_utils.py           # Accuracy verification and Benchmark code for FX2AIT

How to Use Copy Chapter

FX2AIT allows users to directly define a torch model, while fx2ait converter does the conversion for the usage.
Therefore the encoder can be defined normally as

class EncoderBlock(torch.nn.Module):
            def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0):
                super().__init__()
                # Attention layer
                self.attn = torch.nn.MultiheadAttention(
                    embed_dim=input_dim,
                    num_heads=num_heads,
                    batch_first=True,
                )
                # # Two-layer MLP
                self.linear_net = torch.nn.Sequential(
                    torch.nn.Linear(input_dim, dim_feedforward),
                    torch.nn.Dropout(dropout),
                    torch.nn.ReLU(inplace=True),
                    torch.nn.Linear(dim_feedforward, input_dim),
                )
                # Layers to apply in between the main layers
                self.norm1 = torch.nn.LayerNorm(input_dim)
                self.norm2 = torch.nn.LayerNorm(input_dim)
                self.dropout = torch.nn.Dropout(dropout)

            def forward(self, x):
                # Attention part
                attn_out, _ = self.attn(query=x, key=x, value=x)
                # return attn_out
                x = x + self.dropout(attn_out)
                x = self.norm1(x)

                # MLP part
                linear_out = self.linear_net(x)
                x = x + self.dropout(linear_out)
                x = self.norm2(x)

                return x

To run the test and benchmark,

python fx2ait/fx2ait/example/01_transformer_model/test_transformer_encoder.py

Reference Speed vs PyTorch Eager Copy Chapter

A100-40GB / CUDA 11.6.2 Copy Section

PT = PyTorch 1.12 Eager

Note for Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow us to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• Performance results are what we can reproduce and for reference only. It should not be used for other purposes.

Overview (`fx2ait/fx2ait/example/02_vision_model/README.md`) Copy Section

ResNet-18

In this example, we will demo how to use FX2AIT for inference on the ResNet-18 model from torchvision.

Code structure Copy Chapter

test_vision_model.py            # ResNet definition using torch API
../benchmark_utils.py           # Accuracy verification and Benchmark code for FX2AIT

How to Use Copy Chapter

FX2AIT allows users to directly define a torch model, while fx2ait converter does the conversion for the usage.
Therefore the definition of model is as simple as

class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mod = torchvision.models.resnet18()

            def forward(self, x):
                return self.mod(x)

Notice that because AIT supports channel last, while pytorch supports channel first operation, FX2AIT automatically performs this layout conversion for you.

To run the test and benchmark,

python fx2ait/fx2ait/example/02_vision_model/test_vision_model.py

Reference Speed vs PyTorch Eager Copy Chapter

A100-40GB / CUDA 11.6.2 Copy Section

PT = PyTorch 1.12 Eager

Note for Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow us to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• Performance results are what we can reproduce and for reference only. It should not be used for other purposes.

Overview (`fx2ait/fx2ait/example/03_lowering_split/README.md`) Copy Section

AIT Lowerer

Now let's go back the example of encoder. Imagine the cases that

1. You want to test the effect of a particular op, say MultiheadAttention on the entire module.
2. There is some special op that AIT doesn't support.
   AIT actually provide an automatic Lowerer to split the graph into subgraphs and run interpreter,
   so that AIT only run the part it can handle and leave other to AITemplate.

In this example, we will demo how to use AitLowerer for inference on any models.

Code structure Copy Chapter

test_lowerr.py                  # Splited transformer encoder block to illustrate the usage of AitLowerer.
../benchmark_utils.py           # Accuracy verification and Benchmark code for FX2AIT
../lower/
        lower.py                # Lower interface, which integrates lowering passes of Split subgraph and AIT Interpreter
        ait_splitter.py         # Splitter to split graph into submodules
        ait_setting.py          # Lowering settings

How to Use Copy Chapter

To skip an operation can be extremely easy. One just need to register in the method function @torch.fx.wrap

@torch.fx.wrap
def unsupported_attention_op(f, x):
    attn_out, _ = f(x, x, x)
    return attn_out

Then at forward stage, call the function.

class LowerModule(torch.nn.Module):
            def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0):
                super().__init__()
                self.attn = torch.nn.MultiheadAttention(
                    embed_dim=input_dim,
                    num_heads=num_heads,
                    batch_first=True,
                )
                ...

            def forward(self, x):
                # Unsupported op will not be lowered to AIT backend.
                attn_out = unsupported_attention_op(self.attn, x)
                ...

Then AIT won't deal with that part.

lowerer = AitLowerer.create(
            LowerSettings(
                workdir="/tmp",
                name="test_ait_lower",
                min_acc_module_size=0,
            )
        )
        lowered = lowerer(model, inputs)
        lower_output = lowered(*inputs)

The mechanism is that Acc tracer allows user to register wrap function so that Acc won't deal with it.
Then our splitter will split the them into subgraph: _run_on_gpu_0 for pytorch Eager mode and _run_on_acc_1 for AIT,
where _run_on_gpu_0 contains torch.nn.MultiheadAttention and _run_on_acc_1 contains the rest of the model.
Finally, interpreter will be called for the AIT subgraph. (_run_on_acc_1)

Notice that our splitter only split subgraphs with more than 10 ops, since otherwise the subgraph is too small.

To run the test and benchmark,

python fx2ait/fx2ait/example/03_lowering_split/test_lower.py

Reference Speed vs PyTorch Eager Copy Chapter

A100-40GB / CUDA 11.6.2 Copy Section

PT = PyTorch 1.12 Eager

From the example, we learn without AIT's multihead attention module, the speedup will be degraded to 1.2x compared to Pytorch Eager.

Note for Performance Results Copy Section

• For NVIDIA A100, our test cluster doesn't allow us to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency.
• Performance results are what we can reproduce and for reference only. It should not be used for other purposes.

Source Debughints (`docs/source/debughints.rst`) Copy Section

Debug Hints

AITemplate is a new project under active development.
We have a rich test set to avoid bugs but don't be surprised if there is anything unexpected.

Here are some helpful tips we learned during the development of AITemplate:

1. Once the codegen for op which requires profiling is changed, remember to delete old profilers (usually located at workdir), and flush the cache by either deleting ~/.aitemplate or setting the environment variable FLUSH_PROFILE_CACHE=1.

2. Check the pseudo code/visualization generated by each optimization pass if some optimization behaves in unexpected way.

3. Always do the numerical test, from small to large, to make sure the entire model is correct.

4. Try to make the new fusion subgraph work in a manual way, then try to add an automatic pass to rewrite the graph with the fused subgraph.

Source Index (`docs/source/index.rst`) Copy Section

AITemplate Documentation

AITemplate (AIT) is a Python framework that transforms deep neural networks into CUDA (NVIDIA GPU) / HIP (AMD GPU) C++ code for lightning-fast inference serving. AITemplate highlights include:

• High performance: close to roofline fp16 TensorCore (NVIDIA GPU) / MatrixCore (AMD GPU) performance on major models, including ResNet, MaskRCNN, BERT, VisionTransformer, Stable Diffusion, etc.
• Unified, open, and flexible. Seamless fp16 deep neural network models for NVIDIA GPU or AMD GPU. Fully open source, Lego-style easily extendable high-performance primitives for new model support. Supports a significantly more comprehensive range of fusions than existing solutions for both GPU platforms.

.. toctree::
:maxdepth: 1
:caption: Getting Started

install/index

.. toctree::
:maxdepth: 1
:caption: User Guide

tutorial/index
debughints

.. toctree::
:maxdepth: 1
:caption: Runtime Design

runtime/index

.. toctree::
:maxdepth: 1
:caption: Architecture Guide

arch/index

.. toctree::
:maxdepth: 1
:caption: Reference Guide

reference/index
reference/env
genindex

Source Arch Index (`docs/source/arch/index.rst`) Copy Section

Design and Architecture

.. toctree::
:maxdepth: 1

philosophy

Stay tuned for more...

Source Arch Philosophy (`docs/source/arch/philosophy.rst`) Copy Section

Design Philosophy

KISS (Keep it simple and stupid) Copy Chapter

AITemplate avoids deep IR lowering stacks to reduce the system's complexity.
A highly modularized, multiple backend codegen system written in pure Python directly attacks the pain point in high-performance GPU inference.

Pragmatism Copy Chapter

AITemplate provides a PyTorch-style frontend to enable engineers to manually match the PyTorch model & weights to AITemplate for optimization.
Using it is less painful than debugging different lowering IR stacks, especially for complex models such as MaskRCNN.

We believe most of the neural network workload can be decoupled.
For example, most of the network can be decoupled into Encoder, Decoder, and Decoder logics.
For encoder and decoder, it is a computation-bounded problem.
For decoder logic, it may involve more control flows.
By using divide and conquer, we left the decoder logic part to C++ or Python rather than build a unified language / IR stack as a silver bullet.

Source Install Index (`docs/source/install/index.rst`) Copy Section

Installing AITemplate

Using Docker Copy Chapter

The easiest way to get started is to use Docker. Using docker is able to avoid performance regression caused by incorrect version of NVCC and HIPCC.
To use docker, we provide a bash script to build the docker image.

• CUDA:
  .. code-block:: bash

  text Copy

  ./docker/build.sh cuda

• ROCM:
  .. code-block:: bash

  text Copy

  DOCKER_BUILDKIT=1 ./docker/build.sh rocm

This will build a docker image with tag ait:latest.

To launch the docker container

• CUDA:
  .. code-block:: bash

  text Copy

  docker run --gpus all -it ait:latest

• ROCM:
  .. code-block:: bash

  text Copy

  docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined ait:latest

AITemplate will be installed as a Python package in Python 3.8. There will be also a copy of the source code and examples at /AITemplate.

Installing as a Standard Python Package Copy Chapter

Before installing AITemplate, first make sure you have correct hardware and software environment.

• Hardware
  • NVIDIA: AIT is only tested on SM80+ GPUs (Ampere etc).
  • AMD: AIT is only tested on CDNA2 (MI-210/250) GPUs.

.. warning::
- Not all kernels work with old SM75/SM70 (T4/V100) GPUs.
- There may be compiler issues for old CDNA1 (MI-100) GPUs.

• Software
  • NVIDIA: CUDA 11.6
  • AMD: ROCm 5.2, with HIPCC 10736 (commit b0f4678b9058a4ae00200dfb1de0da5f2ea84dcb)

.. warning::
- Incorrect compiler version may lead to performance regression.
- Instruction for building HIPCC 10736 can be founded in docker/Dockerfile.rocm.

When cloning the code, please use the following command to clone the submodules:

.. code-block:: bash

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

.. warning::
Please check that all submodules are cloned correctly before the next step.

Then build the Python wheel package and install it:

.. code-block:: bash

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

Source Reference Backend (`docs/source/reference/backend.rst`) Copy Section

aitemplate.backend

aitemplate.backend.task_runner Copy Chapter

.. automodule:: aitemplate.backend.task_runner
:members:
:imported-members:
:exclude-members: OrderedDict
:autosummary:

aitemplate.backend.builder Copy Chapter

.. automodule:: aitemplate.backend.builder
:members:
:imported-members:
:exclude-members: BaseRunner, Target
:autosummary:

aitemplate.backend.codegen Copy Chapter

.. automodule:: aitemplate.backend.codegen
:members:
:imported-members:
:exclude-members: Tensor, Target
:autosummary:

aitemplate.backend.profiler_cache Copy Chapter

.. automodule:: aitemplate.backend.profiler_cache
:members:
:imported-members:
:exclude-members:
:autosummary:

aitemplate.backend.profiler_runner Copy Chapter

.. automodule:: aitemplate.backend.profiler_runner
:members:
:imported-members:
:exclude-members: Target, Task, namedtuple, BaseRunner
:autosummary:

aitemplate.backend.registry Copy Chapter

.. automodule:: aitemplate.backend.registry
:members:
:imported-members:
:exclude-members:
:autosummary:

aitemplate.backend.target Copy Chapter

.. automodule:: aitemplate.backend.target
:members:
:imported-members:
:exclude-members:
:autosummary:

Source Reference Compiler (`docs/source/reference/compiler.rst`) Copy Section

aitemplate.compiler

base Copy Chapter

.. automodule:: aitemplate.compiler.base
:members:
:imported-members:
:exclude-members: ABC, Enum, abstructmethod, dataclass, pformat, reduce
:autosummary:

tensor_accessor Copy Chapter

.. automodule:: aitemplate.compiler.tensor_accessor
:members:
:imported-members:
:exclude-members: IntImm, IntVar, Tensor, pformat
:autosummary:

compiler Copy Chapter

.. automodule:: aitemplate.compiler.compiler
:members:
:imported-members:
:exclude-members: IntImm, IntVar, Tensor, pformat, DynamicProfileStrategy
:autosummary:

model Copy Chapter

.. automodule:: aitemplate.compiler.model
:members:
:imported-members:
:exclude-members: NamedTuple, TypeVar
:autosummary:

Source Reference Cuda (`docs/source/reference/cuda.rst`) Copy Section

aitemplate.backend.cuda

target_def Copy Chapter

.. automodule:: aitemplate.backend.cuda.target_def
:members:
:imported-members:
:exclude-members: Path, ProfileCacheDB, TargetType
:autosummary:

Source Reference Env (`docs/source/reference/env.rst`) Copy Section

Environment Variables

AITemplate uses environment variables to control the behavior of codegen and profiling.
The environment variables used in AITemplate are listed here.

Codegen Copy Chapter

NUM_BUILDERS: The number of CPU jobs running in parallel during codegen. It controls both the profiler codegen and the final .so codegen. It's set to 12 in NIGHTLY jobs. Internally, it's set to 12 for normal tests and 24 for heavy tests. By default, the builder uses all the available CPUs for building.

AIT_RECOMPILE: If set to "0", it skips compilation for the .so and reuses the previously compiled ones. It is used to speed up local testing. The default value is "1" to always recompile.

AIT_NDEBUG: If set to "1", compile with NDEBUG, disabling debug assertions. Recommended for production builds. "1" by default.

AIT_COMPILER_OPT: The optimization level for a compiler, which is directly passed to the host compiler command line. AITemplate host code may be very light in certain cases, so there is nothing to optimize for a host compiler. Thus, there is no need to make host compiler perform time costly optimizations. It may be very useful to use "-O0" value for debugging GPU kernels. "-O3" by default.

AIT_NVCC_CCBIN: nvcc host compiler (ccbin).

AIT_ENABLE_CUDA_LTO: If set to "1", nvcc will use LTO flags during compilation. Default value is "0".

AIT_TIME_COMPILATION: If set to "1", time each make command at the compilation time. This helps us to do compilation time analysis. Requires to install time <https://man7.org/linux/man-pages/man1/time.1.html>_ package.

AIT_MULTISTREAM_MODE: Controls multi-stream mode. Default mode is "0".

• If set to "0", then no multistreaming is used.
• If set to "1", then a simple multistreaming is used (iteratively track a wavefront of independent operators and execute ones).

AIT_MULTISTREAM_EXTRA_STREAMS: Specifies the number of additional streams used. Default value is "4".

AIT_MULTISTREAM_MAX_MEM_PARALLEL_OPS: Maximum number of parallel operators used in memory planning for simple multi-stream mode. Default value is "99999999" (basically, unlimited).

AIT_USE_CMAKE_COMPILATION: (An experimental feature) If set to "1", then cmake will used instead of make. This allows to build AITemplate using MSVC Compiler + MSBuild on Windows, and it works for linux as well. This builder does not support many features (such as caching) yet. But it allows to generate a cmake project that can be loaded to a modern IDE. Default value is "0".

AIT_ENABLE_STANDALONE: Enable standalone test and benchmark executable generation. Default value is "0" (disabled). If set to "1", this will generate a "test" executable that may be used to run standalone tests and benchmarks. This standalone executable is also well suited for running through debuggers and/or profiling tools, as it does not pull in python and pytorch as dependencies, unlike most python unit tests.

AIT_ENABLE_PTXAS_INFO: Set this to "1" to enable the generation and logging of verbose ( tuning-relevant ) information about CUDA ptx assembly code produced by the CUDA compiler nvcc. Intermediate ptx files, annotated with C++ source info will be written to the build directory. In addition, this flag enables warnings about CUDA register spilling and resource usage.

AIT_CUDA_DEBUG_LEVEL: Configure level of CUDA debug information. Defaults to no debug info. This may either be a string with options passed to nvcc ( for example "-g -G" or "-lineinfo" ) or a CUDA debug level from "0" (default, no debug info), "1" ( "-lineinfo" ) include source code line information. Ideal for profiling with ncu/nsight-compute, "2" full debug information (warning: this disables all optimizations, regardless of other settings)

AIT_ENABLE_CUDA_SOURCE_NAVIGATION_FIX: (Only supported by FBCUDA target so far): When this flag is enabled by setting it to "1" (it is disabled by default), every *.cu file in build dirs into a corresponding *.cu.h file and create a *.cu file which just includes this file. This fixes code navigation issues in some IDE's which don't treat .cu files as C++ files and disable code navigation.

AIT_ENABLE_INCLUDE_FROM_SOURCETREE: (Only supported by FBCUDA target so far) When this flag is enabled by setting it to "1" (it is disabled by default), the target will create an in-place build which tries to directly reference the include paths within the AITemplate source tree. This helps to iterate faster during native Kernel/Operator development and debugging.

Profiling Copy Chapter

CACHE_DIR: The directory for the profiling cache. If unset, it defaults to ~/.aitemplate.

FLUSH_PROFILE_CACHE: If set to "1", it removes the cache file and recreates an empty one.

DISABLE_PROFILER_CODEGEN: Normally in CI we randomly choose two profilers to codegen. If set to "1", this flag disables profiler codegen completely to speed up long running tests so that the tests don't time out. The default value is "0".

CUDA_VISIBLE_DEVICES: This one is from CUDA itself. It's used to set the number of GPU devices available for profiling. Set to "0,1,2,3,4,5,6,7" to speed up profiling. For benchmarking, it's useful to set to a particular device to lower noise.

HIP_VISIBLE_DEVICES: This one is from ROCm itself. It's used to set the number of GPU devices available for profiling. Set to "0,1,2,3,4,5,6,7" to speed up profiling. For benchmarking, it's useful to set to a particular device to lower noise.

FORCE_PROFILE: If set to "1", it will do profiling regardless in_ci_env and disable_profiler_codegen. For non-NIGHTLY CI, we do not do profiling, and we could use FORCE_PROFILE=1 in these CI to do runs with codegen, compile, and profile.

COMBINE_PROFILER_MULTI_SOURCES: Whether to combine multiple profiler sources per target. "0" - Disabled, "1" - Enabled (default).

FORCE_ONE_PROFILER_SOURCE_PER_TARGET: Whether to combine multiple profiler sources per target into one. "0" - Disabled (default), "1" - Enabled.

OSS CI Copy Chapter

CI_FLAG: It is set to "CIRCLECI" in OSS CI to indicate we're in OSS CI environment. The behavior of the profiler and codegen is different in CI to speed up testing. Profiling itself for gemm/conv ops is disabled in CI. But we still compile two random profilers to make sure the profiler codegen is not broken.

AIT_BUILD_DOCS: If set to "1", it will create a fake CUDA target to enable doc building in Github Actions.

Miscellaneous Copy Chapter

LOGLEVEL: It is used to control the logging level in Python. The default value is "INFO". "DEBUG" is useful for debugging.

AIT_PLOT_SHORTEN_TENSOR_NAMES: If set to "1", shorten too long tensor names for a plot of a model graph, thus making a plot much easier to analyze visually. "0" by default.

AIT_USE_FAST_MATH: If set to "0", no fast math option will be used for the device code generation. Default value is "1".

AIT_USE_TANH_FOR_SIGMOID: If set to "1", tanh will be used to approximate sigmoid during device code generation. Default value is "0".

Source Reference Frontend (`docs/source/reference/frontend.rst`) Copy Section

aitemplate.frontend

.. automodule:: aitemplate.frontend.nn
:members:
:imported-members:
:exclude-members:
:autosummary:

.. automodule:: aitemplate.frontend.tensor
:members:
:imported-members:
:exclude-members:
:autosummary:

Source Reference Index (`docs/source/reference/index.rst`) Copy Section

Python API

.. toctree::
:maxdepth: 2

compiler
ops
transform
backend
cuda
rocm
frontend
testing
utils

Source Reference Ops (`docs/source/reference/ops.rst`) Copy Section

aitemplate.compiler.ops

.. automodule:: aitemplate.compiler.ops
:members:
:imported-members:
:exclude-members: Tensor, TensorAccessor, Enum, Operator, IntImm, IntVar, IntVarTensor, wrap_dim
:autosummary:

Source Reference Rocm (`docs/source/reference/rocm.rst`) Copy Section

aitemplate.backend.rocm

target_def Copy Chapter

.. automodule:: aitemplate.backend.rocm.target_def
:members:
:imported-members:
:exclude-members:
:autosummary:

Source Reference Testing (`docs/source/reference/testing.rst`) Copy Section

aitemplate.testing

detect_target Copy Chapter

.. automodule:: aitemplate.testing.detect_target
:members:
:imported-members:
:exclude-members: CUDA, ROCM, Popen
:autosummary:

benchmark_pt Copy Chapter

.. automodule:: aitemplate.testing.benchmark_pt
:members:
:imported-members:
:exclude-members: CUDA, ROCM, Popen
:autosummary:

benchmark_ait Copy Chapter

.. automodule:: aitemplate.testing.benchmark_ait
:members:
:imported-members:
:exclude-members: CUDA, ROCM, Popen
:autosummary:

Source Reference Transform (`docs/source/reference/transform.rst`) Copy Section

aitemplate.compiler.transform

apply_padding Copy Chapter

.. automodule:: aitemplate.compiler.transform.apply_padding
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

bind_constants Copy Chapter

.. automodule:: aitemplate.compiler.transform.bind_constants
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

constant_folding Copy Chapter

.. automodule:: aitemplate.compiler.transform.constant_folding
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor
:autosummary:

fuse_conv_elementwise Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_conv_elementwise
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor
:autosummary:

fuse_group_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_group_ops
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, all_static_dimensions
:autosummary:

fuse_mm_elementwise Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_mm_elementwise
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph
:autosummary:

fuse_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_ops
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul
:autosummary:

fuse_parallel_gemms Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_parallel_gemms
:members:
:imported-members:
:exclude-members: TensorAccessor, is_static_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul
:autosummary:

fuse_permute_bmm Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_permute_bmm
:members:
:imported-members:
:exclude-members: copy_src_op_attributes, remove_tensor_from_sorted_graph, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, gemm_rrr, gemm_rrr_bias, permute021, TensorAccessor, is_static_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul
:autosummary:

fuse_split Copy Chapter

.. automodule:: aitemplate.compiler.transform.fuse_split
:members:
:imported-members: _fuse_split_and_strided_op
:exclude-members: IntVar, copy_src_op_attributes, remove_tensor_from_sorted_graph, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, gemm_rrr, gemm_rrr_bias, permute021, TensorAccessor, is_static_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul
:autosummary:

mark_param_tensor Copy Chapter

.. automodule:: aitemplate.compiler.transform.mark_param_tensor
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

memory_planning Copy Chapter

.. automodule:: aitemplate.compiler.transform.memory_planning
:members:
:imported-members:
:exclude-members: TensorUsageRecord, Workspace, assign_offsets_to_views_and_outputs, greedy_by_size_memory_planning, DimInfo, IntImm, Operator, Source, Tensor, gemm, defaultdict, dataclass
:autosummary:

name_graph Copy Chapter

.. automodule:: aitemplate.compiler.transform.name_graph
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

optimize_graph Copy Chapter

.. automodule:: aitemplate.compiler.transform.optimize_graph
:members:
:imported-members:
:exclude-members: transform_strided_ops, transform_special_ops, transform_odd_alignment, transform_memory_ops, fuse_permute_bmm, fuse_parallel_gemms, fuse_mm_elementwise, apply_padding, fuse_conv_elementwise, fuse_group_ops, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

profile Copy Chapter

.. automodule:: aitemplate.compiler.transform.profile
:members:
:imported-members:
:exclude-members: DynamicProfileStrategy, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

refine_graph Copy Chapter

.. automodule:: aitemplate.compiler.transform.refine_graph
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

remove_no_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.remove_no_ops
:members:
:imported-members:
:exclude-members: IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

remove_unused_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.remove_unused_ops
:members:
:imported-members:
:exclude-members: deque, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

toposort Copy Chapter

.. automodule:: aitemplate.compiler.transform.toposort
:members:
:imported-members:
:exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

transform_memory_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.transform_memory_ops
:members:
:imported-members:
:exclude-members: TensorAccessor, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

transform_odd_alignment Copy Chapter

.. automodule:: aitemplate.compiler.transform.transform_odd_alignment
:members:
:imported-members:
:exclude-members: can_be_constant_folded, copy_src_op_attributes, copy_tensor_attributes, extract_only_one_op, remove_tensor_from_sorted_graph, replace_tensor, sanitize_sorted_graph, toposort, IntVar, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, permute021, unsqueeze, TensorAccessor, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

transform_special_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.transform_special_ops
:members:
:imported-members:
:exclude-members: gemm_rrr, is_singleton_dimension, gemm_rcr, gemm_rrr_small_nk, can_be_constant_folded, copy_src_op_attributes, copy_tensor_attributes, extract_only_one_op, remove_tensor_from_sorted_graph, replace_tensor, sanitize_sorted_graph, toposort, IntVar, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, permute021, unsqueeze, TensorAccessor, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

transform_strided_op_and_view_op Copy Chapter

.. automodule:: aitemplate.compiler.transform.transform_strided_op_and_view_op
:members:
:imported-members:
:exclude-members: IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

transform_strided_ops Copy Chapter

.. automodule:: aitemplate.compiler.transform.transform_strided_ops
:members:
:imported-members:
:exclude-members: get_tensor_index, slice_reshape_scatter, slice_scatter, gen_tensor_index, IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

transform_strided_slice Copy Chapter

.. automodule:: aitemplate.compiler.transform.transform_strided_slice
:members:
:imported-members:
:exclude-members: dynamic_slice, get_tensor_index, slice_reshape_scatter, slice_scatter, gen_tensor_index, IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm
:autosummary:

Source Reference Utils (`docs/source/reference/utils.rst`) Copy Section

aitemplate.utils

visualization.plot Copy Chapter

.. automodule:: aitemplate.utils.visualization.plot
:members:
:imported-members:
:exclude-members: Tensor, Operator
:autosummary:

Source Runtime Cxx Design (`docs/source/runtime/cxx_design.rst`) Copy Section

================ C++ Runtime Note

`Model` v.s. `ModelContainer`

These are the two main classes involved in the C++ runtime implementation:

• The bulk of the runtime implementation is in the Model class.
• The ModelContainer class stores a set of shared constants and a collection of Model instances. Almost all functions in model_interface.h forward to a method in ModelContainer. When Run is invoked, ModelContainer looks for an available Model, or blocks until one becomes available (see the section on asynchronous predictions). It then forwards the run request to the runtime.

Code Structure

Some important files:

1. include/model_interface.h: The interface that we expose in the compiled .so.
2. include/model_container.h: The bulk of the ModelContainer implementation.

Some files are generated at compile time. These include:

• model-generated.h: The implementation of the Model.
• model_container_base.cu: A small part of the implementation for ModelContainer that needs to be generated. ModelContainer inherits from ModelContainerBase, and ModelContainerBase's implementation lives in this file. See model_container.h for more details.

All codegen templates can be found in backend/main_templates.py.
The codegen implementation is in backend/codegen.py.

Note that many of the headers in this directory rely on generated code and thus cannot be #include -d in external projects.
model_interface.h is an exception.

Source Runtime Index (`docs/source/runtime/index.rst`) Copy Section

Runtime Note

.. toctree::
:maxdepth: 1

cxx_design
py_design

Source Runtime Py Design (`docs/source/runtime/py_design.rst`) Copy Section

=================== Python Runtime Note

Python `Model`

Model is a collection of Python bindings to the C++ AIT runtime. This section describes the API.

`AITData` Copy Chapter

This class represents a contiguous blob of memory that AIT will use as a tensor. It is simply a named tuple with these fields:

• data_ptr: int: An unowned pointer to GPU memory. In general, all of the APIs expect that this pointer will be valid for the entire duration of the call.
• shape: List[int]: The shape of the tensor.
• dtype: str: The tensor's dtype; one of "float32", "float16", "int32", "int64". Note that most ops only support float16 at this stage.

When using AITemplate with PyTorch, AITData can be constructed with the torch_to_ait_data utility:

.. code-block:: python

x = torch.randn(3, 3, 3).half().cuda()
  # Equivalent to AITData(x.data_ptr(), [3, 3, 3], "float16")
  x_ait = torch_to_ait_data(x)

If PyTorch is not available, Model provides a set of functions for copying, allocating, and freeing GPU memory. See the docstrings in compiler/model.py for more information.

`run` Copy Chapter

run takes inputs and outputs as collections of AITData instances. Both arguments can be passed as either an ordered list or a dictionary (mapping name to tensor).

.. code-block:: python

# Arguments as a dictionary
module.run(
  {"input0": in0_ait, "input1": in1_ait},
  {"output0": out0_ait, "output1": out1_ait},
)

# Arguments as an ordered list. Note that you might need to query
# the index mapping.
input_name_to_idx = module.get_input_name_to_index_map()
output_name_to_idx = module.get_output_name_to_index_map()

inputs = [None] * len(input_name_to_idx)
outputs = [None] * len(output_name_to_idx)

for name in input_name_to_idx:
  inputs[input_name_to_idx[name]] = ait_inputs[name]

for name in output_name_to_idx:
  outputs[output_name_to_idx[name]] = ait_outputs[name]

module.run(inputs, outputs)

One important caveat is that the output must have the maximum possible size. This is because of dynamic shapes: the size of the output may vary, but its shape is not inferred until inference time. The maximum shape can be queried with the get_output_maximum_shape():

.. code-block:: python

# Can use either name or index.
name_to_idx = module.get_output_name_to_idx()
max_shape = module.get_output_maximum_shape(name_to_idx["output"])
max_shape = module.get_output_maximum_shape("output")

Model.run returns a dictionary of output AITData instances with (possibly dynamic) shapes that inferred in the runtime.

Nullptr Inputs/Outputs Copy Chapter

In general, inputs are allowed to be null if they are size 0 (e.g. at least one dimension is 0). The runtime enforces this with a check before any kernels are launched.

.. code-block:: cpp

If (input_name == nullptr && dim0 * dim1 * … * dimN != 0) {
  throw std::runtime_error(“input_name cannot be null!”);
}

This is convenient since torch.data_ptr() returns null for size zero tensors. The dynamic shape computation is skipped if the lower bound of the tensor’s size is positive.

Constants Copy Chapter

There are two types of constants in AIT; bound and unbound constants. A bound constant is known at compile time and may participate in constant folding. Bound constants are copied into GPU memory at model loading time. Values for bound constants may be provided by passing a dictionary (mapping constant name to AIT tensor) to compile_model.

Unbound constants, on the other hand, do not participate in constant folding and must be provided before running the model. These must be set via Model.set_constant:

.. code-block:: python

module.set_constant("my_constant", AITData(...))
# The pointer in the the tensor must live for the entire duration of run()
module.run(...)

Constants are read-only and shared with all runtimes in the ModelContainer.

`run_with_tensors` Copy Chapter

run_with_tensors is a convenience method with the same interface as run, except it can take lists (or dicts) of torch.Tensor instances:

.. code-block:: python

input0 = torch.randn(input0_shape).cuda().half()
output0 = torch.empty(output0_shape).cuda().half()
# Returns a dictionary of reshaped outputs
result = module.run_with_tensors([input0], [output0])

Streams and Asynchronous Predictions Copy Chapter

A pointer to a stream can optionally be passed to run.
If none is given, the prediction happens on the default stream 0.
If the sync argument is set to True, the stream is synchronized before run() returns.
sync is True by default.

Multiple predictions can happen at the same time (on the same or different streams).
Under the hood, there is a fixed-size pool of runtime objects.
When all the runtimes are used, run() blocks until one becomes available.
The size of this pool can be configured with the num_runtimes option in Model's constructor.

CUDA Graph Copy Chapter

Run also takes a graph_mode option. If set to true, the runtime will try to use CUDA graphs to run the model. graph_mode is not supported on ROCm.

The following is a high level overview of how graph mode works:

1. Each Model has an internal stream used for graph capturing. The model first runs all ops on this stream in capture mode. No kernel launches happen during this stage.
2. If this is the first run, a graph is instantiated via cudaGraphInstantiate.
3. On subsequent runs, we try to avoid the relatively expensive cudaGraphInstantiate call by updating the graph executor (cudaGraphExecUpdate). However, a new graph may still be instantiated if the topology of the graph somehow changed between runs.
4. Once we have the graph executor, we launch a single kernel on the stream that the user provided to run().

Graph mode is mainly beneficial when there are many small kernel launches. A lot of overhead can be avoided since there is only a single kernel launch in graph mode.

Source Tutorial How To Add Op (`docs/source/tutorial/how_to_add_op.rst`) Copy Section

How to add an operator to the AIT codegen

This tutorial will demonstrate how to add a new operator to the AIT codegen.
Full source code can be found at examples/07_how_to_run_pt_model/how_to_run_pt_model.py.

0. Prerequisites

We need to import necessary Python modules:

.. code-block:: python

from typing import Any, Dict, List

import jinja2
import torch

from aitemplate import backend
from aitemplate.backend import registry
from aitemplate.backend.backend_spec import CUDASpec, ROCMSpec
from aitemplate.compiler import compile_model
from aitemplate.compiler.base import IntVar, Operator, Tensor
from aitemplate.testing import detect_target

1. Define the operator graph node

Graph nodes are usually defined at aitemplate/compiler/ops.

.. code-block:: python

class add_one(Operator):
def init(self):
super().init()
# required, unique identity of operator category
self._attrs["op"] = "add_one"
# we can put whatever we want into the op attrs for later use
self._attrs["has_profiler"] = False
self._attrs["nop"] = False

def __call__(self, x: Tensor) -> Tensor:
  # each operator needs to keep a record of input tensors
  self._attrs["inputs"] = [x]
  # optional, to set depth of the op based on inputs' depth, used in DFS
  self._set_depth()
  # infer output shape
  output_shape = self._infer_shape(x)
  # create output Tensor, of which the source op is the current op
  output = Tensor(output_shape, src_ops={self})
  # remember current op's outputs
  self._attrs["outputs"] = [output]
  return output

def _infer_shape(self, x) -> List[IntVar]:
  # infer output shape
  # In case of we need infer shape in C++ side, we will create a jinja2 template
  # for shape inference function, and render to Python code in graph node
  # and render the template into C++ code in codegen
  return x.shape()

def gen_function(self) -> str:
  # this function will be used in codegen
  # here we only need to redirect to backend codegen function
  target = backend.target.Target.current()
  func_key = f"{target.name()}.{self._attrs['op']}.gen_function"
  func = registry.get(func_key)
  return func(self._attrs)

.. note::

• _attrs in Operator is the most important data structure for codegen.
• _attrs["op"] is the identity of operator category, which is used to find the corresponding codegen function in the backend; must be unique.

2. Define the necessary templates for Codegen

In AIT, there are 4 important templates for codegen:

• FUNC_TEMPLATE: the template for generating the function body of the operator, and invoke GPU kernel in the body.
• FUNC_SIGNATURE_TEMPLATE: the template for generating the function signature of the operator. The signature defines the name and arguments of the function.
• FUNC_CALL_TEMPLATE: the template for generating the function call of the operator. The call will be used during inference to invoke the GPU kernel with given arguments.
• FUNC_DECL: the template for forward declaration of the operator function. This is usually an alias of FUNC_SIGNATURE_TEMPLATE.

.. code-block:: python

FUNC_TEMPLATE = jinja2.Template(
"""
{{header_files}}
namespace {
{{kernel}}
} // namespace
{{func_signature}}
{
invoke_add_one(output, input, num_elements, stream);
}
"""
)

FUNC_SIGNATURE = jinja2.Template(
"""
void {{func_name}}(half* output,
const half* input,
const int64_t num_elements,
{{prefix}}Stream_t stream)
"""
)

FUNC_DECL = jinja2.Template(
"""
{{func_signature}};
"""
)

FUNC_CALL_TEMPLATE = jinja2.Template(
"""
{{indent}}int64_t num_elements = 1;
{% for dim_name in dim_names %}
{{indent}}num_elements = {{dim_name}};
{% endfor %}
{{indent}}{{func_name}}(
{{indent}} {{output}}, {{input}}, num_elements, stream / default stream */
{{indent}});
"""
)

3. Create the GPU kernels

In this example we use a simplest add one kernel. The kernel can be written by hand (as what programmer is expected to do), or generated by other tools.

.. code-block:: python

KERNEL_TEMPLATE = jinja2.Template(
"""
global void add_one(half* output, const half* input, const int64_t num_elements) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < num_elements) {
output[idx] = input[idx] + half(1.0);
}
}
void invoke_add_one(half* output, const half* input, int64_t num_elements, {{prefix}}Stream_t stream) {
if (num_elements < 1024) {
dim3 grid(1);
dim3 block(num_elements);
add_one<<<grid, block, 0, stream>>>(output, input, num_elements);
} else {
dim3 grid((num_elements + 1024 - 1) / 1024);
dim3 block(1024);
add_one<<<grid, block, 0, stream>>>(output, input, num_elements);
}
}
"""
)

(Optional) We also provide a helper function to handle CUDA/ROCm float16 data type difference.

.. code-block:: python

FUNC_CALL_FP16_PARAM_TEMPLATE = jinja2.Template(
"""reinterpret_cast<half*>(
{% if is_cuda %}&({% endif %}{{name}}{% if is_cuda %}->raw()){% endif %})"""
)

4. Define the codegen function

The codegen function is the function that renders the templates we defined into valid C++ code string.
The codegen function will take func_attrs from the graph node, and fill in the jinja2 template.

.. code-block:: python

def gen_function_call(func_attrs: Dict[str, Any], indent=" ", is_cuda=False) -> str:
assert len(func_attrs["outputs"]) == 1
assert len(func_attrs["inputs"]) == 1

output_name = FUNC_CALL_FP16_PARAM_TEMPLATE.render(
    name=func_attrs["outputs"][0]._attrs["name"], is_cuda=is_cuda
)
input_name = FUNC_CALL_FP16_PARAM_TEMPLATE.render(
    name=func_attrs["inputs"][0]._attrs["name"], is_cuda=is_cuda
)

dim_names = [dim._attrs["name"] for dim in func_attrs["inputs"][0].shape()]
return FUNC_CALL_TEMPLATE.render(
      func_name=func_attrs["name"],
      output=output_name,
      input=input_name,
      dim_names=dim_names,
      indent=indent,
)

def gen_function(func_attrs: Dict[str, Any], header_files: str, backend_spec) -> str:
prefix = backend_spec.prefix
return FUNC_TEMPLATE.render(
header_files=header_files,
kernel=KERNEL_TEMPLATE.render(prefix=prefix),
func_signature=FUNC_SIGNATURE.render(
func_name=func_attrs["name"], prefix=prefix
),
)

def gen_function_decl(func_attrs: Dict[str, Any], backend_spec) -> str:
return FUNC_DECL.render(
func_signature=FUNC_SIGNATURE.render(
func_name=func_attrs["name"],
prefix=backend_spec.prefix,
).strip()
)

5.1 Register the codegen function in CUDA backend Copy Chapter

CUDA backend functions are usually defined at aitemplate/backend/cuda/.

.. code-block:: python

CUDA_HEADER_FILES = """
#include <cuda_fp16.h>
"""

@registry.reg("cuda.add_one.gen_function")
def cuda_add_one_gen_function(func_attrs: Dict[str, Any]) -> str:
return gen_function(func_attrs, CUDA_HEADER_FILES, CUDASpec())

@registry.reg("cuda.add_one.func_decl")
def cuda_add_one_gen_function_decl(func_attrs: Dict[str, Any]) -> str:
return gen_function_decl(func_attrs, CUDASpec())

@registry.reg("cuda.add_one.func_call")
def cuda_add_one_gen_function_call(func_attrs: Dict[str, Any], indent=" ") -> str:
return gen_function_call(func_attrs, indent, is_cuda=True)

5.2 (Optional) Register the codegen function to ROCm backend Copy Chapter

ROCm backend functions are usually defined at aitemplate/backend/rocm/.

.. code-block:: python

HIP_HEADER_FILES = """
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
"""

@registry.reg("rocm.add_one.gen_function")
def rocm_add_one_gen_function(func_attrs: Dict[str, Any]) -> str:
return gen_function(func_attrs, HIP_HEADER_FILES, ROCMSpec())

@registry.reg("rocm.add_one.func_decl")
def rocm_add_one_gen_function_decl(func_attrs: Dict[str, Any]) -> str:
return gen_function_decl(func_attrs, ROCMSpec())

@registry.reg("rocm.add_one.func_call")
def rocm_add_one_gen_function_call(func_attrs: Dict[str, Any], indent=" ") -> str:
return gen_function_call(func_attrs, indent, is_cuda=False)

6. Compile and verify the results with PyTorch

.. code-block:: python

def create_ait_model(shapes):
X = Tensor(
shape=shapes,
dtype="float16",
name="X",
is_input=True,
)
Y = add_one()(X)
Y._attrs["is_output"] = True
Y._attrs["name"] = "Y"
return Y

def verify_add_one():
shapes = [16, 512]
x = torch.randn(shapes).cuda().half()
y_pt = x + 1.0

Y = create_ait_model([16, 512])
target = detect_target()
with compile_model(Y, target, "./tmp", "add_one") as module:
  y = torch.empty(shapes).cuda().half()
  inputs = {"X": x}
  outputs = {"Y": y}
  module.run_with_tensors(inputs, outputs)
  print(torch.allclose(y, y_pt, atol=1e-2, rtol=1e-2))

Source Tutorial How To Infer Pt (`docs/source/tutorial/how_to_infer_pt.rst`) Copy Section

How to inference a PyTorch model with AIT

This tutorial will demonstrate how to inference a PyTorch model with AIT.
Full source code can be found at examples/07_how_to_run_pt_model/how_to_run_pt_model.py.

0. Prerequisites

We need to import necessary Python modules:

.. code-block:: python

import torch

from aitemplate.compiler import compile_model
from aitemplate.frontend import nn, Tensor
from aitemplate.testing import detect_target
from aitemplate.testing.benchmark_pt import benchmark_torch_function
from aitemplate.utils.graph_utils import sorted_graph_pseudo_code

1. Define a PyTorch module

Here we define a PyTorch model which is commonly seen in Transformers:

.. code-block:: python

class PTSimpleModel(torch.nn.Module):
def init(self, hidden, eps: float = 1e-5):
super().init()
self.dense1 = torch.nn.Linear(hidden, 4 * hidden)
self.act1 = torch.nn.functional.gelu
self.dense2 = torch.nn.Linear(4 * hidden, hidden)
self.layernorm = torch.nn.LayerNorm(hidden, eps=eps)

def forward(self, input):
  hidden_states = self.dense1(input)
  hidden_states = self.act1(hidden_states)
  hidden_states = self.dense2(hidden_states)
  hidden_states = hidden_states + input
  hidden_states = self.layernorm(hidden_states)
  return hidden_states

2. Define an AIT module

We can define a similar AIT module as follows:

.. code-block:: python

class AITSimpleModel(nn.Module):
def init(self, hidden, eps: float = 1e-5):
super().init()
self.dense1 = nn.Linear(hidden, 4 * hidden, specialization="fast_gelu")
self.dense2 = nn.Linear(4 * hidden, hidden)
self.layernorm = nn.LayerNorm(hidden, eps=eps)

def forward(self, input):
  hidden_states = self.dense1(input)
  hidden_states = self.dense2(hidden_states)
  hidden_states = hidden_states + input
  hidden_states = self.layernorm(hidden_states)
  return hidden_states

.. warning::
The nn.Module API in AIT looks similar to PyTorch, but it is not the same.

The fundamental difference is that AIT module is a container to build a graph, while PyTorch module is a container to store parameters for eager.
Which means, each AIT module's forward method can be only called once, and the graph is built during the first call.
If you want to share parameters, you need to use the compiler.ops instead. The compiler.ops is similar to functional in PyTorch.

AITemplate supports automatic fusion of linear followed by other operators. However in many cases, especially for quick iterations, we use manual specialization to specify the fused operator. For example, specialization="fast_gelu" will fuse linear with the fast_gelu operator.

3. Define a helper function to map PyTorch parameters to AIT parameters

In AIT, all names must follow the C variable naming standard, because the names will be used in the codegen process.

.. code-block:: python

def map_pt_params(ait_model, pt_model):
ait_model.name_parameter_tensor()
pt_params = dict(pt_model.named_parameters())
mapped_pt_params = {}
for name, _ in ait_model.named_parameters():
ait_name = name.replace(".", "_")
assert name in pt_params
mapped_pt_params[ait_name] = pt_params[name]
return mapped_pt_params

.. warning::

• Different to PyTorch, it is required to call ait_model .name_parameter_tensor() method to provide each parameter with a name with a direct map to PyTorch.
• Because all names in AIT must follow the C variable naming standard, you can easily replace . by _ or use a regular expression to make sure the name in valid.
• For networks with conv + bn subgraph, we currently don't provide an automatic pass to fold it. Please refer to our ResNet and Detectron2 examples to see how we handle CNN layout transform and BatchNorm folding.

4. Create PyTorch module, inputs/outputs

.. code-block:: python

batch_size=1024
hidden=512

create pt model

pt_model = PTSimpleModel(hidden).cuda().half()

create pt input

x = torch.randn([batch_size, hidden]).cuda().half()

run pt model

pt_model.eval()
y_pt = pt_model(x)

5. Create AIT module, inputs/outputs

.. code-block:: python

batch_size=1024
hidden=512

create AIT model

ait_model = AITSimpleModel(hidden)

create AIT input Tensor

X = Tensor(
shape=[batch_size, hidden],
name="X",
dtype="float16",
is_input=True,
)

run AIT module to generate output tensor

Y = ait_model(X)

mark the output tensor

Y._attrs["is_output"] = True
Y._attrs["name"] = "Y"

.. warning::

• Similar to MetaTensor, LazyTensor and a lot of other lazy evaluation frameworks, AIT's Tensor records the computation graph, and the graph is built when the Tensor is compiled.
• For input tensor, it is required to set the attribute is_input=True.
• For output tensor, it is required to set the attribute Y._attrs["is_output"] = True.
• For input and output tensors, it is better to provide the name attributes to use in runtime.

6. Compile AIT module into runtime and do verification

.. code-block:: python

map pt weights to ait

weights = map_pt_params(ait_model, pt_model)

codegen

target = detect_target()
with compile_model(
Y, target, "./tmp", "simple_model_demo", constants=weights
) as module:
# create storage for output tensor
y = torch.empty([batch_size, hidden]).cuda().half()

# inputs and outputs dict
inputs = {"X": x}
outputs = {"Y": y}

# run
module.run_with_tensors(inputs, outputs, graph_mode=True)

# verify output is correct
print(torch.allclose(y, y_pt, atol=1e-2, rtol=1e-2))

# benchmark ait and pt
count = 1000
ait_t, _, _ = module.benchmark_with_tensors(
    inputs, outputs, graph_mode=True, count=count
)
print(f"AITemplate time: {ait_t} ms/iter")

pt_t = benchmark_torch_function(count, pt_model.forward, x)
print(f"PyTorch eager time: {pt_t} ms/iter")

In this example, AIT will automatically fuse GELU and elementwise addition into the TensorCore/MatrixCore gemm operation. On RTX-3080, in the example AIT is about 1.15X faster than PyTorch Eager.

.. note::

• In this example, we fold the parameters (weights) into AIT runtime. The final dynamic library will contain them as parameters.
• If during the compile time we don't provide the parameters (for example, because the total parameters size is greater than 2GB), we can always call set_constant function in the runtime. Please check the runtime API for the details.

Source Tutorial How To Visualize (`docs/source/tutorial/how_to_visualize.rst`) Copy Section

How to visualize an AIT model

Visualization is important for understanding the behavior of a model optimization.
In AIT, we modify the codegen a little bit, from generating CUDA/HIP C++ code to HTML/Javascript code,
then we can generate a visualization of the model.

The following code will generate a visualization of our first example.

1. Define the AIT Model

.. code-block:: python

from aitemplate import compiler
from aitemplate.frontend import nn, Tensor
from aitemplate.testing import detect_target
from aitemplate.utils.visualization import plot_graph

class AITSimpleModel(nn.Module):
def init(self, hidden, eps: float = 1e-5):
super().init()
self.dense1 = nn.Linear(hidden, 4 * hidden, specialization="fast_gelu")
self.dense2 = nn.Linear(4 * hidden, hidden)
self.layernorm = nn.LayerNorm(hidden, eps=eps)

def forward(self, input):
  hidden_states = self.dense1(input)
  hidden_states = self.dense2(hidden_states)
  hidden_states = hidden_states + input
  hidden_states = self.layernorm(hidden_states)
  return hidden_states

def gen_ait_model():
batch_size = 512
hidden = 1024
ait_model = AITSimpleModel(hidden)
ait_model.name_parameter_tensor()
X = Tensor(
shape=[batch_size, hidden],
name="X",
dtype="float16",
is_input=True,
)
Y = ait_model(X)
Y._attrs["is_output"] = True
Y._attrs["name"] = "Y"
return Y

output_tensor = gen_ait_model()

2. Apply optimizations on the AIT Model

.. code-block:: python

def apply_optimizations(tensors):
target = detect_target()
# first, convert output tensors to graph
with target:
graph = compiler.transform.toposort(tensors)
# second, provide names to the graph
compiler.transform.name_graph(graph)
compiler.transform.mark_param_tensor(graph)
compiler.transform.mark_special_views(graph)
# we can apply optimizations to the graph, or test single optimization pass on the graph
graph = compiler.transform.optimize_graph(graph, "./tmp")
return graph

graph = apply_optimizations(output_tensor)

3. Generate visualization

.. code-block:: python

Plot the graph

plot_graph(graph, file_path="ait_model.html")

The visualization will be generated in the "ait_model.html" file. This file can be opened in Chrome without any web server.

.. raw:: html

Source Tutorial Index (`docs/source/tutorial/index.rst`) Copy Section

Tutorials

.. toctree::
:maxdepth: 1

how_to_infer_pt
how_to_add_op
how_to_visualize

--- METRICS ---

• Files Extracted: 38
• Estimated Token Budget: ~33743 tokens
• Canonical Reference: https://codewiki.google/github.com/facebookincubator/AITemplate