# Past\_key\_value with multiple new tokens

**URL:** <https://discuss.huggingface.co/t/past-key-value-with-multiple-new-tokens/49177>\
**Category:** Intermediate\
**Created:** [August 2, 2023, 4:58pm UTC](https://discuss.huggingface.co/t/past-key-value-with-multiple-new-tokens/49177 "2023-08-02T16:58:13Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Imnimo](https://avatars.discourse-cdn.com/v4/letter/i/ecb155/32.png) [@Imnimo](https://discuss.huggingface.co/u/Imnimo)\
**Post date:** [August 2, 2023, 4:58pm UTC](https://discuss.huggingface.co/t/past-key-value-with-multiple-new-tokens/49177/1 "2023-08-02T16:58:13Z")

</div>

I have a GPT-style model which I’m using to generate text for a set of prompts. Each prompt has a shared prefix, followed by a variable suffix. For example, something like:

Prompt 1: “A B C X Y”  
Prompt 2: “A B C W V”  
Prompt 3: “A B C Q R S”

A simple way of processing Prompt 1 using key-value caching is like this:

```auto
abcxy_tokens = tokenize("A B C X Y")
result = model(abcxy_tokens,use_cache=True)
past_key_values = result['past_key_values']
predicted_token = argmax(result['logits'])
result = model(predicted_token,use_cache=True,past_key_values=past_key_values)
...(repeat)

```

This works great for me, but I would like to save more compute by caching the key-values for the shared prefix, “A B C”. I want to do something like this:

```auto
abc_tokens = tokenize("A B C")
prefix_key_values = model(abcxy_tokens,use_cache=True)['past_key_values']
#prefix_key_values can now be re-used for prompt 1, prompt 2, etc.

prompt_tokens = tokenize("X Y")
result = model(prompt_tokens,use_cache=True,past_key_values=prefix_key_values)
...

```

However, when I try this, I get an error complaining about a shape mismatch. The code only works if there is a single new token.

I don’t think that there’s any fundamental reason why key-value caching would not be possible with multiple new tokens. Does the existing implementation silently assume a single new token? Is there an alternate method I should be using for this scenario?

---

<div class="post-metadata">

**Author:** ![dblakely](https://avatars.discourse-cdn.com/v4/letter/d/54ee81/32.png) [@dblakely](https://discuss.huggingface.co/u/dblakely)\
**Post date:** [August 10, 2023, 4:47pm UTC](https://discuss.huggingface.co/t/past-key-value-with-multiple-new-tokens/49177/2 "2023-08-10T16:47:57Z")

</div>

This should be possible and in fact I’ve done something very similar before. Could you share the error message, stack trace, and which Huggingface model class you’re using?