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"]