[Hardware][AMD][CI/Build][Doc] Upgrade to ROCm 6.1, Dockerfile improvements, test fixes

[mawong-amd]

This PR does the following:

1. Upgrade the default base image to ROCm 6.1.2 with the correct fixes applied, and the default PyTorch version to 2.4.0, bringing it in line with CUDA's update to 2.3.0.
2. Refactor to enable multi-stage builds in Dockerfile.rocm: flash-attn and Triton are now built in separate stages as wheels which are then installed inside the final image. vLLM will eventually also be built in a similar manner once AMD CI is compatible. This improves container build time since component stages can be built in parallel.
3. Restore the use of PYTORCH_ROCM_ARCH as the first point of reference for what ARCHes to build for. This has not been working correctly since the transition to a CMake based system in Cmake based build system #2830, resulting in issues such as that mentioned in [Bugfix][CI/Build][AMD][ROCm]Fixed the cmake build bug which generate garbage on certain devices #5641.
4. Improve miscellaneous bits of Dockerfile logic: for instance, better detection of ROCm version (as opposed to matching BASE_IMAGE exactly), and use of setup.py develop instead of setup.py install which sidesteps the module-shadowing issues that resulted in the need to copy over the vLLM C++ libraries.
5. Other small miscellaneous documentation fixes.
6. Fix a bunch of tests that were broken with upstream changes.

This PR contains the fixes in #5641 in a different form and can be merged after or without it.

---

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

• [Bugfix] for bug fixes.
• [CI/Build] for build or continuous integration improvements.
• [Doc] for documentation fixes and improvements.
• [Model] for adding a new model or improving an existing model. Model name should appear in the title.
• [Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
• [Kernel] for changes affecting CUDA kernels or other compute kernels.
• [Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
• [Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
• [Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

• We adhere to Google Python style guide and Google C++ style guide.
• Pass all linter checks. Please use format.sh to format your code.
• The code need to be well-documented to ensure future contributors can easily understand the code.
• Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
• Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

• After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
• After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
• After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
• Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

[WoosukKwon]

@mawong-amd Thanks for the PR! This is wonderful!

@hongxiayang How would like to merge these two PRs?

[hongxiayang]

Reusing PYTORCH_ROCM_ARCH environment variable in VLLM environment exposes some potential problems. This environment variable was used by pytorch on rocm when building its rocm backend (https://github.com/pytorch/pytorch/blob/0707811286d1846209676435f4f86f2b4b3d1a17/cmake/Dependencies.cmake#L1075). Reuse this will hide this information, and may potentially make the debugging more difficult.

[mawong-amd]

I'm not really sure if this is a concern:

1. This just restores the behavior to what it was prior to Cmake based build system #2830
2. We also now install torch wheels, so PYTORCH_ROCM_ARCH in the base image is decoupled from whatever "Pytorch used when building its ROCm backend"
3. The debugging info provided by this environment variable seems to be very limited.

[hongxiayang]

Yes, I am concerned about the name clash. Can you change it to VLLM_ROCM_SUPPORTED_ARCHS?
Though the dockerfile in this PR installed wheels, the wheels was built using the same environment variable and same logic.

[mawong-amd]

It was working fine prior to #2830 and the workflow people are used to is to specify PYTORCH_ROCM_ARCH to decide the architectures to build vLLM for: this has been the case at ROCm/vllm and here, just that it hasn't been working here in a while.. I haven't yet seen a good reason to change this workflow. If this change is made, it should be made in a separate PR with proper discussion of the change.

Also, if the wheels are all built using the same env var, then it can be easily remembered as a constant.

[mawong-amd]

@WoosukKwon CI is looking all green now other than Documentation, which has been broken for a while, and the AMD Vision Language Models test, which seems to be due to a slow model download hitting the 90 min timeout (machine issue, seems to have originated sometime in the last couple weeks). It passes locally with a faster network, FWIW.

I've also fixed a bunch of tests that were broken due to:

1. [CI/Build] Simplify OpenAI server setup in tests #5100 initializing working directory for ray workers. Reverting this doesn't seem to have affected the tests.
2. Add cuda_device_count_stateless #5473 addition of a test that uses CUDA_VISIBLE_DEVICES. Fixing this requires updating torch to 2.4.0 and installing amdsmi for stateless detection of devices, along with using HIP_VISIBLE_DEVICES.

Can we land this?

[WoosukKwon]

@mawong-amd Thanks for the PR and sorry for the delays in the review. This is big!!

I'm generally good for the changes, but have some questions. Please check out my comments.

Also, I think we can merge this PR faster if you can split the PR into smaller ones. While the PR size is not big, I feel it touches quite many parts of the code. I'd like to integrate some of the changes first, and discuss the rest of the changes more deeply. For example, you could split the PR into 4 smaller PRs:

1. Update PagedAttention hyperparameters
2. Fix the docs
3. Refactor the docker file
4. Update the docker file with torch 2.4.0 & Introduce HIP_VISIBLE_DEVICES

[WoosukKwon]

What about only exposing CUDA_VISIBLE_DEVICES to the vLLM users and we internally convert it to HIP_VISIBLE_DEVICES when using the AMD GPUs?

[mawong-amd]

I've changed places that explicitly set CUDA_VISIBLE_DEVICES in vLLM to use an abstraction layer that copies over changes to HIP_VISIBLE_DEVICES.

We can revisit how we use CUDA_VISIBLE_DEVICES generally in a follow-up PR. Copying over any changes to CUDA_VISIBLE_DEVICES to HIP_VISIBLE_DEVICES is not portable in general, e.g. if it's set outside vLLM, or handling the cases where CUDA_VISIBLE_DEVICES is deleted (we don't need to handle this for now as we never set either env var in vLLM on ROCm).

[mawong-amd]

In general we cannot expect ROCm to recognize CUDA_VISIBLE_DEVICES. Unfortunately we currently don't have a unified interface for setting CUDA_VISIBLE_DEVICES: each use case (Ray executor, MP executor, custom all reduce, tests, etc) does it in its own way, so there will need to be quite a bit of refactoring, and even then, I'm not sure if there's a better solution than just always making sure that changes to CUDA_VISIBLE_DEVICES <=> changes to HIP_VISIBLE_DEVICES.

[mawong-amd]

@mawong-amd Thanks for the PR and sorry for the delays in the review. This is big!!

I'm generally good for the changes, but have some questions. Please check out my comments.

Also, I think we can merge this PR faster if you can split the PR into smaller ones. While the PR size is not big, I feel it touches quite many parts of the code. I'd like to integrate some of the changes first, and discuss the rest of the changes more deeply. For example, you could split the PR into 4 smaller PRs:

1. Update PagedAttention hyperparameters
2. Fix the docs
3. Refactor the docker file
4. Update the docker file with torch 2.4.0 & Introduce HIP_VISIBLE_DEVICES

Hi @WoosukKwon, thanks for your review! The top priority for us is updating the container to ROCm 6.1 (i.e. item 3), so I'll work on that PR first. I think it makes sense to combine 2, 3, and 4. For 2, the doc changes mostly relate to how we build the container. And we have to do 4 before or at the same time as 3, otherwise we will have failing tests of unclear provenance while upgrading to 6.1. Timing is somewhat important here: the longer it takes to do the 6.1 upgrade, the more upstream tests will break on AMD, especially since CI is currently soft-fail which means tests that will obviously fail like #5473 can sneak through. Then the more such tests we have to fix before we can do the upgrade.

[hongxiayang]

1. Update PagedAttention hyperparameters

I agree about pulling out the change for "NUM_THREADS = 1024" to a different pr. As we discussed previously, do you know how this change impact non-mi300x uses?

[DarkLight1337]

@DarkLight1337 It seems runtime_env={"working_dir": VLLM_PATH} was introduced in your PR #5100. Could you confirm that this change is safe and explain the reason if not?

It's unsafe to remove. Otherwise, some vLLM modules may become unavailable to ray workers during testing.

[mawong-amd]

@DarkLight1337 It seems runtime_env={"working_dir": VLLM_PATH} was introduced in your PR #5100. Could you confirm that this change is safe and explain the reason if not?

It's unsafe to remove. Otherwise, some vLLM modules may become unavailable to ray workers during testing.

Can you point out where this happens? Tests were working fine before your PR, and I haven't yet seen any errors from removing it.

[DarkLight1337]

@DarkLight1337 It seems runtime_env={"working_dir": VLLM_PATH} was introduced in your PR #5100. Could you confirm that this change is safe and explain the reason if not?

It's unsafe to remove. Otherwise, some vLLM modules may become unavailable to ray workers during testing.

Can you point out where this happens? Tests were working fine before your PR, and I haven't yet seen any errors from removing it.

I'm not an expert in ray, but I believe that it happened because ray.init is called inside test directory rather than the main vLLM package directory. Not sure why this problem doesn't happen now. I guess you can remove it as long as it doesn't break Ci.

[mawong-amd]

@WoosukKwon Rebased and fixed a couple more conflicts, including adding AMD support for GPU memory clear (amdsmi in place of nvml) to fix distributed tests, now we have an all-green CI again!

[WoosukKwon]

@mawong-amd LGTM! This is big. Huge thanks for your effort!

[WoosukKwon]

@hongxiayang While I understand your concerns mentioned in this PR, I will merge the PR as it includes some critical bug fixes and improvements. However, please feel free to open a new PR to continue the discussion.

[hongxiayang]

@hongxiayang While I understand your concerns mentioned in this PR, I will merge the PR as it includes some critical bug fixes and improvements. However, please feel free to open a new PR to continue the discussion.

Thanks @WoosukKwon for reviewing the PR and merge it.