Perf: overlap DeepSeek V4 MoE combine wait with the scatter

[zhangqi-chen]

Summary

Applies the dispatch-side handshake restructure (#783) to combine, and drops two dependency edges the auto tensormap already provides.

Combine

• Fold the combine_arrived notify into the scatter blocks: each of the N_LOCAL blocks signals every peer after its own puts, and the wait now expects N_LOCAL * moe_epoch. Takes one task launch off the cross-rank critical path. The [1, D] puts are single-shot TPUTs, which drain themselves before the notify issues.
• Make combine_wait wait-only and detach it from the scatter so it spins on the peers' counters while this rank's scatter runs. It anchors on exp_w2_act by reading recv_y_flat, the same view the scatter reads — a wait with no dependency at all is dispatched as soon as the scheduler reaches it and then spins holding a core group, starving the scatters it waits on. Reading the 3-D recv_y instead fails to compile once the expert count changes: it is an exp_w2_act cube output whose inferred layout is nz, while the outlined scalar-read kernel declares nd.
• Restore the plain pl.spmd reduce. pl.spmd_submit + the @pl.jit shared_routed method existed only to pass pl.no_dep(routed_y_buf), and that opt-out is now wrong: with combine_wait off the scatter, the local RAW edge on routed_y_buf is the only thing ordering the reduce after this rank's own dst == my_rank puts. The constant-false specialization branch and the self. noqa go with it.

Dispatch

• Drop the explicit _push_tid edge on dispatch_gather: dispatch_push -> dispatch_gather already carries tensormap edges on recv_x / recv_aux / recv_route.

Measured on a2a3, EP=2, ptoas 0.48, same two cards, --enable-l2-swimlane=1, on the rank that does not stall in the barrier: combine end -> shared_routed start goes from 7.4us to -7.8us (the reduce pre-stages onto a core and starts before the scatter's last block retires), and the tail from exp_w2_act end to hc_post end from 37.4us to 32.5us per MoE layer.

Also trims the verbose rationale comments through dispatch and combine.

[coderabbitai]

📝 Walkthrough

Walkthrough

The MoE dispatch and combine task graph now uses independent synchronization windows, folded payload-arrival notifications, revised wait thresholds, early-resolving tasks, and updated dependency anchors.

Changes

MoE synchronization

Layer / File(s)	Summary
Dispatch barrier and gather dependencies models/deepseek/v4-flash/moe.py	Dispatch documents independent metadata and payload windows, folds payload arrival into the push phase, and makes dispatch_gather depend only on the wait task.
Combine wait and scatter dependencies models/deepseek/v4-flash/moe.py	Combine enables early resolution for scatter and wait, removes scatter gating, anchors receive-buffer dependencies, changes the arrival threshold to moe_epoch * N_LOCAL, and preserves routed-buffer ordering.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• hw-native-sys/pypto-lib#777: Updates the same MoE dispatch/combine barrier logic and folded notifications.
• hw-native-sys/pypto-lib#783: Changes combine synchronization and the moe_epoch * N_LOCAL wait threshold.
• hw-native-sys/pypto-lib#803: Adjusts DSV4 MoE dispatch and combine dependency wiring.

Poem

A rabbit hops through barriers bright,
Folding payloads into flight.
Waits resolve and anchors hold,
Gather follows paths retold.
MoE hops in ordered cheer!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: overlapping the DeepSeek V4 MoE combine wait with scatter.
Description check	✅ Passed	The description is directly related to the changeset and explains the combine and dispatch dependency updates.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}