In assisted decoding, pass model_kwargs to model's forward call (fix prepare_input_for_generation in all models)
Author: sinking-point
Pull request: yes

# What does this PR do?

Previously, assisted decoding would ignore any additional kwargs that it doesn't explicitly handle. This was inconsistent with other generation methods, which pass the model_kwargs through prepare_inputs_for_generation and forward the returned dict to the model's forward call.

The prepare_inputs_for_generation method needs to be amended in all models, as previously it only kept the last input ID when a past_key_values was passed.

Fixes #25020


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@gante 


Comments:
@sinking-point the PR has "WIP" in the title -- is it still under development, or is it ready to review?

Not ready yet. Still have to fix more models and see what's breaking the other test. I've deprioritised this somewhat as it's quite time consuming, but I'll keep chipping away at it whenever I can.

If you need this done quickly, you're welcome to help - lmk and I'll add you as a collaborator on my branch.

Not urgent -- simply double-checking whether it was in need of a review or not :)

@gante This should be ready for review now. Thanks in advance.

@ArthurZucker @LysandreJik This PR is a big one and touches in a core piece of logic for all generative models, so I'm tagging 2 core maintainers.


### Context
Advanced generation techniques (like assisted generation or [medusa](https://github.com/FasterDecoding/Medusa)) may generate more than one token per model forward pass. The original implementation of assisted generation had a lot of custom code, as it breaks one of the assumptions in the models' `prepare_inputs_for_generation` -- that only one token is generated per `forward` pass.

### Solution
@sinking-point has kindly put forward a proposal to get rid of the custom code in assisted generation. After iterating with me, the plan was to remove the assumption of one token per `forward` in `prepare_inputs_for_generation` -- the slicing therein should be done based on how many tokens do not have corresponding past KV values, and not simply taking the last set of inputs. This is the change this PR implements, as well as the removal of some custom assisted generation code. Needless to say, it is fully backwards compatible 😉 

### Postface
To reiterate: this PR gets the green light from me in terms of logic 🟢, and it is a big contribution by @sinking-point. This PR is also important to future-proof our generative techniques -- we will be ready for new types of multiple-token-per-forward-pass strategies as a result of this PR.

I'll be off the next few weeks, but I'm sure this PR will get a quick resolution 🤗 

Thanks @gante , I'll take a look at your comments tomorrow 👍

Hi @sinking-point! Sorry for the delay - I'm taking over this PR from @gante because he's out on a well-deserved rest right now. Is everything ready for review, or are there any other issues you want to discuss with the team before we take a final look at it?

No worries @Rocketknight1 . Thanks for taking this on.

There's one discussion gante opened that I haven't resolved. Could you give your input on this? https://github.com/huggingface/transformers/pull/25242#discussion_r1326154797

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_25242). All of your documentation changes will be reflected on that endpoint.

@sinking-point Replied to the last remaining issue up above!

Thanks @Rocketknight1 , I'll take a look on Monday

Some random tests started failing so I rebased onto main where they're fixed, but it looks like I have some more work to do now.

Ick, yeah. I'm not sure what's causing those test failures, but if you can't figure it out, let me know and I'll dive in!

@Rocketknight1 I think I've sorted it out. The hub test fails seem unrelated to my changes, and might have been a network failure or something. The failing Musicgen pipeline is broken on main as well.

Could you block new generative model merges until this PR is merged? Otherwise I will have to keep applying the changes to any new models that are added.

That's a big ask, unfortunately, but they shouldn't be too frequent! Let's just try to get this PR over the line quickly, so you don't have to keep rebasing and updating. I agree with you on the failing tests, btw, they look like they're unrelated.

Are you happy for me to go ahead and merge this once the doc tests pass, or are there other things left to be cleaned up?

Should be ready to merge if you're happy with it. Thanks!

Looks like doc tests passed @Rocketknight1 , so as you said let's make this a priority before any more models are added.

Understood! It's quite a big PR since it touches so many models, but I'll try to get an internal review in the next few days.

> Could you block new generative model merges until this PR is merged?

Alternatively, could you require that new generative models' `prepare_inputs_for_generation` method follows this PR? That is, instead assuming that if `past_key_values` is provided it covers all but the last position, you should calculate how many positions are remaining after `past_key_values` and keep those.

Hi @Rocketknight1 , any update on this?

Hey @sinking-point 👋 

I'm back from holidays and I'll be doing a quick final check now. Assuming the check comes out positive, we'll tag a core maintainer to greenlight the merge. 

Our apologies for the slow process, it should be quick now 🤗 

@sinking-point regarding the failing test: rebasing the PR should fix it, the bug was fixed last week :)

ping @LysandreJik -- this PR should be ready to be merged after it is rebased. Please read [this comment](https://github.com/huggingface/transformers/pull/25242#issuecomment-1718079624) for context :)

Thanks @gante :)

This seems ok to me but I'd like to ask @patrickvonplaten for his opinion and eventual approval given the experience maintaining this part of the code

Amazing, thank you @LysandreJik and @patrickvonplaten 

And thank you @sinking-point for this big contribution 🔥 💛