Metadata-Version: 2.4
Name: colabfold-kernels
Version: 0.6.0
Summary: Fast fused kernels for ColabFold
Author-email: Milot Mirdita <milot@mirdita.de>
License: MIT
Project-URL: Homepage, https://github.com/mirditalab/colabfold-kernels
Project-URL: Issues, https://github.com/mirditalab/colabfold-kernels/issues
Keywords: alphafold,colabfold,pallas,triton,jax,attention,bioinformatics
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: colabfold-legacy-kernels<0.6,>=0.5.0; sys_platform == "linux" and platform_machine == "x86_64"
Requires-Dist: colabfold-cuda-kernels<0.3,>=0.2.0; sys_platform == "linux"
Dynamic: license-file

# colabfold-kernels

Fast fused kernels for ColabFold, written in
[Pallas](https://docs.jax.dev/en/latest/pallas/index.html) with native CUDA-fallback
for Volta- and Turing-generation GPUs.

```python
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](https://github.com/mirditalab/colabfold-cuda-kernels) megakernel.
