Skip to content

fix(v3_2,qwen3-32b): migrate off retired incore/auto_chunk DSL + tidy naming #671

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(v3_2,qwen3-32b): migrate off retired incore/auto_chunk DSL + tidy naming #671

Show file tree
Hide file tree
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
139 changes: 61 additions & 78 deletions models/deepseek/v3_2/deepseek_v3_2_decode_back.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,89 +32,72 @@
HIDDEN_INV = 1.0 / HIDDEN

# tiling constants.
K_CHUNK = 128
Q_OUT_CHUNK = 64
MLP_OUT_CHUNK = 256
K_TILE = 128
Q_OUT_TILE = 64
MLP_OUT_TILE = 256
BATCH_TILE = 16


def build_deepseek_v3_2_decode_back_program(
batch: int = BATCH,
hidden_size: int = HIDDEN,
intermediate_size: int = INTERMEDIATE,
attn_out_size: int = ATTN_OUT,
ep_nodes: int = EP_NODES,
):
BATCH_SIZE = batch
HIDDEN_SIZE = hidden_size
INTER_SIZE = intermediate_size
ATTN_OUT_SIZE = attn_out_size
EP_NODES_SIZE = ep_nodes

ATTN_BLOCKS = (ATTN_OUT_SIZE + K_CHUNK - 1) // K_CHUNK
HIDDEN_BLOCKS = (HIDDEN_SIZE + K_CHUNK - 1) // K_CHUNK
Q_OUT_BLOCKS = (HIDDEN_SIZE + Q_OUT_CHUNK - 1) // Q_OUT_CHUNK
MLP_OUT_BLOCKS = (INTER_SIZE + MLP_OUT_CHUNK - 1) // MLP_OUT_CHUNK

def build_deepseek_v3_2_decode_back_program():
@pl.program
class DeepSeekV32DecodeBack:
@pl.function(type=pl.FunctionType.Opaque)
def deepseek_v3_2_decode_back_layer(
self,
hidden_states: pl.Tensor[[BATCH_SIZE, HIDDEN_SIZE], pl.BF16],
hidden_states: pl.Tensor[[BATCH, HIDDEN], pl.BF16],
node_id_t: pl.Tensor[[1], pl.INT32],
# combine buffer from cross-node communication
combine_buf: pl.Tensor[[EP_NODES_SIZE, BATCH_SIZE, ATTN_OUT_SIZE], pl.BF16],
wo: pl.Tensor[[ATTN_OUT_SIZE, HIDDEN_SIZE], pl.BF16],
post_rms_weight: pl.Tensor[[1, HIDDEN_SIZE], pl.FP32],
w_gate: pl.Tensor[[HIDDEN_SIZE, INTER_SIZE], pl.BF16],
w_up: pl.Tensor[[HIDDEN_SIZE, INTER_SIZE], pl.BF16],
w_down: pl.Tensor[[INTER_SIZE, HIDDEN_SIZE], pl.BF16],
out: pl.Out[pl.Tensor[[BATCH_SIZE, HIDDEN_SIZE], pl.BF16]],
) -> pl.Tensor[[BATCH_SIZE, HIDDEN_SIZE], pl.BF16]:
combine_buf: pl.Tensor[[EP_NODES, BATCH, ATTN_OUT], pl.BF16],
wo: pl.Tensor[[ATTN_OUT, HIDDEN], pl.BF16],
post_rms_weight: pl.Tensor[[1, HIDDEN], pl.FP32],
w_gate: pl.Tensor[[HIDDEN, INTERMEDIATE], pl.BF16],
w_up: pl.Tensor[[HIDDEN, INTERMEDIATE], pl.BF16],
w_down: pl.Tensor[[INTERMEDIATE, HIDDEN], pl.BF16],
out: pl.Out[pl.Tensor[[BATCH, HIDDEN], pl.BF16]],
) -> pl.Tensor[[BATCH, HIDDEN], pl.BF16]:
Comment on lines +41 to +57

@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

Keep tensor specs fixed with the zero-arg program builder.

The layer signature now hardcodes BATCH, HIDDEN, INTERMEDIATE, ATTN_OUT, and EP_NODES, but build_tensor_specs() still accepts override dimensions. Non-default specs would allocate/golden-validate shapes the compiled program cannot accept. Remove those parameters or assert they match the module constants.

Proposed guard 
 def build_tensor_specs(
     batch: int = BATCH,
     hidden_size: int = HIDDEN,
     intermediate_size: int = INTERMEDIATE,
     attn_out_size: int = ATTN_OUT,
     ep_nodes: int = EP_NODES,
 ):
+    assert batch == BATCH
+    assert hidden_size == HIDDEN
+    assert intermediate_size == INTERMEDIATE
+    assert attn_out_size == ATTN_OUT
+    assert ep_nodes == EP_NODES
🤖 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/v3_2/deepseek_v3_2_decode_back.py` around lines 41 - 57, The
decode-back layer signature in DeepSeekV32DecodeBack is now fixed to the module
constants, but build_tensor_specs still allows caller-provided dimensions that
can drift out of sync. Update build_tensor_specs (and any related spec
construction in build_deepseek_v3_2_decode_back_program) to either remove the
override parameters entirely or validate/assert they exactly match BATCH,
HIDDEN, INTERMEDIATE, ATTN_OUT, and EP_NODES before creating tensors, so the
generated specs always match the compiled program.

# Scope: output projection + residual + post-rms + MLP + residual.
node_id = pl.cast(pl.tensor.read(node_id_t, [0]), pl.INDEX)
for b0 in pl.range(0, BATCH_SIZE, BATCH_TILE):
resid1_tile = pl.create_tensor([BATCH_TILE, HIDDEN_SIZE], dtype=pl.FP32)
for b0 in pl.range(0, BATCH, BATCH_TILE):
resid1_tile = pl.create_tensor([BATCH_TILE, HIDDEN], dtype=pl.FP32)
# Read combine results from this node view.
combined_3d = pl.slice(
combine_buf, [1, BATCH_TILE, ATTN_OUT_SIZE], [node_id, b0, 0]
combine_buf, [1, BATCH_TILE, ATTN_OUT], [node_id, b0, 0]
)
combined = pl.reshape(combined_3d, [BATCH_TILE, ATTN_OUT_SIZE])
combined = pl.reshape(combined_3d, [BATCH_TILE, ATTN_OUT])

# O projection and residual.
for ob in pl.range(Q_OUT_BLOCKS):
o0 = ob * Q_OUT_CHUNK
for ob in pl.range(HIDDEN // Q_OUT_TILE):
o0 = ob * Q_OUT_TILE
with pl.at(level=pl.Level.CORE_GROUP):
a_chunk_0 = pl.slice(combined, [BATCH_TILE, K_CHUNK], [0, 0])
w_chunk_0 = pl.slice(wo, [K_CHUNK, Q_OUT_CHUNK], [0, o0])
a_chunk_0 = pl.slice(combined, [BATCH_TILE, K_TILE], [0, 0])
w_chunk_0 = pl.slice(wo, [K_TILE, Q_OUT_TILE], [0, o0])
o_acc = pl.matmul(a_chunk_0, w_chunk_0, out_dtype=pl.FP32)
for kb in pl.range(1, ATTN_BLOCKS):
k0 = kb * K_CHUNK
a_chunk = pl.slice(combined, [BATCH_TILE, K_CHUNK], [0, k0])
w_chunk = pl.slice(wo, [K_CHUNK, Q_OUT_CHUNK], [k0, o0])
for kb in pl.range(1, ATTN_OUT // K_TILE):
k0 = kb * K_TILE
a_chunk = pl.slice(combined, [BATCH_TILE, K_TILE], [0, k0])
w_chunk = pl.slice(wo, [K_TILE, Q_OUT_TILE], [k0, o0])
o_acc = pl.matmul_acc(o_acc, a_chunk, w_chunk)

with pl.at(level=pl.Level.CORE_GROUP):
resid = pl.cast(
pl.slice(hidden_states, [BATCH_TILE, Q_OUT_CHUNK], [b0, o0]), target_type=pl.FP32
pl.slice(hidden_states, [BATCH_TILE, Q_OUT_TILE], [b0, o0]), target_type=pl.FP32
)
resid1_tile = pl.assemble(resid1_tile, pl.add(o_acc, resid), [0, o0])

# Post RMSNorm.
post_norm_tile = pl.create_tensor([BATCH_TILE, HIDDEN_SIZE], dtype=pl.BF16)
post_norm_tile = pl.create_tensor([BATCH_TILE, HIDDEN], dtype=pl.BF16)
with pl.at(level=pl.Level.CORE_GROUP):
sq_sum = pl.full([1, BATCH_TILE], dtype=pl.FP32, value=0.0)
for kb in pl.range(HIDDEN_BLOCKS):
k0 = kb * K_CHUNK
x_chunk = pl.slice(resid1_tile, [BATCH_TILE, K_CHUNK], [0, k0])
for kb in pl.range(HIDDEN // K_TILE):
k0 = kb * K_TILE
x_chunk = pl.slice(resid1_tile, [BATCH_TILE, K_TILE], [0, k0])
sq_sum = pl.add(sq_sum, pl.reshape(pl.row_sum(pl.mul(x_chunk, x_chunk)), [1, BATCH_TILE]))
inv_rms = pl.recip(pl.sqrt(pl.add(pl.mul(sq_sum, HIDDEN_INV), EPS)))

for kb in pl.range(HIDDEN_BLOCKS):
k0 = kb * K_CHUNK
x_chunk = pl.slice(resid1_tile, [BATCH_TILE, K_CHUNK], [0, k0])
gamma = pl.slice(post_rms_weight, [1, K_CHUNK], [0, k0])
for kb in pl.range(HIDDEN // K_TILE):
k0 = kb * K_TILE
x_chunk = pl.slice(resid1_tile, [BATCH_TILE, K_TILE], [0, k0])
gamma = pl.slice(post_rms_weight, [1, K_TILE], [0, k0])
normed = pl.col_expand_mul(
pl.row_expand_mul(x_chunk, pl.reshape(inv_rms, [BATCH_TILE, 1])), gamma
)
Expand All @@ -123,28 +106,28 @@ def deepseek_v3_2_decode_back_layer(
)

# MLP.
mlp_tile = pl.create_tensor([BATCH_TILE, INTER_SIZE], dtype=pl.BF16)
for ob in pl.range(MLP_OUT_BLOCKS):
o0 = ob * MLP_OUT_CHUNK
mlp_tile = pl.create_tensor([BATCH_TILE, INTERMEDIATE], dtype=pl.BF16)
for ob in pl.range(INTERMEDIATE // MLP_OUT_TILE):
o0 = ob * MLP_OUT_TILE

with pl.at(level=pl.Level.CORE_GROUP):
post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, 0])
wg_0 = pl.slice(w_gate, [K_CHUNK, MLP_OUT_CHUNK], [0, o0])
post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, 0])
wg_0 = pl.slice(w_gate, [K_TILE, MLP_OUT_TILE], [0, o0])
gate_acc = pl.matmul(post_chunk_0, wg_0, out_dtype=pl.FP32)
for kb in pl.range(1, HIDDEN_BLOCKS):
k0 = kb * K_CHUNK
post_chunk = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, k0])
wg = pl.slice(w_gate, [K_CHUNK, MLP_OUT_CHUNK], [k0, o0])
for kb in pl.range(1, HIDDEN // K_TILE):
k0 = kb * K_TILE
post_chunk = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, k0])
wg = pl.slice(w_gate, [K_TILE, MLP_OUT_TILE], [k0, o0])
gate_acc = pl.matmul_acc(gate_acc, post_chunk, wg)

with pl.at(level=pl.Level.CORE_GROUP):
post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, 0])
wu_0 = pl.slice(w_up, [K_CHUNK, MLP_OUT_CHUNK], [0, o0])
post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, 0])
wu_0 = pl.slice(w_up, [K_TILE, MLP_OUT_TILE], [0, o0])
up_acc = pl.matmul(post_chunk_0, wu_0, out_dtype=pl.FP32)
for kb in pl.range(1, HIDDEN_BLOCKS):
k0 = kb * K_CHUNK
post_chunk = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, k0])
wu = pl.slice(w_up, [K_CHUNK, MLP_OUT_CHUNK], [k0, o0])
for kb in pl.range(1, HIDDEN // K_TILE):
k0 = kb * K_TILE
post_chunk = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, k0])
wu = pl.slice(w_up, [K_TILE, MLP_OUT_TILE], [k0, o0])
up_acc = pl.matmul_acc(up_acc, post_chunk, wu)

with pl.at(level=pl.Level.CORE_GROUP):
Expand All @@ -154,22 +137,22 @@ def deepseek_v3_2_decode_back_layer(
mlp_tile = pl.assemble(mlp_tile, mlp_chunk_bf16, [0, o0])

# Down projection + final residual writeback.
for dob in pl.range(HIDDEN_BLOCKS):
d0 = dob * K_CHUNK
for dob in pl.range(HIDDEN // K_TILE):
d0 = dob * K_TILE
with pl.at(level=pl.Level.CORE_GROUP):
mlp_chunk_0 = pl.slice(mlp_tile, [BATCH_TILE, MLP_OUT_CHUNK], [0, 0])
w_down_chunk_0 = pl.slice(w_down, [MLP_OUT_CHUNK, K_CHUNK], [0, d0])
mlp_chunk_0 = pl.slice(mlp_tile, [BATCH_TILE, MLP_OUT_TILE], [0, 0])
w_down_chunk_0 = pl.slice(w_down, [MLP_OUT_TILE, K_TILE], [0, d0])
down_acc = pl.matmul(mlp_chunk_0, w_down_chunk_0, out_dtype=pl.FP32)
for ob in pl.range(1, MLP_OUT_BLOCKS):
o0 = ob * MLP_OUT_CHUNK
down_mlp_chunk_bf16 = pl.slice(mlp_tile, [BATCH_TILE, MLP_OUT_CHUNK], [0, o0])
w_down_chunk = pl.slice(w_down, [MLP_OUT_CHUNK, K_CHUNK], [o0, d0])
for ob in pl.range(1, INTERMEDIATE // MLP_OUT_TILE):
o0 = ob * MLP_OUT_TILE
down_mlp_chunk_bf16 = pl.slice(mlp_tile, [BATCH_TILE, MLP_OUT_TILE], [0, o0])
w_down_chunk = pl.slice(w_down, [MLP_OUT_TILE, K_TILE], [o0, d0])
down_acc = pl.matmul_acc(down_acc, down_mlp_chunk_bf16, w_down_chunk)

with pl.at(level=pl.Level.CORE_GROUP):
out_chunk = pl.add(
down_acc,
pl.slice(resid1_tile, [BATCH_TILE, K_CHUNK], [0, d0]),
pl.slice(resid1_tile, [BATCH_TILE, K_TILE], [0, d0]),
)
out = pl.assemble(out, pl.cast(out_chunk, target_type=pl.BF16), [b0, d0])

Expand Down Expand Up @@ -244,9 +227,9 @@ def golden_deepseek_v3_2_decode_back(tensors):

# Match the chunked accumulation order used by the A2/A3 path to reduce
# BF16/FP32 drift at validation time.
k_chunk = K_CHUNK
q_out_chunk = Q_OUT_CHUNK
mlp_out_chunk = MLP_OUT_CHUNK
k_chunk = K_TILE
q_out_chunk = Q_OUT_TILE
mlp_out_chunk = MLP_OUT_TILE

combined = combine_buf[node_id]
resid1 = torch.zeros(batch, hidden_size, dtype=torch.float32)
Expand Down
Loading
Loading