bugfix: Ensure Loop Termination by Enforcing IEEE-754 Compliance in Sampling Kernels

[yzh119]

This PR addresses issue #769.

As discussed in this comment, the use of the approximate division instruction div.approx.ftz.f32 can break the loop invariant, preventing the loop from terminating. To resolve this, this PR changes the data types of low, high, and mid to double, ensuring that the compiler maintains IEEE-754 compliance and preserves numerical stability.

[dskhudia]

@yzh119 : Thanks a lot for the fix. Wouldn't it still be an issue with loop terminating for some other inputs esp since we have an exact floating-point comparison (} while (min_gt_low != max_le_high);)? Isn't an eps-based solution better here?

[yzh119]

Our initial design uses eps and it might introduce some errors (especially when there are a lot of similar prob values around the top-k boundary) so that some of the values outside of top-k might be incorporated.

Wouldn't it still be an issue with loop terminating for some other inputs esp since we have an exact floating-point comparison

This would not happen if the mid = (low + high) / 2 follows IEEE-754.

[zhyncs]

@yzh119 it breaks the unit test

[zhyncs]

I don't think using double it is a good idea.

[yzh119]

@yzh119 it breaks the unit test

Which unittest you are talking about?

[yzh119]

I don't think using double it is a good idea.

Can you elaborate more? what's your proposal.

[zhyncs]

@yzh119 https://github.com/flashinfer-ai/flashinfer/blob/93e1a2634e22355b0856246b032b285ad1d1da6b/tests/test_sampling.py

[yzh119]

Are you using the main line codebase? We have ci set up for the sampling: https://ci.tlcpack.ai/blue/organizations/jenkins/flashinfer-ci/detail/PR-977/1/pipeline/14

[zhyncs]

This works for me after revert this commit.

[zhyncs]

https://github.com/sgl-project/flashinfer/tree/sgl-kernel