Support mtp run in full graph mode #3903
Conversation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![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 adds support for running MTP models in full graph mode on Ascend hardware, which should improve performance. The changes adapt the MTP model with full graph capture and execution. The implementation introduces a new set of graph parameters and functions for MTP, which are duplicates of existing ones. My main feedback is to refactor this duplicated code to improve maintainability.
vllm_ascend/compilation/acl_graph.py
Outdated
| @dataclass | ||
| class MTPGraphParams: | ||
| events: dict[int, list[torch.npu.ExternalEvent]] | ||
| workspaces: dict[int, torch.Tensor] | ||
| handles: dict[int, list[torch_npu._C._NPUTaskGroupHandle]] | ||
| attn_params: dict[int, list[tuple]] | ||
|
|
||
|
|
||
| _mtp_graph_params: Optional[MTPGraphParams] = None | ||
|
|
||
|
|
||
| def set_mtp_graph_params(aclgraph_capture_sizes: set[int]): | ||
| global _mtp_graph_params | ||
| if _mtp_graph_params is not None: | ||
| raise ValueError("MTPGraph parameters have already been set!") | ||
| _mtp_graph_params = MTPGraphParams( | ||
| {size: [] | ||
| for size in aclgraph_capture_sizes}, | ||
| {size: None | ||
| for size in aclgraph_capture_sizes}, | ||
| {size: [] | ||
| for size in aclgraph_capture_sizes}, | ||
| {size: [] | ||
| for size in aclgraph_capture_sizes}, | ||
| ) | ||
|
|
||
|
|
||
| def update_mtp_graph_params_workspaces(num_tokens: int, workspace: Any): | ||
| global _mtp_graph_params | ||
| if _mtp_graph_params is not None: | ||
| _mtp_graph_params.workspaces[num_tokens] = workspace | ||
|
|
||
|
|
||
| def get_mtp_graph_params(): | ||
| return _mtp_graph_params |
There was a problem hiding this comment.
The MTPGraphParams class and its associated functions (set_mtp_graph_params, update_mtp_graph_params_workspaces, get_mtp_graph_params) are duplicates of GraphParams and its functions. This introduces significant code duplication, which can lead to maintenance issues and potential bugs if one version is updated and the other is forgotten.
Consider refactoring to avoid this duplication. You could, for example, use a single set of functions and a dictionary to manage parameters for different graph types (e.g., DEFAULT and MTP). This would involve:
- Removing
MTPGraphParamsand its related functions.GraphParamscan be used for both. - Using a dictionary to store
GraphParamsinstances for different graph types, keyed by an enum. - Modifying the
set/update/getfunctions to accept agraph_typeparameter to operate on the correctGraphParamsinstance.
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
anon189Ty
force-pushed
the
mtp_dev
branch
7 times, most recently
from
65d1bcc to
d0210e7
Compare
| self.use_sparse = hasattr(vllm_config.model_config.hf_config, | ||
| "index_topk") | ||
|
|
||
| self.query_start_loc = torch.zeros( |
There was a problem hiding this comment.
Maybe these pinned tensors can also reuse corresponding tensors in model runner by reference to self.runner?
|
Plead add UT tests for FULL_DECODE_ONLY aclgraph |
|
has this change be merged into main? which commit |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
anon189Ty
force-pushed
the
mtp_dev
branch
6 times, most recently
from
5817e91 to
9258fa3
Compare
anon189Ty
force-pushed
the
mtp_dev
branch
2 times, most recently
from
cda2170 to
3e7deb1
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Signed-off-by: anon189Ty <[email protected]>
Signed-off-by: anon189Ty <[email protected]>
Signed-off-by: anon189Ty <[email protected]>
Signed-off-by: anon189Ty <[email protected]>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
|
This PR isn’t needed anymore since it’s already implemented on the main branch, and no new features should be added to the v0.11.0-dev release. |
5 participants
What this PR does / why we need it?
Currently, the MTP model still runs in eager in full graph mode. This PR adapts the MTP with the full graph capture and execution. When the graph mode is set to "FULL_DECODE_ONLY", the MTP will run in full-graph to improve the performance.
The change include:
1). Rebuild mtp model with ACLGraphWrapper.
2). Add common attn metadata when start capture in mtp dummy_run.
3). Add common attn metadata update in mtp.
4). Addapted data update when num_speculative_tokens > 1.
Does this PR introduce any user-facing change?
How was this patch tested?