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

Signed-off-by: mgoin <mgoin64@gmail.com>
2025-07-11 01:16:01 +00:00
parent b854321ffe
commit cf75cd2098
1 changed files with 2 additions and 1 deletions
@@ -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