/kraken

Triton-based Symmetric Memory operators and examples

Primary LanguagePythonOtherNOASSERTION

# Kraken

A Triton library of Symmetric Memory operators and examples.

This repository aims to simplify the process of developing distributed AI models using Triton and PyTorch's symmetric memory capabilities. Our initial kernels are adapted from the Symmetric Memory Recipes by Yifu Wang.

Examples

TBD

Requirements

Kraken requires:

  • Triton >= 3.3.0
  • PyTorch >= 2.6.0
  • Python >= 3.10

Installation

git clone https://github.com/meta-pytorch/kraken
cd kraken
pip install -e . -r requirements.txt

License

Source code is made available under a BSD 3 license, however you may have other legal obligations that govern your use of other content linked in this repository.