NOTE: This is a fork of https://github.com/tgale96/grouped_gemm.
Grouped GEMM
A lighweight library exposing grouped GEMM kernels in PyTorch.
Installation
Run pip install grouped-gemm-db to install the package.
Compiling from source
By default, the installed package runs in conservative (cuBLAS) mode:
it launches one GEMM kernel per batch element instead of using a single
grouped GEMM kernel for the whole batch.
To enable using grouped GEMM kernels, you need to switch to the CUTLASS
mode by setting the GROUPED_GEMM_CUTLASS environment variable to 1
when building the library. For example, to build the library in CUTLASS
mode for Ampere (SM 8.0), clone the repository and run the following:
$ TORCH_CUDA_ARCH_LIST=8.0 GROUPED_GEMM_CUTLASS=1 pip install .
See this comment for some performance measurements on A100 and H100.
Upcoming features
- Hopper-optimized grouped GEMM kernels.
1 maintainer
Release files for grouped-gemm-db 0.3.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| grouped_gemm_db-0.3.0.tar.gz | 981.0 kB | Details |
Release files / grouped_gemm_db-0.3.0.tar.gz
| Download URL | grouped_gemm_db-0.3.0.tar.gz |
|---|---|
| Size | 981.0 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
f7a7731d0d4056599b7df868edd8bc1cf436c2276ed56317c645dc67a690f8ec
|
|
BLAKE2b-256 checksum How to use checksums |
46533888df78364044b0ceae5a8dc2a94ed066650222d9fd322f045a787721cd
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.12.10
|
