Fix: compile DeepSeek V4 MoE at EP >= 4, re-anchor dispatch wait

[zhangqi-chen]

Summary

• expert_routed: take the gate/up matmul LHS from a flat [E * RECV_MAX, D] view of recv_x instead of slicing the 3-D parent. The 3-D slice emits an nd tensor_view that ptoas infers as nz once RECV_MAX > 16, so every exp_gate_mm / exp_up_mm failed to compile at EP >= 4. Matches the flat addressing recv_y / gate_i32 / up_i32 already use.
• dispatch_wait: drop deps=[_meta_tid] and anchor on an indices read instead — the same fanin dispatch_push waits on. The wait only spins on the peers' data_arrived counters, so chaining it behind dispatch_meta held its launch until that task's own meta spin and cumsum retired; now both cross-rank spins start together.
• dispatch_gather: take _meta_tid explicitly. recv_meta_local is manual_dep, so the cumsum -> gather edge was riding entirely on the dispatch_wait chain this PR drops (combine reads it transitively).
• Drop the dependency-only seed_dummy task feeding dispatch_meta, left over from the scheduling experiments in perf(dsv4/moe): tune early task resolution #811 — it has no predecessors and never delayed the task it fed.

moe.py x_next PASS at ep2 and ep4 on a2a3.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 4e288eda-9f24-4671-abc7-df6fcd28480c

📥 Commits

Reviewing files that changed from the base of the PR and between c364bbc and 1f1f9a3.

📒 Files selected for processing (2)

• models/deepseek/v4-flash/expert_routed.py
• models/deepseek/v4-flash/moe.py

---

📝 Walkthrough

Walkthrough

The routed-expert kernel now loads tiles from a flattened received-input buffer. The MoE dispatch pipeline removes a dummy dependency, anchors waiting through indices, and requires both wait and metadata tasks before gathering.

Changes

Routed MoE execution

Layer / File(s)	Summary
Flatten routed-expert input tiling models/deepseek/v4-flash/expert_routed.py	recv_x is flattened once, per-tile flat offsets are computed, and gate/up projection loads use the flattened buffer.

Layer / File(s)	Summary
Adjust dispatch task synchronization models/deepseek/v4-flash/moe.py	dispatch_meta no longer depends on a dummy task, dispatch_wait anchors on an indices read, and dispatch_gather depends on both wait and metadata tasks.

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

Sequence Diagram(s)

sequenceDiagram
  participant dispatch_push
  participant indices
  participant dispatch_wait
  participant dispatch_meta
  participant dispatch_gather
  dispatch_push->>indices: produce dispatch indices
  indices->>dispatch_wait: tensor read anchors the wait
  dispatch_meta->>dispatch_gather: provide metadata dependency
  dispatch_wait->>dispatch_gather: provide wait dependency

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#809: Updates related routed-expert indexing and MoE dispatch wiring in another DeepSeek variant.
• hw-native-sys/pypto-lib#811: Adjusts analogous dispatch scheduling and seed dependency behavior.
• hw-native-sys/pypto-lib#820: Changes related dispatch wait and gather synchronization.

Poem

A rabbit watched the tiles run flat,
While dispatch shed its dummy hat.
Metadata waited, gather knew,
The indices anchored signals too.
“Hop, hop!” the routed inputs fly—
Clean paths beneath the model sky.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main DeepSeek V4 MoE compilation and dispatch synchronization fix.
Description check	✅ Passed	The description directly matches the code changes and explains the same MoE compilation and scheduling updates.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

---

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.}