colabfold-kernels
Fast fused kernels for ColabFold, written in Pallas with native CUDA-fallback for Volta- and Turing-generation GPUs.
import jax.numpy as jnp
from colabfold_kernels import attention, dispatch
kernel = attention(dispatch(), jnp.bfloat16, key_dim=96, value_dim=96, layout="seq")
if kernel is not None:
out = kernel(q, k, v, mask_bias, nonbatched_bias, scale)
attention, layer_norm and gated_dual_proj each return None if not supported, as do the other
fused_ops selectors. The FP8 paths are opt-in (global_config["fp8"]) and need compute capability 8.9 or
newer.
diffusion.install() samples AF3-family diffusion in the
colabfold-cuda-kernels megakernel.
1 maintainer
Metadata
1 maintainer
Release files for colabfold-kernels 0.6.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 | |
|---|---|---|---|
| colabfold_kernels-0.6.0.tar.gz | 50.7 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| colabfold_kernels-0.6.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 97.5 kB
Release files / colabfold_kernels-0.6.0.tar.gz
| Download URL | colabfold_kernels-0.6.0.tar.gz |
|---|---|
| Size | 50.7 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
3dd06b2de8d3b3339e6f75c7694ce778f05ccecfa9e3bd433ebcf0159074e1f4
|
|
BLAKE2b-256 checksum How to use checksums |
be10f85875f46700f7af2c427ae45c0e9af02a9e3520914ffac2ad7c201df9b1
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
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.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 16, 2026.
Transparency logRelease files / colabfold_kernels-0.6.0-py3-none-any.whl
| Download URL | colabfold_kernels-0.6.0-py3-none-any.whl |
|---|---|
| Size | 46.8 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
bb8d4ed1285ae1abf72fd88f7ed8a927afec91c0aa4ce1fdae7d9516af0aafdd
|
|
BLAKE2b-256 checksum How to use checksums |
cfec139e20bf5fc93096b61da090abde6897fbca6bca96bd92fe4afa1ecb94d4
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
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.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 16, 2026.
Transparency log
