Support for caching prompt hidden states through multiple calls of `generate()`
Author: offendo
Pull request: no
Labels: name: WIP

### Feature request

Hi there,

I'd like to be able to re-use the hidden states for a common (potentially long) prompt across multiple calls to `model.generate()` in order to reduce redundant computation. Here is how I envision a final API, though I'm sure there are multiple ways to do it.

```python
# Load stuff
model = AutoModel.from_pretrained('huggyllama/llama-7b')
tokenizer = AutoTokenizer.from_pretrained('huggyllama/llama-7b')

# Common prompt that we'd prepend to every example
prompt = "This is a common prompt in every example."
prompt_ids = tokenizer(prompt, return_tensors='pt')

# Examples to pass to generate
examples = ["Ackbar went to", "Billaba enjoys", "Cody is eating some"]

# Generation loop
outputs = []
prompt_hidden_state = None
for ex in examples:
    # Current way of doing things
    out = model.generate(
        **tokenizer(prompt + ex, return_tensors='pt'),   
    )

    # Proposed method to re-use prompt_hidden_state
    out = model.generate(
        **tokenizer(x, return_tensors='pt'),   
        common_prompt_ids=prompt_ids,
        prompt_hidden_state=prompt_hidden_state
    )
    prompt_hidden_state = out.prompt_hidden_state
    outputs.append(out.sequences)
```

Thanks in advance. 

### Motivation

A very common pattern for LLM usage is having a common prompt (e.g., instructions and input/output pairs), a sample input, and asking it to generate the sample output. For example:

```
You are a programmer's assistant which converts English descriptions to Python functions.

English: <example 1 description>
Python: <example 1 function>

English: <example 2 description>
Python: <example 2 function>

English: <example 3 description>
Python: <example 3 function>

English: <input description>
Python: 
```

I'd like to be able to cache the common part of the prompt across inputs, that is, everything before `<input description>` which appears in every example to avoid potentially expensive re-computation.

### Your contribution

The only existing info I could find is the short discussion [here](https://discuss.huggingface.co/t/avoid-recalculating-hidden-states-between-generate-calls/34209). I tried messing around a bit to get this to work but had little luck. I'm not familiar with the inner-workings of `transformers` and ran into numerous errors. One problem is padding, which if we're using left padding, can cause some misalignment with the prompt hidden states, e.g.:
```
<p> <p> <p> common prompt x_1 x_2 x_3
<p> <p> common prompt x_1 x_2 x_3 x_4
<p> <p> <p> <p> common prompt x_1 x_2
```
I don't know the best way to solve this. Do we dynamically pad every tensor in `past_key_values`? That seems slow but I don't know if it actually is.

If someone can suggest a better/easier way or maybe give some more pointers on how to solve padding. I'd be happy to try again myself.

Thanks in advance. 

Comments:
cc @gante 

Hey @offendo 👋 

This is a relevant request, and one that I can't give an exact solution at the moment. In general, `generate()` + prompting is very manual at the moment, and we really want to improve it. As such, the solution to your proposal will depend on our next iteration of prompt-handling!

I'm assigning the issue to me, and I'll keep you posted 🤗 

(cc @patrickvonplaten  -- this is related to our brainstorming session yesterday, about prompting)

This feature would be useful not only for long prompts, but also for incrementally building dialogs / conversations without recomputing the generated parts. And I think there's a simple solution: just expose `past_key_values` in the various `generate` functions, and allow `past_key_values` to be passed in as the parameter for all decoder-only models.

@gante can't we already pass `past_key_values` to `generate` through `model_kwargs` ? 
So we could pre-encode the prompt into `past_key_values` with a single forward pass and then re-use it in generation already I think. 

See this similar (old) issue maybe: https://github.com/huggingface/transformers/issues/4368 should help

@patrickvonplaten Yeah I guess this solves the issue for @offendo , but still `past_key_values` are not being returned with the generated sequences. Hopefully there could be a flag (possibly called `return_past_key_values`) that allows them to be returned.

@lqf96 absolutely. I'll work on adding that, since I've seen others mentioning its usefulness :)

Looks like this is being addressed by #25086... Hopefully it can be accepted soon!

你好，我已经接收到你的邮件！

> @gante can't we already pass `past_key_values` to `generate` through `model_kwargs` ? So we could pre-encode the prompt into `past_key_values` with a single forward pass and then re-use it in generation already I think.
> 
> See this similar (old) issue maybe: #4368 should help

Does this approach work even with the padding issue? If we encode the common prompt (which has no padding) and then call generate on a batch of values with left padding, the hidden states won’t align properly and we’d have to do some dynamic padding of the `past_key_values` or something.

Please correct me if I am misunderstanding something! 

And also, thanks to everyone for their work on this!

@offendo If the inputs are properly passed, then the position ids can be correctly inferred, which will result in the correct output :) Possibly we may need additional logic to ensure this position id inference can happen seamlessly

Hi @gante, thanks for working on this commit https://github.com/huggingface/transformers/pull/25086 ! May I know if this feature request is addressed by this commit or do we need to make additional changes on top of it? 

@louisowen6 definitely needs more work, as it breaks 178 of our tests :)

Looking forward to this!

Hey folks 👋 

If you install from `main` and add `return_dict_in_generate=True` to `generate`, `past_key_values` will be part of the output, assuming your model is configured with `use_cache=True` (the default).

You can then pass `past_key_values` to `generate` to continue generating!

Here's what I did to get it to work with hidden states generated by `model.forward`:

```python
import transformers

model_name = "meta-llama/Llama-2-7b-chat-hf"
m = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
t = transformers.AutoTokenizer.from_pretrained(model_name)

t.padding_side = "left"
t.pad_token = t.eos_token
input_tokens = t(["Lizargrid: Yer a lizard Harry! Your parents", "Lobstargrid: Yer a lobster Harry! Your parents"], return_tensors="pt", padding=True).to(0)


forward_outputs = m(
    **input_tokens,
    return_dict=True,
)

# THIS IS THE ONLY IMPORTANT LINE:
forward_outputs.past_key_values = [[y[:, :, :-1] for y in x] for x in forward_outputs.past_key_values]

new_gen = m.generate(
    input_ids=input_tokens["input_ids"],
    past_key_values=forward_outputs.past_key_values,
    bos_token_id=t.eos_token_id, 
    pad_token_id=t.eos_token_id, 
    max_new_tokens=100, 
    use_cache=True,
)

for text in t.batch_decode(new_gen, skip_special_tokens=False):
    print("\n" + "#" * 80 + "\n\n" + text)

```

Basically, `past_key_values` expects to not have a key and a value for the last token, as it expects `past_key_values` to come from generation, so we drop those key and value associated to the last token of our input sequences.

It would be fun if `generate` was smarter and we didn't have to do that @gante, basically just detect that if the input and the number of key and value are the same as the number of tokens, that they come from `forward` not `generate`

@JulesGM To avoid wasting computations, you have the following options:
1. run the forward pass with `input_tokens["input_ids"][:, :-1]`, which is equivalent to doing the prefill step outside `generate`
2. add a token selection step after your forward pass, and append that token to `input_ids`, which would be somewhat close to running one generation step manually

Automatic slicing would go against our `generate` design philosophy, where we avoid adding transformations except on the most common use cases 🤗 In our experience, adding these sorts of automatic input handling leads to problems in future features.

@gante I agree that silent modifications are generally bad.
I would suggest making it more of a first-grade citizen. The case of having few-shots examples being generated from a million times is just so common. It could be, then, another argument, if it comes from a forward call?

Just in my lab at Mila in Montreal, multiple people have had this issue

I really feel like it should be extremely straightforward to do. In my case, I have like 1000 tokens of few-shot examples, and ~100 tokens or so of text to generate. Being able to, in a straightforward way, not duplicate that memory & work (without having to do tricks like ditch the last token) would be a big deal.

@JulesGM I 100% agree that it is trivial to add from a technical point of view :) 

The catch is that there are 3 possible solutions, and none is good. In the absence of a good solution, it is by far preferable to continue handling the case outside `generate` -- if we stop operating under this principle, `generate` will quickly become bloated and it will be even harder for us to iterate on it.

For completeness, here are the solutions I've considered: (I'm open to new solutions!)
1. Crop the `past_key_valies` outside `generate`, as you suggested in the original snippet (issue = it is a manual task)
2. Automatically crop `past_key_values` if it has the same length as `input_ids` (issue = this condition might be hit in other circumstances, such as when `past_key_values` holds content regarding a past round of a multi-chat conversation, rendering `generate` unusable in those settings)
3. Adding a flag (issue = users have to discover the flag, further bloat on `generate`)

__________________________________________________________

IMO, the underlying issue at the moment is that we allow `past_key_values` and `input_ids` to hold both redundant information and mutually exclusive information, preventing us from making assumptions and/or throwing informative exceptions. 

How about calling them:
`past_key_values_from_generate`
`past_key_values_from_forward`

Otherwise I feel like just having `past_key_values` & silently checking if it's equal in length to input_ids or if it has one fewer of length is honestly probably better than the alternatives

> ### Feature request
> Hi there,
> 
> I'd like to be able to re-use the hidden states for a common (potentially long) prompt across multiple calls to `model.generate()` in order to reduce redundant computation. Here is how I envision a final API, though I'm sure there are multiple ways to do it.
> 
> # Load stuff
> model = AutoModel.from_pretrained('huggyllama/llama-7b')
> tokenizer = AutoTokenizer.from_pretrained('huggyllama/llama-7b')
> 
> # Common prompt that we'd prepend to every example
> prompt = "This is a common prompt in every example."
> prompt_ids = tokenizer(prompt, return_tensors='pt')
> 
> # Examples to pass to generate
> examples = ["Ackbar went to", "Billaba enjoys", "Cody is eating some"]
> 
> # Generation loop
> outputs = []
> prompt_hidden_state = None
> for ex in examples:
>     # Current way of doing things
>     out = model.generate(
>         **tokenizer(prompt + ex, return_tensors='pt'),   
>     )
> 
>     # Proposed method to re-use prompt_hidden_state
>     out = model.generate(
>         **tokenizer(x, return_tensors='pt'),   
>         common_prompt_ids=prompt_ids,
>         prompt_hidden_state=prompt_hidden_state
>     )
>     prompt_hidden_state = out.prompt_hidden_state
>     outputs.append(out.sequences)
> Thanks in advance.
> 
> ### Motivation
> A very common pattern for LLM usage is having a common prompt (e.g., instructions and input/output pairs), a sample input, and asking it to generate the sample output. For example:
> 
> ```
> You are a programmer's assistant which converts English descriptions to Python functions.
> 
> English: <example 1 description>
> Python: <example 1 function>
> 
> English: <example 2 description>
> Python: <example 2 function>
> 
> English: <example 3 description>
> Python: <example 3 function>
> 
> English: <input description>
> Python: 
> ```
> 
> I'd like to be able to cache the common part of the prompt across inputs, that is, everything before `<input description>` which appears in every example to avoid potentially expensive re-computation.
> 
> ### Your contribution
> The only existing info I could find is the short discussion [here](https://discuss.huggingface.co/t/avoid-recalculating-hidden-states-between-generate-calls/34209). I tried messing around a bit to get this to work but had little luck. I'm not familiar with the inner-workings of `transformers` and ran into numerous errors. One problem is padding, which if we're using left padding, can cause some misalignment with the prompt hidden states, e.g.:
> 
> ```
> <p> <p> <p> common prompt x_1 x_2 x_3
> <p> <p> common prompt x_1 x_2 x_3 x_4
> <p> <p> <p> <p> common prompt x_1 x_2
> ```
> 
> I don't know the best way to solve this. Do we dynamically pad every tensor in `past_key_values`? That seems slow but I don't know if it actually is.
> 
> If someone can suggest a better/easier way or maybe give some more pointers on how to solve padding. I'd be happy to try again myself.
> 
> Thanks in advance.

So how will we write the english to python example in the current version? @JulesGM  @offendo 

i want to cache this part only
 ```
> You are a programmer's assistant which converts English descriptions to Python functions.
> 
> English: <example 1 description>
> Python: <example 1 function>
> 
> English: <example 2 description>
> Python: <example 2 function>
> 
> English: <example 3 description>
> Python: <example 3 function> '''

Hi, is there a solution to this now ?

you can do something similar in llama.cpp if your model supports it.

Thanks for the fast reply ! So there isn't an easy implementation for huggingface yet ? Do you have any resources for the llama.cpp implementation ?


so i have done it in this fashion but i believe llama cpp probably has better way to do this through templates (just a  guess)
anyway hope this helps:

```
import json
m1=    {
      "role": "system",
      "content": "You are a helpful AI assistant."
    }

m1 = [m1]
main_message_part_1 = "main overview text "
message = {"role": "user", "content":  main_message_part_1 +"Query : "+str(q1)+"assistant : " +str(a1) + " ,Query : "+str(q2)+"assistant : "+str(a2)+ " ,Query : "+str(q3)+"assistant : "+str(a3) + " ,Query : "+str(q4)+"assistant : "+str(a4)+ " ,Query : "+str(q5)+"assistant : "+str(a5)+ " ,Query : "+str(q6)+"assistant : "+str(a6)+ " ,Query : "+str(q7)+"assistant : "+str(a7)+ " ,Query : "+str(q8)+"assistant : "+str(a8)+ " ,Query : "+str(q9)+"assistant : "+str(a9)+  ", examples are over now. .  return only the correct SQL  query"}
m1.append(message)

text=" whatever is your question"

m1.append({"role": "user", "content": "write the shortest correct sql query for  " + text })

jss =json.dumps({"slot_id":0,  "temperature": 0.1, "n_keep": -1,  "cache_prompt": True,"messages": m1 })
headers = {'Content-Type': 'application/json'}

r = requests.post("LLAMA-SERVER-LINK", data=jss , headers=headers)

```
and u keep a llama-server open hosting the model.
`"cache_prompt": True` is used for caching
and on server side set --gpu-layers 99999 --ctx-size 20000