From e6e6fef41a0b7d8b6bb67501114f0fbca8d500b1 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 15 Jul 2026 18:51:52 +0800 Subject: [PATCH] Skip dense FFN surrogate for MoE profiles --- .../frontier_moe_linear_plan.patch | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_linear_plan.patch diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_linear_plan.patch b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_linear_plan.patch new file mode 100644 index 0000000..1910739 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_linear_plan.patch @@ -0,0 +1,72 @@ +diff --git a/frontier/profiling/linear_op/profiling_plan.py b/frontier/profiling/linear_op/profiling_plan.py +--- a/frontier/profiling/linear_op/profiling_plan.py ++++ b/frontier/profiling/linear_op/profiling_plan.py +@@ -112,5 +112,11 @@ def build_profiling_plan( + + ffn_sharded_enabled = tp_size in ffn_tp_set ++ if is_moe and not _supports_share_expert(model_config): ++ # Routed experts are profiled by the MoE profiler. Keeping the ++ # dense FFN surrogate enabled executes an unprofiled fake weight ++ # path and can reject otherwise-valid attention TP layouts (for ++ # example Qwen3-235B TP8 with 128x128 block FP8 weights). ++ ffn_sharded_enabled = False + + padded_n_embd = model_config.embedding_dim + padded_n_expanded_embd = model_config.mlp_hidden_dim +diff --git a/tests/unit/test_moe_linear_profiling_plan.py b/tests/unit/test_moe_linear_profiling_plan.py +new file mode 100644 +--- /dev/null ++++ b/tests/unit/test_moe_linear_profiling_plan.py +@@ -0,0 +1,52 @@ ++from types import SimpleNamespace ++ ++from frontier.profiling.linear_op.profiling_plan import build_profiling_plan ++ ++ ++def test_routed_moe_plan_does_not_execute_dense_ffn_surrogate() -> None: ++ model_config = SimpleNamespace( ++ no_tensor_parallel=False, ++ embedding_dim=4096, ++ mlp_hidden_dim=1536, ++ num_q_heads=64, ++ num_kv_heads=4, ++ model_type="qwen3_moe", ++ post_attn_norm=True, ++ is_moe=True, ++ is_step2_mini=False, ++ ) ++ ++ plan = build_profiling_plan( ++ model_config=model_config, ++ tp_size=8, ++ attn_tp=[8], ++ ffn_tp=[8], ++ disable_replicated=False, ++ is_moe=True, ++ ) ++ ++ assert plan["attn_sharded_enabled"] is True ++ assert plan["ffn_sharded_enabled"] is False ++ assert plan["ffn_enabled"] is True ++ assert plan["padded_n_expanded_embd"] == 1536 ++ assert "attn_pre_proj" in plan["enabled_ops"] ++ assert "post_attention_layernorm" in plan["enabled_ops"] ++ assert "mlp_up_proj" not in plan["enabled_ops"] ++ assert "mlp_down_proj" not in plan["enabled_ops"] ++ ++ ++def test_dense_plan_still_executes_ffn_surrogate() -> None: ++ model_config = SimpleNamespace( ++ no_tensor_parallel=False, ++ embedding_dim=4096, ++ mlp_hidden_dim=1536, ++ num_q_heads=64, ++ num_kv_heads=4, ++ model_type="qwen2", ++ post_attn_norm=True, ++ is_moe=False, ++ is_step2_mini=False, ++ ) ++ plan = build_profiling_plan(model_config, 8, [8], [8], is_moe=False) ++ assert plan["ffn_sharded_enabled"] is True ++ assert "mlp_up_proj" in plan["enabled_ops"]