Overview

Triton is a programming language designed to simplify the development of highly efficient custom operations for Deep Neural Networks.

You can think of Triton as an alternative to CUDA-C in which compute kernels are single-threaded and automatically parallelized. This is done through a modification of the standard C typesystem which allows for multi-dimensional arrays to be first-class citizen. For example, the Triton-C code for a vector addition would be:

__global__ void add(TYPE* c, TYPE* a, TYPE* b){
   int pid = get_program_id(0);
   float* pa[128] = a + pid*128 + 0 ... 128;
   float* pb[128] = b + pid*128 + 0 ... 128;
   float* pc[128] = c + pid*128 + 0 ... 128;
   *pc = *pa + *pb;
}

The Triton compiler will then look at this code and figure out how to distribute the blocks of pointers p{a,b,c} across different GPU threads so as to maximize memory coalescing, data-reuse and compute efficiency.

The benefits of this approach are most pronounced when writing compute-bound operations. In fact, it is possible to write matrix-multiplication code for tensor cores on par with cuBLAS — without knowing anything about GPUs!

Triton also provides a Python API that automatically generates, compiles and cache PyTorch custom op C code, so that you never have to write any boilerplate code yourself. Check out the tutorials if you want to learn more.