diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index b5fd5d8b98..5eee468bc4 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -3,6 +3,7 @@ import torch import itertools +import gc from contextlib import nullcontext import aiter from aiter import dtypes @@ -1057,6 +1058,9 @@ def _write_bench_csv(rows): os.environ.pop("AITER_BF16_FP8_MOE_BOUND", None) else: os.environ["AITER_BF16_FP8_MOE_BOUND"] = _old_moe_bound + # Reclaim per-case VRAM to avoid OOM under fragmentation. + gc.collect() + torch.cuda.empty_cache() if ret is None: continue ret.update(extras)