Skip to content

Fix DeepSeek V4 routed expert serialization #673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhangqi-chen merged 1 commit into from
Jul 2, 2026
Merged

Fix DeepSeek V4 routed expert serialization #673

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
262 changes: 143 additions & 119 deletions models/deepseek/v4/expert_routed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
D_OUT_TILE_ACT = 512
W2_INNER = 4
W2_ACT_INNER = 8
TILES_PER_EXPERT = RECV_MAX // RECV_TILE

@coderabbitai coderabbitai Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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_TILE

Or, if divisibility is required by the kernel contract:

+assert RECV_MAX % RECV_TILE == 0
 TILES_PER_EXPERT = RECV_MAX // RECV_TILE

Also 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).


assert RECV_MAX % RECV_TILE == 0, "RECV_MAX must be a whole number of RECV_TILE row-tiles"

Expand All @@ -62,131 +63,154 @@ def expert_routed(
recv_y: pl.Tensor[[N_LOCAL_EXPERTS, RECV_MAX, D], pl.BF16],
):
recv_y_flat = pl.reshape(recv_y, [N_LOCAL_EXPERTS * RECV_MAX, D])
w2_act_tids = pl.array.create(N_LOCAL_EXPERTS * TILES_PER_EXPERT, pl.TASK_ID)

# Each local expert owns a disjoint recv_y row slab. Keep the routed stages in
# manual scope so the runtime does not conservatively serialize those dynamic
# writes across experts; the following auto-scope no-ops re-register recv_y
# for downstream combine, only for non-empty expert tiles.
with pl.manual_scope():
for local_i in pl.parallel(N_LOCAL_EXPERTS):
n_rows = pl.read(recv_expert_count, [local_i, 0])
n_tiles = (n_rows + RECV_TILE - 1) // RECV_TILE
flat_base = local_i * RECV_MAX

for t in pl.parallel(n_tiles):
tile_idx = local_i * TILES_PER_EXPERT + t
t0 = t * RECV_TILE
flat_t0 = flat_base + t0
valid_rows = pl.min(RECV_TILE, n_rows - t0)

# gate (w1) and up (w3) cube matmul -> INT32 GM accumulators.
h_tile_fp32 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.FP32, manual_dep=True)
gate_i32 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.INT32, manual_dep=True)
up_i32 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.INT32, manual_dep=True)

with pl.spmd(MOE_INTER // (MM_GATE_INNER * MM_INTER_TILE), name_hint="exp_gate_mm") as gate_tid:
nb_idx = pl.tile.get_block_idx()
n_base = nb_idx * (MM_GATE_INNER * MM_INTER_TILE)
for ng in pl.range(MM_GATE_INNER):
n0 = n_base + ng * MM_INTER_TILE
gate_acc = pl.create_tensor([1, RECV_TILE, MM_INTER_TILE], dtype=pl.INT32)
for k0 in pl.pipeline(0, D, K_TILE, stage=2):
x_k = recv_x[local_i : local_i + 1, t0 : t0 + RECV_TILE, k0 : k0 + K_TILE]
w1_k = routed_w1[local_i : local_i + 1, n0 : n0 + MM_INTER_TILE, k0 : k0 + K_TILE]
if k0 == 0:
gate_acc = pl.matmul(x_k, w1_k, b_trans=True, out_dtype=pl.INT32)
else:
gate_acc = pl.matmul_acc(gate_acc, x_k, w1_k, b_trans=True)
gate_i32[:, n0 : n0 + MM_INTER_TILE] = pl.reshape(gate_acc, [RECV_TILE, MM_INTER_TILE])

with pl.spmd(MOE_INTER // (MM_GATE_INNER * MM_INTER_TILE), name_hint="exp_up_mm") as up_tid:
nb_idx = pl.tile.get_block_idx()
n_base = nb_idx * (MM_GATE_INNER * MM_INTER_TILE)
for ng in pl.range(MM_GATE_INNER):
n0 = n_base + ng * MM_INTER_TILE
up_acc = pl.create_tensor([1, RECV_TILE, MM_INTER_TILE], dtype=pl.INT32)
for k0 in pl.pipeline(0, D, K_TILE, stage=2):
x_k = recv_x[local_i : local_i + 1, t0 : t0 + RECV_TILE, k0 : k0 + K_TILE]
w3_k = routed_w3[local_i : local_i + 1, n0 : n0 + MM_INTER_TILE, k0 : k0 + K_TILE]
if k0 == 0:
up_acc = pl.matmul(x_k, w3_k, b_trans=True, out_dtype=pl.INT32)
else:
up_acc = pl.matmul_acc(up_acc, x_k, w3_k, b_trans=True)
up_i32[:, n0 : n0 + MM_INTER_TILE] = pl.reshape(up_acc, [RECV_TILE, MM_INTER_TILE])

with pl.spmd(
MOE_INTER // (ACT_GATE_INNER * ACT_INTER_TILE),
name_hint="exp_gate_up_act",
deps=[gate_tid, up_tid],
) as act_tid:
nb_idx = pl.tile.get_block_idx()
n_base = nb_idx * (ACT_GATE_INNER * ACT_INTER_TILE)
for ng in pl.pipeline(ACT_GATE_INNER, stage=2):
n0 = n_base + ng * ACT_INTER_TILE
gate_2d_i32 = gate_i32[:, n0 : n0 + ACT_INTER_TILE]
up_2d_i32 = up_i32[:, n0 : n0 + ACT_INTER_TILE]
recv_x_scale_dq = pl.reshape(recv_scale_dq[local_i : local_i + 1, t0 : t0 + RECV_TILE], [RECV_TILE, 1])
w1_scale_chunk = routed_w1_scale[local_i : local_i + 1, n0 : n0 + ACT_INTER_TILE]
w3_scale_chunk = routed_w3_scale[local_i : local_i + 1, n0 : n0 + ACT_INTER_TILE]
gate_2d = pl.cast(gate_2d_i32, target_type=pl.FP32, mode="none")
up_2d = pl.cast(up_2d_i32, target_type=pl.FP32, mode="none")
gate_2d = pl.col_expand_mul(pl.row_expand_mul(gate_2d, recv_x_scale_dq), w1_scale_chunk)
up_2d = pl.col_expand_mul(pl.row_expand_mul(up_2d, recv_x_scale_dq), w3_scale_chunk)
if SWIGLU_LIMIT > 0.0:
gate_2d = pl.minimum(gate_2d, SWIGLU_LIMIT)
up_2d = pl.maximum(pl.minimum(up_2d, SWIGLU_LIMIT), -SWIGLU_LIMIT)
sigmoid = pl.recip(pl.add(pl.exp(pl.neg(gate_2d)), 1.0))
silu = pl.mul(gate_2d, sigmoid)
gated = pl.mul(silu, up_2d)
gated_valid = pl.set_validshape(gated, valid_rows, ACT_INTER_TILE)
gated_masked = pl.fillpad(gated_valid, pad_value=pl.PadValue.zero)
h_tile_fp32[:, n0 : n0 + ACT_INTER_TILE] = gated_masked

h_tile_i8 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.INT8, manual_dep=True)
with pl.at(level=pl.Level.CORE_GROUP, name_hint="exp_h_q", deps=[act_tid]) as hq_tid:
eh_amax = pl.full([1, RECV_TILE], dtype=pl.FP32, value=INT8_AMAX_EPS)
for k0 in pl.pipeline(0, MOE_INTER, QUANT_TILE, stage=2):
eh_a_f32 = h_tile_fp32[:, k0 : k0 + QUANT_TILE]
eh_a_abs = pl.maximum(eh_a_f32, pl.neg(eh_a_f32))
eh_a_max = pl.reshape(pl.row_max(eh_a_abs), [1, RECV_TILE])
eh_amax = pl.maximum(eh_amax, eh_a_max)
eh_sq_row = pl.div(
pl.full([1, RECV_TILE], dtype=pl.FP32, value=INT8_SCALE_MAX), eh_amax
)
h_tile_scale_dq = pl.reshape(pl.recip(eh_sq_row), [RECV_TILE, 1])
eh_sq_col = pl.reshape(eh_sq_row, [RECV_TILE, 1])
for k1 in pl.pipeline(0, MOE_INTER, QUANT_TILE, stage=2):
eh_q_f32 = h_tile_fp32[:, k1 : k1 + QUANT_TILE]
eh_q_scaled = pl.row_expand_mul(eh_q_f32, eh_sq_col)
eh_q_i32 = pl.cast(eh_q_scaled, target_type=pl.INT32, mode="rint")
eh_q_half = pl.cast(eh_q_i32, target_type=pl.FP16, mode="round")
h_tile_i8[:, k1 : k1 + QUANT_TILE] = pl.cast(eh_q_half, target_type=pl.INT8, mode="trunc")

y_i32 = pl.create_tensor([RECV_TILE, D], dtype=pl.INT32, manual_dep=True)
with pl.spmd(D // (W2_INNER * D_OUT_TILE), name_hint="exp_w2_mm", deps=[hq_tid]) as w2_tid:
db_idx = pl.tile.get_block_idx()
d_base = db_idx * (W2_INNER * D_OUT_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]
w2_k = routed_w2[local_i : local_i + 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])

with pl.spmd(D // (W2_ACT_INNER * D_OUT_TILE_ACT), name_hint="exp_w2_act", deps=[w2_tid]) as w2_act_tid:
db_idx = pl.tile.get_block_idx()
d_base = db_idx * (W2_ACT_INNER * D_OUT_TILE_ACT)
w_col_blk = pl.reshape(
recv_weights[local_i : local_i + 1, t0 : t0 + RECV_TILE],
[RECV_TILE, 1],
)
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"
)
Comment on lines +190 to +199

@coderabbitai coderabbitai Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.

Suggested change
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.

w2_act_tids[tile_idx] = w2_act_tid

# Iterate each local expert, then process its rows in tiles of RECV_TILE.
for local_i in pl.parallel(N_LOCAL_EXPERTS):
n_rows = pl.read(recv_expert_count, [local_i, 0])
n_tiles = (n_rows + RECV_TILE - 1) // RECV_TILE
flat_base = local_i * RECV_MAX # row offset in recv_y_flat
flat_base = local_i * RECV_MAX

for t in pl.parallel(n_tiles):
t0 = t * RECV_TILE
flat_t0 = flat_base + t0

valid_rows = pl.min(RECV_TILE, n_rows - t0)

# gate (w1) and up (w3) cube matmul -> INT32 GM accumulators.
h_tile_fp32 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.FP32)
gate_i32 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.INT32)
up_i32 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.INT32)

for nb_idx in pl.spmd(MOE_INTER // (MM_GATE_INNER * MM_INTER_TILE), name_hint="exp_gate_mm"):
n_base = nb_idx * (MM_GATE_INNER * MM_INTER_TILE)
for ng in pl.range(MM_GATE_INNER):
n0 = n_base + ng * MM_INTER_TILE
gate_acc = pl.create_tensor([1, RECV_TILE, MM_INTER_TILE], dtype=pl.INT32)
for k0 in pl.pipeline(0, D, K_TILE, stage=2):
x_k = recv_x[local_i : local_i + 1, t0 : t0 + RECV_TILE, k0 : k0 + K_TILE]
w1_k = routed_w1[local_i : local_i + 1, n0 : n0 + MM_INTER_TILE, k0 : k0 + K_TILE]
if k0 == 0:
gate_acc = pl.matmul(x_k, w1_k, b_trans=True, out_dtype=pl.INT32)
else:
gate_acc = pl.matmul_acc(gate_acc, x_k, w1_k, b_trans=True)
gate_i32[:, n0 : n0 + MM_INTER_TILE] = pl.reshape(gate_acc, [RECV_TILE, MM_INTER_TILE])

for nb_idx in pl.spmd(MOE_INTER // (MM_GATE_INNER * MM_INTER_TILE), name_hint="exp_up_mm"):
n_base = nb_idx * (MM_GATE_INNER * MM_INTER_TILE)
for ng in pl.range(MM_GATE_INNER):
n0 = n_base + ng * MM_INTER_TILE
up_acc = pl.create_tensor([1, RECV_TILE, MM_INTER_TILE], dtype=pl.INT32)
for k0 in pl.pipeline(0, D, K_TILE, stage=2):
x_k = recv_x[local_i : local_i + 1, t0 : t0 + RECV_TILE, k0 : k0 + K_TILE]
w3_k = routed_w3[local_i : local_i + 1, n0 : n0 + MM_INTER_TILE, k0 : k0 + K_TILE]
if k0 == 0:
up_acc = pl.matmul(x_k, w3_k, b_trans=True, out_dtype=pl.INT32)
else:
up_acc = pl.matmul_acc(up_acc, x_k, w3_k, b_trans=True)
up_i32[:, n0 : n0 + MM_INTER_TILE] = pl.reshape(up_acc, [RECV_TILE, MM_INTER_TILE])

for nb_idx in pl.spmd(MOE_INTER // (ACT_GATE_INNER * ACT_INTER_TILE), name_hint="exp_gate_up_act"):
n_base = nb_idx * (ACT_GATE_INNER * ACT_INTER_TILE)
for ng in pl.pipeline(ACT_GATE_INNER, stage=2):
n0 = n_base + ng * ACT_INTER_TILE
gate_2d_i32 = gate_i32[:, n0 : n0 + ACT_INTER_TILE]
up_2d_i32 = up_i32[:, n0 : n0 + ACT_INTER_TILE]
recv_x_scale_dq = pl.reshape(recv_scale_dq[local_i : local_i + 1, t0 : t0 + RECV_TILE], [RECV_TILE, 1])
w1_scale_chunk = routed_w1_scale[local_i : local_i + 1, n0 : n0 + ACT_INTER_TILE]
w3_scale_chunk = routed_w3_scale[local_i : local_i + 1, n0 : n0 + ACT_INTER_TILE]
gate_2d = pl.cast(gate_2d_i32, target_type=pl.FP32, mode="none")
up_2d = pl.cast(up_2d_i32, target_type=pl.FP32, mode="none")
gate_2d = pl.col_expand_mul(pl.row_expand_mul(gate_2d, recv_x_scale_dq), w1_scale_chunk)
up_2d = pl.col_expand_mul(pl.row_expand_mul(up_2d, recv_x_scale_dq), w3_scale_chunk)
if SWIGLU_LIMIT > 0.0:
gate_2d = pl.minimum(gate_2d, SWIGLU_LIMIT)
up_2d = pl.maximum(pl.minimum(up_2d, SWIGLU_LIMIT), -SWIGLU_LIMIT)
sigmoid = pl.recip(pl.add(pl.exp(pl.neg(gate_2d)), 1.0))
silu = pl.mul(gate_2d, sigmoid)
gated = pl.mul(silu, up_2d)
# Zero rows >= valid_rows.
gated_valid = pl.set_validshape(gated, valid_rows, ACT_INTER_TILE)
gated_masked = pl.fillpad(gated_valid, pad_value=pl.PadValue.zero)
h_tile_fp32[:, n0 : n0 + ACT_INTER_TILE] = gated_masked

# Per-row A8 requant of h_tile (amax across full MOE_INTER row).
h_tile_i8 = pl.create_tensor([RECV_TILE, MOE_INTER], dtype=pl.INT8)
with pl.at(level=pl.Level.CORE_GROUP, name_hint="exp_h_q"):
eh_amax = pl.full([1, RECV_TILE], dtype=pl.FP32, value=INT8_AMAX_EPS)
for k0 in pl.pipeline(0, MOE_INTER, QUANT_TILE, stage=2):
eh_a_f32 = h_tile_fp32[:, k0 : k0 + QUANT_TILE]
eh_a_abs = pl.maximum(eh_a_f32, pl.neg(eh_a_f32))
eh_a_max = pl.reshape(pl.row_max(eh_a_abs), [1, RECV_TILE])
eh_amax = pl.maximum(eh_amax, eh_a_max)
eh_sq_row = pl.div(
pl.full([1, RECV_TILE], dtype=pl.FP32, value=INT8_SCALE_MAX), eh_amax
)
h_tile_scale_dq = pl.reshape(pl.recip(eh_sq_row), [RECV_TILE, 1])
eh_sq_col = pl.reshape(eh_sq_row, [RECV_TILE, 1])
for k1 in pl.pipeline(0, MOE_INTER, QUANT_TILE, stage=2):
eh_q_f32 = h_tile_fp32[:, k1 : k1 + QUANT_TILE]
eh_q_scaled = pl.row_expand_mul(eh_q_f32, eh_sq_col)
eh_q_i32 = pl.cast(eh_q_scaled, target_type=pl.INT32, mode="rint")
eh_q_half = pl.cast(eh_q_i32, target_type=pl.FP16, mode="round")
h_tile_i8[:, k1 : k1 + QUANT_TILE] = pl.cast(eh_q_half, target_type=pl.INT8, mode="trunc")

# w2 cube matmul -> INT32 GM accumulators.
y_i32 = pl.create_tensor([RECV_TILE, D], dtype=pl.INT32)

for db_idx in pl.spmd(D // (W2_INNER * D_OUT_TILE), name_hint="exp_w2_mm"):
d_base = db_idx * (W2_INNER * D_OUT_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]
w2_k = routed_w2[local_i : local_i + 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])

for db_idx in pl.spmd(D // (W2_ACT_INNER * D_OUT_TILE_ACT), name_hint="exp_w2_act"):
d_base = db_idx * (W2_ACT_INNER * D_OUT_TILE_ACT)
# Per-row scale = h-dequant scale * routing weight.
w_col_blk = pl.reshape(
recv_weights[local_i : local_i + 1, t0 : t0 + RECV_TILE],
[RECV_TILE, 1],
)
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"
)
tile_idx = local_i * TILES_PER_EXPERT + t
flat_t0 = flat_base + t * RECV_TILE
with pl.at(level=pl.Level.CORE_GROUP, name_hint="exp_routed_tile_done", deps=[w2_act_tids[tile_idx]]):
recv_y_flat[flat_t0 : flat_t0 + RECV_TILE, 0:16] = recv_y_flat[
flat_t0 : flat_t0 + RECV_TILE, 0:16
]

return recv_y

Expand Down
Loading