Bumped base image version; reverted Triton to 3.5.1

Dockerfile

@@ -6,7 +6,7 @@ ARG BUILD_JOBS=16
 # =========================================================
 # STAGE 1: Base Image (Installs Dependencies)
 # =========================================================
-FROM nvidia/cuda:13.1.0-devel-ubuntu24.04 AS base
+FROM nvidia/cuda:13.1.1-devel-ubuntu24.04 AS base
 
 # Build parallemism
 ARG BUILD_JOBS
@@ -102,7 +102,7 @@ RUN git clone https://github.com/triton-lang/triton.git
 
 # We expect TRITON_REF to be passed from the command line to break the cache
 # Set to v3.5.1 tag by default
-ARG TRITON_REF=v3.6.0
+ARG TRITON_REF=v3.5.1
 
 WORKDIR $VLLM_BASE_DIR/triton
 
@@ -193,7 +193,7 @@ RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
 # =========================================================
 # STAGE 4: Runner (Transfers only necessary artifacts)
 # =========================================================
-FROM nvidia/cuda:13.1.0-devel-ubuntu24.04 AS runner
+FROM nvidia/cuda:13.1.1-devel-ubuntu24.04 AS runner
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PIP_BREAK_SYSTEM_PACKAGES=1

Dockerfile.mxfp4

@@ -258,7 +258,8 @@ RUN mkdir -p tiktoken_encodings && \
 # No need to copy source code, as it's already in the site-packages
 COPY --from=builder /workspace/wheels /workspace/wheels
 RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
-    uv pip install /workspace/wheels/*.whl
+    uv pip install /workspace/wheels/*.whl && \
+    rm -rf /workspace/wheels
 
 # Setup Env for Runtime
 ENV TORCH_CUDA_ARCH_LIST="12.0;12.1"

Dockerfile.wheels

@@ -1,6 +1,6 @@
 # syntax=docker/dockerfile:1.6
 
-FROM nvidia/cuda:13.1.0-devel-ubuntu24.04
+FROM nvidia/cuda:13.1.1-devel-ubuntu24.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PIP_BREAK_SYSTEM_PACKAGES=1

build-and-copy.sh

@@ -11,7 +11,7 @@ REBUILD_VLLM=false
 COPY_HOSTS=()
 SSH_USER="$USER"
 NO_BUILD=false
-TRITON_REF="v3.6.0"
+TRITON_REF="v3.5.1"
 VLLM_REF="main"
 TMP_IMAGE=""
 PARALLEL_COPY=false