PoC: fold dsv4/moe dispatch/combine wait into syncall barriers#768
PoC: fold dsv4/moe dispatch/combine wait into syncall barriers#768lwDavid wants to merge 2 commits into
Conversation
Experimental: eliminate the dispatch_wait / combine_wait runtime tasks by folding their notify/wait into the neighbouring SPMD launch, rendezvousing with a full-occupancy hard (FFTS) pl.system.syncall instead of a separate task boundary. Also merges dispatch_push + dispatch_gather into one dispatch_exchange task (removes the push->gather boundary + its scheduling gaps). Mechanism: each cross-rank TPUT self-drains (codegen emits PIPE_ALL after every put), so a barrier that rendezvouses all blocks is sufficient for block 0 to notify peers without racing the data — no separate landing fence needed. Status: NOT ready to merge. Correctness PASSes on moe.py's own golden, and the scheduling gaps between push/wait/gather are removed as intended, but: - Net perf is a wash-to-negative: dispatch end-to-end (incl. gaps) 33.6us -> 41.8us on the compute rank; decode_fwd overall ~unchanged/slightly slower. The full-occupancy (48-core) hard barriers cost more than the folded task boundaries save. - Intermittent spin observed once on moe.py (349ms, input-dependent). - AIV_CORES=48 is hardcoded for a2a3; a5 would deadlock (507018). Opened as a draft to document the approach. A real win needs a lightweight partial-occupancy barrier (the soft GM-polling form currently spins ~1.2s).
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes the DeepSeek-V4 MoE single-layer decode implementation by folding the separate dispatch/combine wait and gather tasks into full-occupancy SPMD launches using fast hardware barriers, which eliminates task-boundary stalls. Feedback points out a potential correctness issue in the folded tasks if N_LOCAL exceeds AIV_CORES, recommending a grid-stride loop to safely handle all configurations.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…w-native-sys#768) Address gemini review: the combine scatter and dispatch gather guarded expert work with `if e < N_LOCAL`, which silently drops experts e >= AIV_CORES when a config has N_LOCAL > AIV_CORES. Replace with a grid-stride loop `pl.range(block_idx, N_LOCAL, AIV_CORES)` so every expert is covered for any N_LOCAL/AIV_CORES ratio (no-op when N_LOCAL <= AIV_CORES, the current ep2 case). Note: the separate intermittent-correctness issue on sim (folded syncall barrier cannot reproduce the task-boundary drain/visibility the removed wait task provided) is not addressed here — it is a framework limitation, tracked in the PR description.
1 participant
Summary
Experiment: remove the
dispatch_waitandcombine_waitruntime tasks indsv4/moeby folding their cross-ranknotify/waitinto the neighbouring SPMDlaunch, using a full-occupancy hard (FFTS)
pl.system.syncallto rendezvous theblocks instead of a separate task boundary. Also merges
dispatch_push+dispatch_gatherinto a singledispatch_exchangetask, removing thepush -> gatherboundary and its scheduling gaps.Mechanism (why the fold is correct): every cross-rank
TPUTself-drains —codegen emits a
PIPE_ALLbarrier after each put (pto_ops_distributed.cpp) —so once a barrier rendezvouses all blocks, every outgoing put has landed and
block 0 can
notifypeers without racing the data. No separate landing fence isneeded; the former
pl.at(...wait...)task only existed to provide thatcross-block rendezvous, which the barrier now supplies.
Structural result: dispatch
meta + push + wait + gather(4 tasks) →meta + exchange(2 tasks); combinescatter + wait + reduce(3) →scatter + reduce(2). The
waittasks and thepush -> gatherscheduling gaps are gone.Status — NOT ready to merge
Correctness PASSes on
moe.py's own golden (fixed input, ep2), and thescheduling gaps are removed as intended, but the net result does not justify
merging:
(straggler-free), dispatch end-to-end incl. scheduling gaps: 33.6 µs →
41.8 µs. The removed gaps (~4.7 µs) are real, but the full-occupancy
(48-core) hard barriers + 32→48 block launch cost more than the folded task
boundaries save.
decode_fwdoverall is unchanged-to-slightly-slower(82.5 → 83.0 ms, same-scale run).
moe.pyrun (aroute_hash/barrier lane spun ~349 ms). Input-dependent; root cause is thefull-occupancy barrier's sensitivity to per-rank arrival skew.
AIV_CORES = 48is hardcoded to the a2a3 physical AIV count;a5 has a different count and would deadlock (507018).
moe.pyis shared by 6forward drivers (decode/prefill × fwd/layer/mtp), so this is a broad blast
radius.
Added
# ci: no-dep-genbecause the full-occupancy hardsyncalltrips 507018under dep_gen/DFX (same constraint
hc_prealready carries, pypto#1931).What would make it a real win
A lightweight partial-occupancy barrier: keep the fold at
pl.spmd(N_LOCAL)(32 blocks, no full-occupancy drain gap / idle chaperones) with a fast barrier.
The soft GM-polling
syncallis partial-occupancy but currently spins ~1.2 s, soit is not usable today — fixing that (framework side) is the real unlock. With a
cheap partial-occupancy barrier the removed scheduling gaps would turn net
positive.
Related Issues
None (references pypto#1931 for the dep_gen / hard-syncall constraint).