diff --git a/Dockerfile.rocm.ubi b/Dockerfile.rocm.ubi
index b9010c6afa217..17909b97f50b2 100644
--- a/Dockerfile.rocm.ubi
+++ b/Dockerfile.rocm.ubi
@@ -1,6 +1,6 @@
 ## Global Args #################################################################
 ARG BASE_UBI_IMAGE_TAG=9.4
-ARG PYTHON_VERSION=3.11
+ARG PYTHON_VERSION=3.12
 # Default ROCm ARCHes to build vLLM for.
 ARG PYTORCH_ROCM_ARCH="gfx908;gfx90a;gfx942;gfx1100"
 ARG MAX_JOBS=12
@@ -132,7 +132,6 @@ FROM rocm_devel AS build_triton
 # Whether to build triton on rocm
 ARG BUILD_TRITON="0"
 ARG TRITON_BRANCH="e0fc12c"
-ARG TRITON_WHEEL_SHA=21eae954ef
 
 # Build triton wheel if `BUILD_TRITON` is set to 1
 RUN --mount=type=cache,target=/root/.cache/ccache \
@@ -148,7 +147,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
         python3 setup.py bdist_wheel --dist-dir=/install; \
     else \
         mkdir -p /install && \
-        wget -P /install "https://download.pytorch.org/whl/nightly/pytorch_triton_rocm-3.0.0%2B${TRITON_WHEEL_SHA}-cp311-cp311-linux_x86_64.whl"; \
+        wget -P /install "https://download.pytorch.org/whl/pytorch_triton_rocm-3.0.0-cp312-cp312-linux_x86_64.whl#sha256=650ce7a6624ecbbe90c30fccb472a57772c08bc6f12312246f7af584f275b7ca"; \
     fi
 
 ##################################################################################################
@@ -179,7 +178,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
 
 FROM rocm_base AS vllm-openai
 ARG MAX_JOBS
-ARG PYTHON_VERSION
 
 WORKDIR /workspace
 
diff --git a/Dockerfile.ubi b/Dockerfile.ubi
index 3edd65fb1d368..ee9aaf4dd9636 100644
--- a/Dockerfile.ubi
+++ b/Dockerfile.ubi
@@ -1,15 +1,13 @@
 ## Global Args #################################################################
 ARG BASE_UBI_IMAGE_TAG=9.4
-ARG PYTHON_VERSION=3.11
+ARG PYTHON_VERSION=3.12
 
 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
 
 ## Base Layer ##################################################################
 FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} as base
-ARG PYTHON_VERSION
-
-RUN microdnf -y update && \
-    microdnf install -y \
+ENV PYTHON_VERSION=$PYTHON_VERSION
+ARG PYTHON_VERSION RUN microdnf -y update && \ microdnf install -y \
     python${PYTHON_VERSION}-pip python${PYTHON_VERSION}-wheel \
     && microdnf clean all
 
@@ -26,13 +24,13 @@ RUN microdnf install -y \
 
 ## Python Installer ############################################################
 FROM base as python-install
-
 ARG PYTHON_VERSION
 
 ENV VIRTUAL_ENV=/opt/vllm
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+ENV PYTHON_VERSION=$PYTHON_VERSION
 RUN microdnf install -y \
-    python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip python${PYTHON_VERSION}-wheel && \
+    python${PYTHON_VERSION}-devel  && \
     python${PYTHON_VERSION} -m venv $VIRTUAL_ENV && pip install --no-cache -U pip wheel uv && microdnf clean all
 
 
@@ -151,7 +149,6 @@ RUN CFLAGS="-O3 -Wall -Werror=format-security -Wno-unused-function -Wp,-D_GLIBCX
 
 ## Release #####################################################################
 FROM python-install AS vllm-openai
-ARG PYTHON_VERSION
 
 WORKDIR /workspace
 
@@ -180,7 +177,8 @@ RUN --mount=type=bind,from=libsodium-builder,src=/usr/src/libsodium,target=/usr/
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=cache,target=/root/.cache/uv \
-    uv pip install https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.2/flashinfer-0.1.2+cu121torch2.4-cp311-cp311-linux_x86_64.whl
+    uv pip install \
+        "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.6/flashinfer-0.1.6+cu124torch2.4-cp312-cp312-linux_x86_64.whl"
 
 ENV HF_HUB_OFFLINE=1 \
     HOME=/home/vllm \