@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dcfe95234c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 09:45 UTC

Update Dockerfile to build for Blackwell (#18095)

Michael Goin committed May 17, 2025 at 03:23 UTC dcfe95234c1154abec9144266c66d137b12d4689
1 file changed +3 -3
docker/Dockerfile
+3 -3
@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
77 # can be useful for both `dev` and `test`
78 # explicitly set the list to avoid issues with torch 2.2
79 # see https://github.com/pytorch/pytorch/pull/123243
80 -ARG torch_cuda_arch_list='7.0 7.5 8.0 8.6 8.9 9.0+PTX'
80 +ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0+PTX'
81 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
82 # Override the arch list for flash-attn to reduce the binary size
83 ARG vllm_fa_cmake_gpu_arches='80-real;90-real'
@@ -257,8 +257,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
257 if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \
258 # uv pip install --system https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.4/flashinfer_python-0.2.4+cu124torch2.6-cp38-abi3-linux_x86_64.whl ; \
259 # TESTING: install FlashInfer from source to test 2.7.0 final RC
260 - FLASHINFER_ENABLE_AOT=1 TORCH_CUDA_ARCH_LIST='7.5 8.0 8.6 8.9 9.0+PTX' \
261 - uv pip install --system --no-build-isolation "git+https://github.com/flashinfer-ai/flashinfer@v0.2.4" ; \
260 + FLASHINFER_ENABLE_AOT=1 TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0 10.0+PTX' \
261 + uv pip install --system --no-build-isolation "git+https://github.com/flashinfer-ai/flashinfer@e00e8cedbfcb220f328fd36aa8f529f869b01e6b" ; \
262 fi
263 COPY examples examples
264 COPY benchmarks benchmarks