Fix DeepSeek V4 routed expert serialization#673
Conversation
Move routed expert tile execution into manual scope with explicit task dependencies so independent expert tiles are not serialized by conservative dynamic recv_y writes. Re-register non-empty tile outputs for downstream combine with lightweight marker tasks.
📝 WalkthroughWalkthroughThe expert_routed kernel in models/deepseek/v4/expert_routed.py is restructured to use explicit manual dependency scoping. A new TILES_PER_EXPERT constant sizes a per-expert tile completion tracking array (w2_act_tids), and the loop uses pl.manual_scope(), manual_dep=True tensors, and pl.at deps-based barriers before writing to recv_y_flat, replacing prior conservative serialization. ChangesExpert Routed Kernel Dependency Rework
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant OuterLoop as expert_routed loop
participant TileCompute as Per-expert tile compute
participant DepArray as w2_act_tids
participant RecvY as recv_y_flat
OuterLoop->>TileCompute: enter pl.manual_scope() for expert tile
TileCompute->>TileCompute: compute gate/up, SwiGLU, INT8 requant, w2 GEMM
TileCompute->>DepArray: mark tile completion (manual_dep=True)
OuterLoop->>DepArray: pl.at(deps=[...]) check tile-done barrier
DepArray-->>OuterLoop: dependency resolved
OuterLoop->>RecvY: write/update recv_y_flat slice for tile
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}
There was a problem hiding this comment.
Code Review
This pull request introduces manual scoping and explicit task dependency tracking in the expert_routed function to optimize parallel execution and prevent conservative serialization of dynamic writes across experts. Specifically, it wraps the expert routing stages in a pl.manual_scope() block, tracks task IDs for the activation stages, and uses a subsequent loop with pl.at to synchronize on the completed tiles. There are no review comments, so we have no additional feedback to provide.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@models/deepseek/v4/expert_routed.py`:
- Line 46: `TILES_PER_EXPERT` is using floor division while the routing code
computes `n_tiles` with ceil division, which can make `w2_act_tids` too small
when `RECV_MAX` is not divisible by `RECV_TILE`. Update the `TILES_PER_EXPERT`
definition in `expert_routed.py` to use the same ceil-based tile count as the
loop, or add a hard divisibility assertion near the `RECV_MAX`/`RECV_TILE`
constants and ensure the same contract is enforced in the related routing logic
(`w2_act_tids`, `n_tiles`, and the tile loop sites).
- Around line 190-199: The final write in the W2 path can reintroduce nonzero
values in padded rows because `recv_weights` is applied after invalid rows were
already zeroed; update the `expert_routed.py` logic around `row_scale_blk`,
`w2_scale_chunk`, and the `recv_y_flat` assignment so `y_2d` is re-masked after
the last scaling step and before casting/writing. Keep the zeroing aligned with
the final tile write in the `pl.pipeline(W2_ACT_INNER, stage=2)` loop so
`recv_y[e, cnt[e]:, :]` remains zero even if padding-row weights are
uninitialized or NaN.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 4f045e2a-138f-4c26-bd82-dc96f1070344
📒 Files selected for processing (1)
models/deepseek/v4/expert_routed.py
| D_OUT_TILE_ACT = 512 | ||
| W2_INNER = 4 | ||
| W2_ACT_INNER = 8 | ||
| TILES_PER_EXPERT = RECV_MAX // RECV_TILE |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Make TILES_PER_EXPERT match the ceil-tile loop.
n_tiles is computed with ceiling division, but TILES_PER_EXPERT uses floor division. If RECV_MAX is ever not divisible by RECV_TILE, w2_act_tids[tile_idx] can index past the allocated task-id array. Use ceiling division or add a static divisibility guard.
Proposed fix
-TILES_PER_EXPERT = RECV_MAX // RECV_TILE
+TILES_PER_EXPERT = (RECV_MAX + RECV_TILE - 1) // RECV_TILEOr, if divisibility is required by the kernel contract:
+assert RECV_MAX % RECV_TILE == 0
TILES_PER_EXPERT = RECV_MAX // RECV_TILEAlso applies to: 75-79, 207-210
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@models/deepseek/v4/expert_routed.py` at line 46, `TILES_PER_EXPERT` is using
floor division while the routing code computes `n_tiles` with ceil division,
which can make `w2_act_tids` too small when `RECV_MAX` is not divisible by
`RECV_TILE`. Update the `TILES_PER_EXPERT` definition in `expert_routed.py` to
use the same ceil-based tile count as the loop, or add a hard divisibility
assertion near the `RECV_MAX`/`RECV_TILE` constants and ensure the same contract
is enforced in the related routing logic (`w2_act_tids`, `n_tiles`, and the tile
loop sites).
| row_scale_blk = pl.mul(h_tile_scale_dq, w_col_blk) | ||
| for dg in pl.pipeline(W2_ACT_INNER, stage=2): | ||
| d0 = d_base + dg * D_OUT_TILE_ACT | ||
| y_2d_i32 = y_i32[:, d0 : d0 + D_OUT_TILE_ACT] | ||
| w2_scale_chunk = routed_w2_scale[local_i : local_i + 1, d0 : 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_t0 : flat_t0 + RECV_TILE, d0 : d0 + D_OUT_TILE_ACT] = pl.cast( | ||
| y_2d, target_type=pl.BF16, mode="rint" | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Re-mask the final tile before writing recv_y_flat.
Invalid rows are zeroed before W2, but recv_weights is applied afterward. If padding-row weights are uninitialized/NaN, the full-tile write can violate the reference contract that recv_y[e, cnt[e]:, :] stays zero. Mask y_2d after the final scaling step.
Proposed fix
for dg in pl.pipeline(W2_ACT_INNER, stage=2):
d0 = d_base + dg * D_OUT_TILE_ACT
y_2d_i32 = y_i32[:, d0 : d0 + D_OUT_TILE_ACT]
w2_scale_chunk = routed_w2_scale[local_i : local_i + 1, d0 : 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)
+ y_2d = pl.fillpad(
+ pl.set_validshape(y_2d, valid_rows, D_OUT_TILE_ACT),
+ pad_value=pl.PadValue.zero,
+ )
recv_y_flat[flat_t0 : flat_t0 + RECV_TILE, d0 : d0 + D_OUT_TILE_ACT] = pl.cast(
y_2d, target_type=pl.BF16, mode="rint"
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| row_scale_blk = pl.mul(h_tile_scale_dq, w_col_blk) | |
| for dg in pl.pipeline(W2_ACT_INNER, stage=2): | |
| d0 = d_base + dg * D_OUT_TILE_ACT | |
| y_2d_i32 = y_i32[:, d0 : d0 + D_OUT_TILE_ACT] | |
| w2_scale_chunk = routed_w2_scale[local_i : local_i + 1, d0 : 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_t0 : flat_t0 + RECV_TILE, d0 : d0 + D_OUT_TILE_ACT] = pl.cast( | |
| y_2d, target_type=pl.BF16, mode="rint" | |
| ) | |
| row_scale_blk = pl.mul(h_tile_scale_dq, w_col_blk) | |
| for dg in pl.pipeline(W2_ACT_INNER, stage=2): | |
| d0 = d_base + dg * D_OUT_TILE_ACT | |
| y_2d_i32 = y_i32[:, d0 : d0 + D_OUT_TILE_ACT] | |
| w2_scale_chunk = routed_w2_scale[local_i_i : local_i + 1, d0 : 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) | |
| y_2d = pl.fillpad( | |
| pl.set_validshape(y_2d, valid_rows, D_OUT_TILE_ACT), | |
| pad_value=pl.PadValue.zero, | |
| ) | |
| recv_y_flat[flat_t0 : flat_t0 + RECV_TILE, d0 : d0 + D_OUT_TILE_ACT] = pl.cast( | |
| y_2d, target_type=pl.BF16, mode="rint" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@models/deepseek/v4/expert_routed.py` around lines 190 - 199, The final write
in the W2 path can reintroduce nonzero values in padded rows because
`recv_weights` is applied after invalid rows were already zeroed; update the
`expert_routed.py` logic around `row_scale_blk`, `w2_scale_chunk`, and the
`recv_y_flat` assignment so `y_2d` is re-masked after the last scaling step and
before casting/writing. Keep the zeroing aligned with the final tile write in
the `pl.pipeline(W2_ACT_INNER, stage=2)` loop so `recv_y[e, cnt[e]:, :]` remains
zero even if padding-row weights are uninitialized or NaN.
## Summary - replace the routed expert marker-task bridge with auto-scope W2 epilogue block assembly - keep gate/up/activation/W2 matmul in manual scope with explicit dependencies - write static W2 epilogue channel blocks through pl.assemble to avoid exp_w2_act_1/2 self-dependency chains ## Validation - python -m py_compile models/deepseek/v4/expert_routed.py - ruff check models/deepseek/v4/expert_routed.py - python models/deepseek/v4/decode_fwd.py -p a2a3 --ep 2 -d 0,1 --enable-l2-swimlane --compile-only - Full device run before rebasing onto current main: build_output/_jit_l3_decode_fwd_20260703_111753, PASS; exp_w2_act_2 self-chain edges dropped to 0 on rank0/rank1 swimlane traces. PR #673 is already merged, so this follow-up PR carries only the additional W2 epilogue assembly change on top of current main.
2 participants
Summary
expert_routedtile execution intomanual_scopeand wire explicit task dependencies for gate/up/act/quant/w2 stages, avoiding conservative serialization across independent local expert tiles.n_tilesloop so empty experts do not launch routed work; only non-empty tiles emit downstream marker tasks to makerecv_yvisible to the following combine stage.decode_fwdrank1 improved from 11057.26 us / 10779.18 us prior normal runs to 10240.10 us, and avoids the previously observed 38670.98 us / 46190.78 us serialized cases.expert_routed.pywith the same 25 non-empty tiles is 521.30 us vs 515.96 us before because it adds 25exp_routed_tile_donemarker tasks; this is the expected small marker overhead that buys correct downstream dependencies in the full graph.python -m py_compile models/deepseek/v4/expert_routed.py;ruff check --config ruff.toml models/deepseek/v4/expert_routed.py;python tests/lint/check_headers.py;python tests/lint/check_english_only.py; standaloneexpert_routed.py -p a2a3 -d 0 --enable-l2-swimlanePASS withrecv_yvalidation.Related Issues
None