[Inference] Expert tensor parallelism is rejected for inference-optimized MoE and has no regression coverage

Summary

Current main (d564dd01de4549a4884be26b1d74e169798fd877) rejects expert_tensor_parallel_size > 1 for transformer_impl == "inference_optimized":

# megatron/core/transformer/transformer_config.py:1856
if self.transformer_impl == "inference_optimized" and self.num_moe_experts is not None:
    ...
    if self.expert_tensor_parallel_size > 1:
        raise ValueError(
            "Inference-optimized MoE layers does not support expert tensor parallelism."
        )

So inference MoE has no ETP path at all, and there is no regression coverage for how the dispatcher behaves under ragged or empty-rank inputs in that domain.

Measured capability on that commit

Produced by constructing the real configs, not by reading documentation:

One trap worth recording: the inference_optimized path also requires normalization="RMSNorm". A config that omits it fails inside __post_init__ for a reason unrelated to ETP, which makes a capability matrix report a false rejection.

Path contract relevant to any ETP work

dynamic_context sets

NCCLAllGatherDispatcher._use_allgather_v = not self.using_cuda_graph_this_step()

so the two dispatcher paths are not interchangeable:

• _use_allgather_v=False is the CUDA-graph (decode) path, where equal token counts across ranks hold by construction;
• _use_allgather_v=True is the eager/ragged (prefill) path, where counts may differ.

Feeding unequal counts to the equal-count path is not a reachable state and hangs rather than returning a wrong answer, because that gather assumes a shared row count. Any test or future ETP change has to keep each count profile with its legal path.

Relation to #4743

#4743 ("Inference: Add support for expert tensor parallelism") proposed this but is an unmerged draft with a stale base (main @ 6486d52c, mergeable_state=dirty), and its tp_ep plumbing is not in this base. This issue is therefore not a request to revive that diff as-is.

Proposal

Two separable pieces:

1. Regression coverage that does not depend on an ETP implementation (available now, see the linked PR): zero-local-token ranks on the ragged path, the compact row-order contract asserted from per-source-rank tags, equal/ragged/zero-local transitions on one live dispatcher, and a pinned capability test for the ETP boundary above. This is what the dispatcher currently lacks independent of ETP.

2. Inference ETP itself, which needs design agreement before code:

• which axis the fused expert weights are sharded along, and how the shards reconstruct the unsharded expert;
• where the combine-side partial-sum reduction over the ETP ranks happens, performed exactly once (a duplicated or missing reduction is the failure mode to test for);
• which of equal / ragged / graph / zero-local combinations are in scope for a first change;
• whether the existing guard should be narrowed or replaced, and the explicit, tested rejection kept for combinations that stay unsupported.

First-cut scope that seems natural: BF16, no FP4/FP8 dependence, and ETP=1 behaviour bit-identical to today.

Question for maintainers: is inference ETP still wanted on the inference-optimized path, and if so along which sharding axis? If the answer is that ETP is not planned there, the capability test in the linked PR should be kept as the documented boundary and the guard left in place.

What is deliberately not claimed

No ETP implementation is proposed as already working, and no performance benefit is claimed. The linked PR contains tests and a capability boundary only; it does not touch the guard or add an ETP path.