feat: add vLLM ROCm variant for AMD GPUs

[doringeman]

make docker-run-vllm-rocm

MODEL_RUNNER_HOST=http://localhost:8080 docker model smollm2-vllm

[gemini-code-assist]

Summary of Changes

Hello @doringeman, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the model runner's compatibility by integrating vLLM support for AMD GPUs via ROCm. By introducing a new Docker variant and corresponding build/run commands, users with AMD hardware can now easily deploy and utilize vLLM, broadening the accessibility and utility of the project across different GPU architectures.

Highlights

• New vLLM ROCm Variant: Introduced a new Docker build variant for vLLM specifically tailored for AMD GPUs using ROCm, enabling broader hardware compatibility.
• Dockerfile Modifications: Modified the Dockerfile to include a dedicated vllm-rocm build stage, which installs vLLM with ROCm support from its specific wheel index.
• Makefile Updates: Updated the Makefile to define new variables for the ROCm base image and Docker image, and added new targets (docker-build-vllm-rocm, docker-run-vllm-rocm) for building and running this variant.
• Documentation for New Commands: Expanded the Makefile's help message to include instructions for the new ROCm-enabled vLLM build and run commands.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• Dockerfile
  • Added a new vllm-rocm build stage, inheriting from llamacpp, to specifically build vLLM with ROCm support.
  • Configured the vllm-rocm stage to install necessary Python packages and vllm version 0.15.1 from the ROCm-specific extra index URL (https://wheels.vllm.ai/rocm/).
  • Introduced a final-vllm-rocm stage to package the model-runner binary with the ROCm-enabled vLLM environment.
• Makefile
  • Declared new variables: VLLM_ROCM_BASE_IMAGE for the ROCm PyTorch base image and DOCKER_IMAGE_VLLM_ROCM for the new Docker image tag.
  • Extended the .PHONY target list to include docker-build-vllm-rocm and docker-run-vllm-rocm.
  • Added new docker-build-vllm-rocm and docker-run-vllm-rocm targets, enabling the build and execution of the vLLM ROCm Docker image.
  • Updated the help target to document the newly added ROCm-specific build and run commands.

Activity

• No human activity (comments, reviews, or progress updates) has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a vLLM ROCm variant for AMD GPU support by adding a new build stage in the Dockerfile and corresponding targets in the Makefile. The changes are generally sound, but I've identified a critical issue in the Makefile that will likely cause the build to fail due to an incorrect variable. Additionally, I've suggested a refactoring in the Dockerfile to improve maintainability by reducing code duplication. Addressing these points will make the implementation more robust and easier to manage.

[gemini-code-assist]

This LLAMA_SERVER_VARIANT value will likely cause the build to fail. The Dockerfile uses this variable to construct the path to the llama.cpp binary inside the llama-server image. Based on the llamacpp/native/rocm.Dockerfile provided in the context, the path includes the ROCm version (e.g., ...rocm7.0.2...), but here you are setting the variant to just rocm.

This will result in a path mismatch and a COPY failure in the Dockerfile.

To fix this, the LLAMA_SERVER_VARIANT should include the ROCm version. Assuming the version is 7.0.2 as per rocm.Dockerfile, this line should be updated.

		LLAMA_SERVER_VARIANT=rocm7.0.2 \

[gemini-code-assist]

There's significant duplication between this new vllm-rocm stage and the existing vllm stage (lines 86-112). Many of the setup steps, such as installing system dependencies, creating directories, and installing uv, are identical.

To improve maintainability and reduce code duplication, consider creating a common base stage (e.g., vllm-base) from which both vllm and vllm-rocm can inherit. This base stage would handle the common setup, and the specific vllm and vllm-rocm stages would only contain the logic for installing their respective versions of the vllm package.

For example:

# --- vLLM base ---
FROM llamacpp AS vllm-base
USER root
RUN apt update && apt install -y python3 python3-venv python3-dev curl ca-certificates build-essential && rm -rf /var/lib/apt/lists/*
RUN mkdir -p /opt/vllm-env && chown -R modelrunner:modelrunner /opt/vllm-env
USER modelrunner
RUN curl -LsSf https://astral.sh/uv/install.sh | sh

# --- vLLM variant ---
FROM vllm-base AS vllm
# ... vllm specific install
...

# --- vLLM ROCm variant ---
FROM vllm-base AS vllm-rocm
# ... vllm-rocm specific install
...

This would make the Dockerfile cleaner and easier to manage as more variants are added.

[ericcurtin]

It's a fair point about the duplication here

[ericcurtin]

vllm tends to pull it's own version of everything on install in the venv, etc. We might be able to use plain-old ubuntu for all vLLM images. Llama.cpp is a bit different, need to install the dependancies ourselves

[ericcurtin]

llama.cpp and vLLM can be on different ROCm versions if needs be. Hard to avoid that, we are moving to the one container image per backend model anyway.