top_p_renorm_prob is not numerically stable for small probability values

[merrymercy]

When top_p is 1.0, we expect top_p_renorm_prob to be an identity function. However, it will make small values zero.

Reproduce

The following example constructs some value small probability values (1e-9). After top_p_renorm_prob, all of them become zero.

import torch
import flashinfer

pre_norm_prob = torch.tensor([[1 - 1e-9 * 20] + [1e-9] * 20]).to("cuda")
probs = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
print(f"sum={probs.sum(dim=1)}")

renormed_probs = flashinfer.sampling.top_p_renorm_prob(probs, 1.0)
print(f"{probs=}")
print(f"{renormed_probs=}")

print(f"{probs.min()=}")
print(f"{renormed_probs.min()=}")

assert torch.all(probs == probs)
assert torch.all(probs == renormed_probs)

Output

sum=tensor([1.], device='cuda:0')
probs=tensor([[1.0000e+00, 1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09,
         1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09,
         1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09, 1.0000e-09,
         1.0000e-09, 1.0000e-09, 1.0000e-09]], device='cuda:0')
renormed_probs=tensor([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
         0., 0., 0.]], device='cuda:0')
probs.min()=tensor(1.0000e-09, device='cuda:0')
renormed_probs.min()=tensor(0., device='cuda:0')
Traceback (most recent call last):
  File "/root/test_top_p_renorm.py", line 16, in <module>
    assert torch.all(probs == renormed_probs)
AssertionError