@@ -3,24 +3,29 @@ set -xeo pipefail
33export DEBIAN_FRONTEND=noninteractive
44
55TORCH_CUDA_ARCH_LIST=' '
6- FILTER_ARCHES=' '
7- BUILD_TRITON=' '
86
9- while getopts ' a:ft ' ; do
7+ while getopts ' a:' ; do
108 case " ${OPT} " in
119 a) TORCH_CUDA_ARCH_LIST=" ${OPTARG} "
12- f) FILTER_ARCHES=' 1'
13- t) BUILD_TRITON=' 1'
1410 * ) exit 92 ;;
1511 esac
1612done
1713
18- export NVCC_APPEND_FLAGS=' -gencode=arch=compute_100,code=[sm_100,compute_100] -gencode=arch=compute_100a,code=sm_100a --diag-suppress 174'
19- export TORCH_CUDA_ARCH_LIST=" ${TORCH_CUDA_ARCH_LIST:- 9.0 10.0+PTX} "
14+ printf ' Using %s=%s\n'
15+ FLASHINFER_COMMIT " ${FLASHINFER_COMMIT:- <None>} "
16+ SGLANG_COMMIT " ${SGLANG_COMMIT:- <None>} "
17+
18+ export NVCC_APPEND_FLAGS=' --diag-suppress 174,177,2361'
19+ export TORCH_CUDA_ARCH_LIST=" ${TORCH_CUDA_ARCH_LIST:- 9.0a 10.0a 12.0+PTX} "
2020
2121mkdir -p /wheels/logs
2222
23- _BUILD () { python3 -m build -w -n -v -o /wheels " ${1:- .} " ; }
23+ _CLONE () {
24+ git clone --filter=tree:0 --no-single-branch --no-checkout " ${1:? } " " ${2:? } " && \
25+ git -C " ${2:? } " " ${3:? } " && \
26+ git -C " ${2:? } " ;
27+ }
28+ _BUILD () { python3 -m build -w -n -v -o /wheels " ${@:- .} " ; }
2429_LOG () { tee -a " /wheels/logs/${1:? } " ; }
2530_CONSTRAINTS=" $( python3 -m pip list | sed -En ' s@^(torch(vision|audio)?)\s+(\S+)$@\1==\3@p' ) "
2631_PIP_INSTALL () {
@@ -29,122 +34,55 @@ _PIP_INSTALL() {
2934 " $@ "
3035}
3136
32- _PIP_INSTALL -U pip setuptools wheel build pybind11 ninja cmake
33-
34- # triton (not compatible with torch 2.6)
35- if [ " ${BUILD_TRITON} " = 1 ]; then (
36- : " ${TRITON_COMMIT:? } "
37- echo ' Building triton-lang/triton'
38- git clone --recursive --filter=blob:none https://github.com/triton-lang/triton
39- cd triton
40- git checkout " ${TRITON_COMMIT} "
41- _BUILD python | & _LOG triton.log
42- ); fi
37+ _PIP_INSTALL -U pip setuptools wheel build pybind11 ninja ' cmake<4.0.0' ' scikit-build-core>=0.10' ' setuptools-scm>=8.0'
4338
4439# flashinfer
4540: " ${FLASHINFER_COMMIT:? } "
46- : " ${CUTLASS_COMMIT:? } "
4741(
48- echo ' Building flashinfer-ai/flashinfer'
49- git clone --recursive --filter=blob:none https://github.com/flashinfer-ai/flashinfer
42+ echo " Building flashinfer-ai/flashinfer @ ${FLASHINFER_COMMIT} "
43+ _CLONE https://github.com/flashinfer-ai/flashinfer flashinfer " ${FLASHINFER_COMMIT} "
5044cd flashinfer
51- git checkout " ${FLASHINFER_COMMIT} "
52- sed -i ' s/name = "flashinfer-python"/name = "flashinfer"/'
53- git -C 3rdparty/cutlass checkout " ${CUTLASS_COMMIT} "
54- _PIP_INSTALL -U optree
45+ # flashinfer v0.6+ uses TVM for AOT kernel compilation
46+ _PIP_INSTALL -U optree ' apache-tvm-ffi>=0.1.5,<0.2'
47+ # Convert TORCH_CUDA_ARCH_LIST to FlashInfer's format:
48+ # - Add 'a' suffix to >=9.0 arches for architecture-specific instructions
49+ # - Strip +PTX (FlashInfer AOT only generates native SASS, not PTX)
50+ # See vllm-tensorizer/Dockerfile for the reference pattern.
51+ FLASHINFER_ARCH_LIST=" $( echo " ${TORCH_CUDA_ARCH_LIST} " | sed -E ' s@\b(9|10|12)\.0\b@\1.0a@g; s@\+PTX\b@@g' | xargs) "
52+ FLASHINFER_CUDA_ARCH_LIST=" ${FLASHINFER_ARCH_LIST} "
5553NVCC_APPEND_FLAGS=" ${NVCC_APPEND_FLAGS: +$NVCC_APPEND_FLAGS } --diag-suppress 20281,174"
56- FLASHINFER_ENABLE_AOT=1 _BUILD . | & _LOG flashinfer.log
57- )
58-
59- # Setup cutlass repo for vLLM to use
60- git clone --recursive --filter=blob:none https://github.com/NVIDIA/cutlass
61- git -C cutlass checkout " ${CUTLASS_COMMIT} "
62-
63- # vLLM
64- : " ${VLLM_COMMIT:? } "
65- (
66- echo ' Building vllm-project/vllm'
67- export VLLM_CUTLASS_SRC_DIR=" ${PWD} /cutlass"
68- test -d " ${VLLM_CUTLASS_SRC_DIR} "
69- git clone --recursive --filter=blob:none https://github.com/vllm-project/vllm
70- cd vllm
71- git checkout " ${VLLM_COMMIT} "
72- # For lsmod
73- apt-get -qq update && apt-get -qq install --no-install-recommends -y kmod
74- python3 use_existing_torch.py
75- _PIP_INSTALL -r requirements-build.txt
76- USE_CUDNN=1 USE_CUSPARSELT=1 _BUILD . | & _LOG vllm.log
54+ _BUILD . \
55+ | & _LOG flashinfer.log \
56+ | sed ' /^Killed$/{p; Q1}'
7757)
7858
79- # sglang
59+ # sglang (sgl-kernel CUDA extension + Python package)
8060: " ${SGLANG_COMMIT:? } "
8161(
82- echo ' Building sglang'
83- git clone --recursive --filter=blob:none https://github.com/sgl-project/sglang
62+ echo " Building sgl-project/ sglang @ ${SGLANG_COMMIT} "
63+ _CLONE https://github.com/sgl-project/sglang sglang " ${SGLANG_COMMIT} "
8464cd sglang
85- git checkout " ${SGLANG_COMMIT} "
65+
66+ # sgl-kernel (CUDA extension, compiled via scikit-build-core/CMake)
8667(
8768cd sgl-kernel
88- git -C 3rdparty/cutlass checkout " ${CUTLASS_COMMIT} "
89- git -C 3rdparty/flashinfer/3rdparty/cutlass checkout " ${CUTLASS_COMMIT} "
90-
91- ARCH_TRIPLE=" $( gcc -print-multiarch) "
92- LIB_DIR=" /usr/lib/${ARCH_TRIPLE:? } "
93- test -d " ${LIB_DIR:? } "
94- PYTHON_API_VER=" $(
95- python3 --version | sed -En ' s@Python ([0-9])\.([0-9]+)\..*@cp\1\2@p'
96- ) "
97- ARCH_FILTER=()
98- if [ " ${FILTER_ARCHES} " = 1 ]; then
99- ARCH_FILTER=(-e ' s@"-gencode=arch=compute_[78][0-9],code=sm_[78][0-9]",@#\0@'
100- fi
10169
102- sed -Ei \
103- " ${ARCH_FILTER[@]} "
104- -e ' s@/usr/lib/x86_64-linux-gnu@' " ${LIB_DIR} " ' @'
105- -e ' s@(\s+)(\w.+manylinux2014_x86_64.+)@\1pass # \2@'
106- -e ' s@\{"py_limited_api": "cp39"}@{"py_limited_api": "' " ${PYTHON_API_VER:- cp310} " ' "}@'
107- setup.py
108- SGL_KERNEL_ENABLE_BF16=1 SGL_KERNEL_ENABLE_FP8=1 SGL_KERNEL_ENABLE_SM90A=1 \
109- _BUILD . | & _LOG sglang.log
70+ _BUILD \
71+ -Cbuild-dir=build \
72+ -Ccmake.define.SGL_KERNEL_ENABLE_SM100A=1 \
73+ -Ccmake.define.SGL_KERNEL_ENABLE_SM90A=1 \
74+ -Ccmake.define.SGL_KERNEL_ENABLE_BF16=1 \
75+ -Ccmake.define.SGL_KERNEL_ENABLE_FP8=1 \
76+ -Ccmake.define.SGL_KERNEL_ENABLE_FP4=1 \
77+ . \
78+ | & _LOG sgl-kernel.log \
79+ | sed ' /^Killed$/{p; Q1}'
11080)
81+
82+ # sglang Python package (no CUDA compilation)
83+ # Remove torch/torchaudio pins — the base image provides these.
84+ sed -i -E ' /torch(audio)?[><=~]/d'
11185_BUILD python | & _LOG sglang.log
11286)
11387
114- # decord and xgrammar aren't available on PyPI for ARM64
115-
116- if [ ! " $( uname -m) " = ' x86_64' ; then
117- # xgrammar (for sglang)
118- (
119- git clone --recursive --filter=blob:none -b v0.1.11 https://github.com/mlc-ai/xgrammar && \
120- cd xgrammar
121- (
122- mkdir build && cd build
123- cmake -S.. -B. -DCMAKE_BUILD_TYPE=Release -GNinja | & _LOG xgrammar.log
124- cmake --build . | & _LOG xgrammar.log
125- )
126- _BUILD python | & _LOG xgrammar.log
127- )
128-
129- # decord (for sglang)
130- : " ${DECORD_COMMIT:? } "
131- (
132- apt-get -qq update && apt-get -q install --no-install-recommends -y \
133- build-essential python3-dev python3-setuptools \
134- make cmake ffmpeg \
135- libavcodec-dev libavfilter-dev libavformat-dev libavutil-dev
136- git clone --recursive --filter=blob:none https://github.com/dmlc/decord
137- cd decord
138- git checkout " ${DECORD_COMMIT} "
139- (
140- mkdir build && cd build
141- cmake -S.. -B. -DUSE_CUDA=0 -DCMAKE_BUILD_TYPE=Release -GNinja | & _LOG decord.log
142- cmake --build . | & _LOG decord.log
143- cp libdecord.so /wheels/libdecord.so
144- )
145- cd python
146- _BUILD . | & _LOG decord.log
147- )
148- fi
149-
15088apt-get clean
0 commit comments