diff --git a/sgl-kernel/tests/test_moe_topk_sigmoid.py b/sgl-kernel/tests/test_moe_topk_sigmoid.py index e9a9240ee..6b2cba858 100644 --- a/sgl-kernel/tests/test_moe_topk_sigmoid.py +++ b/sgl-kernel/tests/test_moe_topk_sigmoid.py @@ -5,18 +5,14 @@ import pytest import torch from sgl_kernel import topk_sigmoid -_IS_HIP = torch.version.hip is not None - @pytest.fixture(autouse=True) def _deterministic_seed(): - # AMD/ROCm only: pin RNG so torch.randn produces identical gating scores - # across runs. atol=0 indices comparison is otherwise tripped by near-tied - # sigmoid scores where hipCUB's tie-break inside torch.topk disagrees with - # sgl_kernel.topk_sigmoid. Not observed on CUDA, so leave CUDA behavior - # unchanged. - if _IS_HIP: - torch.manual_seed(0) + # Pin RNG on every backend so torch.randn produces identical gating scores + # across runs. The exact index comparison can otherwise be tripped by + # near-tied sigmoid scores whose rounded values have different top-k + # tie-break ordering between torch.topk and sgl_kernel.topk_sigmoid. + torch.manual_seed(0) @pytest.mark.parametrize(