[AMD] test(sgl-kernel): seed RNG on ROCm in test_moe_topk_sigmoid to fix tie-break flake (#25356)
This commit is contained in:
@@ -5,6 +5,19 @@ 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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens, num_experts, topk",
|
||||
|
||||
Reference in New Issue
Block a user