Triton Kernels

Triton kernels for Stable Diffusion 3 and Flux

Installation

pip install -e .

# for tests
pip install -e .[testing]

# for benchmarks
pip install matplotlib pandas

Benchmarks

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

Tests

python -m pytest