[MoE] Use CPU split-size sum for EP permute output size

[sanketpurandare]

Stacked PRs:

• ->[MoE] Use CPU split-size sum for EP permute output size #3627
• [graph_trainer] Add EP overlap scheduling pass #3328
• [graph_trainer] Add graph EP chunking pass #3325
• [graph_trainer] Add EP overlap eager chunking scaffolding #3363
• [graph_trainer] Support hinted symbolic input dims in tracing #3362

---

[MoE] Use CPU split-size sum for EP permute output size

Compute the EP permute output size from the CPU output-splits tensor and pass that size into repeat_interleave(output_size=...) and the matching arange path. This keeps graph capture from constructing a long symbolic expression over the split-size list for the permute size.

This depends on PyTorch support for unbacked SymInt output sizes in repeat_interleave/arange:

pytorch/pytorch#186573

Related TorchTitan issue:

#3336

[tianyu-l]

does this only not work?

[sanketpurandare]

It works but num_global_tokens_per_local_expert_E is a GPU tensor and summing it up gives a GPU scalar value. if we pass this to repeat_interleave or arange it will cause another cpu-gpu sync.

[tianyu-l]

why if-else

[sanketpurandare]

I kept the output size passing as optional, so that users of the _permute function may or may not pass it.

[tianyu-l]

do you want this to be on CPU or GPU

[sanketpurandare]

total should be a python int on cpu

[tianyu-l]

can you use routed_input_RD.shape[0]? Here's what claude potentially hallucinates:

---

Why routed_input_RD.shape[0] is known without adding a sync — provenance

This is the part that makes it free rather than just relocating the sync. Trace where the shape
came from in dispatch():

  output_splits_list = output_splits.tolist()      # line 298 — the ONE existing D2H sync
  ...
  routed_input_RD = all_to_all_single_autograd(
      routed_input_ND, output_splits_list, input_splits_list, self.ep_mesh,   # line 301-306
  )

all_to_all_single is given the output split sizes as a Python list of host ints. The output
tensor's row count is sum(output_splits_list) — computed by host-side arithmetic at allocation
time, before any device work. So routed_input_RD.shape[0] is a host int that was fixed by the
.tolist() sync that already happened at line 298 (the unavoidable one, needed because
all_to_all_single requires list splits).

Therefore reading routed_input_RD.shape[0] in _permute reuses a value the host already knows. It
adds zero new transfers, while deleting the repeat_interleave sync from 2a. Net: two D2H syncs
→ one.

(Under CooR/compile the same holds symbolically: output_splits.tolist() produces unbacked
SymInts, sum(...) is a host-side symbolic int, the a2a output shape[0] is that SymInt, and
repeat_interleave accepts a SymInt output_size — so no .item() graph break is inserted either.)

[tianyu-l]

do you want this on CPU or GPU

[sanketpurandare]

total should be a python int on cpu

[tianyu-l]

Could you explain a bit more on before vs. after? I couldn't follow all the compiler language

[sanketpurandare]

When we convert a tensor to a list each item of that list is a sym-int. If I sum the elements of the list I get a sym-int expression. For instance tensor t = [a, b, c, d], when I call list_t = t.to_list(), list_t = [u0, u1, u2, u3], where each ui is a sym-int. And then if I sum it up I will get a expression u0+u1+u2+u3. When the compiler gets such an expression for size allocation in repeat_interleave or arange, it will try to calculate the upper bound which ends up being large. Instead if we first call sum on the tensor and then call .item I will get one sym-int

[tianyu-l]

the sentence doesn't finish:

Instead if we first call sum on the tensor and then call .item I will get one sym-int

so that the largest bound is 65536? compared with the sum of a list of 65536's?