Iterative generation using Input embeds and `past_key_values`
Author: yaswanth19
Pull request: yes

# What does this PR do?
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #34678 #35707
**Logic**: If cache is present along with `inputs_embeds` then use `inputs_embeds` to generate first token for every prompt rather than only for the first token of the cache

## 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#create-a-pull-request),
      Pull Request section?
- [ ] 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

@zucchini-nlp 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Code which I am using to check the feature branch
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

model_id = "microsoft/phi-2"
model = AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

model.generation_config.max_new_tokens = 30

prompt_cache = StaticCache(config=model.config, batch_size=1, max_cache_len=1000)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt")

inputs_embeds = model.get_input_embeddings()(inputs_initial_prompt.input_ids)
outputs = model.generate(inputs_embeds=inputs_embeds, past_key_values=prompt_cache)

response = tokenizer.batch_decode(outputs)[0]
print(response)

prompts = ["Help me to write a blogpost about travelling.", "Write a short note on AI"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(prompt, return_tensors="pt")
    new_input_ids = torch.cat([outputs, new_inputs.input_ids], dim=1)

    inputs_embeds = torch.cat([inputs_embeds,model.get_input_embeddings()(new_input_ids)],dim=1) # Necessary to align with cache

    outputs = model.generate(inputs_embeds=inputs_embeds, past_key_values=prompt_cache)
    response = tokenizer.batch_decode(outputs)[0]
    print(response)
    responses.append(response)
```

@zucchini-nlp I have written a test case to check the continuation using input embeds. Could you check if it's a valid test case or not? This test case is failing for some models, and I am looking into them (most likely need to overwrite in the model class). Also, should we make this functionality very specific to the static cache or apply it to all types of caches? I'm not sure whether this will break with other types of caches. I will test with different caches once the test cases are complete.

@yaswanth19 thanks a lot! Yes, the test looks good, but one small nit. I'd recommend to not use `input ids` from first and second batches and then flat-concat them. You can take a look at how it is done in `continue_generate_from_input_ids` test

And using only default cache is enough imo, if it works then StaticCache should also work

@zucchini-nlp I have modified the testcase and it's failing on few models like `idefics`,`zamba2`, `moshi`. It's either because they have their own custom prepare_input_for_generation function, due to non-standard kv cache etc.. Can you help in running the `generate` CI for all the tests coz rn the CI stops once a single test fails. Running the CI on all the cases will help me in identifying all the models the testcase is failing on and then I can either overwrite or skip them.

AFAIK this feature is breaking existing testcase for only `Moshi` . Will need to debug it.

@yaswanth19 sorry, completely forgot about this PR. Let me quick run the tests. For the models you mentioned, we can skip them with a reason why (different cache type, needs pixels etc) 

Btw, you should be able to run tests locally with `pytest -k test_generate_continue_from_inputs_embeds tests/models/`, they don't require any GPU 

UPDATE: failing models are Bloom, Chameleon, CLVP, Cohere2, Fuyu, Ideifcs, Moshi, Qwen2-VL, Zamba. For some models it might be fixable be applying the fix in model code, we override `prepare_inputs_for_generation` for multimodal models usually

@zucchini-nlp I think the PR is ready for review; Only catch is as I said before, this feature is breaking a existing functionality of `Moshi` as there instead of overwriting, it uses the general `prepare_inputs_for_generation` func and performs some post-processing on top of that and calculates input_embeds. Hence I am not sure how to fix/skip it.

cc: @ylacombe 

@yaswanth19 for `Moshi` if adding this new condition breaks generation, we should override the method in Moshi (same way as Qwen2-VL does it) and add a comment on top why it is overriden

If it actually doesn't break anything and you mean Moshi can't do `continue_from_input_embeds`, feel free just to skip the test. Noone would want to continue from embeds in audio models imo

@zucchini-nlp The CI is green :smiley: and the failing tests are unrelated to this PR. Ready for Review :hugs: 

> Code which I am using to check the feature branch
> 
> ```
> import torch
> from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache
> 
> model_id = "microsoft/phi-2"
> model = AutoModelForCausalLM.from_pretrained(model_id)
> tokenizer = AutoTokenizer.from_pretrained(model_id)
> 
> model.generation_config.max_new_tokens = 30
> 
> prompt_cache = StaticCache(config=model.config, batch_size=1, max_cache_len=1000)
> 
> INITIAL_PROMPT = "You are a helpful assistant. "
> inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt")
> 
> inputs_embeds = model.get_input_embeddings()(inputs_initial_prompt.input_ids)
> outputs = model.generate(inputs_embeds=inputs_embeds, past_key_values=prompt_cache)
> 
> response = tokenizer.batch_decode(outputs)[0]
> print(response)
> 
> prompts = ["Help me to write a blogpost about travelling.", "Write a short note on AI"]
> responses = []
> for prompt in prompts:
>     new_inputs = tokenizer(prompt, return_tensors="pt")
>     new_input_ids = torch.cat([outputs, new_inputs.input_ids], dim=1)
> 
>     inputs_embeds = torch.cat([inputs_embeds,model.get_input_embeddings()(new_input_ids)],dim=1) # Necessary to align with cache
> 
>     outputs = model.generate(inputs_embeds=inputs_embeds, past_key_values=prompt_cache)
>     response = tokenizer.batch_decode(outputs)[0]
>     print(response)
>     responses.append(response)
> ```

Thanks for your outstanding contribution! Does this PR support batch inference like this by pre-computing a attention_mask_real
`model_outputs = self.llm.generate(
                inputs_embeds=inputs_embeds_audio,
                # max_length=kwargs.get("max_length", 200),
                max_new_tokens=kwargs.get("max_new_tokens", 200),
                num_beams=kwargs.get("num_beams", 4),
                do_sample=kwargs.get("do_sample", False),
                min_length=kwargs.get("min_length", 1),
                top_p=kwargs.get("top_p", 1.0),
                repetition_penalty=kwargs.get("repetition_penalty", 1.0),
                length_penalty=kwargs.get("length_penalty", 1.0),
                temperature=kwargs.get("temperature", 1.0),
                attention_mask=attention_mask_real,
                bos_token_id=self.tokenizer.bos_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                pad_token_id=self.tokenizer.pad_token_id,
                past_key_values=kv_cache
            )`