perf(dsv4/moe): remove redundant dependency edges

[Little-oil]

Summary

• exclude every redundant edge whose predecessor or successor is one of the 33 alloc tasks
• handle the corrected 83-edge deletion set: 8 cross-stage, 50 dispatch-to-expert, and 25 expert-local edges
• opt selected MoE intermediates out of automatic dependency tracking while preserving alloc creator retention and required transitive ordering
• keep the explicit combine wait while marking the routed-window read as no-dependency

Testing

• validated 33 alloc nodes, 83 unique deletable edges, and 0 alloc-touching deletable edges
• python -m ruff check models/deepseek/v4/moe.py models/deepseek/v4/expert_shared.py models/deepseek/v4/expert_routed.py
• python models/deepseek/v4/moe.py -p a2a3 --ep 2 -d 0,1 --compile-only
• device test passed as reported by the requester

[coderabbitai]

📝 Walkthrough

Walkthrough

DeepSeek-V4 MoE decode now buffers dispatch metadata locally, updates gather/combine ordering, extracts shared-routed accumulation into a submitted helper, and applies explicit manual dependencies to intermediate and expert quantization tensors.

Changes

DeepSeek-V4 decode updates

Layer / File(s)	Summary
Expert quantization tensor dependencies models/deepseek/v4/expert_routed.py, models/deepseek/v4/expert_shared.py	Dequantization scale tensors are pre-created with manual_dep=True and populated with reciprocal-derived values.
Dispatch metadata buffering and gather ordering models/deepseek/v4/moe.py	Dispatch writes per-source expert counts to recv_meta_local; gather reads that buffer with updated task ordering and manual-dependency-controlled outputs.
Shared-routed and combine integration models/deepseek/v4/moe.py	Shared-routed accumulation uses a new submitted helper, and combine receives recv_meta_local instead of distributed metadata.
Estimated code review effort: 3 (Moderate)	~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant moe
  participant dispatch
  participant recv_meta_local
  participant dispatch_gather
  participant combine
  moe->>dispatch: pass recv_meta_local and dispatch buffers
  dispatch->>recv_meta_local: store per-source expert counts
  dispatch->>dispatch_gather: signal metadata and local push completion
  dispatch_gather->>recv_meta_local: read local counts
  dispatch_gather->>combine: provide gathered expert outputs
  combine->>recv_meta_local: read local counts

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#777: Changes related dispatch metadata ordering in the same MoE pipeline.
• hw-native-sys/pypto-lib#783: Changes related combine and shared-routed dependency wiring.

Poem

A bunny found counts in a local little tray,
While gathers and combines aligned on their way.
Scales waited clearly, then hopped into place,
Shared routes submitted with nimble-foot grace.
MoE’s tidy trail now flows bright and true—
One carrot for ordering, and one for review! 🐇

🚥 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 accurately summarizes the main change: removing redundant dependency edges in DSV4 MoE execution.
Description check	✅ Passed	The description clearly matches the changeset and testing focus, especially dependency tracking and MoE ordering 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.}

[gemini-code-assist]

Code Review

This pull request optimizes the DeepSeek v4 MoE implementation by introducing a local metadata tensor recv_meta_local to reduce distributed reads, refactoring the shared_routed logic into an incore JIT function submitted via pl.spmd_submit, and marking several intermediate tensors with manual_dep=True to avoid automatic dependency tracking overhead. The reviewer feedback suggests further performance optimizations by also marking comb_ffn and recv_count_out with manual_dep=True, as their ordering is already transitively guaranteed.

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.

[Little-oil]

Method: removing dependency edges

The code-level mechanism depends on where the target edge comes from. In PyPTO, the effective fanin is:

auto-tracked dependencies U explicit deps

Therefore, if the same structural edge is produced by both mechanisms, both sources must be removed.

1. Remove an explicit TaskId edge

For edges introduced by deps=[...], remove the corresponding producer TaskId from the list.

Before:

with pl.spmd(
    N_LOCAL,
    name_hint="dispatch_gather",
    deps=[_meta_tid, _wait_tid, _push_tid],
) as _gather_tid:
    ...

After:

with pl.spmd(
    N_LOCAL,
    name_hint="dispatch_gather",
    deps=[_wait_tid, _push_tid],
) as _gather_tid:
    ...

This removes the direct dispatch_meta -> dispatch_gather explicit edge. The ordering remains available through _meta_tid -> _wait_tid -> dispatch_gather. The _push_tid edge is retained because the remote-arrival wait does not cover same-rank local puts.

The same rule applies to pl.submit, pl.spmd_submit, pl.at, pl.spmd, and pl.system.task_dummy dependency lists.

2. Disable automatic tracking for a tensor lifetime

For automatic edges created by reads/writes of a tensor, mark the allocation with manual_dep=True:

recv_scale_out = pl.create_tensor(
    [N_LOCAL, RECV_MAX],
    dtype=pl.FP32,
    manual_dep=True,
)

Every task that reads or writes this tensor skips OverlapMap lookup and producer insertion for its entire lifetime. The original allocation creator/owner retention still applies.

Examples in this PR include:

• x_norm_scale
• recv_scale_out
• recv_w_out
• recv_r_route_out
• recv_meta_local
• shared/routed expert h_tile_scale_dq

Use this only when every automatic edge carried by that tensor is redundant or is covered by another data path or an explicit dependency. manual_dep=True does not add replacement edges automatically.

3. Disable tracking for one call argument

When only one consumer should ignore a tensor dependency, use the narrower call-site form:

ffn_out, _reduce_tid = pl.spmd_submit(
    self.shared_routed,
    sh,
    pl.no_dep(routed_y_buf),
    ffn_out,
    num_tokens,
    core_num=T,
    deps=[_cwait_tid],
)

pl.no_dep(routed_y_buf) removes automatic lookup/insertion only for this argument of this submission. Other uses of routed_y_buf remain auto-tracked.

Here the local automatic combine -> shared_routed edge is suppressed, while deps=[_cwait_tid] retains the cross-rank arrival guarantee:

combine scatter -> combine_wait -> shared_routed

For an outlined pl.at block, the analogous form is no_dep_args=[tensor].

4. Change the data path before disabling tracking

If the original tensor still has other consumers that need automatic tracking, introduce a dedicated local snapshot rather than marking the original tensor manual.

This PR snapshots distributed metadata once:

recv_meta_local = pl.create_tensor(
    [N_RANKS, N_LOCAL],
    dtype=pl.INT32,
    manual_dep=True,
)

count = pl.read(recv_meta, [src, e])
pl.write(recv_meta_local, [src, e], count)

Gather and combine then read recv_meta_local instead of the distributed recv_meta. This avoids recreating direct metadata edges while leaving the original distributed-window behavior intact.

5. Preallocate expression results that need manual tracking

If a value is created as an expression result inside a task, first create explicit storage with manual_dep=True, then write the expression into it.

Before:

h_tile_scale_dq = pl.reshape(pl.recip(eh_sq_row), [RECV_TILE, 1])

After:

h_tile_scale_dq = pl.create_tensor(
    [RECV_TILE, 1],
    dtype=pl.FP32,
    manual_dep=True,
)

with pl.at(level=pl.Level.CORE_GROUP, name_hint="exp_h_q"):
    ...
    h_tile_scale_dq[:, :] = pl.reshape(
        pl.recip(eh_sq_row),
        [RECV_TILE, 1],
    )

This allows the scale path to opt out while the main h_tile_i8 -> exp_w2_mm -> exp_w2_act chain continues to provide ordering.

Practical rule

For each target structural edge A -> B, inspect every source of that edge:

• explicit deps: remove the producer TaskId;
• tensor-lifetime automatic tracking: use manual_dep=True;
• one call-site argument: use pl.no_dep(...);
• mixed/over-broad dataflow: split or snapshot the tensor first.

A structural edge disappears only after all of its sources have been removed. For example, removing dispatch_meta -> dispatch_gather requires both removing _meta_tid from gather's explicit deps and preventing the metadata tensor from recreating the automatic edge.

Host-read control tensors and cross-rank synchronization require extra care. In this PR, recv_count_out remains auto-tracked because orchestration reads it before routed-expert submission, and the explicit combine wait remains because local tensor tracking cannot prove peer remote writes have arrived. Follow-up #879 also illustrates why a lifetime-wide manual_dep=True opt-out must be validated against the full scheduling path.