-
-
Notifications
You must be signed in to change notification settings - Fork 8.7k
[WIP] Draft to remove torch.cuda calls and use Platform APIs. #20721
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Summary of Changes
Hello @jikunshang, I'm Gemini Code Assist1! 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 initiates a significant refactoring effort to abstract direct torch.cuda
calls behind a new Platform
API. The primary goal is to enhance multi-platform compatibility, particularly for XPU devices, by centralizing device-specific operations. This change makes the codebase more modular and easier to extend to new hardware backends.
Highlights
- Platform Abstraction: Introduced new abstract methods (
empty_cache
,reset_peak_memory_stats
,mem_get_info
,memory_stats
,memory_reserved
,synchronize
) to thePlatform
interface (vllm/platforms/interface.py
) to standardize device-specific operations across different hardware backends. - CUDA and XPU Implementations: Provided concrete implementations for these new platform methods in
vllm/platforms/cuda.py
(wrappingtorch.cuda
calls) andvllm/platforms/xpu.py
(wrappingtorch.xpu
calls), enabling consistent interaction with device functionalities. - Migration to Platform API: Replaced numerous direct
torch.cuda
calls with calls to thecurrent_platform
API acrossvllm/utils/__init__.py
,vllm/v1/worker/gpu_model_runner.py
, andvllm/v1/worker/gpu_worker.py
for memory management, device synchronization, and device property retrieval, centralizing device interactions. - XPU Worker Unification: Updated
vllm/platforms/xpu.py
to use the genericvllm.v1.worker.gpu_worker.Worker
class, streamlining the worker implementation for XPU devices and reducing code duplication. - Platform-Agnostic Device Initialization: Modified
vllm/v1/worker/gpu_worker.py
to usecurrent_platform.device_name
for creatingtorch.device
instances, further decoupling device initialization from specific CUDA calls and improving multi-platform compatibility.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command>
or @gemini-code-assist <command>
. Below is a summary of the supported commands.
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 issue 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 is currently in preview and 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 to provide feedback.
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
-
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. ↩
This pull request has merge conflicts that must be resolved before it can be |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request effectively refactors device-specific torch.cuda
calls to use a new Platform
API, which is a great step towards supporting multiple hardware backends. The changes are consistent and well-abstracted in most places.
I've identified one high-severity issue in vllm/v1/worker/gpu_worker.py
where platform-specific setup logic is hardcoded for CUDA and XPU, which could be incorrect for XPU and is not maintainable. I've suggested moving this logic into the respective Platform
classes for better abstraction and correctness.
if self.device_config.device.type == "cuda" or \ | ||
self.device_config.device.type == "xpu": |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This condition hardcodes device types 'cuda' and 'xpu', which is not easily extensible to other platforms like ROCm. The environment variables being set within this block (TORCH_NCCL_AVOID_RECORD_STREAMS
and NCCL_ASYNC_ERROR_HANDLING
) are specific to NCCL, which is used by the CUDA backend. It's not clear if these are applicable or correct for the XPU backend, which uses ccl
or xccl
. Applying NCCL-specific workarounds to other platforms could lead to unexpected behavior or bugs.
A better approach would be to abstract this platform-specific setup into the Platform
classes. For example, you could add an on_worker_init()
method to the Platform
interface and implement it in CudaPlatform
and XpuPlatform
with their respective environment variable settings. This would make the code more modular, maintainable, and less prone to errors when adding new hardware backends.
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Essential Elements of an Effective PR Description Checklist
supported_models.md
andexamples
for a new model.Purpose
Test Plan
Test Result
(Optional) Documentation Update