Triton Kernels

Triton kernels for Flux

Installation

pip install -e .

# for tests
pip install -e .[testing]

# for benchmarks
pip install matplotlib pandas

Tests

python -m pytest

Benchmarks

• LayerNorm + Modulation Kernel
  • python ./benchmarks/layer_norm.py
• RMSNorm Kernel
  • python ./benchmarks/rms_norm.py
• RoPE Kernel
  • python ./benchmarks/rope_test.py