[V1][Prototype] MTP Support #17683
[V1][Prototype] MTP Support #17683
Conversation
Signed-off-by: Rui Qiao <ruisearch42@gmail.com>
ruisearch42
requested review from
njhill,
LiuXiaoxuanPKU,
WoosukKwon,
robertgshaw2-redhat,
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 🚀 |
|
@ruisearch42 Any update on this PR? I think the overall direction is good, while we can probably share more code with the eagle proposer. |
|
This pull request has merge conflicts that must be resolved before it can be |
| ) | ||
|
|
||
| with set_forward_context(attn_metadata, self.vllm_config): | ||
| hidden_states = self.model( |
There was a problem hiding this comment.
is this for single layer prototype?
There was a problem hiding this comment.
yes. I think DeepSeek only released one layer of weights for MTP. This can be extended to run multiple times as a follow up.
| num_tokens_per_req = query_len_per_req - num_rejected_tokens | ||
|
|
||
| cu_num_tokens = torch.empty_like(cu_target_query_lens) | ||
| torch.cumsum(num_tokens_per_req, dim=0, out=cu_num_tokens[1:]) |
There was a problem hiding this comment.
shall we move this into triton as well?
|
|
||
| # TODO(woosuk): Consider seeds. | ||
| q = torch.empty_like(probs) | ||
| q.exponential_() |
There was a problem hiding this comment.
curious why we have separate sampling logic here for MTP?
There was a problem hiding this comment.
needs to be cleaned up
| # SPDX-License-Identifier: Apache-2.0 | ||
| import torch | ||
| import torch.nn as nn | ||
| import triton |
|
Done in #18435 |
I'd like to first address the main logic and correctness, and then the following TODOs:
The main challenge is that we need to prepare MLA attention metadata, instead of flash attention metadata, and may need to refactor related code.