Cache position incorrectly inferred for generation
Author: killjoygoose
Pull request: no
Labels: name: bug

### System Info

In https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py line 369 you check whether the cache class is supported, which works, if cache is None basically (as far as I could reproduce). However, if one adds `past_key_values` to the input in generate, at some point the `cache_position` becomes an empty torch.tensor, leading to an `IndexError` on line 388 `cache_position[-1]` . Maybe there is some hidden way to fix this, I could definitely create the `cache_position`s myself, but I think it is rather counterintuitive that something starts as None and gets converted to tensor in some cases and in others not. 
```
        if self._supports_cache_class: <- here : it is once None, once torch.tensor([]), which I find counterintuitive and also errorprone
            model_inputs["cache_position"] = cache_position
        # - `cache_position` was not a mandatory input in `prepare_inputs_for_generation` for those models, and this
        #   function may be called outside of `generate`. Handle most use cases by creating `cache_position` on the fly
        #   (this alternative is not as robust as calling `generate` and letting it create `cache_position`)
        elif cache_position is None:
            past_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
            cache_position = torch.arange(past_length, input_ids.shape[1], dtype=torch.long, device=input_ids.device)

        # 2. Generic cache-dependent input preparation
        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        # Exception 1: when passing input_embeds, input_ids may be missing entries
        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
        #              (we can't check exception 3 while compiling)
        if past_key_values is not None:
            model_inputs["past_key_values"] = past_key_values
            if (
                inputs_embeds is not None  # Exception 1
                or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1])  # Exception 3
            ):
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]
```


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Run 2 consecutive .generate calls on (in my case `LlavaForConditionalGeneration` model but it should be doable with any model which inherits from `GenerationMixin`).
In the 1st call set `return_dict_in_generate=True`, use the `past_key_values=output_of_the_first_call.past_key_values` in the second call -> then you get the `IndexError`

### Expected behavior

I would expect that the case positions are always None if not provided explicitly and that they are correctly inferred from the input_ids length + the currently available cache length. 

Comments:
Hey! Do you mean continuing generation with prev past key values? Can you provide a reproducer?

```
from transformers.models.llava import LlavaForConditionalGeneration
from transformers import AutoProcessor

MODEL_PATH = "./models/llava-1.5-7b-hf"
model = LlavaForConditionalGeneration.from_pretrained(
    MODEL_PATH,
    device_map="auto",
    torch_dtype=torch.float16,
)
processor = AutoProcessor.from_pretrained(MODEL_PATH)
inputs = processor(text="USER: tell me something nice! ASSISTANT:", return_tensors="pt")
generate_ids = model.generate(**inputs, max_new_tokens=1, return_dict_in_generate=True)
decoded = processor.batch_decode(
    generate_ids.sequences,
)[0]
print("First generation!")
print(decoded)

inputs = processor(text="USER: How do you say hello! in german? ASSISTANT", return_tensors="pt")
generate_ids = model.generate(**inputs,
                              max_new_tokens=500,
                              return_dict_in_generate=True,
                              past_key_values=generate_ids.past_key_values) <- here it crashes with IndexError
decoded = processor.batch_decode(
    generate_ids.sequences,
)[0]



```

Btw, if you aggree, I would also change the name of the variable `cache_positions` because it is misleading - I would imagine that in means positions that are cached, it means actually positions to be cached. How about `ids_to_cache` or sth in that sense. 

I see, you are trying to continue generation using a cache of initial system prompt. In that case you have to pass in the total inputs (prev cached ids + new ids). So it has to be `model.generate(generate_ids, past_key_values=past_key_values)`, no need to pass images as those were cached already in prev run

Take a look at cache docs here https://huggingface.co/docs/transformers/v4.48.0/kv_cache#iterative-generation-with-cache 🤗   

hmm, ok, I understand. It is still suboptimal to have to re-tokenize though :     `new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to("cuda")`, and that is just in order to compute the correct `torch.arange`, when it can be much more user friendly by allowing to only pass the cache. Also you are also checking if `cache_position is None` in the function which throws the error, so why don't you just allow it to be None until that point and there assign it? Btw how is this code even reachable, when the `cache_position` gets updated in `_get_initial_cache_position`?

Passing `cache_position` is not necessary if you pass the correct input ids, generate will take care of it. The error here is thrown because the inputs are of shorter length than we expect. Try with `model.generate(generate_ids, past_key_values=past_key_values)`

OK, so here is the problem - it is not in the input ids, but it the attention mask. In `_sample` we keep the attention mask to (Batch, seqlength+N_generated). So I will close the issue, but please look into the `_get_initial_cache_position` as it can be easily improved.

Also- it is generally bad practice not to inform the users why a particular exception is raised - here the problem is, as I said, with the attention masks, otherwise, after small touches you can allow for just adding the next batch of ids + the cache. However, you should either validate that the attentions shape == B, len_cache+len_current and raise some Exception - informing the user that the attention mask must be of a particular shape (this is also logical as we don't know if the batch processing did not involve padding) OR you should simply prepend 1s, as if no attention mask was passed - what you actually already do if no attention mask is passed.

I ran into this issue while implementing a custom generation loop. I'm working with a multi-modal model, so I can't exactly reconstruct the input IDs for the cached keys and values. As a workaround, I just simply prepended 1s to my new prompt, and it seems to work. Is this an acceptable workaround @killjoygoose, @zucchini-nlp? It'd be nice to support this use case natively. 

@yukw777 yep, acceptable as the prefix key/values will be cropped internally and will not be forwarded to the model

This issue stems from legacy code, and actually not asking users to pass the whole context is the solution. However it will be quite breaking, so we decided not to fix it yet. 

Also cc @gante with whom we discussed it. I see this issue coming up way too more than expected, do you have any idea what we can do (except fro having it in the documentation)?