Switch to vllm0.8.0 and torch2.6.0+cu126

- Only targeting Hopper for now
- Add vllm/cuda versions to image tags for easier tracking

Author: Apsu

.github/workflows/build-vllm.yaml

@@ -2,18 +2,16 @@ name: Build vLLM Docker Image Matrix
 
 env:
   PARALLELISM: 1
-  MAX_JOBS: 32 # Limit for building multiple archs
-  NVCC_THREADS: 2
-  TORCH_CUDA_ARCH_LIST: 9.0a;10.0a
-  VLLM_FA_CMAKE_GPU_ARCHES: 90a-real;100a-real
-  TRITON_REF: release/3.3.x
-  TRITON_BUILD_VERSION: 3.3.0
-  XFORMERS_REF: v0.0.29.post3
-  XFORMERS_BUILD_VERSION: 0.0.29.post3
+  TORCH_CUDA_ARCH_LIST: 9.0a
+  VLLM_FA_CMAKE_GPU_ARCHES: 90a-real
+  TRITON_REF: release/3.2.x
+  TRITON_BUILD_VERSION: 3.2.0
+  XFORMERS_REF: v0.0.29.post2
+  XFORMERS_BUILD_VERSION: 0.0.29.post2
   FLASHINFER_REF: v0.2.2.post1
   FLASHINFER_BUILD_VERSION: 0.2.2.post1
-  VLLM_REF: d47807ba
-  VLLM_BUILD_VERSION: 0.7.4
+  VLLM_REF: v0.8.0
+  VLLM_BUILD_VERSION: 0.8.0
 
 on:
   push:
@@ -25,13 +23,16 @@ jobs:
     strategy:
       matrix:
         arch: [amd64, arm64]
-        cuda_version: [12.8.0]
+        cuda_version: [12.6.3]
         image_distro: [ubuntu24.04]
     runs-on: [self-hosted, "${{ matrix.arch }}"]
     steps:
-      - name: Generate image name
+      - name: Prepare some env vars
         run: |
           echo "GHCR_IMAGE=ghcr.io/${GITHUB_REPOSITORY@L}" >> ${GITHUB_ENV}
+          CUDA_VERSION=${{ matrix.cuda_version }}
+          CUDA_SHORT=${CUDA_VERSION%.*}
+          echo "CUDA_TAG=${CUDA_SHORT//./}" >> ${GITHUB_ENV}
 
       - name: Login to GHCR
         uses: docker/login-action@v3
@@ -65,21 +66,23 @@ jobs:
             FLASHINFER_BUILD_VERSION=${{ env.FLASHINFER_BUILD_VERSION }}
             VLLM_REF=${{ env.VLLM_REF }}
             VLLM_BUILD_VERSION=${{ env.VLLM_BUILD_VERSION }}
-          cache-from: type=registry,ref=${{ env.GHCR_IMAGE }}:cache-${{ matrix.arch }}
-          cache-to: type=registry,ref=${{ env.GHCR_IMAGE }}:cache-${{ matrix.arch }},mode=max
+          cache-from: type=registry,ref=${{ env.GHCR_IMAGE }}:cache-cu${{ env.CUDA_TAG }}-${{ matrix.arch }}
+          cache-to: type=registry,ref=${{ env.GHCR_IMAGE }}:cache-cu${{ env.CUDA_TAG }}-${{ matrix.arch }},mode=max
           context: .
           file: Dockerfile
           platforms: linux/${{ matrix.arch }}
           push: true
-          tags: ${{ env.GHCR_IMAGE }}:${{ matrix.arch }}
+          tags: ${{ env.GHCR_IMAGE }}:${{ env.VLLM_BUILD_VERSION }}-cu${{ env.CUDA_TAG }}-${{ matrix.arch }}
 
+  # Fix this to use matrix and handle imagetools create --append
   ghcr:
     needs: build
     runs-on: self-hosted
     steps:
-      - name: Generate image name
+      - name: Prepare some env vars
         run: |
           echo "GHCR_IMAGE=ghcr.io/${GITHUB_REPOSITORY@L}" >> ${GITHUB_ENV}
+          echo "CUDA_TAG=126" >> ${GITHUB_ENV}
 
       - name: Login to GHCR
         uses: docker/login-action@v3
@@ -90,5 +93,5 @@ jobs:
 
       - name: Append images
         run: |
-          ARCHS=(amd64 arm64)
-          docker buildx imagetools create -t ${GHCR_IMAGE}:latest ${ARCHS[@]/#/${GHCR_IMAGE}:}
+          TAGS=(${VLLM_BUILD_VERSION}-cu${CUDA_TAG}-{amd,arm}64)
+          docker buildx imagetools create -t ${GHCR_IMAGE}:latest ${TAGS[@]/#/${GHCR_IMAGE}:}

Dockerfile

@@ -1,16 +1,10 @@
-ARG CUDA_VERSION=12.8.0
+ARG CUDA_VERSION=12.6.3
 ARG IMAGE_DISTRO=ubuntu24.04
 ARG PYTHON_VERSION=3.12
 
 # ---------- Builder Base ----------
 FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-devel-${IMAGE_DISTRO} AS base
 
-# Set build scaling
-ARG MAX_JOBS=32
-ENV MAX_JOBS=${MAX_JOBS}
-ARG NVCC_THREADS=2
-ENV NVCC_THREADS=${NVCC_THREADS}
-
 # Set arch lists for all targets
 # 'a' suffix is not forward compatible but enables all optimizations
 ARG TORCH_CUDA_ARCH_LIST="9.0a"
@@ -52,7 +46,7 @@ ENV CUDA_HOME=/usr/local/cuda
 ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
 
 # Install pytorch nightly
-RUN uv pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu128
+RUN uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu126
 
 FROM base AS build-base
 RUN mkdir /wheels
@@ -62,8 +56,8 @@ RUN mkdir /wheels
 RUN uv pip install -U build cmake ninja pybind11 setuptools wheel
 
 FROM build-base AS build-triton
-ARG TRITON_REF=release/3.3.x
-ARG TRITON_BUILD_VERSION=3.3.0
+ARG TRITON_REF=release/3.2.x
+ARG TRITON_BUILD_VERSION=3.2.0
 ENV BUILD_VERSION=${TRITON_BUILD_VERSION:-${TRITON_REF#v}}
 RUN git clone https://github.com/triton-lang/triton.git
 RUN cd triton && \
@@ -73,8 +67,8 @@ RUN cd triton && \
     uv build python --wheel --no-build-isolation -o /wheels
 
 FROM build-base AS build-xformers
-ARG XFORMERS_REF=v0.0.29.post3
-ARG XFORMERS_BUILD_VERSION=0.0.29.post3
+ARG XFORMERS_REF=v0.0.29.post2
+ARG XFORMERS_BUILD_VERSION=0.0.29.post2
 ENV BUILD_VERSION=${XFORMERS_BUILD_VERSION:-${XFORMERS_REF#v}}
 RUN git clone  https://github.com/facebookresearch/xformers.git
 RUN cd xformers && \
@@ -96,15 +90,14 @@ RUN cd flashinfer && \
     uv build --wheel --no-build-isolation -o /wheels
 
 FROM build-base AS build-vllm
-ARG VLLM_REF=53be4a86
-ARG VLLM_BUILD_VERSION=0.7.4
+ARG VLLM_REF=v0.8.0
+ARG VLLM_BUILD_VERSION=0.8.0
 ENV BUILD_VERSION=${VLLM_BUILD_VERSION:-${VLLM_REF#v}}
 RUN git clone https://github.com/vllm-project/vllm.git
 RUN cd vllm && \
     git checkout ${VLLM_REF} && \
     git submodule sync && \
     git submodule update --init --recursive -j 8 && \
-    python use_existing_torch.py && \
     uv pip install -r requirements/build.txt && \
     uv build --wheel --no-build-isolation -o /wheels