BobMcDear/attorch

A subset of PyTorch's neural network modules, written in Python using OpenAI's Triton.

What it solves

attorch provides a readable, self-contained, and hackable alternative to PyTorch's nn module. It allows developers to create custom deep learning operations with the efficiency of Triton-based kernels without needing the technical expertise required to write raw CUDA kernels.

How it works

The project is written purely in Python using OpenAI's Triton. It implements a wide variety of neural network layers (such as convolutions, attention, and various activation functions) and mathematical functions. It supports both forward and backward passes, making it suitable for both training and inference. To simplify integration, it includes a fallback mechanism (attorch.nn) that automatically uses PyTorch layers when a specific attorch module is unavailable or when performance-critical layers like convolutions and pooling are better handled by PyTorch.

Who it’s for

It is designed for researchers and developers who want to develop custom deep learning operations and need a performance boost over pure PyTorch implementations but cannot write CUDA kernels.

Highlights

  • Triton-powered: Uses OpenAI's Triton for efficient GPU computations.
  • Cuda-free: Enables high-performance kernels without writing CUDA code.
  • Broad Layer Support: Includes a wide range of layers including MultiheadAttention, BatchNorm, and various activation functions with optional kernel fusion.
  • Training and Inference: Fully supports both forward and backward passes.
  • PyTorch Integration: Offers a fallback interface to ensure compatibility with existing PyTorch workflows.

Related

  • Project
    triton-inference-server/backend

    A framework and set of utilities for creating backends that allow Triton Inference Server to execute models from various deep-learning frameworks or custom C/C++ logic.

  • Project
    pytorch/pytorch

    A Python-based machine learning framework that provides GPU-accelerated tensor computation and a dynamic autograd system for building deep neural networks.

  • Project
    triton-lang/triton

    Triton is a language and compiler for writing highly efficient custom deep-learning primitives, offering higher productivity than CUDA and more flexibility than other DSLs.

  • Dispatch
    OpenAI Triton 1.0 Release

    OpenAI has released Triton 1.0, an open-source Python-like language and compiler that allows researchers to write highly efficient GPU kernels without deep CUDA expertise.

  • Project
    Ascend/pytorch

    A PyTorch adapter plugin that enables the PyTorch framework to run on Ascend AI processors, providing hardware acceleration for AI model training and inference.