From ef0f996df62c50992079237ea63eab29e91cbdaa Mon Sep 17 00:00:00 2001
From: Eugene Rakhmatulin <eugr@eugr.com>
Date: Thu, 29 Jan 2026 23:14:43 -0800
Subject: [PATCH] Bumped base image version; reverted Triton to 3.5.1

---
 Dockerfile        | 6 +++---
 Dockerfile.mxfp4  | 3 ++-
 Dockerfile.wheels | 2 +-
 build-and-copy.sh | 2 +-
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index cda048c..5dc3f1e 100644
--- a/Dockerfile
+++ b/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
diff --git a/Dockerfile.mxfp4 b/Dockerfile.mxfp4
index 11af230..b528f8d 100644
--- a/Dockerfile.mxfp4
+++ b/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"
diff --git a/Dockerfile.wheels b/Dockerfile.wheels
index 77dd45f..e287f1c 100644
--- a/Dockerfile.wheels
+++ b/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
diff --git a/build-and-copy.sh b/build-and-copy.sh
index 2030909..0b45eb2 100755
--- a/build-and-copy.sh
+++ b/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