[core / attention] Fix fused attention generation with newest transformers version

[younesbelkada]

What does this PR do?

Currently in the latest transformers release, using AutoAWQ + fused attention with cache is broken
In huggingface/transformers#25242 the logic of caching has changed a bit, now when using transformers cache + a past key value length of 1 (as done here), the input ids will be sliced as such:

input_ids = input_ids[:, 1:]

Meaning the assumption if seqlen == 1: to deal with the transformers cache case needs now to be adapted, one can just check if past_key_values is present in kwargs and contains valid tensors, and slice out only the last token if that's the case

cc @casper-hansen

[younesbelkada]

I also checked out to this commit in transformers: huggingface/transformers#26162 (before huggingface/transformers#25242) and can confirm it works in both cases.

[casper-hansen]

Tested and looks good. No performance regression on my end.

[casper-hansen]

I take the remark about performance regression back. I tested using my benchmark.py script found in examples and saw no difference. But using the .generate() function, it is 50% slower.

Slicing the hidden states in every attention layer for every token is a lot of overhead. We should instead slice it at a higher level, e.g. in the model. However, that requires implementing a LlamaModel, MistralModel, AquilaModel. This is probably the right solution but requires a bit of work, which I will look into.

[younesbelkada]

Thanks for benchmarking ! yes slicing only once at the model level makes sense!