Cách build Docker image cho vLLM hỗ trợ CUDA 12.8 và GPU RTX 5090 (SM120)

Phiên bản nightly v0.8.3rc2.dev172+g3cdc57669

Ban đầu thì mình xài Ollama cho AI Inference kết hợp với Open WebUI. Em này cài là chạy, không phải suy nghĩ nhiều. Nhưng giờ do cần chạy dual RTX 5090 nên phải kiếm giải pháp thay thế. Hỏi em ChatGPT Lê Ky (tên mình đặt cho ẻm) thì em ấy gợi ý xài hàng top vLLM. Mình sẽ có bài so sánh 2 em này sau.

Ngặt nỗi vLLM chưa support chính thức cho dòng Blackwell. Thế là mất gần 1 tuần mò mẫm tâm sự với ChatGPT để build thành công docker (phiên bản nightly: v0.8.3rc2.dev172+g3cdc57669) với:

• CUDA 12.8

• PyTorch nightly 2.8.0.dev20250413+cu128

• Hỗ trợ kiến trúc SM120 cho RTX 5090

🚩 Chuẩn bị môi trường

Bước 1: Pull Docker container nền được chính NVIDA recommend

Xem thêm: https://github.com/vllm-project/vllm/issues/14452

docker pull nvcr.io/nvidia/pytorch:25.02-py3

Bước 2: Khởi động container để build

docker run -it --gpus all -v D:\Projects\vLLM:/vllm-workspace --shm-size=64g --name vllm-cu128-build nvcr.io/nvidia/pytorch:25.02-py3

🚧 Các bước build bên trong Docker container

Bước 3: Kiểm tra và cập nhật PyTorch nightly (CUDA 12.8)

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 --force-reinstall --no-deps --break-system-packages

• pip install --pre torch

  • --pre (pre-release): Cho phép cài đặt phiên bản nightly (bản thử nghiệm trước khi phát hành chính thức) của PyTorch. Các phiên bản nightly chứa những cập nhật và tính năng mới nhất chưa có trong phiên bản stable.

• --index-url https://download.pytorch.org/whl/nightly/cu128

  • Chỉ định nguồn tải package PyTorch nightly build, đặc biệt hỗ trợ CUDA 12.8 (cu128).

  • Sử dụng đường dẫn này nhằm đảm bảo PyTorch cài đặt sẽ tương thích hoàn toàn với phiên bản CUDA 12.8, tối ưu hóa cho GPU NVIDIA RTX 5090 (kiến trúc SM120).

• --force-reinstall

  • Buộc cài đặt lại PyTorch ngay cả khi phiên bản tương tự đã tồn tại trên hệ thống, nhằm đảm bảo chắc chắn bạn có phiên bản mới nhất (nightly) mà không bị xung đột với bản trước đó.

• --no-deps

  • Không tự động cài các package phụ thuộc khác đi kèm với PyTorch, giúp bạn chủ động kiểm soát phiên bản các thư viện đi kèm như CUDA Toolkit, torchvision, torchaudio, v.v. Điều này đặc biệt hữu ích trong môi trường đang build tùy chỉnh hoặc môi trường đang sử dụng các phiên bản đặc biệt.

• --break-system-packages

  • Cho phép pip cài đặt package vào thư mục hệ thống (system-wide) ngay cả khi điều này có thể làm hỏng hoặc xung đột với các package hệ thống đã có trước đó.

  • Tham số này được sử dụng khi bạn cài đặt vào môi trường gốc của container hoặc hệ thống Linux, không phải môi trường ảo hóa như virtualenv hay conda.

  • Lưu ý quan trọng: Đây là một tùy chọn mạnh, chỉ nên sử dụng khi bạn đã biết chắc chắn mình đang làm gì và chấp nhận các rủi ro về tính ổn định của hệ thống. Đối với môi trường Docker thì lựa chọn này ít rủi ro hơn, do container dễ dàng khởi tạo lại nếu xảy ra vấn đề.

• Sau khi thực hiện lệnh này, bạn sẽ có phiên bản PyTorch nightly mới nhất với khả năng hỗ trợ đầy đủ CUDA 12.8 (cu128), sẵn sàng để build và chạy các mô hình AI (ví dụ vLLM) trên GPU RTX 5090 một cách tối ưu nhất.

Bước 4: Thiết lập biến môi trường để build đúng kiến trúc RTX 5090 (SM120)

export TORCH_CUDA_ARCH_LIST="12.0"
export CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=120 -DVLLM_TARGET_DEVICE=cuda"
export CMAKE_GENERATOR=Ninja
export PATH="/usr/local/bin:$PATH"
export CMAKE_BUILD_PARALLEL_LEVEL=4

Giải thích các biến môi trường:

• TORCH_CUDA_ARCH_LIST="12.0"

  • Chỉ định phiên bản kiến trúc CUDA (CUDA Compute Capability) cụ thể của GPU cần hỗ trợ. RTX 5090 sử dụng kiến trúc CUDA Compute Capability là SM120, nên giá trị này được đặt là 12.0.

• CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=120 -DVLLM_TARGET_DEVICE=cuda"

  • Biến này cung cấp các tham số cho quá trình build bằng CMake:

    • -DCMAKE_CUDA_ARCHITECTURES=120: Hướng dẫn CMake build riêng cho kiến trúc SM120, tối ưu cho RTX 5090.

    • -DVLLM_TARGET_DEVICE=cuda: Chỉ định thiết bị mục tiêu là GPU NVIDIA CUDA.

• CMAKE_GENERATOR=Ninja

  • Ninja là trình build nhanh hơn và hiệu quả hơn so với trình mặc định của CMake (Make). Thiết lập này yêu cầu hệ thống có Ninja đã cài đặt.

• PATH="/usr/local/bin:$PATH"

  • Đảm bảo trình biên dịch và các công cụ phụ trợ nằm trong thư mục /usr/local/bin được sử dụng ưu tiên hơn các thư mục khác trong biến PATH.

• CMAKE_BUILD_PARALLEL_LEVEL=4

  • Điều chỉnh số lượng luồng sử dụng trong quá trình biên dịch CMake. Bạn có thể tăng hoặc giảm con số này tùy thuộc vào số lượng CPU cores mà hệ thống của bạn có để tối ưu thời gian build.

Bước 5: Xoá cache build cũ (nếu có)

cd /vllm-workspace
rm -rf build dist vllm.egg-info

Bước 6: Build vLLM từ source

CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=120" \
python3 -m pip install . --force-reinstall --no-deps --no-build-isolation

Giải thích lệnh cài đặt:

• CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=120"

  • Đảm bảo pip sử dụng tham số này khi build package, giúp biên dịch thư viện (trong trường hợp này là vLLM) chính xác với kiến trúc CUDA của RTX 5090.

• python3 -m pip install .

  • Cài đặt package từ thư mục hiện tại (đã clone source code của vLLM vào đây).

• --force-reinstall

  • Cài đặt lại package ngay cả khi đã có phiên bản cũ được cài đặt trước đó. Điều này đảm bảo phiên bản build mới nhất được cài đặt thành công.

• --no-deps

  • Không tự động cài đặt các dependencies khác. Việc này giúp bạn quản lý riêng các dependencies, tránh các lỗi phát sinh do pip cài sai phiên bản hoặc không tương thích.

• --no-build-isolation

  • Tắt chế độ build isolation, cho phép sử dụng các gói Python đã được cài sẵn trên hệ thống (đặc biệt hữu ích khi dùng với các phiên bản nightly build hoặc tự build như PyTorch nightly 2.8.0 cu128).

Bước 7: Cài đặt thêm các dependency (ví dụ openai)

pip install openai

Bước 8: Commit Docker container vừa build thành image để sử dụng lại

docker commit <container_id> vllm-cu128-ready

🚀 Chạy Docker vLLM với multi-GPU

Bước 9: Chạy container vLLM Docker (ví dụ 2 GPU RTX 5090)

docker run -it --gpus all `
  -p 8000:8000 `
  --shm-size=64g `
  --name vllm-cu128-multi-gpu `
  vllm-cu128-ready:latest `
  python3 -m vllm.entrypoints.api_server `
    --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 `
    --host 0.0.0.0 `
    --port 8000 `
    --tensor-parallel-size 2

Bước 10: Kiểm tra inference

curl.exe -X POST "http://localhost:8000/generate" `
  -H "Content-Type: application/json" `
  -d "{\"prompt\":\"Hello, who are you?\",\"max_tokens\":10}"

✅Kết quả:

Sau khi thực hiện các bước trên, chúng ta sẽ có một bản build vLLM Docker container tối ưu hóa hoàn hảo cho GPU NVIDIA RTX 5090 với CUDA 12.8 và PyTorch 2.8.0 nightly build. Bạn có thể chạy thử nghiệm hoặc deploy model ngay lập tức.

Hướng dẫn này đã được thử nghiệm thực tế và hoạt động ổn định với hệ thống của mình sử dụng 2 card RTX 5090.