[Feature] Speculative Decoding

[josephrocca]

Motivation

Speculative decoding can speed up generation more than 2x. This degree of speedup is an important feature for a production-grade LM deployment library, and it seems the methods are starting to mature enough to make their way into frameworks like TGI and vLLM, so might be a good time for LMDeploy to consider adding support for a popular/established speculative decoding method.

Related resources

• TGI (supports Medusa and MLPSpeculator as of writing):
  • https://huggingface.co/docs/text-generation-inference/basic_tutorials/train_medusa
  • MLPSpeculator. huggingface/text-generation-inference#1865
• vLLM (groundwork for several speculation methods in progress as of writing):
  • [WIP] Speculative decoding using a draft model vllm-project/vllm#2188
  • [Model] MLPSpeculator speculative decoding support vllm-project/vllm#4947
  • [Speculative Decoding] Medusa Implementation with Top-1 proposer vllm-project/vllm#4978
  • [Speculative Decoding] Enable arbitrary model inputs vllm-project/vllm#5101
  • [Speculative Decoding 1/2 ] Add typical acceptance sampling as one of the sampling techniques in the verifier vllm-project/vllm#5131
• MLC-LLM (supports only EAGLE as of writing):
  • [SpecDecode] Support Eagle in speculative decoding mlc-ai/mlc-llm#2080
  • [Eagle] Make eagle disco compatible mlc-ai/mlc-llm#2197
  • [Eagle] Avoid worker - engine transfer for hidden states mlc-ai/mlc-llm#2256
  • [Eagle] Fix token shifting for prefill step mlc-ai/mlc-llm#2266
  • [Eagle] Run additional decode for draft model when all proposals are accepted mlc-ai/mlc-llm#2294
  • [Eagle] Fix the requests for additional decode in eagle verify mlc-ai/mlc-llm#2336

Below is a copy-paste from a neat project called Spec-Bench. The ranking when running 33B models is similar. Please see the linked repo for latest data.

• Device: a single NVIDIA GeForce RTX 3090 GPU (24GB) with 12 CPU cores
• Testing environment: Pytorch 2.0.1, under CUDA 11.8
• Experimental Settings: Vicuna-7B-v1.3, greedy decoding, FP16 precision, batch size = 1

Models	Multi-turn Conversation	Translation	Summa-rization	Question Answering	Mathematical Reasoning	Retrieval-aug. Generation	#Mean Accepted Tokens	Overall
EAGLE🏅	2.44x	1.81x	2.13x	2.11x	2.54x	1.82x	3.57	2.16x
SpS🥈	1.98x	1.37x	2.00x	1.95x	1.89x	1.76x	2.29	1.83x
Hydra🥉	2.04x	1.67x	1.56x	1.81x	2.16x	1.48x	3.26	1.80x
PLD	1.57x	1.07x	2.31x	1.25x	1.62x	1.56x	1.74	1.55x
Medusa	1.60x	1.38x	1.28x	1.46x	1.64x	1.22x	2.32	1.44x
REST	1.49x	1.18x	1.21x	1.46x	1.35x	1.27x	1.63	1.32x
Lookahead	1.13x	0.97x	1.05x	1.07x	1.29x	0.98x	1.65	1.08x

Note that MLPSpeculator is not included in the benchmark since it is newer. Another new method that isn't included in Spec-Bench as of writing:

• https://github.com/apple/ml-recurrent-drafter

[zhyncs]

In fact, we have already implemented the Medusa TreeMask version in LMDeploy. When batch=1, the acceleration ratio and RPS improvement relative to the main branch are consistent with those in the blog.

And when the batch size increases, the overhead of Medusa prefill is greater than the benefit of generating multiple tokens at each iteration. We are currently working on solving this problem. Please stay tuned.

[zhyncs]

EAGLE also has plans to support open source in the future.

[Dbxwz]

@zhyncs I implemented EAGLE in vllm and met the same probelm when the batch size increases.
Here is a simple analysis (bs is batch size, k is proposal length, the batch size bottleneck of target model is 3):

Because the calculation of rejected tokens wastes GPU resources, so skipping speculative decoding is the best choice sometimes.

Meituan's solution introduce a novel sampling mechanism that leverages Thompson Sampling to regulate the generation processes. And someone use a trained control module (I forgot the source).

Or, similar to VLLM's current approach, we can simply skip speculative decoding when the batch size exceeds a certain threshold. It's simple and effective, additional judgment conditions is useful for future enhancements.

[zhyncs]

we can simply skip speculative decoding when the batch size exceeds a certain threshold

Thank you for sharing. In fact, this is currently how we do it internally as well, but this approach is still a bit rough. If we want speculative decoding to take effect by default without burdening the user's mind when they are not using it, we also need to dynamically adjust the threshold based on actual workloads, which introduces a certain level of complexity.

In actual usage, the reception rate of Eagle is slightly higher than that of Medusa.

Thompson Sampling Control Mechanism Currently not implemented in actual production environment.

[coolhok]

EAGLE also has plans to support open source in the future.

Can you reveal the schedule。Or share the development of the branch together，thanks!!

[snippetzero]

In fact, we have already implemented the Medusa TreeMask version in LMDeploy. When batch=1, the acceleration ratio and RPS improvement relative to the main branch are consistent with those in the blog.

And when the batch size increases, the overhead of Medusa prefill is greater than the benefit of generating multiple tokens at each iteration. We are currently working on solving this problem. Please stay tuned.

@zhyncs Hey, could you let me know how things are going right now? Maybe there's something I can do to lend a hand? Appreciate it.

[zhyncs]

I will split the internal implementation of the TreeMask version into multiple PRs and then submit them.

[snippetzero]

I will split the internal implementation of the TreeMask version into multiple PRs and then submit them.

Thank you, could you share the methods to solve the performance degradation after the batch size increases

[zhyncs]

The overall design and detailed implementation have been discussed with @lzhangzz before, and there was improvement in small batch sizes, but it didn't work well in large batch sizes. As far as I know, the performance achieved on vLLM is also similar.

[zhyncs]

EAGLE has a higher computational load than Medusa, but it has a higher acceptance rate. It performs better in large batches compared to Medusa. However, this is just a temporary solution. The reason this approach works, which involves trading more computation for reduced latency, is because in small batches, computational resources are not fully utilized.

[snippetzero]

EAGLE has a higher computational load than Medusa, but it has a higher acceptance rate. It performs better in large batches compared to Medusa. However, this is just a temporary solution. The reason this approach works, which involves trading more computation for reduced latency, is because in small batches, computational resources are not fully utilized.

how is the attention kernel chosen during the verification stage? As mentioned in the FlashInfer blog, the computational intensity of the append/verification stage is between that of decode and prefill. It doesn't seem optimal to use either the decode or prefill kernel from the lmdeploy engine directly. It should be noted that when Q is relatively long, using the current prefill kernel for verification might not be the optimal approach

[zhyncs]

how is the attention kernel chosen during the verification stage?

@snippetzero The current implementation uses the prefill kernel in TurboMind. cc @lzhangzz

[GxjGit]

In fact, we have already implemented the Medusa TreeMask version in LMDeploy. When batch=1, the acceleration ratio and RPS improvement relative to the main branch are consistent with those in the blog.

And when the batch size increases, the overhead of Medusa prefill is greater than the benefit of generating multiple tokens at each iteration. We are currently working on solving this problem. Please stay tuned.

Hi, It seems that I can not find any code related to Speculative Decoding in LMdeploy. Has it not to been pushed in the sources? If done, could you provide me commit id or simply some key words?

[josephrocca]

Has it not to been pushed in the sources?

I don't think it's available yet, and I'm not sure if this can be prioritised right now. Maybe @lvhan028 can comment? It certainly seems very exciting based on vLLM's findings:

• [Speculative decoding] [Help wanted] [Performance] Optimize draft-model speculative decoding vllm-project/vllm#4630

In the Anyscale fork we saw a 50% speedup on bs=8 with a 68m-sized draft model on TP1/70B target model on TP8 and a 7B draft model on TP(1|8)/70B target model on TP8. This was with the optimizations listed above as "P0".

and also based on together.ai's findings:

• https://www.together.ai/blog/speculative-decoding-for-high-throughput-long-context-inference

Conventional wisdom (e.g., Chen et al., 2023; Li et al., 2024; Liu et al., 2024) is that in the high-throughput regime (i.e., large batch sizes), speculative decoding—which leverages underutilized GPU compute during memory-bound decoding—does not make sense, because decoding will be compute-bound and the GPUs will thus be fully utilized. Surprisingly, we show analytically and empirically that for large batch sizes, if the input sequences are long enough, decoding once again becomes memory-bound due to the large size of the KV cache. Building on this key observation, we demonstrate that speculative decoding can improve throughput and latency by up to 2x on 8 A100s in this large-batch, long-context setting.

[lvhan028]

No, it hasn't