Speed up TopKLogitsWarper and TopPLogitsWarper (pytorch)
Author: LSinev
Pull request: yes

# What does this PR do?

Speeds up TopKLogitsWarper and TopPLogitsWarper using torch filling functions.
Here's a minimal example to reproduce the slow behavior (and test speed of improvements):
```
import torch
from transformers import TopPLogitsWarper, TopKLogitsWarper, LogitsWarper
import timeit


class TopKLogitsWarperNew(LogitsWarper):
    r"""
    :class:`transformers.LogitsWarper` that performs top-k, i.e. restricting to the k highest probability elements.

    Args:
        top_k (:obj:`int`):
            The number of highest probability vocabulary tokens to keep for top-k-filtering.
        filter_value (:obj:`float`, `optional`, defaults to :obj:`-float("Inf")`):
            All filtered values will be set to this float value.
        min_tokens_to_keep (:obj:`int`, `optional`, defaults to 1):
            Minimum number of tokens that cannot be filtered.
    """

    def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not isinstance(top_k, int) or top_k <= 0:
            raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")

        self.top_k = top_k
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        top_k = min(max(self.top_k, self.min_tokens_to_keep), scores.size(-1))  # Safety check
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
        scores = scores.masked_fill(indices_to_remove, self.filter_value)  # changed here
        return scores


class TopPLogitsWarperNew(LogitsWarper):
    """
    :class:`transformers.LogitsWarper` that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <=
    prob_cut_off.

    Args:
        top_p (:obj:`float`):
            If set to < 1, only the most probable tokens with probabilities that add up to :obj:`top_p` or higher are
            kept for generation.
        filter_value (:obj:`float`, `optional`, defaults to :obj:`-float("Inf")`):
            All filtered values will be set to this float value.
        min_tokens_to_keep (:obj:`int`, `optional`, defaults to 1):
            Minimum number of tokens that cannot be filtered.
    """

    def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0):
            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")

        self.top_p = top_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        sorted_logits, sorted_indices = torch.sort(scores, descending=True)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)  # changed here

        # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs > self.top_p
        if self.min_tokens_to_keep > 1:
            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
            sorted_indices_to_remove[..., : self.min_tokens_to_keep - 1] = 0
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # scatter sorted tensors to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        scores = scores.masked_fill(indices_to_remove, self.filter_value)  # changed here
        return scores


input_ids = torch.randint(0, 10000, (256, 256))
scores = torch.randn(256, 10000)

top_k_lw = TopKLogitsWarper(100)
top_p_lw = TopPLogitsWarper(0.95)

top_k_lw_new = TopKLogitsWarperNew(100)
top_p_lw_new = TopPLogitsWarperNew(0.95)

print(f"Existing top_k impl time for 100 iterations on CPU = {timeit.timeit(lambda: top_k_lw(input_ids, scores), number=100)}")
print(f"Proposed top_k impl time for 100 iterations on CPU = {timeit.timeit(lambda: top_k_lw_new(input_ids, scores), number=100)}")

print(f"Existing top_p impl time for 100 iterations on CPU = {timeit.timeit(lambda: top_p_lw(input_ids, scores), number=100)}")
print(f"Proposed top_p impl time for 100 iterations on CPU = {timeit.timeit(lambda: top_p_lw_new(input_ids, scores), number=100)}")

if torch.cuda.is_available():
    input_ids = input_ids.cuda()
    scores = scores.cuda()

    print(f"Existing top_k impl time for 100 iterations on GPU = {timeit.timeit(lambda: top_k_lw(input_ids, scores), number=100)}")
    print(f"Proposed top_k impl time for 100 iterations on GPU = {timeit.timeit(lambda: top_k_lw_new(input_ids, scores), number=100)}")

    print(f"Existing top_p impl time for 100 iterations on GPU = {timeit.timeit(lambda: top_p_lw(input_ids, scores), number=100)}")
    print(f"Proposed top_p impl time for 100 iterations on GPU = {timeit.timeit(lambda: top_p_lw_new(input_ids, scores), number=100)}")

```
Timings reported:
```
Existing top_k impl time for 100 iterations on CPU = 2.5527561419994527
Proposed top_k impl time for 100 iterations on CPU = 0.36601612999947974

Existing top_p impl time for 100 iterations on CPU = 6.4072540179995485
Proposed top_p impl time for 100 iterations on CPU = 4.1470332960007

Existing top_k impl time for 100 iterations on GPU = 0.09082965299967327
Proposed top_k impl time for 100 iterations on GPU = 0.008193381999262783

Existing top_p impl time for 100 iterations on GPU = 1.1027910299999348
Proposed top_p impl time for 100 iterations on GPU = 0.9008321309993335
```

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@patrickvonplaten

 -->


Comments:
This looks great @LSinev! 

Ok, looks good to merge. I checked that your implementation works with Pytorch 1.4 as well