From b116adfb550d8ea4842232921c5746aac5147287 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Mon, 2 Nov 2026 17:38:01 +0800 Subject: [PATCH 1/4] fix(moe): restore TP MoE num_local_tokens --- atom/model_ops/moe.py | 6 ++---- tests/model_ops/test_dp_gather_routed_rows.py | 10 ++++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 230e0a5365..fdb9c2eda2 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -524,7 +524,7 @@ def _select_logical( fused_shared_experts_scoring_func=layer.shared_expert_scoring_func, routed_scaling_factor=layer.routed_scaling_factor, # TODO: a -1 route faults the TP a16w4 MoE; until that consumer is - # found, pure TP routes and runs its pad rows like real ones. + # found, pure TP routes its pad rows like real ones. mask_pad_rows=layer.use_ep, ) @@ -5405,9 +5405,7 @@ def _run_experts( prefix=f"{self.prefix}.fused_moe", x_scale=x_scale, scheduled_tokens_gpu=( - None - if gathered or not self.use_ep - else pass_scheduled_tokens_gpu(topk_ids.shape[0]) + None if gathered else pass_scheduled_tokens_gpu(topk_ids.shape[0]) ), ) if split: diff --git a/tests/model_ops/test_dp_gather_routed_rows.py b/tests/model_ops/test_dp_gather_routed_rows.py index 409fff9741..754bc1c64c 100644 --- a/tests/model_ops/test_dp_gather_routed_rows.py +++ b/tests/model_ops/test_dp_gather_routed_rows.py @@ -6,6 +6,7 @@ from types import SimpleNamespace +import pytest import torch from import_guard import skip_if_dependency_missing @@ -138,9 +139,11 @@ def prefill(rank, rows): assert runs[0] == {0, 2} and runs[1] == {1, 3} -def test_gathered_rows_are_placed_after_the_gather_and_uncounted(monkeypatch): +@pytest.mark.parametrize("use_ep", [False, True], ids=["tp", "ep"]) +def test_gathered_rows_are_placed_after_the_gather_and_uncounted(monkeypatch, use_ep): """Placement (the EPLB remap, which shifts the shared column) follows the - gather and the shared-id re-mark, and gathered rows carry no count.""" + gather and the shared-id re-mark. Only ungathered rows carry the pass's + count, under both TP and EP.""" calls = [] method = SimpleNamespace( place_experts=lambda layer, w, ids, gathered: ( @@ -156,8 +159,7 @@ def test_gathered_rows_are_placed_after_the_gather_and_uncounted(monkeypatch): activation=None, apply_router_weight_on_input=False, prefix="t", - # Only an expert-parallel layer hands its experts the row count. - use_ep=True, + use_ep=use_ep, ) count = torch.tensor([2], dtype=torch.int32) context = SimpleNamespace( From 544ca781fadece44f87404333040e8e73cbeb8c3 Mon Sep 17 00:00:00 2001 From: "1229775764@qq.com" Date: Tue, 3 Nov 2026 18:04:18 +0800 Subject: [PATCH 2/4] test(moe): exercise TP token count on CPU --- atom/model_ops/moe.py | 4 +++ tests/model_ops/test_dp_gather_routed_rows.py | 25 ++++++++++--------- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 25dbd90d9a..1c754df382 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -5401,6 +5401,10 @@ def _run_experts( apply_router_weight_on_input=self.apply_router_weight_on_input, prefix=f"{self.prefix}.fused_moe", x_scale=x_scale, + # Pure-TP padded graphs require ROCm/aiter#7294: its FlyDSL, + # FlyDSL-v2, and native-MXFP4 stage2 paths reject the sorting + # sentinel at the real token count instead of scattering every + # padded expert slot into that row. scheduled_tokens_gpu=( None if gathered else pass_scheduled_tokens_gpu(topk_ids.shape[0]) ), diff --git a/tests/model_ops/test_dp_gather_routed_rows.py b/tests/model_ops/test_dp_gather_routed_rows.py index 754bc1c64c..3ec0ce3fb0 100644 --- a/tests/model_ops/test_dp_gather_routed_rows.py +++ b/tests/model_ops/test_dp_gather_routed_rows.py @@ -8,13 +8,19 @@ import pytest import torch -from import_guard import skip_if_dependency_missing - -try: - # aiter/triton absent under bare non-GPU pytest - from atom.model_ops import moe -except ImportError as exc: - skip_if_dependency_missing(exc, "requires full atom import env") +from aiter_stub import stubbed_aiter + +with stubbed_aiter( + "aiter.fused_moe", + "aiter.jit.utils.chip_info", + "aiter.jit.utils.torch_guard", + "aiter.ops.flydsl.moe_common", + "aiter.ops.shuffle", + "aiter.ops.triton.quant", + "aiter.utility.fp4_utils", +): + from atom.model_ops import eplb, moe, topK + from atom.model_ops.fused_moe.expert_layout import SharedExpertMode class _TwoRankGroup: @@ -110,9 +116,6 @@ def test_gathered_shared_ids_split_the_rows_between_ep_ranks(monkeypatch): only where i % ep == r. Routed on their own ranks, the rows carry their owners' marks through the gather -- every rank would run the same rows' shared expert. Re-marked by gathered position, the ranks split the rows.""" - from atom.model_ops import topK - from atom.model_ops.fused_moe.expert_layout import SharedExpertMode - ep, top_k, shared, fake, local = 2, 2, 8, 9, 2 layer = SimpleNamespace( expert_layout=SimpleNamespace(mode=SharedExpertMode.LEGACY_AITER), @@ -190,8 +193,6 @@ def test_gathered_rows_are_placed_after_the_gather_and_uncounted(monkeypatch, us def test_gathered_rows_pick_one_replica_on_every_rank(monkeypatch): """The forced-remote replica hash mixes in the DP rank only for a rank's own rows; gathered rows are hashed by the row alone.""" - from atom.model_ops import eplb - monkeypatch.setattr(eplb, "_MANAGER", SimpleNamespace(replica_spread=(4, 3))) assert eplb._replica_spread() == (4, 3) assert eplb._replica_spread(gathered=True) == (1, 0) From bad5598b06a56a7ea55411fdf31700e09ae26312 Mon Sep 17 00:00:00 2001 From: "1229775764@qq.com" Date: Tue, 3 Nov 2026 18:08:19 +0800 Subject: [PATCH 3/4] test(moe): stub Triton for CPU collection --- tests/model_ops/test_dp_gather_routed_rows.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/model_ops/test_dp_gather_routed_rows.py b/tests/model_ops/test_dp_gather_routed_rows.py index 3ec0ce3fb0..2240b68793 100644 --- a/tests/model_ops/test_dp_gather_routed_rows.py +++ b/tests/model_ops/test_dp_gather_routed_rows.py @@ -18,6 +18,8 @@ "aiter.ops.shuffle", "aiter.ops.triton.quant", "aiter.utility.fp4_utils", + "triton", + "triton.language", ): from atom.model_ops import eplb, moe, topK from atom.model_ops.fused_moe.expert_layout import SharedExpertMode From 9705a5e7662a6571f9e27d43da0f23941a9d725b Mon Sep 17 00:00:00 2001 From: "1229775764@qq.com" Date: Tue, 3 Nov 2026 20:24:06 +0800 Subject: [PATCH 4/4] test(moe): isolate CUDA registration in CPU contracts --- tests/model_ops/test_dp_gather_routed_rows.py | 45 +++++++++++-------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/tests/model_ops/test_dp_gather_routed_rows.py b/tests/model_ops/test_dp_gather_routed_rows.py index 2240b68793..80f687b981 100644 --- a/tests/model_ops/test_dp_gather_routed_rows.py +++ b/tests/model_ops/test_dp_gather_routed_rows.py @@ -5,21 +5,30 @@ `peer_rows` rows of zeros.""" from types import SimpleNamespace +from unittest.mock import patch import pytest import torch from aiter_stub import stubbed_aiter -with stubbed_aiter( - "aiter.fused_moe", - "aiter.jit.utils.chip_info", - "aiter.jit.utils.torch_guard", - "aiter.ops.flydsl.moe_common", - "aiter.ops.shuffle", - "aiter.ops.triton.quant", - "aiter.utility.fp4_utils", - "triton", - "triton.language", +from atom.utils import forward_context as fc + +with ( + stubbed_aiter( + "aiter.fused_moe", + "aiter.jit.utils.chip_info", + "aiter.jit.utils.torch_guard", + "aiter.ops.flydsl.moe_common", + "aiter.ops.shuffle", + "aiter.ops.triton.quant", + "aiter.utility.fp4_utils", + "triton", + "triton.language", + ), + # These CPU contracts never call CUDA custom ops. Registering them would + # infer schemas from the fake AITER enum defaults and retain stub kernels + # in torch's process-wide registry after stubbed_aiter removes its modules. + patch("atom.utils.custom_register.direct_register_custom_op"), ): from atom.model_ops import eplb, moe, topK from atom.model_ops.fused_moe.expert_layout import SharedExpertMode @@ -147,8 +156,8 @@ def prefill(rank, rows): @pytest.mark.parametrize("use_ep", [False, True], ids=["tp", "ep"]) def test_gathered_rows_are_placed_after_the_gather_and_uncounted(monkeypatch, use_ep): """Placement (the EPLB remap, which shifts the shared column) follows the - gather and the shared-id re-mark. Only ungathered rows carry the pass's - count, under both TP and EP.""" + gather and the shared-id re-mark. Only the pass's whole ungathered rows + carry its count, under both TP and EP.""" calls = [] method = SimpleNamespace( place_experts=lambda layer, w, ids, gathered: ( @@ -173,22 +182,22 @@ def test_gathered_rows_are_placed_after_the_gather_and_uncounted(monkeypatch, us running_tokens=3, scheduled_tokens_gpu=count, ) - monkeypatch.setattr( - moe, "get_forward_context", lambda: SimpleNamespace(context=context) - ) - monkeypatch.setattr( - moe, "pass_scheduled_tokens_gpu", lambda rows: count if rows == 3 else None - ) + forward = SimpleNamespace(context=context) + monkeypatch.setattr(moe, "get_forward_context", lambda: forward) + monkeypatch.setattr(fc, "get_forward_context", lambda: forward) w, ids = torch.rand(3, 2), torch.zeros(3, 2, dtype=torch.int32) moe.FusedMoE._run_experts(layer, torch.zeros(3, 8), w, ids, gathered=True) moe.FusedMoE._run_experts(layer, torch.zeros(3, 8), w, ids) + moe.FusedMoE._run_experts(layer, torch.zeros(1, 8), w[:1], ids[:1]) assert calls == [ ("place", True), ("apply", None), ("place", False), ("apply", count), + ("place", False), + ("apply", None), ]