[feature] support pcp + mtp in full graph #4572
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 introduces support for PCP (Prefill Context Parallelism) and MTP (Multi-Token Prediction) in full graph mode, along with several related bug fixes. The changes correctly generalize PCP-only logic to accommodate DCP (Decode Context Parallelism) as well. A notable improvement is the handling of variable query lengths in speculative decoding batches, which replaces assumptions of fixed lengths with more robust logic. However, I've identified one critical issue in the implementation that needs to be addressed.
zhangsicheng5
force-pushed
the
mtp
branch
3 times, most recently
from
b08104c to
b330b75
Compare
weijinqian0
added
ready
| # prefill target_hidden_states: pcp split | ||
| num_tokens_d = num_decode_reqs * self.decode_threshold | ||
| query_lens_d = self.runner.query_lens[:num_decode_reqs] | ||
| num_tokens_d = query_lens_d.sum().item() |
There was a problem hiding this comment.
is query_lens_d a device tensor? if so, you call query_lens_d.sum().item() will incur cpu blocking, please fix it.
There was a problem hiding this comment.
This is a host tensor, refer to model_runner_v1.py: self.query_lens = torch.from_numpy(num_scheduled_tokens), so it will not influence host device sync. The _d is abbreviation of _decode.
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
zhangsicheng5
force-pushed
the
mtp
branch
7 times, most recently
from
bf5139c to
7cc6d7b
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
zhangsicheng5
force-pushed
the
mtp
branch
2 times, most recently
from
64c24c1 to
df5a7d2
Compare
Signed-off-by: zhangsicheng5 <zhangsicheng5@huawei.com>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
4 participants
Uh oh!
There was an error while loading. Please reload this page.