From e35b14e7081d5702466c2a797e4436c98a285c88 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniele=20Trifir=C3=B2?= <dtrifiro@redhat.com>
Date: Thu, 12 Sep 2024 19:02:02 +0200
Subject: [PATCH] Dockerfile.ubi: bump flashinfer, use correct version

---
 Dockerfile.ubi | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/Dockerfile.ubi b/Dockerfile.ubi
index f22f6b3520b62..0819c08164e59 100644
--- a/Dockerfile.ubi
+++ b/Dockerfile.ubi
@@ -180,7 +180,11 @@ 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 \
+        env \
+        python_tag=$(echo ${PYTHON_VERSION} | sed 's/\.//') \
+        flashinfer_version="0.1.6" \
+        https://github.com/flashinfer-ai/flashinfer/releases/download/v${flashinfer_version}/flashinfer-${flashinfer_version}+cu124torch2.4-${python_tag}-${python_tag}-linux_x86_64.whl
 
 ENV HF_HUB_OFFLINE=1 \
     HOME=/home/vllm \