Derive AutoEP rank splits from the per-expert count exchange

[yh0903]

Redundancy

compute_split_plan (deepspeed/module_inject/auto_ep_layer.py) runs two all-to-all exchanges over overlapping metadata:

# exchange 1: [ep_size] per-rank totals -> output_splits
local_counts_tensor = count_matrix.sum(dim=1).clone()
dist.all_to_all_single(remote_counts_tensor, local_counts_tensor, group=ep_group)
output_splits = remote_counts_tensor.cpu().tolist()

# exchange 2: [ep_size, E_local] per-expert matrix -> local_counts
dist.all_to_all_single(received_counts_flat, local_expert_counts_flat, group=ep_group)
received_counts = received_counts_flat.view(ep_size, num_local_experts)

The second exchange already contains the first. received_counts[s, e] is the number of tokens source rank s sends to local expert e, so

received_counts.sum(dim=1)[s] == (tokens rank s sends to this rank) == output_splits[s]

element-wise and exactly — the first exchange transmits a strict projection of what the second one transmits.

The planner also issued two separate .cpu() calls, each of which is a device-to-host sync.

Fix

Drop the rank-total exchange and derive output_splits from the detailed one:

received_counts = received_counts_flat.view(ep_size, num_local_experts)
input_splits, output_splits = torch.stack((
    count_matrix.sum(dim=1),
    received_counts.sum(dim=1),
)).cpu().tolist()

Per split-plan call this changes:

	before	after
all-to-all collectives	2	1
device-to-host syncs	2	1

compute_split_plan_from_expert_indices (folded-TP path) now shares the same helper instead of duplicating the logic.

Only metadata is affected. The dispatch/combine payload _AllToAllV calls are untouched.

Measurements

EP=8 (2 nodes x 8 H100), 256 experts, 1024 tokens/rank, top_k=8, 50 trials, p50 of split-plan time per layer:

token distribution	before	after	saved
balanced	0.366 ms	0.190 ms	0.176 ms (-48%)
mild skew	0.352 ms	0.182 ms	0.171 ms (-49%)
severe skew	0.360 ms	0.183 ms	0.176 ms (-49%)

Collective time alone goes 0.120 ms -> 0.058 ms, consistent with halving the number of exchanges. The saving is flat across skew because what is removed is fixed per-call overhead, not payload-dependent work. For a 20-layer MoE model this is ~3.5 ms per forward pass.

Tests

tests/unit/v1/moe/test_autoep_unit.py::TestSplitPlan adds:

• an oracle check that the derived output_splits equal the per-rank totals, for every rank and three token distributions;
• a differential test against a local reimplementation of the old two-collective planner, asserting the old path issues 2 collectives, the new one issues 1, and all four SplitPlan fields are identical;
• a test that compute_split_plan_from_expert_indices produces an identical plan for the same routing assignment.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 58cfe55c7b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".