break execute_model in gpu_model_runner into sub-functions for custom scopes #24265
Conversation
bangshengtang
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
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 You ask your reviewers to trigger select CI tests on top of 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 If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. 🚀 |
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
![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 refactors the execute_model function by breaking it down into smaller, more focused sub-functions, each wrapped in a profiling scope. This is a great improvement for code readability, maintainability, and performance profiling. The refactoring appears to be logically sound. I've identified a couple of issues with incorrect type hints in the new function signatures that should be addressed to ensure code correctness from a static analysis perspective.
vllm/v1/worker/gpu_model_runner.py
Outdated
There was a problem hiding this comment.
The return type hint for _preprocess is incorrect. It specifies 9 elements in the tuple, but the function returns 8. Additionally, several of the types are mismatched with the actual returned values (e.g., the 3rd element is Optional[torch.Tensor] but hinted as int). This should be corrected to match the returned tuple for type consistency and to help static analysis tools.
) -> tuple[
int,
int,
Optional[torch.Tensor],
Optional[torch.Tensor],
Optional[torch.Tensor],
torch.Tensor,
Optional[IntermediateTensors],
dict[str, Any],
]:
vllm/v1/worker/gpu_model_runner.py
Outdated
There was a problem hiding this comment.
The type hint for logprobs_lists in the return tuple is LogprobsLists, but the value can be None if sampler_output.logprobs_tensors is None. The type hint should be Optional[LogprobsLists] to accurately reflect this.
Optional[LogprobsLists],
bangshengtang
force-pushed
the
export-D81009244
branch
from
41d00e1 to
757588e
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
bangshengtang
force-pushed
the
export-D81009244
branch
from
757588e to
c18d76a
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
bangshengtang
force-pushed
the
export-D81009244
branch
from
61d77aa to
e823d15
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
bangshengtang
force-pushed
the
export-D81009244
branch
from
e823d15 to
2e9ed8d
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
bangshengtang
force-pushed
the
export-D81009244
branch
from
2e9ed8d to
e0f8874
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
bangshengtang
force-pushed
the
export-D81009244
branch
from
e0f8874 to
994669b
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
houseroad
left a comment
houseroad
left a comment
There was a problem hiding this comment.
I like this move. It will make the profiling result easier to parse. As by product, it also break long function into smaller pieces.
vllm/v1/worker/gpu_model_runner.py
Outdated
There was a problem hiding this comment.
later we can move this into something like _post_processing.
vllm/v1/utils.py
Outdated
There was a problem hiding this comment.
is there a reason not to turn it on by default? how much is the overhead? (given it's just nvtx context?)
There was a problem hiding this comment.
For use case for extreme perf, we would like to avoid additional overhead. Although, this should be light.
|
This pull request has merge conflicts that must be resolved before it can be |
… scopes (vllm-project#24265) Summary: break down the main execute_model() function into multiple parts so that we can have better latency breakdown in profiles: - preprocess - forward - postprocess - bookkeep (includes sync) - draft (if spec decoding is enabled) this is meant to be a refactor + some new function scopes, there's no functionality difference Test Plan: CI run the vLLM predictor runbook and take a trace {F1981500506} Reviewed By: houseroad, frank-wei Differential Revision: D81009244
bangshengtang
force-pushed
the
export-D81009244
branch
from
3f21393 to
9610aee
Compare
|
This pull request was exported from Phabricator. Differential Revision: D81009244 |
… scopes (vllm-project#24265) Co-authored-by: Bangsheng Tang <bangsheng@meta.com>
… scopes (vllm-project#24265) Co-authored-by: Bangsheng Tang <bangsheng@meta.com>
… scopes (vllm-project#24265) Co-authored-by: Bangsheng Tang <bangsheng@meta.com>
… scopes (vllm-project#24265) Co-authored-by: Bangsheng Tang <bangsheng@meta.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>
… scopes (vllm-project#24265) Co-authored-by: Bangsheng Tang <bangsheng@meta.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>
5 participants
Summary:
break down the main execute_model() function into multiple parts so that we can have better latency breakdown in profiles:
this is meant to be a refactor + some new function scopes, there's no functionality difference
Test Plan:
CI
run the vLLM predictor runbook and take a trace
Differential Revision: D81009244