diff --git a/test/registered/ops/test_aiter_allreduce_fusion_amd.py b/test/registered/ops/test_aiter_allreduce_fusion_amd.py index 69ed9a194..7b99f4b6c 100755 --- a/test/registered/ops/test_aiter_allreduce_fusion_amd.py +++ b/test/registered/ops/test_aiter_allreduce_fusion_amd.py @@ -409,11 +409,15 @@ class TestAiterAllreduceFusionGate(CustomTestCase): ) ) stack.enter_context(mock.patch.object(comm, "_use_aiter", use_aiter)) + # moe_ep_size/moe_tp_size of 1 keep the hybrid EP+TP guard inactive + # so the aiter branch is what decides. stack.enter_context( mock.patch.object( comm, "get_parallel", - lambda: types.SimpleNamespace(tp_size=tp_world_size), + lambda: types.SimpleNamespace( + tp_size=tp_world_size, moe_ep_size=1, moe_tp_size=1 + ), ) ) # the gate reads get_exec().comm.enable_aiter_allreduce_fusion