Assistant Model With Falcon Fails 
Author: yrapop01
Pull request: no

### System Info

transformers version: 4.30.0.dev0 (and also 4.29.2)
python version: 3.9
platform: sagemaker notebook on aws running on g5.12xlarge

### Who can help?

@ArthurZucker @younesbelkada @gante 

### Information

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

### Tasks

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

### Reproduction

Code:

```
import transformers
import torch

bnb_config = transformers.BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = "tiiuae/falcon-40b-instruct"

tokenizer = transformers.AutoTokenizer.from_pretrained(model)
model = transformers.AutoModelForCausalLM.from_pretrained(model,
                                                          quantization_config=bnb_config,
                                                          device_map="auto",
                                                          trust_remote_code=True)

assistant = "tiiuae/falcon-7b-instruct"
assistant = transformers.AutoModelForCausalLM.from_pretrained(assistant,
                                                              quantization_config=bnb_config,
                                                              device_map="auto",
                                                              trust_remote_code=True)

text = "Girafatron is"

inputs = tokenizer(text, return_tensors="pt", return_token_type_ids=False)
inputs = {k: v.cuda() for k, v in inputs.items()}

model.eval()
with torch.no_grad():
    outputs = model.generate(**inputs,
                             max_new_tokens=16,
                             assistant_model=assistant,
                             do_sample=False)
```

Exception:
```
IndexError                                Traceback (most recent call last)
/tmp/ipykernel_38536/3021020828.py in <cell line: 7>()
      6 model.eval()
      7 with torch.no_grad():
----> 8     outputs = model.generate(**inputs,
      9                              max_new_tokens=16,
     10                              assistant_model=assistant,

~/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/torch/utils/_contextlib.py in decorate_context(*args, **kwargs)
    113     def decorate_context(*args, **kwargs):
    114         with ctx_factory():
--> 115             return func(*args, **kwargs)
    116 
    117     return decorate_context

~/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/transformers/generation/utils.py in generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, **kwargs)
   1490 
   1491             # 12. run assisted generate
-> 1492             return self.assisted_decoding(
   1493                 input_ids,
   1494                 assistant_model=assistant_model,

~/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/transformers/generation/utils.py in assisted_decoding(self, input_ids, assistant_model, do_sample, logits_processor, logits_warper, stopping_criteria, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, streamer, **model_kwargs)
   4380             # 5.3. Discard past key values relative to unused assistant tokens
   4381             new_cache_size = new_cur_len - 1
-> 4382             outputs.past_key_values = _crop_past_key_values(self, outputs.past_key_values, new_cache_size)
   4383             model_kwargs["assistant_past_key_values"] = _crop_past_key_values(
   4384                 assistant_model, model_kwargs["assistant_past_key_values"], new_cache_size - 1

~/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/transformers/generation/utils.py in _crop_past_key_values(model, past_key_values, maximum_length)
   4522             new_past.append(
   4523                 (
-> 4524                     past_key_values[idx][0][:, :, :maximum_length, :],
   4525                     past_key_values[idx][1][:, :, :maximum_length, :],
   4526                 )

IndexError: too many indices for tensor of dimension 3```

### Expected behavior

expected to run without errors

Comments:
Hey thanks for reporting! Will have a look asap! Though the model is on the hub we should try to make this run smoothly! 

Hey @yrapop01 👋 

I've had a look at assisted generation, and here's what I found:
1. The immediate error you see can be fixed -- Falcon has a custom cache structure, so it needs custom code for cache slicing. Easy peasy.
2. We then hit a more complex wall -- the modeling code does not handle the case where there is a cache and the input ids' length is larger than 1 (a bit of a special case that is needed for the assistant [here](https://github.com/huggingface/transformers/blob/0675600a60b260d6bdb9c8ad91d932d690672bf0/src/transformers/generation/utils.py#L4268)).

This means it requires modelling changes, but the model is not yet in `transformers`. I'm going to discuss internally, and let you know of our next steps :)

@yrapop01 we are adding Falcon to `transformers` (as opposed to hub-loaded model code), I'll make sure assisted generation works in the transformers version!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Do you have ETA for adding Falcon to transformers?

@yrapop01 You can follow the PR here: #24523 

Thank you!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.