[Feature]: Support for a draft model that takes inputs from base model (to support Medusa/EAGLE/Hydra)

[abhigoyal1997]

🚀 The feature, motivation and pitch

In approaches like Medusa/EAGLE/Hydra, the speculative model uses the last hidden states from the base model to propose candidates. This feature will allow any such approaches to be implemented with ease. One idea is to store the required base model's outputs along with the sequence and then use that while generating candidates for the next iteration.

Alternatives

No response

Additional context

No response

[abhigoyal1997]

I have implemented Medusa using this. If this makes sense and can be accepted as a contribution, I would love to create a PR (including the implementation of Medusa). I am also working on implementing the EAGLE approach.

[KexinFeng]

@abhigoyal1997 This is indeed an important feature that people have been looking for. It's also within my exploration radar, and I look forward to its implementation in vllm.
Here is some detailed question. I know for Medusa, tree-draft-tokens play an essential role; for Eagle, it is also important. In your implementation, did you enable tree-draft-tokens, or is it still the single-sequence draft-tokens?

I'm asking this because I'm developing this tree-style speculation. And it will be a perfect match with the Medusa/Eagle/Hydra here. We can maybe combine the effort and see how the performance boost when the two technieques are put together. #4565 (comment)

[abhigoyal1997]

Hi @KexinFeng
Currently what I've implemented only takes top-1 predictions to get single-sequence draft tokens. I agree that tree-style speculation is essential to get a significant acceleration. I've observed it in a torch.compile based implementation I worked on (based on gpt-fast), but I've not tried implementing that in vllm yet as it looked more complicated at the time and I knew it is already being worked on.

As for the current implementation of Medusa and EAGLE using a single sequence, I'll create a PR as soon as I've tested it a bit more and have company approvals.

[youkaichao]

cc @cadedaniel @LiuXiaoxuanPKU for visibility.

[Siegfried-qgf]

@abhigoyal1997 This is indeed an important feature that people have been looking for. It's also within my exploration radar, and I look forward to its implementation in vllm. Here is some detailed question. I know for Medusa, tree-draft-tokens play an essential role; for Eagle, it is also important. In your implementation, did you enable tree-draft-tokens, or is it still the single-sequence draft-tokens?

I'm asking this because I'm developing this tree-style speculation. And it will be a perfect match with the Medusa/Eagle/Hydra here. We can maybe combine the effort and see how the performance boost when the two technieques are put together. #4565 (comment)

I'm excited that you're working on this. I'm also considering adding tree attention to vllm to adapt it to EAGLE. How is your work going now？ I'd like to ask where I should make efforts to modify it, and do you have an open source plan?