FlexGEMM Wheels

Pre-built wheels for FlexGEMM - Efficient sparse convolution based on Triton.

Available Configurations

• CUDA 12.4 + PyTorch 2.5.1 (Python 3.10-3.12)
• CUDA 12.6 + PyTorch 2.6.0 (Python 3.10-3.13)
• CUDA 12.6 + PyTorch 2.8.0 (Python 3.10-3.13)
• CUDA 12.8 + PyTorch 2.8.0 (Python 3.10-3.13)
• CUDA 12.8 + PyTorch 2.9.1 (Python 3.10-3.13)

Installation

Choose the configuration matching your CUDA and PyTorch versions:

pip install flexgemm --find-links https://pozzettiandrea.github.io/flexgemm-wheels/cu128-torch291/

Supported GPUs

• CUDA 12.4/12.6: RTX 20/30/40 series, A100
• CUDA 12.8: RTX 20/30/40/50 series, A100, B100/B200