flux.triton

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

• Linear Kernel
  • python ./benchmarks/linear.py
• Attention Kernel
  • python ./benchmarks/attention.py
• LayerNorm + Modulation Kernel
  • python ./benchmarks/layer_norm.py
• RMSNorm Kernel
  • python ./benchmarks/rms_norm.py
• RoPE Kernel
  • python ./benchmarks/rope.py