[CI Bugfix] Specify same TORCH_CUDA_ARCH_LIST for flashinfer aot and install (#20772)

Signed-off-by: mgoin <mgoin64@gmail.com>

[CI Bugfix] Specify same TORCH_CUDA_ARCH_LIST for flashinfer aot and install (#20772)
Signed-off-by: mgoin <mgoin64@gmail.com>
cf75cd20 · Michael Goin · GitHub · b854321f · cf75cd20
Unverified Commit cf75cd20 authored Jul 11, 2025 by Michael Goin Committed by GitHub Jul 11, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

docker/Dockerfile docker/Dockerfile +2 -1

No files found.
--- 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