Investigate top-p performance

[mattdangerw]

Currently our top-p (nucleus) sampler has the worst performance of our entire sampling offering, likely because of the sort operation we need to do for our softmax probabilities. See these rough numbers for example, where top-p is slower than beam, despite needing much less computation on the model itself.

We should see if there is anything we can do to speed up top-p sampling with XLA (e.g. a more XLA friendly sort op).

[mattdangerw]

One thing we could try here is adding a top-k heuristic cutoff, before applying top-p. This would mean you only need to sort a much smaller range of vocabulary terms to find the correct top-p cutoff point.

We could expose this as an option fairly easily. Less clear if we could expose some sort of defaults for this, but if we could it could give us a much more performant top-p out of the box.