[Discussion] Will vLLM consider using Speculative Sampling to accelerating LLM decoding?

[gesanqiu]

Sampling is an already known bottleneck of vLLM(see #421 and #670 ). Last weekend I saw a project named Medusa, in it's blog, it introduce a new simple decoding way to accelerate LLM generation and reach a good performance. As far as I known, lepton.ai is alreay use this method.
Adopting Medusa Heads is not difficult, since there is no seperate model. But tree attention and typical acceptance scheme is not a standard process for most LLM inference framework and should take a huge effort.

Any advice or comments?

[void-main]

I'm actually trying to integrate Medusa to vLLM.

After some quick and dirty experiments, I found out what really makes the integration hard is PagedAttention itself, which, by the way, is one of the core features of vLLM.

The whole vLLM system is built with an assumption that in decoder phase, there will be ONE newly generated token which ties to ONE kv cache block for each sequence. With speculative decoding methods, like Medusa, this assumption won't hold anymore.

Though it's hard, I believe it's totally doable. Here are some places we might want to tweak:

• should take candidates into consideration when you _prepare_inputs in the worker
• update block tables when you accept one candidate and release other blocks
• change the semantics of forward, make it return a list of tokens in SequenceOutputs
• set medusa tree mask with custom_attention_mask, and take mask into consideration in single_query_cached_kv_attention

Working on a PoC demo (only supports single running query), hopefully could share some new thoughts and findings later.

[Data-drone]

How is your testing going @void-main ?

[void-main]

Still working on it. Needs some modification to PagedAttention, which is the core to vLLM. @Data-drone

[Data-drone]

Is there a branch I can have a look at?

[void-main]

Currently the code change is in a private repo in my company. Later we'd like to release the working version.

[InkdyeHuang]

Currently the code change is in a private repo in my company. Later we'd like to release the working version.

when input_token_len = 450 ouput_token_len = 150, the first prompt step time and the second generate time is 1:1, so when the accurate of draftmodel is 96%, the speedtime is 20%, only the accurate is 30% can cover additional consumption

[zhaoyang-star]

Thank @void-main for the sharing the progress on porting Medusa.

I am porting Speculative Decoding into vLLM. I also found the main blocker is PagedAttention after some quick and dirty codes. More clearly, in speculative decoding mode, more than ONE tokens need to be taken as input when KV Cache has already existed. However, PagedAttention only supports two situations:

• Prefill stage: NO KV Cache exist. Calling xformers.xops.memory_efficient_attention_forward function internally.
• Decoding stage: KV Cache has already existed and ONLY ONE generated new tokens per prompt. Calling handwritten CUDA kernel paged_attention_v1/v2 internally.

The main blocker for the porting work happend at decoding stage. I should create a new kernel which could take more than one generated new tokens and existing KV Cache as input. The new kernel may have large difference to the paged_attention_v1/v2 kernel.

By now HuggingFace Transformers and llama.cpp has supported Speculative Decoding or similar method.

I have found more and more vLLMers consider using Speculative Decoding to accelerating LLM inference. Hope free discussion here and any suggestions are welcome. @WoosukKwon @zhuohan123 @casper-hansen @Yard1

[void-main]

@zhaoyang-star

You're right on the decoding stage, but one more thing, when doing medusa SpS, the m dimension for a single sequence is no longer 1, which means you should consider using Tensor Core for performance. I've implemented a kernel called PagedFlashAttention to do exactly that.

[void-main]

Recent update: I got a working demo version with vLLM + Medusa, for a single sentence, the average accept length is 1.8 (which means you get 1.8 tokens per forward on average), the performance boots 30%+ compared to vllm.

Here's the demo video:
https://github.com/vllm-project/vllm/assets/552990/0f22d79c-262c-4bfb-ab67-4c6ccab80dd5

Left part is vLLM + Medusa, right side part is pure vLLM.

The result is pretty interesting, vLLM + Medusa is slower on first token, because Medusa doesn't produce any token on first forward pass. But Medusa catches up soon, and become faster on later tokens.

[void-main]

Notes on the PagedFlashAttention kernel, as you can tell from the name, it combines PagedAttention with FlashAttention.

PagedAttention only works for decoding stages where you generate 1 token for each sequence, so you could use CUDA core to calculate attention score for each sequence. But with tree candidates from medusa, for each sequence, you need to process ~7-30 candidates. Sticking to CUDA core would make it too slow to get any benefits.

But, wait a sec, processing many tokens at once, that's what FlashAttention does. The only blocker is vLLM pages KV caches that needs to be used during generation phase. So I wrote a OpenAI triton kernel to do that. The code is based off lightllm's tokenattention, and added some page-related loading.

I'd like to say it's a pretty fun journey to implement PagedFlashAttention with triton, but I ran into several issues(triton-lang/triton#2488, triton-lang/triton#2522 and triton-lang/triton#2637), and the issues are not resolved yet, so hopefully we could optimize the kernel later.

[void-main]

I have found more and more vLLMers consider using Speculative Decoding to accelerating LLM inference. Hope free discussion here and any suggestions are welcome. @WoosukKwon @zhuohan123 @casper-hansen @Yard1

And I totally agree with @zhaoyang-star , vLLM is a great framework, but the whole framework is based on the assumption that each forward pass generates 1 token. Maybe later we should propose an RFC (maybe named VFC?) to extend this.

[AlvL1225]

I have found more and more vLLMers consider using Speculative Decoding to accelerating LLM inference. Hope free discussion here and any suggestions are welcome. @WoosukKwon @zhuohan123 @casper-hansen @Yard1

And I totally agree with @zhaoyang-star , vLLM is a great framework, but the whole framework is based on the assumption that each forward pass generates 1 token. Maybe later we should propose an RFC (maybe named VFC?) to extend this.

Hi, does medusa or speculative decoding support topp or topk sampling?

[zhaoyang-star]

I have found more and more vLLMers consider using Speculative Decoding to accelerating LLM inference. Hope free discussion here and any suggestions are welcome. @WoosukKwon @zhuohan123 @casper-hansen @Yard1

And I totally agree with @zhaoyang-star , vLLM is a great framework, but the whole framework is based on the assumption that each forward pass generates 1 token. Maybe later we should propose an RFC (maybe named VFC?) to extend this.

Hi, does medusa or speculative decoding support topp or topk sampling?

Speculative decoding supports temperature/topk/topp.