From cf75cd2098f6a3f0bc38d92d1669810c084dab9b Mon Sep 17 00:00:00 2001
From: Michael Goin <mgoin64@gmail.com>
Date: Fri, 11 Jul 2025 10:16:01 +0900
Subject: [PATCH] [CI Bugfix] Specify same TORCH_CUDA_ARCH_LIST for flashinfer
 aot and install (#20772)

Signed-off-by: mgoin <mgoin64@gmail.com>
---
 docker/Dockerfile | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 9ef021687a3c8..6ae4f789f0593 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -408,7 +408,8 @@ RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
 
           # Needed to build AOT kernels
           pushd flashinfer
-            python3 -m flashinfer.aot
+            TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
+              python3 -m flashinfer.aot
             TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
               uv pip install --system --no-build-isolation .
           popd