how does vllm handle wrong tokens in speculative decoding?

[Tomorrowdawn]

The core problem is, accepted tokens in different sentences of same batch are different.
I tried to search a lot about batched speculative decoding and I found most inference frameworks ignore this problem by simply setting batch size == 1, except vllm. I read the code in vllm/spec_decode, and it seems vllm fill all wrong tokens with -1. I am still unclear about:

1. if there is other techniques to remove those -1s?
2. I know a trivial method is to mask these -1s off, if vllm does the same, how does it affects the performance?

thanks.

[cadedaniel]

Before the engine appends token ids to sequences, it removes -1 tokens. The logic is here:

vllm/vllm/engine/output_processor/multi_step.py

Line 68 in 34128a6

# -1 means the output token is not valid (eg. due to spec decode

[cadedaniel]

This allows speculative decoding to fit within continuous batching scheduling.

[Tomorrowdawn]

Thanks! I found this PR #3103 which introduces the "batch expansion" used for verification and I also reviewed the relevant code in spec_decode/batch_expansion.py and worker/worker. From the image you provided, it appears that batch expansion splits the output sentence into a group of sentences, but tokens "1", "2", and "3" are not present in the kvcache because the target model hasn't executed the forward pass on them yet.

I was wondering if you could further explain the underlying mechanics of this process? I'm quite interested in gaining a deeper understanding of how it works. Any additional insights you could provide would be greatly appreciated.

[cadedaniel]

Your description is correct! The KV for tokens 1,2,3 is populated in the fwd pass of the target model, which occurs after batch expansion and before batch contraction.

For more details, check out this doc.
It doesn't have any treatment of the draft model KV, but explains how the target model KV is populated using lookahead allocation. https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a

[Tomorrowdawn]

Thanks! I found this PR #3103 which introduces the "batch expansion" used for verification and I also reviewed the relevant code in spec_decode/batch_expansion.py and worker/worker. From the image you provided, it appears that batch expansion splits the output sentence into a group of sentences, but tokens "1", "2", and "3" are not present in the kvcache because the target model hasn't executed the forward pass on them yet. I was wondering if you could further explain the underlying mechanics of this process? I'm quite interested in gaining a deeper understanding of how it works. Any additional insights you could provide would be greatly appreciated.

these lines suggests only the last token id would be appended into "input_ids". For example, in the last sentence, I don't understand how target model compute the key & value for 1, 2, 3.

[cadedaniel]

In batch expansion we create a new vLLM sequence for each token to be scored. So a single sequence with spec tokens 1,2,3 becomes three sequences. This means that there is still one input id per sequence, namely, 1, 2, and 3.

Because QKV projection happens before saving to KV cache in each layer, which happens before each attention computation, we can safely do this without impacting the computation result.