What does this PR do?

This PR aims at solving issue #32165.

I've started adapting code to enable speculative decoding with batch_size >1. I've reused some of the work in former PR #26875.

Main steps of the solution:

When batch size > 1:

1. Compute the number of similar tokens between candidate tokens and tokens obtained after doing a forward pass with the main model. This results in a tensor n_matches with the number of matches for each sequence in the batch.

2. We keep the matching tokens. For that, we keep all tokens from the output to the main model with a sequence position inferior to n_matches.max() + 1. In doing so, we also retain some potential mismatched tokens, which we will deal with in the next steps using padding tokens. The resulting tensor, input_ids, is thus in the form (batch_size, n_matches + 1).

3. We shift each sequence i in input_ids by n_matches.max() - n_matches[i]. The matching tokens are displaced to the right of input_ids and n_matches.max() - n_matches[i] padding tokens are added to the left.

4. Left cut: We cut all columns that contain only padding tokens. By design, these columns are to the left of the input_ids. In this way we keep the smallest possible `input_ids' that contain all the information needed to continue assisted generation.

Steps 1 to 4 are the main addition to the the original speculative decoding loop described in detail in this blog to enable assisted generation with BS > 1.

To make this work, we also need to adapt the computation of the attention_masks, past_key_values and position_ids to take into account the shifted positions of the generated tokens.

To do:

For now, I want to make this work with Whisper using this snippet.

I've implemented steps 1 to 4 and adapted the computation of the attention_masks and past_key_values to handle the new padding tokens.

I still need to make some adaptations with the position_ids to make this work properly. From what I can see:

• For the main model (here WhisperForConditionalGeneration), position_ids are inferred directly from the attention_mask as we can see here. So if we pass the right attention mask to generate we should be good.

• For the assistant model (WhisperForCausalLM in our example), position_ids are currently not computed nor passed to prepare_inputs_for_generation, which I'm not sure exactly why. I've done a first attempt at solving this with no success so far.

cc @sanchit-gandhi @gante