v1/offloading: Add worker-side CPU support #21448
Conversation
orozery
requested review from
WoosukKwon,
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
as code owners
|
👋 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 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces worker-side support for CPU offloading, a significant feature enhancement. The changes include new abstractions for offloading, a worker queue manager for handling asynchronous transfers, and the necessary CPU-specific logic for tensor creation and data movement. The implementation is accompanied by a comprehensive set of tests that cover both the transfer logic and the asynchronous worker management.
My review has identified a couple of high-severity issues. One is a misleading docstring in the OffloadingManager abstract class that could lead to incorrect implementations. The other is the use of __del__ for resource cleanup in OffloadingQueueManager, which is unreliable and could lead to resource leaks. Addressing these points will improve the robustness and maintainability of the new offloading framework. Overall, this is a well-structured contribution.
This commit adds a new offloading component, composed of: 1. A scheduler side OffloadingManager (abstract) which kicks-off KV data transfers and keeps track of offloaded data. 2. A worker side OffloadingQueueManager which asynchronously manages KV transfers. Signed-off-by: Or Ozeri <oro@il.ibm.com>
orozery
force-pushed
the
cpu-offloading-worker
branch
from
4a19c04 to
a625bee
Compare
This commit adds worker-side support for CPU offloading. It uses the swap_blocks function to perform the actual copying between GPU and CPU. Supports any CPU block size which is divided by the GPU block size. Signed-off-by: Or Ozeri <oro@il.ibm.com>
orozery
force-pushed
the
cpu-offloading-worker
branch
from
a625bee to
86a96a4
Compare
This PR adds worker-side support for CPU offloading.
It uses the
swap_blocksfunction to perform the actual copying between CPU and GPU.Supports any
cpu_block_sizewhich is divided bygpu_block_size.Part of the work described in RFC #19854
Depends on #19848