NVIDIA/TileGym
GitHub: NVIDIA/TileGym
TileGym 是 NVIDIA 推出的基于 Tile 的 GPU 编程 kernel 教程与示例库,帮助开发者学习高效 CUDA kernel 编写并探索其在 LLM 推理中的集成。
Stars: 770 | Forks: 79
English | [简体中文](README_chs.md) | [繁體中文](README_cht.md) | [日本語](README_ja.md) | [Français](README_fr.md)
# TileGym
TileGym 是一个 CUDA Tile kernel 库,为基于 Tile 的 GPU 编程提供了丰富的 kernel 教程和示例。
[**概述**](#overview) |
[**功能**](#features) |
[**安装说明**](#installation) |
[**快速入门**](#quick-start) |
[**贡献指南**](#contributing) |
[**许可证**](#license-and-third-party-notices)
## 概述
本仓库旨在为基于 Tile 的 GPU 编程提供实用的 kernel 教程和示例。TileGym 是一个用于试验 CUDA Tile 的平台,您可以在这里学习如何构建高效的 GPU kernel,并探索它们在实际的大型语言模型(如 Llama 3.1 和 DeepSeek V2)中的集成。无论您是正在学习基于 Tile 的 GPU 编程,还是希望优化您的 LLM 实现,TileGym 都能提供实用的示例和全面的指导。
## 功能
- 丰富的 CUDA Tile kernel 示例集合
- 针对常见深度学习算子的实用 kernel 实现
- 用于评估 kernel 效率的性能基准测试
- 与主流 LLM(Llama 3.1、DeepSeek V2)的端到端集成示例
## 安装说明
### 前置条件
- PyTorch(版本 2.9.1 或兼容版本)
- **[CUDA 13.1+](https://developer.nvidia.com/cuda-downloads)**(必需 - TileGym 仅在 CUDA 13.1+ 上构建和测试)
- Triton(随 PyTorch 安装一同包含)
### 安装步骤
#### 1. 准备 `torch` 和 `triton` 环境
如果您已经有 `torch` 和 `triton`,请跳过此步骤。
```
pip install --pre torch --index-url https://download.pytorch.org/whl/cu130
```
我们已验证 `torch==2.9.1` 可以正常工作。在安装 `torch` 时,您也可以获得 `triton` 包。
#### 2. 安装 TileGym
TileGym 使用 [`cuda-tile`](https://github.com/nvidia/cutile-python)(≥ 1.3.0)进行 GPU kernel 编程,它在运行时依赖于 `tileiras` 编译器。
##### 从 PyPI 安装(推荐)
```
pip install tilegym[tileiras]
```
这将安装 TileGym 和所有运行时依赖,包括 `cuda-tile[tileiras]`,它会将 `tileiras` 编译器直接打包到您的 Python 环境中。
如果您的系统上已经有可用的 `tileiras`(例如,来自 [CUDA Toolkit 13.1+](https://developer.nvidia.com/cuda-downloads)),您可以省略该附加项:
```
pip install tilegym
```
##### 从源码安装
```
git clone https://github.com/NVIDIA/TileGym.git
cd TileGym
pip install .[tileiras] # or: pip install . (if you have system tileiras)
```
对于可编辑(开发)模式,请使用 `pip install -e .` 或 `pip install -e .[tileiras]`。
所有运行时依赖都在 [`requirements.txt`](requirements.txt) 中声明,并由 `pip install tilegym` 和 `pip install .` 自动安装。
我们还提供了 Dockerfile,您可以参考 [modeling/transformers/README.md](modeling/transformers/README.md)。
### 后端
TileGym 为以下后端提供 kernel,每个后端都位于 `src/tilegym/ops/` 下各自的文件夹中:
- **cuTile**(默认)— [`src/tilegym/ops/cutile`](src/tilegym/ops/cutile),详情请见 [cutile-python](https://github.com/nvidia/cutile-python)。
- **CUDA Tile C++** — [`src/tilegym/ops/tilecpp`](src/tilegym/ops/tilecpp),详情请见 [README.tilecpp.md](README.tilecpp.md)。
- **Triton CUDA Tile IR** — [`src/tilegym/ops/triton`](src/tilegym/ops/triton),详情请见 [Triton-to-tile-IR](https://github.com/triton-lang/Triton-to-tile-IR)。
要使用 Triton CUDA Tile IR 后端,请将其 wheel 包安装到一个单独的目录中,并在运行时使用 `ENABLE_TILE=1` 进行选择。适用于 CPython 3.12 和 3.13 的 wheel 包可在 [releases 页面](https://github.com/triton-lang/Triton-to-tile-IR/releases) 获取:
```
# 安装到单独的目录中,与默认环境分开
pip install --target /opt/nvtriton .whl
# 在运行时选择 Triton CUDA Tile IR backend
PYTHONPATH=/opt/nvtriton ENABLE_TILE=1 python your_script.py
```
## 快速入门
使用 TileGym 主要有三种方式:
### 1. 探索 Kernel 示例
所有 kernel 实现都位于 `src/tilegym/ops/` 目录中。您可以使用最简脚本测试单个操作。各个操作的函数级用法和最简脚本记录在 [tests/ops/README.md](tests/ops/README.md) 中。
### 2. 运行基准测试
通过微基准测试评估 kernel 性能:
```
cd tests/benchmark
bash run_all.sh
```
完整的基准测试指南可在 [tests/benchmark/README.md](tests/benchmark/README.md) 中找到。
### 3. 运行 LLM Transformer 示例
在端到端推理场景中使用 TileGym kernel。我们为使用 TileGym kernel 加速的 Transformer 语言模型(例如 Llama 3.1-8B)提供了可运行的脚本和说明。
首先,安装附加依赖:
```
pip install accelerate==1.13.0 --no-deps
```
**容器化设置 (Docker)**:
```
docker build -t tilegym-transformers -f modeling/transformers/Dockerfile .
docker run --gpus all -it tilegym-transformers bash
```
更多详情请见 [modeling/transformers/README.md](modeling/transformers/README.md)
### 4. Julia (cuTile.jl) Kernel(可选)
TileGym 还在 Julia 中包含了实验性的 [cuTile.jl](https://github.com/JuliaGPU/cuTile.jl) kernel 实现。这些内容独立存在于 `julia/` 目录中,不需要 Python TileGym 包。
**前置条件**:[Julia 1.12+](https://julialang.org/downloads/)、CUDA 13.1、Blackwell GPU
```
# 安装 Julia(如果尚未安装)
curl -fsSL https://install.julialang.org | sh
# 安装依赖项
julia --project=julia/ -e 'using Pkg; Pkg.instantiate()'
# 运行测试
julia --project=julia/ julia/test/runtests.jl
```
完整依赖列表请参见 `julia/Project.toml`。
### 5. 启用 cuTile-rs (Rust) 后端(可选)
部分操作在以下路径附带了额外的 **cuTile-rs** 后端
[`src/tilegym/ops/cutile_rs`](src/tilegym/ops/cutile_rs) — 这些 kernel 使用 Rust 和
[`cutile-rs`](https://github.com/NVlabs/cutile-rs) 编写,并通过
C-ABI `libcutile_kernels.so` 加载。它是可选启用的,并且只能从源码检出版本中使用。
**前置条件**(除上述基础安装外),需与
[cuTile-rs](https://github.com/NVlabs/cutile-rs) 匹配:
- **Rust 1.89+** — `cargo` 和 `rustc` 需在 `PATH` 中:
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
rustup default stable
- **带有 headers 的 CUDA toolkit** — Rust 构建会针对
`cuda.h` 运行 `bindgen`。将 `CUDA_TOOLKIT_PATH` 设置为您的安装路径;如果未设置,cuTile-rs 将
回退到 `/usr/local/cuda`:
export CUDA_TOOLKIT_PATH=/usr/local/cuda # 必须包含 include/cuda.h
**使用它。** 后端加载器会在首次使用时惰性构建共享库
(`cargo build --release`),因此无需手动构建步骤:
```
import tilegym
tilegym.set_backend("cutile-rs")
from tilegym.backend.selector import get_available_backends
print(get_available_backends()) # should include "cutile-rs"
from tilegym.ops import bmm # backend-agnostic import
# ... bmm(...) 现在会分发到 cuTile-rs kernel
```
**可选的环境变量配置:**
```
export CUTILE_RS_AUTOBUILD=0 # skip the lazy rebuild; use a prebuilt .so
export CUTILE_RS_KERNELS_DIR=/abs/path/to/cutile_kernels # override the crate location
```
**对 cuTile-rs 进行基准测试。** 在将 cuTile-rs 性能与
cuTile-Python 基准进行比较时,请使用 **`CUPTI=1`** 运行性能测试(使用 CUPTI /
`torch.profiler` 设备时间,而不是 CUDA 事件)。cuTile-rs kernel 通常具有与参考实现不同的
主机/启动开销,而 CUDA 事件挂钟计时对于小型(亚微秒级)kernel 的计算会偏大;CUPTI 测量的是纯 GPU kernel 时间,能提供稳定且公平的对比比例:
```
CUPTI=1 pytest tests/ops/test_bmm.py -k "test_perf and cutile_rs" --print-record
```
## 贡献指南
我们欢迎各种形式的贡献。请阅读我们的 [CONTRIBUTING.md](CONTRIBUTING.md) 了解指南,包括贡献者许可协议 (CLA) 流程。
## 许可证和第三方声明
- 项目许可证:MIT
- [LICENSE](LICENSE)
- 第三方归属和许可证文本:
- [LICENSES/ATTRIBUTIONS.md](LICENSES/ATTRIBUTIONS.md)
## 功能
- 丰富的 CUDA Tile kernel 示例集合
- 针对常见深度学习算子的实用 kernel 实现
- 用于评估 kernel 效率的性能基准测试
- 与主流 LLM(Llama 3.1、DeepSeek V2)的端到端集成示例
## 安装说明
### 前置条件
- PyTorch(版本 2.9.1 或兼容版本)
- **[CUDA 13.1+](https://developer.nvidia.com/cuda-downloads)**(必需 - TileGym 仅在 CUDA 13.1+ 上构建和测试)
- Triton(随 PyTorch 安装一同包含)
### 安装步骤
#### 1. 准备 `torch` 和 `triton` 环境
如果您已经有 `torch` 和 `triton`,请跳过此步骤。
```
pip install --pre torch --index-url https://download.pytorch.org/whl/cu130
```
我们已验证 `torch==2.9.1` 可以正常工作。在安装 `torch` 时,您也可以获得 `triton` 包。
#### 2. 安装 TileGym
TileGym 使用 [`cuda-tile`](https://github.com/nvidia/cutile-python)(≥ 1.3.0)进行 GPU kernel 编程,它在运行时依赖于 `tileiras` 编译器。
##### 从 PyPI 安装(推荐)
```
pip install tilegym[tileiras]
```
这将安装 TileGym 和所有运行时依赖,包括 `cuda-tile[tileiras]`,它会将 `tileiras` 编译器直接打包到您的 Python 环境中。
如果您的系统上已经有可用的 `tileiras`(例如,来自 [CUDA Toolkit 13.1+](https://developer.nvidia.com/cuda-downloads)),您可以省略该附加项:
```
pip install tilegym
```
##### 从源码安装
```
git clone https://github.com/NVIDIA/TileGym.git
cd TileGym
pip install .[tileiras] # or: pip install . (if you have system tileiras)
```
对于可编辑(开发)模式,请使用 `pip install -e .` 或 `pip install -e .[tileiras]`。
所有运行时依赖都在 [`requirements.txt`](requirements.txt) 中声明,并由 `pip install tilegym` 和 `pip install .` 自动安装。
我们还提供了 Dockerfile,您可以参考 [modeling/transformers/README.md](modeling/transformers/README.md)。
### 后端
TileGym 为以下后端提供 kernel,每个后端都位于 `src/tilegym/ops/` 下各自的文件夹中:
- **cuTile**(默认)— [`src/tilegym/ops/cutile`](src/tilegym/ops/cutile),详情请见 [cutile-python](https://github.com/nvidia/cutile-python)。
- **CUDA Tile C++** — [`src/tilegym/ops/tilecpp`](src/tilegym/ops/tilecpp),详情请见 [README.tilecpp.md](README.tilecpp.md)。
- **Triton CUDA Tile IR** — [`src/tilegym/ops/triton`](src/tilegym/ops/triton),详情请见 [Triton-to-tile-IR](https://github.com/triton-lang/Triton-to-tile-IR)。
要使用 Triton CUDA Tile IR 后端,请将其 wheel 包安装到一个单独的目录中,并在运行时使用 `ENABLE_TILE=1` 进行选择。适用于 CPython 3.12 和 3.13 的 wheel 包可在 [releases 页面](https://github.com/triton-lang/Triton-to-tile-IR/releases) 获取:
```
# 安装到单独的目录中,与默认环境分开
pip install --target /opt/nvtriton 标签:CUDA, DLL 劫持, GPU编程, Vectored Exception Handling, 凭据扫描, 可视化界面, 大语言模型, 深度学习, 算子开发, 请求拦截, 逆向工具, 高性能计算