[Reformer] - Cache hidden states and buckets to speed up inference
Author: patrickvonplaten
Pull request: yes

As discussed with the authors in Reformer we cache the hidden states and buckets to speed up inference for language generation. Caching only the hidden states can save at least twice the memory versus caching the key and value output vectors (often more since the `hidden_size` can be smaller than key and query projections). 

The idea is to only recompute key and value projections within the same chunk so that the output will be equal.

- [x] Implement caching mechanism for Local Attention
- [x] Implement caching mechanism for LSH Attention
- [x] Add test
- [x] Add logic to generation
- [x] Refactoring and better naming

=> This results in a 10x speed up when generating up to 1000 tokens.

**Some fun trials trying to make Reformer generate 16000 tokens for its own Wikipedia article [here](https://colab.research.google.com/drive/1Oao8vBtDkz6v1E1efUhTlug5uuxC_BnE?usp=sharing).**

*Review*:
Not urgent. Merging this PR can wait a bit. Added a couple of tests to make sure caching gives same output as non-caching.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=h1) Report
> Merging [#5578](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/8ab565a4be5a7fd96b19ef88d474037ef31f27e5&el=desc) will **decrease** coverage by `0.10%`.
> The diff coverage is `98.47%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/5578/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #5578      +/-   ##
==========================================
- Coverage   77.32%   77.21%   -0.11%     
==========================================
  Files         146      146              
  Lines       26047    26198     +151     
==========================================
+ Hits        20141    20230      +89     
- Misses       5906     5968      +62     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_xlnet.py](https://codecov.io/gh/huggingface/transformers/pull/5578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ194bG5ldC5weQ==) | `81.49% <ø> (ø)` | |
| [src/transformers/modeling\_reformer.py](https://codecov.io/gh/huggingface/transformers/pull/5578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yZWZvcm1lci5weQ==) | `90.56% <98.47%> (+2.70%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_electra.py](https://codecov.io/gh/huggingface/transformers/pull/5578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9lbGVjdHJhLnB5) | `26.02% <0.00%> (-69.52%)` | :arrow_down: |
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `85.46% <0.00%> (-0.76%)` | :arrow_down: |
| [src/transformers/generation\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3V0aWxzLnB5) | `97.11% <0.00%> (+0.28%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/5578/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9yb2JlcnRhLnB5) | `93.36% <0.00%> (+49.37%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=footer). Last update [8ab565a...7cd6a8f](https://codecov.io/gh/huggingface/transformers/pull/5578?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


@LysandreJik @sgugger - when you guys are back, feel free to merge to PR or we can wait until I'm back - not too urgent this PR. LGTM to merge.

Okey merging this. 

@sgugger @LysandreJik @sshleifer -> checked caching extensively and added multiple tests for equality. 
Would be nice if you can tag me for possible future issues related to this caching mechanism.