The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Currently, model.tp_plan returns _ep_plan instead of _tp_plan when enable_expert_parallel=True (see modeling_utils.py ). This means the EP plan fully replaces the TP plan.

This creates a tension between attn layer TP and EP :

If we include attention entries (like Llama4 does):

base_model_ep_plan = {
  "layers.*.self_attn.q_proj": "colwise",  # ← attention TP
  ...
  "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",  # ← expert EP
}

• TP+EP on a 2D mesh works (our test: TP=2, DP=4 ✓)
• But EP size is constrained to divide num_kv_heads if I understand correctly (Qwen3Attention.forward infers the shape (*input_shape, -1, self.head_dim). ). For Qwen3-30B-A3B (num_kv_heads=4) means the max EP=4. With EP=64 on a 1D mesh attention forward would probably crashes because q_proj output can't be reshaped into whole heads.

If we exclude attention entries (like gpt_oss does): Pure EP works at any scale (EP=16, 32, 64) BUT TP+EP on a 2D mesh loses attention sharding :/

What's the preferred approach? Should we go expert-only in the EP plan for maximum flexibility, or include attention for combined TP+EP at the cost of constraining EP size?

I think for now, we should create a new mesh dim ep that will use the base_model_ep_plan without any TP. And then use the base_model_tp_plan for attention layers only on the tp mesh dim

But EP size is constrained to divide num_kv_heads

It would be nice if we add a check on this some way or document this on ep plan / tp plan combination.

Recommended approach is to go with the best defaults / what makes most sense. pure EP would probably be "faster" as less coms, EP + FSDP is easier to get and probably makes more sense? but we want to allow people to have EP + TP if they want -> error early if simple!

It would be nice if we add a check on this some way or document this on ep plan / tp plan combination.
Great idea ! For now, I removed the attn layers from base_ep_plan because they limited the sharding scheme. I'll see where we can assert this post init to make if attn layer are in the base_ep_plan we assert that n_kv % ep_size ==0

pure EP would probably be "faster" as less coms, EP + FSDP is easier to get and probably makes more sense

Here are some numbers for Qwen3-30B-A3B (MoE, 128 experts, 8 active):

For comparison FSDP2 without EP:

I also ran EP=8 2D mesh vs EP=16 flat mesh on 2 nodes, just to see if the layout matters. so EP =9 in 2 nodes, means 8 experts (intra node) replicated across 2 nodes.

Flat and 2D meshes give identical performance, because I think the current EP implementation uses all-reduce (96% inter-node bandwidth), not all-to-all.

So EP seems to be slower than pure FSDP2 or FSDP2+CP for long ctx. 🤔

Yeah, this could just be because we do send 100% of all hidden states to all experts instead of sending just the ones allocated to that expert, same for when we reduce?

In any case that's good to have we can investigate later on perf issues / bottlenecks !

run-slow: qwen3_moe, qwen3_omni_moe, qwen3_vl_moe

is the fastest no?

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/qwen3_moe", "models/qwen3_omni_moe", "models/qwen3_vl_moe"]
quantizations: []

CI Results

Workflow Run ⚙️

Commit Info

✅ No failing test specific to this PR 🎉 👏 !

16k 8 64 64 1 1 64 21.4% 230,500 3,602 51.1 GB (64%)
is the fastest no?

for MFU and TPS we lose perf by doubling EP size each time :/

In any case that's good to have we can investigate later on perf issues / bottlenecks !

Yes I profiled the Qwen3 30B MoE briefly to get a sense of what was slow. It seems to be the all-reduce in MoE layers 🤔

The all_reduce_forward and it takes a huge slice of the forward comm

happy to merge once green!

@bot /style

Style fix bot fixed some files and pushed the changes.

[For maintainers] Suggested jobs to run (before merge)

run-slow: qwen3_moe, qwen3_omni_moe, qwen3_vl_moe

Finally CI is 🟢 !