Dao-AILab/quack
GitHub: Dao-AILab/quack
QuACK 是一组使用 CuTe-DSL 编写的高性能 CUDA 核函数库，为深度学习提供经高度优化的归一化、激活和矩阵乘法等算子。
Stars: 1056 | Forks: 142
# 🦆 QuACK：一组奇特可爱的 CuTe Kernel 🦆
Kernel 使用 [CuTe-DSL](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.html) 编写。
## 安装
```
# 对于 CUDA 12.9：
pip install quack-kernels
# 对于 CUDA 13.x：
pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
# 暂时不要将 uv 用于 CUDA 13.x 安装：它可能会竞争/安装
# nvidia-cutlass-dsl[cu13]，导致顺序错误 (NVIDIA/cutlass#3259)：
# https://github.com/NVIDIA/cutlass/issues/3259
# 可选：安装 NVIDIA matmul heuristics 以获得更好的未调优 GEMM 配置
pip install 'quack-kernels[heuristics]'
# 可选：JAX 绑定（会拉取 jax 和 jax-tvm-ffi）
pip install 'quack-kernels[jax]'
```
## 环境要求
- H100、B200/B300 或 RTX 50 GPU
- CUDA toolkit 12.9+
- Python 3.12
## Kernel 🐥
- 🦆 RMSNorm 前向 + 反向
- 🦆 Softmax 前向 + 反向
- 🦆 交叉熵前向 + 反向
- 🦆 Layernorm 前向 + 反向
- 🦆 Hopper gemm + epilogue
- 🦆 Blackwell gemm + epilogue
- 🦆 Blackwell GeForce gemm + epilogue
## 用法
```
from quack import rmsnorm, softmax, cross_entropy
```
部分 Kernel 也提供了 JAX 绑定（参见 [docs/jax.md](docs/jax.md)）：
```
from quack.softmax_jax import softmax
```
## 文档
- [JAX 接口](docs/jax.md) — 可选的 `jax` + `jax-tvm-ffi` 绑定，示例请参见 `quack/softmax_jax.py`。
[2025-07-10] 我们发布了一篇详尽的
[博客文章](media/2025-07-10-membound-sol.md)，介绍了如何借助 [CuTe-DSL](https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.html)，
在 Python 的舒适环境中将受限于内存带宽的 kernel 提速到极限。
## 性能
详情请参阅我们的[博客文章](media/2025-07-10-membound-sol.md)。
## 开发
要设置开发环境：
```
pip install -e '.[dev]'
pre-commit install
# 对于 CUDA 13.x：
pip install 'quack-kernels[dev,cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
# 暂时不要将 uv 用于 CUDA 13.x 安装；请改用 pip。
# 参见 https://github.com/NVIDIA/cutlass/issues/3259
```
标签：CUDA, NVIDIA CuTe, Vectored Exception Handling, 凭据扫描, 深度学习, 算子库, 逆向工具, 高性能计算