[Speculative Decoding] Fixing hidden states handling in batch expansion

[abhigoyal1997]

This PR fixes the handling of hidden_states in batch expansion.

Fix #7505

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[abhigoyal1997]

/ready

[cadedaniel]

• can you add a test?
• what exactly does this fix? were the hidden states applied incorrectly?

[abhigoyal1997]

• can you add a test?

Would an E2E test be sufficient? I can add the example I mentioned in the issue where we see the error.

• what exactly does this fix? were the hidden states applied incorrectly?

There was an error in a specific scenario. Consider sequences for which spec. decode is disabled (e.g., when num_tokens + spec_tokens > max_model_len). For such sequences, proposal length is set to 0 and they are handled separately from other sequences. Then when contracting the batch (in BatchExpansionTop1Scorer._contract_batch), we merge spec and non_spec sequences into shape (batch_size, spec_length + 1). While this was done correctly for other tensors, hidden_states were kept unchanged before and only reshaped later on:

vllm/vllm/spec_decode/spec_decode_worker.py

Line 650 in fc93e56

hidden_states = hidden_states.reshape(-1, max_proposal_len + 1,

This reshaping works fine if all sequences have proposal length > 0, but if for any sequence spec decode is disabled, it doesn't work. This PR changes batch expansion (mainly BatchExpansionTop1Scorer._contract_batch) such that hidden_states are manipulated and arranged like other tensors like tokens, probs and logprobs. This ensures hidden_states are correctly updated in all cases.

[cadedaniel]

Thanks for the writeup.

Would an E2E test be sufficient? I can add the example I mentioned in the issue where we see the error.

An E2E test is good. I think the issue here is that we should validate an accuracy level, not just that it doesn't crash. For this you can build on this PR #6454 which adds an assertion that draft acceptance rate is 100%.

[abhigoyal1997]

An E2E test is good. I think the issue here is that we should validate an accuracy level, not just that it doesn't crash. For this you can build on this PR #6454 which adds an assertion that draft acceptance rate is 100%.

Makes sense, but any draft model which uses hidden states can't be identical to the target model. In that case I don't think we can get draft acceptance rate of 100%. What do you suggest then?

[cadedaniel]

You can run it for the test prompts and record the accuracy, then assert it doesn’t go below that fixed value plus some epsilon.

[abhigoyal1997]

@cadedaniel I've added the test along with other MLPSpeculator tests.

[njhill]

Thanks @abhigoyal1997!

[njhill]

@cadedaniel you good with this being merged?

[cadedaniel]

yep, good to merge! thanks for the great testing here

[cadedaniel]

@njhill do you know what value we should expect here?

[cadedaniel]

(like, is this ballpark correct)

[njhill]

@cadedaniel the value looks reasonable but I’m actually not sure what’s expected, will try to find out when I get a chance.