[ROCm][CI] Stabilize MI355 FlyDSL MoE accuracy test (#53024)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 21, 2026 at 15:17 UTC
a0af854f5d9647637a1f8c9ab9edd3722fe9956c
1 file changed
+49
-1
tests/kernels/moe/test_flydsl_moe.py
+49
-1
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
15
)
16
from vllm.platforms import current_platform
17
from vllm.platforms.rocm import on_gfx950
18
+from vllm.utils.torch_utils import set_random_seed
19
20
if not (current_platform.is_rocm() and on_gfx950()):
21
pytest.skip("This test can only run on ROCm and gfx950.", allow_module_level=True)
@@ -41,11 +42,58 @@ RoutingBuffers = tuple[
42
]
43
44
45
+def _assert_flydsl_matches_reference(
46
+ actual: torch.Tensor,
47
+ expected: torch.Tensor,
48
+ *,
49
+ atol: float = 0.5,
50
+ rtol: float = 0.1,
51
+ max_mismatch_fraction: float = 1e-5,
52
+ max_error_ratio: float = 3.0,
53
+) -> None:
54
+ """Compare kernels without making accuracy depend on output size.
55
+
56
+ The FlyDSL and reference kernels accumulate BF16 values in different
57
+ orders. A strict allclose makes one expected rounding outlier fail an
58
+ otherwise accurate output, which becomes increasingly likely for the
59
+ largest token counts (up to 117 million output elements here).
60
+
61
+ Keep the original elementwise tolerance for 99.999% of values, while also
62
+ bounding every tolerated outlier to three times its elementwise tolerance.
63
+ """
64
+ assert actual.shape == expected.shape
65
+
66
+ mismatch = ~torch.isclose(actual, expected, atol=atol, rtol=rtol)
67
+ mismatch_count = int(mismatch.sum().item())
68
+ total = actual.numel()
69
+ allowed_mismatches = int(total * max_mismatch_fraction)
70
+ mismatch_msg = (
71
+ f"FlyDSL/reference mismatch: {mismatch_count}/{total} values "
72
+ f"({mismatch_count / total:.6%}) exceed atol={atol}, rtol={rtol}; "
73
+ f"allowed <= {allowed_mismatches}/{total} "
74
+ f"({max_mismatch_fraction:.6%})"
75
+ )
76
+ assert mismatch_count <= allowed_mismatches, mismatch_msg
77
+
78
+ if mismatch_count:
79
+ actual_mismatch = actual[mismatch].float()
80
+ expected_mismatch = expected[mismatch].float()
81
+ abs_error = (actual_mismatch - expected_mismatch).abs()
82
+ tolerance = atol + rtol * expected_mismatch.abs()
83
+ worst_error_ratio = (abs_error / tolerance).max().item()
84
+ assert worst_error_ratio <= max_error_ratio, (
85
+ f"{mismatch_msg}; worst error is {worst_error_ratio:.4f}x its "
86
+ f"elementwise tolerance (allowed <= {max_error_ratio:.4f}x)"
87
+ )
88
+
89
+
90
@pytest.mark.parametrize(
91
"num_tokens", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
92
)
93
@pytest.mark.parametrize("inter_dim", [256, 512])
94
def test_flydsl_moe(num_tokens: int, inter_dim: int):
95
+ set_random_seed(0)
96
+
97
device = "cuda"
98
topk = 8
99
num_experts = 384
@@ -172,7 +220,7 @@ def test_flydsl_moe(num_tokens: int, inter_dim: int):
220
scale_is_bf16=True,
221
)
222
175
- assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1)
223
+ _assert_flydsl_matches_reference(out, out_ref)
224
225
226
if __name__ == "__main__":