perf(dsv4/moe): predicate empty routed-expert stages - #809
Open
wangqin1723-max wants to merge 4 commits into
Open
perf(dsv4/moe): predicate empty routed-expert stages#809wangqin1723-max wants to merge 4 commits into
wangqin1723-max wants to merge 4 commits into
Conversation
Each routed expert stage (gate mm, up mm, SwiGLU act, h requant, w2 mm) used to sit inside a `pl.parallel(n_tiles)` row-tile loop, so every stage was launched once per row tile. Push the tile loop inside each stage instead: a stage now launches once per local expert and iterates its valid row tiles internally, cutting the runtime task count. - Widen the per-expert intermediates (gate/up/h) from [RECV_TILE, ...] to [RECV_MAX, ...] so data flows across stages. - Promote y_i32 and h_tile_scale_dq to flat expert-major [N_LOCAL_EXPERTS * RECV_MAX, ...] tensors, since they cross loops. - Split the w2 dequant + routing-weight epilogue into its own expert loop. recv_y's nz layout forces the writeback through pl.assemble, which is frame-level only, so that loop keeps a frame-level tile loop and assembles each tile right after the spmd produces it.
The per-expert valid row count (recv_expert_count) was read at the pl.parallel frame level, so a single AICPU scalar task fed every routed compute stage. Move the pl.read into each stage's own scope (gate mm, up mm, SwiGLU act, h requant, w2 mm) so the count scalar lives on the stage's cores. moe ep2 PASS on a2a3.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughDeepSeek-V4 FLASH now routes through 128 experts. Dispatch task handles flow into ChangesDeepSeek-V4 routed expert execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant moe
participant dispatch
participant expert_routed
participant SPMDStages
moe->>dispatch: Produce count_tid and gather_tid
dispatch-->>moe: Return task handles
moe->>expert_routed: Pass dispatch dependencies
expert_routed->>SPMDStages: Run gated, activation, quantization, and W2 stages
SPMDStages-->>expert_routed: Store routed output accumulators
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the DeepSeek-V4 MoE single-layer decode implementation, reducing the number of routed experts per rank from 32 to 16 and introducing explicit task-dependency tracking (using task IDs) to optimize SPMD execution and predication. The feedback suggests further optimizing the expert_routed kernel by carrying and saving the task ID for the exp_w2_mm stage (w2_tids), and hoisting the exp_w2_act stage outside the tile loop to reduce launch overhead, eliminate temporary tensors, and avoid unnecessary assembly operations.
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.
Comment on lines
+79
to
+80
| gate_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) | ||
| up_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) |
Contributor
There was a problem hiding this comment.
To support hoisting and predication of the exp_w2_act stage, we should carry its dependency task ID w2_tids across the passes, similar to gate_tids and up_tids.
Suggested change
| gate_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) | |
| up_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) | |
| gate_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) | |
| up_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) | |
| w2_tids = pl.array.create(N_LOCAL_EXPERTS, pl.TASK_ID) |
Comment on lines
+239
to
+264
| with pl.spmd( | ||
| D // (W2_INNER * D_OUT_TILE), | ||
| name_hint="exp_w2_mm", | ||
| deps=[count_tid, quant_tid], | ||
| predicate=(recv_expert_count[local_e, 0] > 0), | ||
| ) as _w2_tid: | ||
| w2_rows = pl.read(recv_expert_count, [local_e, 0]) | ||
| w2_tiles = (w2_rows + RECV_TILE - 1) // RECV_TILE | ||
| wb_idx = pl.tile.get_block_idx() | ||
| d_base = wb_idx * (W2_INNER * D_OUT_TILE) | ||
| for w2_t in pl.range(w2_tiles): | ||
| w2_t0 = w2_t * RECV_TILE | ||
| for dg in pl.range(W2_INNER): | ||
| d0 = d_base + dg * D_OUT_TILE | ||
| y_acc = pl.create_tensor([1, RECV_TILE, D_OUT_TILE], dtype=pl.INT32) | ||
| for k0 in pl.pipeline(0, MOE_INTER, INTER_K, stage=2): | ||
| h_k = h_tile_i8[:, k0 : k0 + INTER_K] | ||
| h_k = h_tile_i8[w2_t0 : w2_t0 + RECV_TILE, k0 : k0 + INTER_K] | ||
| w2_k = routed_w2[local_e : local_e + 1, d0 : d0 + D_OUT_TILE, k0 : k0 + INTER_K] | ||
| if k0 == 0: | ||
| y_acc = pl.matmul(h_k, w2_k, b_trans=True, out_dtype=pl.INT32) | ||
| else: | ||
| y_acc = pl.matmul_acc(y_acc, h_k, w2_k, b_trans=True) | ||
| y_i32[:, d0 : d0 + D_OUT_TILE] = pl.reshape(y_acc, [RECV_TILE, D_OUT_TILE]) | ||
|
|
||
| y_i32[ | ||
| e_flat_base + w2_t0 : e_flat_base + w2_t0 + RECV_TILE, | ||
| d0 : d0 + D_OUT_TILE, | ||
| ] = pl.reshape(y_acc, [RECV_TILE, D_OUT_TILE]) |
Contributor
There was a problem hiding this comment.
Save the task ID of exp_w2_mm into w2_tids so that the hoisted exp_w2_act stage can explicitly depend on it for correct chaining of predicated tasks.
with pl.spmd(
D // (W2_INNER * D_OUT_TILE),
name_hint="exp_w2_mm",
deps=[count_tid, quant_tid],
predicate=(recv_expert_count[local_e, 0] > 0),
) as w2_tid:
w2_rows = pl.read(recv_expert_count, [local_e, 0])
w2_tiles = (w2_rows + RECV_TILE - 1) // RECV_TILE
wb_idx = pl.tile.get_block_idx()
d_base = wb_idx * (W2_INNER * D_OUT_TILE)
for w2_t in pl.range(w2_tiles):
w2_t0 = w2_t * RECV_TILE
for dg in pl.range(W2_INNER):
d0 = d_base + dg * D_OUT_TILE
y_acc = pl.create_tensor([1, RECV_TILE, D_OUT_TILE], dtype=pl.INT32)
for k0 in pl.pipeline(0, MOE_INTER, INTER_K, stage=2):
h_k = h_tile_i8[w2_t0 : w2_t0 + RECV_TILE, k0 : k0 + INTER_K]
w2_k = routed_w2[local_e : local_e + 1, d0 : d0 + D_OUT_TILE, k0 : k0 + INTER_K]
if k0 == 0:
y_acc = pl.matmul(h_k, w2_k, b_trans=True, out_dtype=pl.INT32)
else:
y_acc = pl.matmul_acc(y_acc, h_k, w2_k, b_trans=True)
y_i32[
e_flat_base + w2_t0 : e_flat_base + w2_t0 + RECV_TILE,
d0 : d0 + D_OUT_TILE,
] = pl.reshape(y_acc, [RECV_TILE, D_OUT_TILE])
w2_tids[local_e] = w2_tidReferences
- In PyPTO on CANN/Ascend hardware, keeping a conditional check (e.g.,
if db == 0) inside apl.pipelineloop can be preferred over peeling the first iteration. This allows the first chunk's load to overlap with the rest of the pipeline rather than running as an un-pipelined prologue, provided the compiler successfully pipelines the loop-index branch.
Comment on lines
+266
to
+297
| for output_e in pl.parallel(N_LOCAL_EXPERTS): | ||
| output_rows = pl.read(recv_expert_count, [output_e, 0]) | ||
| output_tiles = (output_rows + RECV_TILE - 1) // RECV_TILE | ||
| output_flat_base = output_e * RECV_MAX | ||
| for output_t in pl.range(output_tiles): | ||
| output_t0 = output_t * RECV_TILE | ||
| flat_output_t0 = output_flat_base + output_t0 | ||
| recv_y_tile = pl.create_tensor([RECV_TILE, D], dtype=pl.BF16) | ||
| with pl.spmd(D // (W2_ACT_INNER * D_OUT_TILE_ACT), name_hint="exp_w2_act"): | ||
| db_idx = pl.tile.get_block_idx() | ||
| act_d_base = db_idx * (W2_ACT_INNER * D_OUT_TILE_ACT) | ||
| w_col_blk = pl.reshape( | ||
| recv_weights[local_e : local_e + 1, tt0 : tt0 + RECV_TILE], | ||
| recv_weights[output_e : output_e + 1, output_t0 : output_t0 + RECV_TILE], | ||
| [RECV_TILE, 1], | ||
| ) | ||
| row_scale_blk = pl.mul(h_tile_scale_dq, w_col_blk) | ||
| row_scale_blk = pl.mul( | ||
| h_tile_scale_dq[flat_output_t0 : flat_output_t0 + RECV_TILE, :], w_col_blk | ||
| ) | ||
| for dg in pl.pipeline(W2_ACT_INNER, stage=2): | ||
| act_d0 = act_d_base + dg * D_OUT_TILE_ACT | ||
| y_2d_i32 = y_i32[:, act_d0 : act_d0 + D_OUT_TILE_ACT] | ||
| w2_scale_chunk = routed_w2_scale[local_e : local_e + 1, act_d0 : act_d0 + D_OUT_TILE_ACT] | ||
| y_2d_i32 = y_i32[ | ||
| flat_output_t0 : flat_output_t0 + RECV_TILE, act_d0 : act_d0 + D_OUT_TILE_ACT | ||
| ] | ||
| w2_scale_chunk = routed_w2_scale[ | ||
| output_e : output_e + 1, act_d0 : act_d0 + D_OUT_TILE_ACT | ||
| ] | ||
| y_2d = pl.cast(y_2d_i32, target_type=pl.FP32, mode="none") | ||
| y_2d = pl.col_expand_mul(pl.row_expand_mul(y_2d, row_scale_blk), w2_scale_chunk) | ||
| recv_y_tile[:, act_d0 : act_d0 + D_OUT_TILE_ACT] = pl.cast( | ||
| y_2d, target_type=pl.BF16, mode="rint" | ||
| ) | ||
| recv_y_flat = pl.assemble(recv_y_flat, recv_y_tile, [flat_tt0, 0]) | ||
| recv_y_flat = pl.assemble(recv_y_flat, recv_y_tile, [flat_output_t0, 0]) |
Contributor
There was a problem hiding this comment.
The exp_w2_act stage is currently launched inside the for output_t in pl.range(output_tiles): loop, which incurs significant launch overhead for every tile.
We can optimize this by hoisting exp_w2_act outside the loop (similar to the other stages) and looping inside the SPMD block. This also allows us to predicate exp_w2_act on recv_expert_count[output_e, 0] > 0 and write directly to recv_y_flat, completely eliminating the temporary recv_y_tile tensor and the pl.assemble call.
for output_e in pl.parallel(N_LOCAL_EXPERTS):
output_rows = pl.read(recv_expert_count, [output_e, 0])
output_tiles = (output_rows + RECV_TILE - 1) // RECV_TILE
output_flat_base = output_e * RECV_MAX
with pl.spmd(
D // (W2_ACT_INNER * D_OUT_TILE_ACT),
name_hint="exp_w2_act",
deps=[count_tid, w2_tids[output_e]],
predicate=(recv_expert_count[output_e, 0] > 0),
):
db_idx = pl.tile.get_block_idx()
act_d_base = db_idx * (W2_ACT_INNER * D_OUT_TILE_ACT)
for output_t in pl.range(output_tiles):
output_t0 = output_t * RECV_TILE
flat_output_t0 = output_flat_base + output_t0
w_col_blk = pl.reshape(
recv_weights[output_e : output_e + 1, output_t0 : output_t0 + RECV_TILE],
[RECV_TILE, 1],
)
row_scale_blk = pl.mul(
h_tile_scale_dq[flat_output_t0 : flat_output_t0 + RECV_TILE, :], w_col_blk
)
for dg in pl.pipeline(W2_ACT_INNER, stage=2):
act_d0 = act_d_base + dg * D_OUT_TILE_ACT
y_2d_i32 = y_i32[
flat_output_t0 : flat_output_t0 + RECV_TILE, act_d0 : act_d0 + D_OUT_TILE_ACT
]
w2_scale_chunk = routed_w2_scale[
output_e : output_e + 1, act_d0 : act_d0 + D_OUT_TILE_ACT
]
y_2d = pl.cast(y_2d_i32, target_type=pl.FP32, mode="none")
y_2d = pl.col_expand_mul(pl.row_expand_mul(y_2d, row_scale_blk), w2_scale_chunk)
recv_y_flat[
flat_output_t0 : flat_output_t0 + RECV_TILE, act_d0 : act_d0 + D_OUT_TILE_ACT
] = pl.cast(y_2d, target_type=pl.BF16, mode="rint")
This was referenced Jul 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
recv_expert_count[local, 0] > 0so experts with zero tokens skip all compute.gate_tids/up_tids/act_tid/quant_tid) across the split compute passes and wire them as explicitdeps=[count_tid, gather_tid, ...]edges, so predicated tasks chain correctly while auto dependency tracking still covers tensor reads/writes.dispatch()returns its_meta_tid/_gather_tidand threads them intoexpert_routedas the count/gather producers.recv_xto a 2D ND view (recv_x_flat) to stop PTOAS inferring NZ on the 3-D view whenRECV_MAX == 32(EP=4), which conflicted with the row-major matmul input layout.exp_h_qfrompl.at(CORE_GROUP)topl.spmd(1)for a uniform predicated scope; lowern_routed_experts256 -> 128 (16/rank) so more experts hit the empty path.This branch also carries the routed-expert stage-hoisting from #738 (launch each stage once per expert; read the routed row count inside each stage), on top of which the predicate work is layered.
Related Issues
Stacks on #738.