Skip to main content

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.

Metadata

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)

Source distribution for grouped-gemm-db 0.3.0
File Size Uploaded Provenance Actions
grouped_gemm_db-0.3.0.tar.gz 981.0 kB Without provenance Details for grouped_gemm_db-0.3.0.tar.gz

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

Security for grouped-gemm-db 0.3.0

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

No provenance found

None of the files in this release were uploaded with provenance metadata.

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.

To learn more about protecting yourself from malicious actors, visit our package security documentation.

Release history Release notifications | RSS feed

This release

0.3.0 This release

1 release file No provenance

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page