[ROCm] [CI] Optimize Dockerfilerocm and reduce build time on CI (#885)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

.buildkite/test-template-amd-omni.j2

@@ -14,7 +14,7 @@
         soft_fail: false
         commands:
           - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --progress plain ."
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
         env:

docker/Dockerfile.rocm

@@ -1,9 +1,7 @@
-ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205
-FROM ${BASE_IMAGE}
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
+FROM ${BASE_IMAGE} AS final
 
 ARG COMMON_WORKDIR=/app
-ARG VLLM_VERSION=v0.14.0
-ARG PYTORCH_ROCM_ARCH="gfx942;gfx950"
 
 WORKDIR ${COMMON_WORKDIR}
 
@@ -13,30 +11,18 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Step 2: Reinstall vllm from source
-RUN python3 -m pip uninstall -y vllm && rm -rf vllm &&\
-    git clone https://github.com/vllm-project/vllm.git && \
-    cd vllm && \
-    git checkout ${VLLM_VERSION} && \
-    python3 -m pip install -r requirements/rocm.txt && \
-    python3 setup.py clean --all && \
-    PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
-    cd ../ && \
-    rm -rf vllm/.git
-
 RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 
-# Step 3: Copy vllm-omni code and install without uv
+# Step 2: Copy vllm-omni code and install without uv
 COPY . ${COMMON_WORKDIR}/vllm-omni
-RUN cd ${COMMON_WORKDIR}/vllm-omni && python3 -m pip install --no-cache-dir ".[dev]"
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
 
-# Create python symlink
-# `GPU_ARCHS` is an environment variable that is used to set the GPU archs for the AITER.
-# This is needed to prevent the AITER automatic GPU arch detection from failing on MI325X.
-# The AITER version used in this dockerfile has issues with handling
-# the GPU archs of MI325X (CI machine) correctly. So we manually set the GPU archs here.
-# We reuse AITER_ROCM_ARCH from the base image to avoid duplication.
-ENV GPU_ARCHS=${AITER_ROCM_ARCH}
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 
+CMD ["/bin/bash"]
+
 ENTRYPOINT []
+
+#Set entrypoint for vllm-openai official images
+FROM final AS vllm-openai
+ENTRYPOINT ["vllm", "serve"]

docs/getting_started/installation/gpu/rocm.inc.md

@@ -34,8 +34,28 @@ DOCKER_BUILDKIT=1 docker build \
 
 #### Launch the docker image
 
+##### Launch with OpenAI API Server
+
 ```
-docker run -it \
+docker run --rm \
+--group-add=video \
+--ipc=host \
+--cap-add=SYS_PTRACE \
+--security-opt seccomp=unconfined \
+--device /dev/kfd \
+--device /dev/dri \
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--env "HF_TOKEN=$HF_TOKEN" \
+-p 8091:8091 \
+--ipc=host \
+vllm-omni-rocm \
+--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+```
+
+##### Launch with interactive session for development
+
+```
+docker run --rm -it \
 --network=host \
 --group-add=video \
 --ipc=host \
@@ -44,8 +64,9 @@ docker run -it \
 --device /dev/kfd \
 --device /dev/dri \
 -v <path/to/model>:/app/model \
-vllm-omni-rocm \
-bash
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--entrypoint bash \
+vllm-omni-rocm
 ```
 
 # --8<-- [end:build-docker]
@@ -57,8 +78,7 @@ vLLM-Omni offers an official docker image for deployment. These images are built
 #### Launch vLLM-Omni Server
 Here's an example deployment command that has been verified on 2 x MI300's:
 ```bash
-docker run -it \
-  --network=host \
+docker run --rm \
   --group-add=video \
   --ipc=host \
   --cap-add=SYS_PTRACE \
@@ -68,14 +88,15 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
+  -p 8091:8091 \
   vllm/vllm-omni-rocm:v0.14.0 \
-  vllm serve --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
 #### Launch an interactive terminal with prebuilt docker image.
 If you want to run in dev environment you can launch the docker image as follows:
 ```bash
-docker run -it \
+docker run --rm -it \
   --network=host \
   --group-add=video \
   --ipc=host \
@@ -86,8 +107,8 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.14.0 \
-  bash
+  --entrypoint bash \
+  vllm/vllm-omni-rocm:v0.14.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -17,7 +17,13 @@ For installation on GPU from source:
 ```bash
 uv venv --python 3.12 --seed
 source .venv/bin/activate
+
+# On CUDA
 uv pip install vllm==0.14.0 --torch-backend=auto
+
+# On ROCm
+uv pip install vllm==0.14.0 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .