TileLang 容器化环境搭建：Docker 镜像构建与 GPU 容器运行全解

TileLang 容器化环境搭建：Docker 镜像构建与 GPU 容器运行全解

【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang

TileLang 的构建依赖一条较长的工具链：定制版 TVM 子模块、CUTLASS 头文件、CUDA/HIP 工具链以及一套 C++/CMake 构建环境。本文以仓库 docker/README.md 为主体，完整讲解如何用docker/目录下的官方 Dockerfile 构建 TileLang 开发镜像、如何在 NVIDIA 与 AMD 两类 GPU 主机上以正确参数运行容器，并结合各 Dockerfile 的逐行内容说明镜像内部实际安装了什么、为什么要这样装，使你能在不破坏宿主环境的前提下获得一个可直接编译和调试 TileLang 核函数的可复现环境。

一、为什么用 Docker：一次构建，消除依赖差异

安装指南中给出的源码构建流程要求宿主机具备 Python >= 3.10、CMake >= 3.26.1（见 pyproject.toml 中tool.scikit-build.cmake.version的约束）、protobuf/z3/LLVM 等系统依赖，并且 TileLang 依赖一个以 git submodule 形式内嵌在3rdparty/tvm下的定制 TVM——任何一个环节版本不对，构建都可能失败。

docker/目录的做法是把这些依赖固化进镜像：

• 基础镜像统一选用 NVIDIA 官方的 PyTorch 容器（自带对应版本 CUDA 与 PyTorch），AMD 侧选用 ROCm 官方 PyTorch 容器；
• 系统级编译依赖通过apt-get一次性安装；
• 在镜像构建阶段直接pip install -e .以可编辑模式编译 TileLang，容器启动后即可直接import tilelang。

按 docker/README.md 的说法，镜像基于 Ubuntu 20.04 生态（实际基础镜像为 NVIDIA/ROCm 官方 PyTorch 容器，均为 Ubuntu 系），包含运行实验所需的全部依赖；官方目前仅维护 NVIDIA GPU 的 Dockerfile 为主流路径，AMD 的 Dockerfile 同样已经提供在仓库中（docker/Dockerfile.rocm）。

二、可用的 Dockerfile 与基础镜像对照

docker/目录下共有 9 个 Dockerfile，覆盖 CUDA 11.8 到 12.8 以及 ROCm 7.2。CUDA 系列 Dockerfile 之间除基础镜像版本外结构完全一致，差异仅体现在第一行FROM，因此下表可以直接作为选型依据：

选择原则：以你宿主机驱动所支持的最高 CUDA 版本为准，取不超过该版本的最接近的Dockerfile.cu*（NVIDIA 驱动向后兼容，较新驱动可运行较旧 CUDA 容器）。

三、构建 CUDA 镜像：README 官方流程

docker/README.md 给出的标准流程如下（git clone使用官方仓库地址，-f指定的 Dockerfile 换成你要的 CUDA 版本）：

git clone --recursive https://github.com/tile-ai/tilelang TileLang cd TileLang/docker # 构建镜像，可能需要一段时间（官方测试机上约 10 分钟以上） # 将版本号 cu124 替换为你想用的版本 # AMD GPU 用户将 .cu** 替换为 .rocm docker build -t tilelang_workspace -f Dockerfile.cu124 .

几个关键点：

1. --recursive不可省略。TileLang 依赖3rdparty/tvm定制 TVM 子模块与3rdparty/cutlass，[安装指南](https://link.gitcode.com/i/6b7620ebfa49625c54fcde9acccfb47a)在源码构建一节中也明确强调了这一点。
2. 构建上下文是docker/目录（CUDA 系列），因为 Dockerfile 内通过git clone自行拉取源码，上下文内容无关；而 ROCm 系列 Dockerfile 用COPY .拷贝源码，构建方式不同（见第六节）。
3. 构建耗时主要来自 TVM 相关 C++ 库与 TileLang 本体（libtilelang）的编译，CPU 核数多的机器上会更快。

镜像内部实际安装了什么

以 Dockerfile.cu124 为例逐段解读，各 CUDA 版本 Dockerfile 的结构与此一致：

FROM nvcr.io/nvidia/pytorch:24.05-py3 # 自带 CUDA 12.4 + PyTorch 的官方镜像 WORKDIR /root # 1) 系统编译依赖：构建工具 + TileLang/TVM 需要的库 RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential git wget \ libgtest-dev libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev llvm-dev \ && apt-get clean autoclean && ... # 2) 安装 Miniconda (py310) 到 /opt/conda，并把它加入 PATH RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.5.2-0-Linux-x86_64.sh -O install_miniconda.sh && \ bash install_miniconda.sh -b -p /opt/conda && rm install_miniconda.sh ENV PATH="/opt/conda/bin:${PATH}" # 3) conda 安装 pip、cmake，并用 conda-forge 的 libstdcxx-ng=12 保证 C++ 运行时 RUN conda install pip cmake && conda install -c conda-forge libstdcxx-ng=12 && conda clean --all # 4) 再补一轮 apt 的 Python 头文件与构建工具 RUN apt-get install -y python3 python3-dev python3-setuptools gcc libtinfo-dev zlib1g-dev build-essential cmake libedit-dev libxml2-dev # 5) 拉取源码（含子模块）并以可编辑模式安装 RUN git clone https://github.com/tile-ai/tilelang.git --recursive -b main TileLang \ && cd TileLang && USE_CUDA=1 pip install -e . -v CMD bash

两点值得注意：

• USE_CUDA=1是构建开关。TileLang 的 CMake 入口会读取该环境变量决定后端：在 CMakeLists.txt 中，USE_CUDA/USE_ROCM环境变量分别被映射为 CMake 的USE_CUDA/USE_ROCM开关；未显式设置时，Linux 上默认按能否找到 CUDA 头文件自动判断。
• **pip install -e .（可编辑模式）**意味着源码目录/root/TileLang就是 Python 包的安装位置。进入容器后修改tilelang/下的 Python 文件立即生效；若改动 C++ 代码，则需在build/目录下重新make/ninja（这一"开发者快速重编"路径在 安装指南 的Faster Rebuild for Developers一节有详述）。
• Dockerfile.cu128相比其他版本额外安装了cython3并执行pip install cython，对应构建系统对 Cython 的依赖（pyproject.toml中build-system.requires声明了cython>=3.1.0）。

四、运行 NVIDIA GPU 容器

镜像构建完成后，按 docker/README.md 启动：

docker run -it --cap-add=SYS_ADMIN --network=host --gpus all --cap-add=SYS_PTRACE --shm-size=4G --security-opt seccomp=unconfined --security-opt apparmor=unconfined --name tilelang_test tilelang_workspace bash

各参数作用：

进入容器后，验证安装（与 安装指南 中的验证命令一致）：

python -c "import tilelang; print(tilelang.__version__)"

如果后续要在容器内跑性能敏感的实验，也可以参考 安装指南 Docker 一节的建议，把--shm-size调大（例如32g）并用-v挂载宿主目录做代码/数据共享。

五、ROCm 镜像：不同的构建约定

Dockerfile.rocm 与 CUDA 系列有两点本质区别：

1. 构建上下文是仓库根目录，因为该文件用COPY . /root/tilelang直接拷贝本地源码而不是git clone。文件内注释已写明构建方式：

   # 在 tilelang 仓库根目录执行 docker build -f docker/Dockerfile.rocm -t tilelang_rocm:latest .

   这意味着你本地改过的源码会直接进入镜像，适合定制构建。

2. 安装时临时移除 torch 依赖。TileLang 的 pyproject.toml 声明了torch为运行时依赖，而基础镜像已内置 ROCm 版 PyTorch 2.10.0。为避免 pip 从 PyPI 重新拉取 CUDA 版 torch 覆盖容器内的 ROCm 版，Dockerfile 在安装前把pyproject.toml中以torch开头的依赖行临时删掉、安装完成后再恢复：

   RUN cd /root/tilelang && \ cp pyproject.toml pyproject.toml.bak && \ sed -i '/^[[:space:]]*\"torch/d' pyproject.toml && \ USE_ROCM=1 USE_CUDA=0 pip install -e . -v && \ mv pyproject.toml.bak pyproject.toml

此外，ROCm Dockerfile 相比 CUDA 系列额外安装了rocm-dev rocm-libs hip-dev hipblas-dev rocblas-dev，并设置了一组固定环境变量：

ENV USE_ROCM=1 ENV USE_CUDA=0 ENV ROCM_HOME=/opt/rocm ENV HIP_PLATFORM=amd ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1201;gfx1100"

其中USE_ROCM/USE_CUDA对应 CMakeLists.txt 中的环境变量开关；PYTORCH_ROCM_ARCH列出了预编译 PyTorch 所覆盖的 GPU 架构（如 MI300X 的gfx942）。

提示：安装指南 同时说明，仅仅运行TileLang 在 AMD GPU 上并不需要源码构建——PyPI 的 Linux x86_64 wheel 是 CUDA+ROCm 的 fat 构建，pip install tilelang配合 ROCm 版 torch 即可，只需宿主机提供hipcc（JIT 用）。docker/Dockerfile.rocm面向的是"开发 TileLang 本身或需要定制构建"的场景；该文档还附有一节在既有 ROCm 容器（如 sglang 镜像）内就地重编 TileLang 的完整步骤（-DUSE_CUDA=OFF -DUSE_ROCM=ON、--no-deps安装等），遇到容器内构建问题可对照阅读。

六、运行 AMD GPU 容器

AMD 主机的启动命令与 NVIDIA 版只有设备暴露方式不同，来自 docker/README.md：

docker run -it --cap-add=SYS_ADMIN --network=host --device=/dev/kfd --device=/dev/dri \ --cap-add=SYS_PTRACE --shm-size=4G --security-opt seccomp=unconfined --security-opt apparmor=unconfined \ --name tilelang_test tilelang_workspace bash

关键差异是用--device=/dev/kfd --device=/dev/dri直接把 ROCm 的设备节点传入容器，替代--gpus all（ROCm 不使用 nvidia-container-toolkit）。容器内验证目标检测可执行：

python -c "import tilelang; from tilelang.backend.target import determine_target; print(tilelang.__version__, determine_target(return_object=True))"

正常应输出hiptarget 及你的 GPU 架构（例如mcpu=gfx942）；若 ROCm 不在/opt/rocm或 PATH 上存在多个 HIP 工具链，需按 安装指南 的说明设置ROCM_PATH指向正确的 ROCm 前缀。

七、实用建议与常见注意事项

1. 构建不需要 GPU 宿主。按 安装指南 的说法，Docker 构建阶段只是编译，宿主机没有 GPU 也可完成docker build，然后把镜像docker save/load到有 GPU 的机器上使用。
2. 版本对应关系：tilelang.__version__默认会嵌入 SDK 与 git 信息（形如<sdk>.git<git_hash>，<sdk>取cuda/rocm/metal，CUDA 侧表现为cu124这类标签），可用它确认容器内构建所用的工具链；如需去掉该标签，构建时设置NO_VERSION_LABEL=ON（详见 安装指南 的Install Configs一节，构建期环境变量全集如USE_CUDA、USE_ROCM、USE_LLVM、TVM_ROOT、WITH_PIP_CUDA_TOOLCHAIN等也集中说明在该处）。
3. 容器内二次安装依赖：由于容器使用宿主网络且 conda 在 PATH 最前，pip install的行为与宿主机环境完全隔离，不会互相污染，这正是容器化方案相对裸机安装的最大优势。
4. 镜像选择核对：如果你只看到nvcc/CUDA 头文件相关报错，先确认docker build -f指定的 Dockerfile 与你宿主驱动支持的 CUDA 版本匹配，再确认启动命令带上了--gpus all（NVIDIA）或/dev/kfd、/dev/dri设备（AMD）。

通过上述流程，你可以在 docker/README.md 提供的官方 Dockerfile 基础上，为 TileLang 搭建一个从编译到调试都开箱即用的 GPU 容器环境：NVIDIA 侧选定Dockerfile.cu*构建后以--gpus all启动，AMD 侧在仓库根目录以Dockerfile.rocm构建后以 KFD/DRI 设备节点启动，并用import tilelang及 target 检测命令完成闭环验证。

【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang