Skip to content

Ahead of Time (AOT) Triton Math Library

License

Notifications You must be signed in to change notification settings

cms-externals/aotriton

 
 

Repository files navigation

Build Instructions

mkdir build
cd build
cmake .. -DCMAKE_INSTALL_PREFIX=./install_dir -DCMAKE_BUILD_TYPE=Release -G Ninja
# Use ccmake to tweak options
ninja install

The library and the header file can be found under build/install_dir afterwards.

Note: do not run make separately, due to the limit of the current build system, make install will run the whole build process unconditionally.

Prerequisites

  • hipcc in /opt/rocm/bin, as a part of ROCm
  • cmake
  • ninja

Generation

The kernel definition for generation is done in rules.py. Edits to this file are needed for each new kernel, but it is extensible and generic.

Include files can be added in this directory.

The final build output is an archive object file any new project may link against.

The archive file and header files are installed in the path specified by CMAKE_INSTALL_PREFIX.

Kernel Support

Currently the first kernel supported is FlashAttention as based on the algorithm from Tri Dao.

PyTorch Consumption

PyTorch recently expanded AOTriton support for FlashAttention. AOTriton is consumed in PyTorch through the SDPA kernels. The Triton kernels and bundled archive are built at PyTorch build time.

About

Ahead of Time (AOT) Triton Math Library

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 87.0%
  • C++ 10.3%
  • CMake 2.7%