[MISC] Move bind_kv_cache to worker module #20900
Conversation
Signed-off-by: wangxiyuan <[email protected]>
wangxiyuan
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @wangxiyuan, 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 focuses on optimizing the vLLM initialization process by strategically relocating the bind_kv_cache function. By moving this function and its associated model_executor dependencies from a general utility module to a worker-specific one, the PR aims to prevent the loading of unnecessary modules during the initial LLM import, thereby improving startup efficiency.
Highlights
- Code Relocation: The
bind_kv_cachefunction, along with its related imports, has been moved fromvllm/v1/utils.pytovllm/v1/worker/utils.py. - Import Optimization: This refactoring allows for the removal of
model_executorrelated imports (such asextract_layer_indexandAttentiontype hints) from the generalvllm.v1.utilsmodule. These dependencies are now only loaded within the worker context wherebind_kv_cacheis actually utilized. - Startup Performance Improvement: The primary benefit of this change is to reduce unnecessary module imports during the initial
LLMsetup, which is expected to lead to a leaner and potentially faster initialization ofprocess0.
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. ↩
There was a problem hiding this comment.
Code Review
This pull request refactors the bind_kv_cache function by moving it from vllm.v1.utils to vllm.v1.worker.utils, improving import efficiency. I've identified one area for improvement in the newly moved function to enhance its API and efficiency for one of its callers. My feedback is focused on making runner_kv_caches an optional parameter.
| def bind_kv_cache( | ||
| kv_caches: dict[str, torch.Tensor], | ||
| forward_context: dict[str, "Attention"], | ||
| runner_kv_caches: list[torch.Tensor], | ||
| ) -> None: | ||
| """ | ||
| Bind the allocated KV cache to both ModelRunner and forward context so | ||
| that the KV cache can be used in the forward pass. | ||
| This function: | ||
| 1) Fills the ModelRunner's kv cache list (`runner_kv_caches`) with | ||
| kv_caches. | ||
| 2) Associates each attention layer in the `forward_context` with its | ||
| corresponding KV cache in kv_caches. | ||
| Args: | ||
| kv_caches: The allocated kv_caches with layer names as keys. | ||
| forward_context: The global forward context containing all Attention | ||
| layers with layer names as keys. | ||
| runner_kv_caches: The kv_cache declared by ModelRunner. | ||
| """ | ||
| # Bind kv_caches to ModelRunner | ||
| assert len(runner_kv_caches) == 0 | ||
|
|
||
| # Convert kv_caches dict to a list of tensors in the order of layer_index. | ||
| index2name = defaultdict(list) | ||
| for layer_name in kv_caches: | ||
| index2name[extract_layer_index(layer_name)].append(layer_name) | ||
|
|
||
| for layer_index in sorted(index2name.keys()): | ||
| layer_names = index2name[layer_index] | ||
| if len(layer_names) > 1: | ||
| # One typical case is encoder-decoder model, e.g., bart. | ||
| # The cross attention and self attention in the same decoder layer | ||
| # has different layer_name but the same layer_index. | ||
| raise NotImplementedError | ||
| layer_name = layer_names[0] | ||
| runner_kv_caches.append(kv_caches[layer_name]) | ||
|
|
||
| # Bind kv_caches to forward context | ||
| for layer_name, kv_cache in kv_caches.items(): | ||
| # NOTE: Use list because of v0 PP virtual engine. | ||
| forward_context[layer_name].kv_cache = [kv_cache] |
There was a problem hiding this comment.
Consider making runner_kv_caches an optional parameter to avoid unnecessary computation when the caller doesn't need it. This improves the function's cohesion and efficiency.
def bind_kv_cache(
kv_caches: dict[str, torch.Tensor],
forward_context: dict[str, "Attention"],
runner_kv_caches: Optional[list[torch.Tensor]] = None,
) -> None:
"""
Bind the allocated KV cache to both ModelRunner and forward context so
that the KV cache can be used in the forward pass.
This function:
1) Fills the ModelRunner's kv cache list (`runner_kv_caches`) with
kv_caches.
2) Associates each attention layer in the `forward_context` with its
corresponding KV cache in kv_caches.
Args:
kv_caches: The allocated kv_caches with layer names as keys.
forward_context: The global forward context containing all Attention
layers with layer names as keys.
runner_kv_caches: The kv_cache declared by ModelRunner.
"""
if runner_kv_caches is not None:
# Bind kv_caches to ModelRunner
assert len(runner_kv_caches) == 0
# Convert kv_caches dict to a list of tensors in the order of layer_index.
index2name = defaultdict(list)
for layer_name in kv_caches:
index2name[extract_layer_index(layer_name)].append(layer_name)
for layer_index in sorted(index2name.keys()):
layer_names = index2name[layer_index]
if len(layer_names) > 1:
# One typical case is encoder-decoder model, e.g., bart.
# The cross attention and self attention in the same decoder layer
# has different layer_name but the same layer_index.
raise NotImplementedError
layer_name = layer_names[0]
runner_kv_caches.append(kv_caches[layer_name])
# Bind kv_caches to forward context
for layer_name, kv_cache in kv_caches.items():
# NOTE: Use list because of v0 PP virtual engine.
forward_context[layer_name].kv_cache = [kv_cache]There was a problem hiding this comment.
This suggested change is out of the PR scope, we can do it in the next PR.
|
👋 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.
LGTM, just move bind_kv_cache from utils to worker.utils to avoid useless import, this is reasonable.
This PR also resolve the vllm-ascend downstream bug: vllm-project/vllm-ascend#1773 , the broken happened after #20061 (but it's a reasonable PR just had some impact on import order which impacts on vllm ascend patch order and make the patch no longer working indirectly.)
This issue will be resolved on vllm-ascend after torch deps upgrade to 2.7.1 vllm-project/vllm-ascend#1562 , but for now we need this PR.
This PR just move utils to reasonable postion, so I'm +1 on it.
|
cc @heheda12345 Would you mind taking a look? |
heheda12345
left a comment
heheda12345
left a comment
There was a problem hiding this comment.
LGTM! Thanks for finding this problem.
|
Does v0 bind_kv_cache cause similar problem? |
github-actions
bot
added
the
ready
|
@heheda12345 I assume v0 code will be removed soon. So I did not change it. |
This PR fixed the broken CI. It require vllm-project/vllm#20900 merged first. - vLLM version: v0.9.2 - vLLM main: vllm-project/vllm@e8cc53a Signed-off-by: wangxiyuan <[email protected]>
Signed-off-by: wangxiyuan <[email protected]> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: wangxiyuan <[email protected]>
Signed-off-by: wangxiyuan <[email protected]>
Signed-off-by: Konrad Zawora <[email protected]>
Signed-off-by: wangxiyuan <[email protected]> Signed-off-by: Jinzhen Lin <[email protected]>
Signed-off-by: wangxiyuan <[email protected]> Signed-off-by: Paul Pak <[email protected]>
Signed-off-by: wangxiyuan <[email protected]> Signed-off-by: Diego-Castan <[email protected]>
Signed-off-by: wangxiyuan <[email protected]>
This PR fixed the broken CI. It require vllm-project/vllm#20900 merged first. - vLLM version: v0.9.2 - vLLM main: vllm-project/vllm@e8cc53a Signed-off-by: wangxiyuan <[email protected]>
This PR fixed the broken CI. It require vllm-project/vllm#20900 merged first. - vLLM version: v0.9.2 - vLLM main: vllm-project/vllm@e8cc53a Signed-off-by: wangxiyuan <[email protected]>
3 participants
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
When vLLM is setup by called
from vllm import LLM, SamplingParams, a lots of uselessimportis called. One main step isfrom vllm.model_executor.models.utils import extract_layer_indexinvllm.v1.utils. It will import much files inmodel_executorwhich won't be used in process0.This PR move
bind_kv_cachetovllm.v1.worker.utilsto skip uselessimportwhen initializeLLM.Test Plan
Just a code cleanup. No new test is needed.
Test Result
All test should pass
(Optional) Documentation Update