[CI][MoE] Moe kernels test cleanup (#54954)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Sep 2, 2026 at 20:31 UTC
9b38e3ad53967856df6dd18dbaf91f8f5e3a5f32
2 files changed
+54
-4
tests/kernels/moe/conftest.py
+46
@@ -2,6 +2,32 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
5
+from vllm.distributed import cleanup_dist_env_and_memory
6
+
7
+NEEDS_CLEAN_ENTRY = frozenset(
8
+ {
9
+ # Fail unless they start from a clean distributed and memory state,
10
+ # which they used to get only as a side effect of the preceding test
11
+ # cleaning up after itself.
12
+ "test_moe_layer.py",
13
+ "test_ocp_mx_moe.py",
14
+ "test_zero_expert_moe.py",
15
+ # Execute no test on ROCm, where this was measured, so nothing is known
16
+ # about how they behave without the per-test cleanup. Cleaned on entry
17
+ # as a precaution; drop once they have been seen green without it.
18
+ "test_batched_deepgemm.py",
19
+ "test_deepep_deepgemm_moe.py",
20
+ "test_deepep_moe.py",
21
+ "test_deepep_v2_moe.py",
22
+ "test_deepgemm.py",
23
+ "test_grouped_topk.py",
24
+ "test_mxfp4_moe.py",
25
+ "test_shared_fused_moe_routed_transform.py",
26
+ "test_silu_mul_per_token_group_quant_fp8_colmajor.py",
27
+ "test_situ_mul_fp8_quant.py",
28
+ }
29
+)
30
+
31
32
def pytest_addoption(parser):
33
parser.addoption(
@@ -12,3 +38,23 @@ def pytest_addoption(parser):
38
@pytest.fixture
39
def subtests(request):
40
return request.config.getoption("--subtests")
41
+
42
+
43
+@pytest.fixture()
44
+def should_do_global_cleanup_after_test() -> bool:
45
+ """Drop the per-test global cleanup for this directory.
46
+
47
+ A full ``gc.collect()`` plus an emptied caching allocator after each of the
48
+ ~1450 small kernel tests per shard costs more than the tests themselves.
49
+ The tests that depend on isolation get a clean state on entry instead, in
50
+ the hook below.
51
+ """
52
+ return False
53
+
54
+
55
+def pytest_runtest_setup(item):
56
+ # Not tryfirst: the skipping plugin evaluates skip marks in a tryfirst
57
+ # hook, so tests about to be skipped never reach this, while fixture setup
58
+ # still runs afterwards.
59
+ if item.path.name in NEEDS_CLEAN_ENTRY:
60
+ cleanup_dist_env_and_memory()
tests/kernels/moe/test_fused_topk.py
+8
-4
@@ -45,7 +45,8 @@ def torch_topk(
45
46
47
@pytest.mark.skipif(
48
- not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
48
+ not current_platform.is_cuda_alike(),
49
+ reason="This test is skipped on non-CUDA platform.",
50
)
51
@pytest.mark.parametrize("num_tokens", [1, 33, 56])
52
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@@ -89,7 +90,8 @@ def test_fused_topk(
90
91
92
@pytest.mark.skipif(
92
- not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
93
+ not current_platform.is_cuda_alike(),
94
+ reason="This test is skipped on non-CUDA platform.",
95
)
96
@pytest.mark.parametrize("num_tokens", [1, 33, 56])
97
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@@ -138,7 +140,8 @@ def test_fused_topk_bias(
140
141
142
@pytest.mark.skipif(
141
- not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
143
+ not current_platform.is_cuda_alike(),
144
+ reason="This test is skipped on non-CUDA platform.",
145
)
146
@pytest.mark.parametrize("num_experts", [6, 8, 16])
147
@pytest.mark.parametrize("topk", [3, 4])
@@ -205,7 +208,8 @@ def test_fused_topk_nan_inf_clamp(
208
209
210
@pytest.mark.skipif(
208
- not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
211
+ not current_platform.is_cuda_alike(),
212
+ reason="This test is skipped on non-CUDA platform.",
213
)
214
@pytest.mark.parametrize("num_experts", [6, 8, 16])
215
@pytest.mark.parametrize("topk", [3, 4])