From 68df160342a3a46ef3913a3c0f26165d40115a28 Mon Sep 17 00:00:00 2001 From: zhangqi-chen Date: Thu, 2 Jul 2026 11:36:44 +0800 Subject: [PATCH] fix(v3_2,qwen3-32b): migrate off retired incore/auto_chunk DSL + tidy naming pypto retired pl.incore / pl.auto_incore / pl.auto_chunk and the loop chunk= kwarg (#1819, #1895), which broke the v3.2 decode_front and prefill_back CI cases at compile time. Migrate to pl.at(level=CORE_GROUP) and manual loop tiling (parallel step + pl.at + pl.range(chunk)). - prefill_back: also fixes a runtime task-ring deadlock. Stage 4 nested the down-proj loop inside the MLP loop (~16560 core scopes, over the 16384 task ring). Restructured to the decode_back shape: materialize mlp_tile, then a separate down-proj loop that merges the final residual. BATCH 4 -> 1 to keep golden fast. - qwen3_32b_decode: drop the now-redundant pl.auto_chunk, keep pl.split(pl.SplitMode.UP_DOWN). - Naming cleanup across v3.2 and qwen3_32b: _CHUNK -> _TILE, flatten the program builders to use module-level constants (drop _CFG/_SIZE aliases and builder params), inline _BLOCKS. - Remove deepseek_v3_2_prefill_front_draft.py. Validated on a2a3 (--ptoas 0.46): decode_front, decode_back, prefill_back, qwen3_32b_decode all PASS. --- .../v3_2/deepseek_v3_2_decode_back.py | 139 ++--- .../v3_2/deepseek_v3_2_decode_front.py | 113 ++-- .../v3_2/deepseek_v3_2_prefill_back.py | 227 ++++--- .../v3_2/deepseek_v3_2_prefill_front_draft.py | 584 ------------------ models/qwen3/32b/qwen3_32b_decode.py | 244 ++++---- 5 files changed, 339 insertions(+), 968 deletions(-) delete mode 100644 models/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.py diff --git a/models/deepseek/v3_2/deepseek_v3_2_decode_back.py b/models/deepseek/v3_2/deepseek_v3_2_decode_back.py index cf3eef40..3a8c1b01 100644 --- a/models/deepseek/v3_2/deepseek_v3_2_decode_back.py +++ b/models/deepseek/v3_2/deepseek_v3_2_decode_back.py @@ -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]: # 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 ) @@ -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): @@ -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]) @@ -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) diff --git a/models/deepseek/v3_2/deepseek_v3_2_decode_front.py b/models/deepseek/v3_2/deepseek_v3_2_decode_front.py index 7c9c6069..764492d6 100644 --- a/models/deepseek/v3_2/deepseek_v3_2_decode_front.py +++ b/models/deepseek/v3_2/deepseek_v3_2_decode_front.py @@ -53,21 +53,20 @@ # Scope1 tiles RMSNORM_K = 512 PROJ_K = 512 -Q_OUT_CHUNK = 64 -KV_OUT_CHUNK = 64 -LORA_CHUNK = 64 +Q_OUT_TILE = 64 +KV_OUT_TILE = 64 +LORA_TILE = 64 # Scope2 tiles -K_CHUNK = 128 -IDX_OUT_CHUNK = 128 -KIDX_OUT_CHUNK = 64 -QREDUCE_OUT_CHUNK = 64 +K_TILE = 128 +IDX_OUT_TILE = 128 +KIDX_OUT_TILE = 64 +QREDUCE_OUT_TILE = 64 QREDUCE_BATCH_TILE = 16 -WEIGHTS_OUT_CHUNK = 16 +WEIGHTS_OUT_TILE = 16 # Scope3 tiles SEQ_TILE = 64 -MAX_SEQ_BLOCKS = (MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE Q_VALID = 1 Q_PAD = 16 SORT_LEN = 8192 @@ -75,8 +74,8 @@ # Scope4 tiles ATTN_SCALE = 1.0 / (QK_HEAD_DIM**0.5) -Q_LATENT_CHUNK = 128 -V_OUT_CHUNK = 32 +Q_LATENT_TILE = 128 +V_OUT_TILE = 32 MATMUL_ROW_PAD = 16 @@ -136,14 +135,14 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 1.2: Project qr = normed @ wq_a. qr_fp32 = pl.create_tensor([BATCH, Q_LORA_RANK], dtype=pl.FP32) - for q0 in pl.parallel(0, Q_LORA_RANK, LORA_CHUNK): + for q0 in pl.parallel(0, Q_LORA_RANK, LORA_TILE): with pl.at(level=pl.Level.CORE_GROUP, name_hint="q_lora_proj"): q_tile_a = normed_states[:, 0 : PROJ_K] - q_tile_b = wq_a[0 : PROJ_K, q0 : q0 + LORA_CHUNK] + q_tile_b = wq_a[0 : PROJ_K, q0 : q0 + LORA_TILE] q_acc = pl.matmul(q_tile_a, q_tile_b, out_dtype=pl.FP32) for k0 in pl.range(PROJ_K, HIDDEN, PROJ_K): q_tile_a_i = normed_states[:, k0 : k0 + PROJ_K] - q_tile_b_i = wq_a[k0 : k0 + PROJ_K, q0 : q0 + LORA_CHUNK] + q_tile_b_i = wq_a[k0 : k0 + PROJ_K, q0 : q0 + LORA_TILE] q_acc = pl.matmul_acc(q_acc, q_tile_a_i, q_tile_b_i) qr_fp32 = pl.assemble(qr_fp32, q_acc, [0, q0]) @@ -151,29 +150,29 @@ def deepseek_v3_2_decode_front_scope1234( qr_out = pl.create_tensor([BATCH, Q_LORA_RANK], dtype=pl.BF16) with pl.at(level=pl.Level.CORE_GROUP, name_hint="q_lora_rmsnorm"): q_partial_sq = pl.full([1, BATCH], dtype=pl.FP32, value=0.0) - for k0 in pl.range(0, Q_LORA_RANK, LORA_CHUNK): - qr_chunk_fp32 = qr_fp32[:, k0 : k0 + LORA_CHUNK] + for k0 in pl.range(0, Q_LORA_RANK, LORA_TILE): + qr_chunk_fp32 = qr_fp32[:, k0 : k0 + LORA_TILE] q_partial = pl.reshape(pl.row_sum(pl.mul(qr_chunk_fp32, qr_chunk_fp32)), [1, BATCH]) q_partial_sq = pl.add(q_partial_sq, q_partial) q_variance = pl.reshape(pl.add(pl.mul(q_partial_sq, Q_LORA_INV), EPS), [BATCH, 1]) q_inv_rms = pl.recip(pl.sqrt(q_variance)) - for k0 in pl.range(0, Q_LORA_RANK, LORA_CHUNK): - qr_chunk_bf16 = pl.cast(qr_fp32[:, k0 : k0 + LORA_CHUNK], target_type=pl.BF16) + for k0 in pl.range(0, Q_LORA_RANK, LORA_TILE): + qr_chunk_bf16 = pl.cast(qr_fp32[:, k0 : k0 + LORA_TILE], target_type=pl.BF16) qr_chunk_fp32 = pl.cast(qr_chunk_bf16, target_type=pl.FP32) - q_gamma = q_norm_weight[:, k0 : k0 + LORA_CHUNK] + q_gamma = q_norm_weight[:, k0 : k0 + LORA_TILE] q_normed = pl.col_expand_mul(pl.row_expand_mul(qr_chunk_fp32, q_inv_rms), q_gamma) qr_out = pl.assemble(qr_out, pl.cast(q_normed, target_type=pl.BF16), [0, k0]) # Stage 1.4: Project q_proj = qr @ wq_b. q_proj = pl.create_tensor([BATCH, NUM_HEADS * QK_HEAD_DIM], dtype=pl.BF16) - for q0 in pl.parallel(0, NUM_HEADS * QK_HEAD_DIM, Q_OUT_CHUNK): + for q0 in pl.parallel(0, NUM_HEADS * QK_HEAD_DIM, Q_OUT_TILE): with pl.at(level=pl.Level.CORE_GROUP, name_hint="q_head_proj"): - q_chunk_init = qr_out[:, 0 : LORA_CHUNK] - wq_b_init = wq_b[0 : LORA_CHUNK, q0 : q0 + Q_OUT_CHUNK] + q_chunk_init = qr_out[:, 0 : LORA_TILE] + wq_b_init = wq_b[0 : LORA_TILE, q0 : q0 + Q_OUT_TILE] q_out_acc = pl.matmul(q_chunk_init, wq_b_init, out_dtype=pl.FP32) - for k0 in pl.range(LORA_CHUNK, Q_LORA_RANK, LORA_CHUNK): - q_chunk = qr_out[:, k0 : k0 + LORA_CHUNK] - wq_b_chunk = wq_b[k0 : k0 + LORA_CHUNK, q0 : q0 + Q_OUT_CHUNK] + for k0 in pl.range(LORA_TILE, Q_LORA_RANK, LORA_TILE): + q_chunk = qr_out[:, k0 : k0 + LORA_TILE] + wq_b_chunk = wq_b[k0 : k0 + LORA_TILE, q0 : q0 + Q_OUT_TILE] q_out_acc = pl.matmul_acc(q_out_acc, q_chunk, wq_b_chunk) with pl.at(level=pl.Level.CORE_GROUP, name_hint="q_head_proj_write"): @@ -181,14 +180,14 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 1.5: Project kv_a = normed @ wkv_a. kv_a_out = pl.create_tensor([BATCH, KV_A_OUT], dtype=pl.BF16) - for kv0 in pl.parallel(0, KV_A_OUT, KV_OUT_CHUNK): + for kv0 in pl.parallel(0, KV_A_OUT, KV_OUT_TILE): with pl.at(level=pl.Level.CORE_GROUP, name_hint="kv_a_proj"): kv_tile_a = normed_states[:, 0 : PROJ_K] - kv_tile_b = wkv_a[0 : PROJ_K, kv0 : kv0 + KV_OUT_CHUNK] + kv_tile_b = wkv_a[0 : PROJ_K, kv0 : kv0 + KV_OUT_TILE] kv_acc = pl.matmul(kv_tile_a, kv_tile_b, out_dtype=pl.FP32) for k0 in pl.range(PROJ_K, HIDDEN, PROJ_K): kv_tile_a_i = normed_states[:, k0 : k0 + PROJ_K] - kv_tile_b_i = wkv_a[k0 : k0 + PROJ_K, kv0 : kv0 + KV_OUT_CHUNK] + kv_tile_b_i = wkv_a[k0 : k0 + PROJ_K, kv0 : kv0 + KV_OUT_TILE] kv_acc = pl.matmul_acc(kv_acc, kv_tile_a_i, kv_tile_b_i) # Stage 1.6: Final KV output cast. @@ -257,14 +256,14 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 2.1: q_idx_full = wq_b_idx(qr_out). q_idx_full = pl.create_tensor([BATCH, INDEX_Q_OUT], dtype=pl.BF16) - for q0 in pl.parallel(0, INDEX_Q_OUT, IDX_OUT_CHUNK): + for q0 in pl.parallel(0, INDEX_Q_OUT, IDX_OUT_TILE): with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_q_idx_proj"): - s2_q_chunk_init = qr_out[:, 0 : LORA_CHUNK] - s2_wq_chunk_init = wq_b_idx[0 : LORA_CHUNK, q0 : q0 + IDX_OUT_CHUNK] + s2_q_chunk_init = qr_out[:, 0 : LORA_TILE] + s2_wq_chunk_init = wq_b_idx[0 : LORA_TILE, q0 : q0 + IDX_OUT_TILE] s2_q_acc = pl.matmul(s2_q_chunk_init, s2_wq_chunk_init, out_dtype=pl.FP32) - for k0 in pl.range(LORA_CHUNK, Q_LORA_RANK, LORA_CHUNK): - qr_chunk = qr_out[:, k0 : k0 + LORA_CHUNK] - wq_chunk = wq_b_idx[k0 : k0 + LORA_CHUNK, q0 : q0 + IDX_OUT_CHUNK] + for k0 in pl.range(LORA_TILE, Q_LORA_RANK, LORA_TILE): + qr_chunk = qr_out[:, k0 : k0 + LORA_TILE] + wq_chunk = wq_b_idx[k0 : k0 + LORA_TILE, q0 : q0 + IDX_OUT_TILE] s2_q_acc = pl.matmul_acc(s2_q_acc, qr_chunk, wq_chunk) with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_q_idx_proj_write"): @@ -272,14 +271,14 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 2.2: k_idx = wk_idx(hidden_states). k_idx = pl.create_tensor([BATCH, INDEX_HEAD_DIM], dtype=pl.BF16) - for k1 in pl.parallel(0, INDEX_HEAD_DIM, KIDX_OUT_CHUNK): + for k1 in pl.parallel(0, INDEX_HEAD_DIM, KIDX_OUT_TILE): with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_k_idx_proj"): - s2_x_init = hidden_states[:, 0 : K_CHUNK] - wk_init = wk_idx[0 : K_CHUNK, k1 : k1 + KIDX_OUT_CHUNK] + s2_x_init = hidden_states[:, 0 : K_TILE] + wk_init = wk_idx[0 : K_TILE, k1 : k1 + KIDX_OUT_TILE] s2_k_acc = pl.matmul(s2_x_init, wk_init, out_dtype=pl.FP32) - for k0 in pl.range(K_CHUNK, HIDDEN, K_CHUNK): - s2_x_chunk = hidden_states[:, k0 : k0 + K_CHUNK] - wk_chunk = wk_idx[k0 : k0 + K_CHUNK, k1 : k1 + KIDX_OUT_CHUNK] + for k0 in pl.range(K_TILE, HIDDEN, K_TILE): + s2_x_chunk = hidden_states[:, k0 : k0 + K_TILE] + wk_chunk = wk_idx[k0 : k0 + K_TILE, k1 : k1 + KIDX_OUT_TILE] s2_k_acc = pl.matmul_acc(s2_k_acc, s2_x_chunk, wk_chunk) with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_k_idx_proj_write"): @@ -343,19 +342,19 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 2.6: weights = weights_proj(hidden_states) * n_heads^-0.5 * head_dim^-0.5. weights_proj_bf16 = pl.create_tensor([HIDDEN, INDEX_HEADS], dtype=pl.BF16) with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_weights_proj_bf16"): - for k0 in pl.range(0, HIDDEN, K_CHUNK): - wp_tile = pl.cast(weights_proj[k0 : k0 + K_CHUNK, :], target_type=pl.BF16) + for k0 in pl.range(0, HIDDEN, K_TILE): + wp_tile = pl.cast(weights_proj[k0 : k0 + K_TILE, :], target_type=pl.BF16) weights_proj_bf16 = pl.assemble(weights_proj_bf16, wp_tile, [k0, 0]) weights = pl.create_tensor([BATCH, INDEX_HEADS], dtype=pl.FP32) - for w0 in pl.parallel(0, INDEX_HEADS, WEIGHTS_OUT_CHUNK): + for w0 in pl.parallel(0, INDEX_HEADS, WEIGHTS_OUT_TILE): with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_weights_matmul"): - s2_x_init = hidden_states[:, 0 : K_CHUNK] - wp_init = weights_proj_bf16[0 : K_CHUNK, w0 : w0 + WEIGHTS_OUT_CHUNK] + s2_x_init = hidden_states[:, 0 : K_TILE] + wp_init = weights_proj_bf16[0 : K_TILE, w0 : w0 + WEIGHTS_OUT_TILE] s2_w_acc = pl.matmul(s2_x_init, wp_init, out_dtype=pl.FP32) - for k0 in pl.range(K_CHUNK, HIDDEN, K_CHUNK): - s2_x_chunk = hidden_states[:, k0 : k0 + K_CHUNK] - wp_chunk = weights_proj_bf16[k0 : k0 + K_CHUNK, w0 : w0 + WEIGHTS_OUT_CHUNK] + for k0 in pl.range(K_TILE, HIDDEN, K_TILE): + s2_x_chunk = hidden_states[:, k0 : k0 + K_TILE] + wp_chunk = weights_proj_bf16[k0 : k0 + K_TILE, w0 : w0 + WEIGHTS_OUT_TILE] s2_w_acc = pl.matmul_acc(s2_w_acc, s2_x_chunk, wp_chunk) weights = pl.assemble(weights, s2_w_acc, [0, w0]) @@ -372,13 +371,13 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 2.8: Reduce q_idx_full across heads with weights to get q_idx_out. q_idx_out = pl.create_tensor([BATCH, INDEX_HEAD_DIM], dtype=pl.BF16) weights_flat = pl.reshape(weights, [BATCH * INDEX_HEADS]) - with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.auto_chunk], name_hint="s2_q_reduce"): - for d0 in pl.parallel(0, INDEX_HEAD_DIM, QREDUCE_OUT_CHUNK): + with pl.at(level=pl.Level.CORE_GROUP, name_hint="s2_q_reduce"): + for d0 in pl.parallel(0, INDEX_HEAD_DIM, QREDUCE_OUT_TILE): for b in pl.range(BATCH): - s2_acc_b = pl.full([1, QREDUCE_OUT_CHUNK], dtype=pl.FP32, value=0.0) + s2_acc_b = pl.full([1, QREDUCE_OUT_TILE], dtype=pl.FP32, value=0.0) for h in pl.range(INDEX_HEADS): s2_q_h_b = pl.cast( - pl.slice(q_idx_full, [1, QREDUCE_OUT_CHUNK], [b, h * INDEX_HEAD_DIM + d0]), + pl.slice(q_idx_full, [1, QREDUCE_OUT_TILE], [b, h * INDEX_HEAD_DIM + d0]), target_type=pl.FP32, ) s2_w_h_b = pl.read(weights_flat, [b * INDEX_HEADS + h]) @@ -399,7 +398,7 @@ def deepseek_v3_2_decode_front_scope1234( s3_q_padded = pl.create_tensor([BATCH * Q_PAD, INDEX_HEAD_DIM], dtype=pl.BF16) for b in pl.parallel(BATCH): - s3_score_tiles = pl.create_tensor([MAX_SEQ_BLOCKS * Q_PAD, SEQ_TILE], dtype=pl.FP32) + s3_score_tiles = pl.create_tensor([(MAX_SEQ // SEQ_TILE) * Q_PAD, SEQ_TILE], dtype=pl.FP32) # Stage 3.0: Pad q_idx_out and pre-fill scores. with pl.at(level=pl.Level.CORE_GROUP, name_hint="s3_init"): @@ -488,8 +487,8 @@ def deepseek_v3_2_decode_front_scope1234( # Stage 4.2: Project q_nope to latent space chunk-by-chunk. with pl.at(level=pl.Level.CORE_GROUP, name_hint="s4_q_nope_latent_proj"): q_w_row = h * QK_NOPE_HEAD_DIM - for q0 in pl.range(0, KV_LORA_RANK, Q_LATENT_CHUNK): - w_qn_h_chunk = w_q_nope_to_latent_2d[q_w_row : q_w_row + QK_NOPE_HEAD_DIM, q0 : q0 + Q_LATENT_CHUNK] + for q0 in pl.range(0, KV_LORA_RANK, Q_LATENT_TILE): + w_qn_h_chunk = w_q_nope_to_latent_2d[q_w_row : q_w_row + QK_NOPE_HEAD_DIM, q0 : q0 + Q_LATENT_TILE] q_nope_latent_part = pl.matmul(q_nope_padded, w_qn_h_chunk, out_dtype=pl.FP32) q_nope_latent_batch = pl.assemble(q_nope_latent_batch, q_nope_latent_part, [0, q0]) @@ -537,8 +536,8 @@ def deepseek_v3_2_decode_front_scope1234( with pl.at(level=pl.Level.CORE_GROUP, name_hint="s4_v_proj"): ctx_v_batch = pl.full([MATMUL_ROW_PAD, V_HEAD_DIM], dtype=pl.FP32, value=0.0) v_w_row = h * KV_LORA_RANK - for v0 in pl.range(0, V_HEAD_DIM, V_OUT_CHUNK): - wv_tile_chunk = w_latent_to_v_2d[v_w_row : v_w_row + KV_LORA_RANK, v0 : v0 + V_OUT_CHUNK] + for v0 in pl.range(0, V_HEAD_DIM, V_OUT_TILE): + wv_tile_chunk = w_latent_to_v_2d[v_w_row : v_w_row + KV_LORA_RANK, v0 : v0 + V_OUT_TILE] v_part_batch = pl.matmul(ctx_latent_bf16_batch, wv_tile_chunk, out_dtype=pl.FP32) ctx_v_batch = pl.assemble(ctx_v_batch, v_part_batch, [0, v0]) diff --git a/models/deepseek/v3_2/deepseek_v3_2_prefill_back.py b/models/deepseek/v3_2/deepseek_v3_2_prefill_back.py index 7edf23fc..2854ae52 100644 --- a/models/deepseek/v3_2/deepseek_v3_2_prefill_back.py +++ b/models/deepseek/v3_2/deepseek_v3_2_prefill_back.py @@ -18,7 +18,7 @@ import pypto.language as pl -BATCH = 4 +BATCH = 1 MAX_SEQ = 128 HIDDEN = 7168 INTERMEDIATE = 18432 @@ -30,48 +30,32 @@ EPS = 1e-6 HIDDEN_INV = 1.0 / HIDDEN -K_CHUNK = 128 -Q_OUT_CHUNK = 64 -MLP_OUT_CHUNK = 128 +K_TILE = 128 +Q_OUT_TILE = 64 +MLP_OUT_TILE = 128 TOK_TILE = 64 +# RMSNorm block-chunking: each core scope normalizes NORM_TILE contiguous +# HIDDEN // K_TILE blocks (= 56 blocks; 56 = 7 * 8). +NORM_TILE = 8 -def build_deepseek_v3_2_prefill_back_program( - batch: int = BATCH, - max_seq_len: int = MAX_SEQ, - hidden_size: int = HIDDEN, - intermediate_size: int = INTERMEDIATE, - attn_out_size: int = ATTN_OUT, - ep_nodes: int = EP_NODES, -): - BATCH_CFG = batch - MAX_SEQ_CFG = max_seq_len - HIDDEN_CFG = hidden_size - INTER_CFG = intermediate_size - ATTN_OUT_CFG = attn_out_size - EP_NODES_CFG = ep_nodes - - ATTN_BLOCKS = (ATTN_OUT_CFG + K_CHUNK - 1) // K_CHUNK - HIDDEN_BLOCKS = (HIDDEN_CFG + K_CHUNK - 1) // K_CHUNK - Q_OUT_BLOCKS = (HIDDEN_CFG + Q_OUT_CHUNK - 1) // Q_OUT_CHUNK - MLP_OUT_BLOCKS = (INTER_CFG + MLP_OUT_CHUNK - 1) // MLP_OUT_CHUNK - +def build_deepseek_v3_2_prefill_back_program(): @pl.program class DeepSeekV32PrefillBack: @pl.function(type=pl.FunctionType.Opaque) def deepseek_v3_2_prefill_back_layer( self, - hidden_states: pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16], - seq_lens: pl.Tensor[[BATCH_CFG], pl.INT32], - combine_buf: pl.Tensor[[EP_NODES_CFG, BATCH_CFG, MAX_SEQ_CFG, ATTN_OUT_CFG], pl.BF16], - wo: pl.Tensor[[ATTN_OUT_CFG, HIDDEN_CFG], pl.BF16], - post_rms_weight: pl.Tensor[[1, HIDDEN_CFG], pl.FP32], - w_gate: pl.Tensor[[HIDDEN_CFG, INTER_CFG], pl.BF16], - w_up: pl.Tensor[[HIDDEN_CFG, INTER_CFG], pl.BF16], - w_down: pl.Tensor[[INTER_CFG, HIDDEN_CFG], pl.BF16], - out: pl.Out[pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16]], - ) -> pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16]: - for b in pl.parallel(0, BATCH_CFG, 1): + hidden_states: pl.Tensor[[BATCH, MAX_SEQ, HIDDEN], pl.BF16], + seq_lens: pl.Tensor[[BATCH], pl.INT32], + combine_buf: pl.Tensor[[EP_NODES, BATCH, MAX_SEQ, 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, MAX_SEQ, HIDDEN], pl.BF16]], + ) -> pl.Tensor[[BATCH, MAX_SEQ, HIDDEN], pl.BF16]: + for b in pl.parallel(0, BATCH, 1): seq_len_b = pl.tensor.read(seq_lens, [b]) tok_blocks = (seq_len_b + TOK_TILE - 1) // TOK_TILE for p0_idx in pl.range(tok_blocks): @@ -79,141 +63,138 @@ def deepseek_v3_2_prefill_back_layer( valid_tok = pl.min(TOK_TILE, seq_len_b - p0) # GM intermediate tensors. - resid1_tile = pl.create_tensor([TOK_TILE, HIDDEN_CFG], dtype=pl.FP32) - attn_tile = pl.create_tensor([TOK_TILE, ATTN_OUT_CFG], dtype=pl.BF16) + resid1_tile = pl.create_tensor([TOK_TILE, HIDDEN], dtype=pl.FP32) + attn_tile = pl.create_tensor([TOK_TILE, ATTN_OUT], dtype=pl.BF16) # Stage 1: Copy combine_buf 4D -> attn_tile 2D. - with pl.incore(): - for kb in pl.range(ATTN_BLOCKS): - k0 = kb * K_CHUNK + with pl.at(level=pl.Level.CORE_GROUP): + for kb in pl.range(ATTN_OUT // K_TILE): + k0 = kb * K_TILE a_chunk_fp32 = pl.reshape( pl.cast( - pl.slice(combine_buf, [1, 1, TOK_TILE, K_CHUNK], [0, b, p0, k0], - valid_shape=[1, 1, valid_tok, K_CHUNK]), + pl.slice(combine_buf, [1, 1, TOK_TILE, K_TILE], [0, b, p0, k0], + valid_shape=[1, 1, valid_tok, K_TILE]), target_type=pl.FP32, ), - [TOK_TILE, K_CHUNK], + [TOK_TILE, K_TILE], ) a_chunk_bf16 = pl.cast(a_chunk_fp32, target_type=pl.BF16) attn_tile = pl.assemble(attn_tile, a_chunk_bf16, [0, k0]) # Stage 2: Output projection + first 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 # Cube: chained matmul. - with pl.incore(): - tile_a = pl.slice(attn_tile, [TOK_TILE, K_CHUNK], [0, 0]) - tile_w = pl.slice(wo, [K_CHUNK, Q_OUT_CHUNK], [0, o0]) + with pl.at(level=pl.Level.CORE_GROUP): + tile_a = pl.slice(attn_tile, [TOK_TILE, K_TILE], [0, 0]) + tile_w = pl.slice(wo, [K_TILE, Q_OUT_TILE], [0, o0]) o_acc = pl.matmul(tile_a, tile_w, out_dtype=pl.FP32) - for kb in pl.range(1, ATTN_BLOCKS): - k0 = kb * K_CHUNK - tile_a_i = pl.slice(attn_tile, [TOK_TILE, K_CHUNK], [0, k0]) - tile_w_i = pl.slice(wo, [K_CHUNK, Q_OUT_CHUNK], [k0, o0]) + for kb in pl.range(1, ATTN_OUT // K_TILE): + k0 = kb * K_TILE + tile_a_i = pl.slice(attn_tile, [TOK_TILE, K_TILE], [0, k0]) + tile_w_i = pl.slice(wo, [K_TILE, Q_OUT_TILE], [k0, o0]) o_acc = pl.matmul_acc(o_acc, tile_a_i, tile_w_i) resid1_tile = pl.assemble(resid1_tile, o_acc, [0, o0]) # Vector: add residual. - with pl.incore(): + with pl.at(level=pl.Level.CORE_GROUP): resid_chunk = pl.reshape( pl.cast( - pl.slice(hidden_states, [1, TOK_TILE, Q_OUT_CHUNK], [b, p0, o0], - valid_shape=[1, valid_tok, Q_OUT_CHUNK]), + pl.slice(hidden_states, [1, TOK_TILE, Q_OUT_TILE], [b, p0, o0], + valid_shape=[1, valid_tok, Q_OUT_TILE]), target_type=pl.FP32, ), - [TOK_TILE, Q_OUT_CHUNK], + [TOK_TILE, Q_OUT_TILE], ) - mm_out = pl.slice(resid1_tile, [TOK_TILE, Q_OUT_CHUNK], [0, o0]) + mm_out = pl.slice(resid1_tile, [TOK_TILE, Q_OUT_TILE], [0, o0]) resid_sum = pl.add(mm_out, resid_chunk) resid1_tile = pl.assemble(resid1_tile, resid_sum, [0, o0]) # Stage 3: Post-attention RMSNorm. - post_norm_tile = pl.create_tensor([TOK_TILE, HIDDEN_CFG], dtype=pl.BF16) - down_fp32_tile = pl.create_tensor([TOK_TILE, HIDDEN_CFG], dtype=pl.FP32) + post_norm_tile = pl.create_tensor([TOK_TILE, HIDDEN], dtype=pl.BF16) # 3a: Compute inv_rms (reduction — sequential). - with pl.auto_incore(): + with pl.at(level=pl.Level.CORE_GROUP): sq_sum = pl.full([1, TOK_TILE], dtype=pl.FP32, value=0.0) - for kb in pl.range(HIDDEN_BLOCKS): - k0 = kb * K_CHUNK - x_chunk = pl.slice(resid1_tile, [TOK_TILE, K_CHUNK], [0, k0]) + for kb in pl.range(HIDDEN // K_TILE): + k0 = kb * K_TILE + x_chunk = pl.slice(resid1_tile, [TOK_TILE, K_TILE], [0, k0]) sq_sum = pl.add( sq_sum, pl.reshape(pl.row_sum(pl.mul(x_chunk, x_chunk)), [1, TOK_TILE]), ) inv_rms = pl.rsqrt(pl.add(pl.mul(sq_sum, HIDDEN_INV), EPS)) - # 3b: Normalize + gamma + zero-init down_proj (parallel — independent offsets). - with pl.auto_incore(): - for kb in pl.parallel(0, HIDDEN_BLOCKS, chunk=8): - k0 = kb * K_CHUNK - x_chunk = pl.slice(resid1_tile, [TOK_TILE, K_CHUNK], [0, k0]) - gamma = pl.slice(post_rms_weight, [1, K_CHUNK], [0, k0]) - normed = pl.col_expand_mul( - pl.row_expand_mul(x_chunk, pl.reshape(inv_rms, [TOK_TILE, 1])), - gamma, - ) - normed_bf16 = pl.cast(normed, target_type=pl.BF16) - post_norm_tile = pl.assemble(post_norm_tile, normed_bf16, [0, k0]) - down_zero_chunk = pl.full([TOK_TILE, K_CHUNK], dtype=pl.FP32, value=0.0) - down_fp32_tile = pl.assemble(down_fp32_tile, down_zero_chunk, [0, k0]) - - # Stage 4: MLP gate/up + SiLU + down projection. - for ob in pl.range(MLP_OUT_BLOCKS): - o0 = ob * MLP_OUT_CHUNK + # 3b: Normalize + gamma. Chunk HIDDEN // K_TILE across cores; each + # core scope handles NORM_TILE contiguous blocks. + for kb0 in pl.parallel(0, HIDDEN // K_TILE, NORM_TILE): + with pl.at(level=pl.Level.CORE_GROUP): + for kb in pl.range(kb0, kb0 + NORM_TILE): + k0 = kb * K_TILE + x_chunk = pl.slice(resid1_tile, [TOK_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, [TOK_TILE, 1])), + gamma, + ) + normed_bf16 = pl.cast(normed, target_type=pl.BF16) + post_norm_tile = pl.assemble(post_norm_tile, normed_bf16, [0, k0]) + + # Stage 4: MLP gate/up + SiLU -> mlp_tile. + mlp_tile = pl.create_tensor([TOK_TILE, INTERMEDIATE], dtype=pl.BF16) + for ob in pl.range(INTERMEDIATE // MLP_OUT_TILE): + o0 = ob * MLP_OUT_TILE # Gate matmul chain. - with pl.incore(): - pc0 = pl.slice(post_norm_tile, [TOK_TILE, K_CHUNK], [0, 0]) - wg0 = pl.slice(w_gate, [K_CHUNK, MLP_OUT_CHUNK], [0, o0]) + with pl.at(level=pl.Level.CORE_GROUP): + pc0 = pl.slice(post_norm_tile, [TOK_TILE, K_TILE], [0, 0]) + wg0 = pl.slice(w_gate, [K_TILE, MLP_OUT_TILE], [0, o0]) gate_acc = pl.matmul(pc0, wg0, out_dtype=pl.FP32) - for kb in pl.range(1, HIDDEN_BLOCKS): - k0 = kb * K_CHUNK - pci = pl.slice(post_norm_tile, [TOK_TILE, K_CHUNK], [0, k0]) - wgi = pl.slice(w_gate, [K_CHUNK, MLP_OUT_CHUNK], [k0, o0]) + for kb in pl.range(1, HIDDEN // K_TILE): + k0 = kb * K_TILE + pci = pl.slice(post_norm_tile, [TOK_TILE, K_TILE], [0, k0]) + wgi = pl.slice(w_gate, [K_TILE, MLP_OUT_TILE], [k0, o0]) gate_acc = pl.matmul_acc(gate_acc, pci, wgi) # Up matmul chain. - with pl.incore(): - pc0 = pl.slice(post_norm_tile, [TOK_TILE, K_CHUNK], [0, 0]) - wu0 = pl.slice(w_up, [K_CHUNK, MLP_OUT_CHUNK], [0, o0]) + with pl.at(level=pl.Level.CORE_GROUP): + pc0 = pl.slice(post_norm_tile, [TOK_TILE, K_TILE], [0, 0]) + wu0 = pl.slice(w_up, [K_TILE, MLP_OUT_TILE], [0, o0]) up_acc = pl.matmul(pc0, wu0, out_dtype=pl.FP32) - for kb in pl.range(1, HIDDEN_BLOCKS): - k0 = kb * K_CHUNK - pci = pl.slice(post_norm_tile, [TOK_TILE, K_CHUNK], [0, k0]) - wui = pl.slice(w_up, [K_CHUNK, MLP_OUT_CHUNK], [k0, o0]) + for kb in pl.range(1, HIDDEN // K_TILE): + k0 = kb * K_TILE + pci = pl.slice(post_norm_tile, [TOK_TILE, K_TILE], [0, k0]) + wui = pl.slice(w_up, [K_TILE, MLP_OUT_TILE], [k0, o0]) up_acc = pl.matmul_acc(up_acc, pci, wui) - # SiLU activation. - with pl.auto_incore(): + # SiLU activation -> mlp_tile. + with pl.at(level=pl.Level.CORE_GROUP): sigmoid = pl.recip(pl.add(pl.exp(pl.neg(gate_acc)), 1.0)) mlp_chunk = pl.mul(pl.mul(gate_acc, sigmoid), up_acc) - mlp_chunk_bf16 = pl.cast(mlp_chunk, target_type=pl.BF16) - - # Down projection: cube matmul + vector accumulate. - for dob in pl.range(HIDDEN_BLOCKS): - d0 = dob * K_CHUNK - - with pl.incore(): - w_down_chunk = pl.slice(w_down, [MLP_OUT_CHUNK, K_CHUNK], [o0, d0]) - down_next = pl.matmul(mlp_chunk_bf16, w_down_chunk, out_dtype=pl.FP32) - - with pl.incore(): - down_prev = pl.slice(down_fp32_tile, [TOK_TILE, K_CHUNK], [0, d0]) - accum = pl.add(down_prev, down_next) - down_fp32_tile = pl.assemble(down_fp32_tile, accum, [0, d0]) - - # Stage 5: Final residual add -> BF16 output (parallel — independent offsets). - with pl.auto_incore(): - for ob in pl.parallel(0, HIDDEN_BLOCKS, chunk=8): - o0 = ob * K_CHUNK - final_sum = pl.add( - pl.slice(down_fp32_tile, [TOK_TILE, K_CHUNK], [0, o0]), - pl.slice(resid1_tile, [TOK_TILE, K_CHUNK], [0, o0]), + mlp_tile = pl.assemble(mlp_tile, pl.cast(mlp_chunk, target_type=pl.BF16), [0, o0]) + + # Stage 5: Down projection + final residual -> BF16 output. + for dob in pl.range(HIDDEN // K_TILE): + d0 = dob * K_TILE + + with pl.at(level=pl.Level.CORE_GROUP): + mlp0 = pl.slice(mlp_tile, [TOK_TILE, MLP_OUT_TILE], [0, 0]) + wd0 = pl.slice(w_down, [MLP_OUT_TILE, K_TILE], [0, d0]) + down_acc = pl.matmul(mlp0, wd0, out_dtype=pl.FP32) + for ob in pl.range(1, INTERMEDIATE // MLP_OUT_TILE): + o0 = ob * MLP_OUT_TILE + mlpi = pl.slice(mlp_tile, [TOK_TILE, MLP_OUT_TILE], [0, o0]) + wdi = pl.slice(w_down, [MLP_OUT_TILE, K_TILE], [o0, d0]) + down_acc = pl.matmul_acc(down_acc, mlpi, wdi) + + with pl.at(level=pl.Level.CORE_GROUP): + out_chunk = pl.add( + down_acc, + pl.slice(resid1_tile, [TOK_TILE, K_TILE], [0, d0]), ) - final_bf16 = pl.cast(final_sum, target_type=pl.BF16) - out = pl.assemble(out, final_bf16, [b, p0, o0]) + out = pl.assemble(out, pl.cast(out_chunk, target_type=pl.BF16), [b, p0, d0]) return out @@ -298,7 +279,7 @@ def golden_prefill_back(tensors): eps = EPS post_rms_f = post_rms_weight.float() - def chunked_bf16_matmul(a_bf16, b_bf16, k_chunk=K_CHUNK): + def chunked_bf16_matmul(a_bf16, b_bf16, k_chunk=K_TILE): """BF16 x BF16 -> FP32 with chunked K-dim accumulation.""" K = a_bf16.shape[-1] acc = torch.zeros(*a_bf16.shape[:-1], b_bf16.shape[-1], dtype=torch.float32) @@ -328,7 +309,7 @@ def chunked_bf16_matmul(a_bf16, b_bf16, k_chunk=K_CHUNK): gate = chunked_bf16_matmul(normed_bf16, w_gate) up = chunked_bf16_matmul(normed_bf16, w_up) mlp_bf16 = (gate * torch.sigmoid(gate) * up).bfloat16() - down = chunked_bf16_matmul(mlp_bf16, w_down, k_chunk=MLP_OUT_CHUNK) + down = chunked_bf16_matmul(mlp_bf16, w_down, k_chunk=MLP_OUT_TILE) # 5. Final residual -> BF16. out_t[b, sl, :] = (down + resid1).bfloat16() diff --git a/models/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.py b/models/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.py deleted file mode 100644 index 65b6ec6c..00000000 --- a/models/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.py +++ /dev/null @@ -1,584 +0,0 @@ -# Copyright (c) PyPTO Contributors. -# This program is free software, you can redistribute it and/or modify it under the terms and conditions of -# CANN Open Software License Agreement Version 2.0 (the "License"). -# Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -# See LICENSE in the root of the software repository for the full text of the License. -# ----------------------------------------------------------------------------------------------------------- -""" -DeepSeek V3.2-EXP single-layer prefill FRONT part (batch=16, max_seq=4096). - -Aligned to official v3.2-exp MLA/DSA shapes. Indexer routing is externalized as -`index_topk_pos`, and the kernel applies sparse attention on those indices. -""" - -import os - -import pypto.language as pl - - -BATCH = 16 -MAX_SEQ = 4096 -HIDDEN = 7168 -NUM_HEADS = 128 -Q_LORA_RANK = 1536 -KV_LORA_RANK = 512 -QK_NOPE_HEAD_DIM = 128 -QK_ROPE_HEAD_DIM = 64 -QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM -V_HEAD_DIM = 128 -ATTN_OUT = NUM_HEADS * V_HEAD_DIM -INDEX_HEADS = 64 -INDEX_TOPK = 2048 -EP_NODES = 128 - -EPS = 1e-6 -ATTN_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5) -HIDDEN_INV = 1.0 / HIDDEN - -K_CHUNK = 256 -K_CHUNK = 512 -Q_OUT_CHUNK = 256 -KV_OUT_CHUNK = 128 -LORA_CHUNK = 128 -V_OUT_CHUNK = 64 -SEQ_TILE = 120 -TOK_TILE = 4 -# Extra local pad tensor width to raise explicit Vec occupancy in memory report. -LOCAL_PAD_WIDTH = 16384 - -# Conservative software guard for AIV Vec/UB working set (bytes). This does -# not reflect final backend private buffers, but helps avoid overly aggressive -# source-side tile settings. -UB_SOFT_LIMIT_BYTES = 160 * 1024 - - -def build_deepseek_v3_2_prefill_front_program( - batch: int = BATCH, - max_seq_len: int = MAX_SEQ, - hidden_size: int = HIDDEN, - num_heads: int = NUM_HEADS, - q_lora_rank: int = Q_LORA_RANK, - kv_lora_rank: int = KV_LORA_RANK, - qk_nope_head_dim: int = QK_NOPE_HEAD_DIM, - qk_rope_head_dim: int = QK_ROPE_HEAD_DIM, - v_head_dim: int = V_HEAD_DIM, - index_heads: int = INDEX_HEADS, - index_topk: int = INDEX_TOPK, - ep_nodes: int = EP_NODES, -): - BATCH_CFG = batch - MAX_SEQ_CFG = max_seq_len - HIDDEN_CFG = hidden_size - NUM_HEADS_CFG = num_heads - Q_LORA_RANK_CFG = q_lora_rank - KV_LORA_RANK_CFG = kv_lora_rank - QK_NOPE_HEAD_DIM_CFG = qk_nope_head_dim - QK_ROPE_HEAD_DIM_CFG = qk_rope_head_dim - QK_HEAD_DIM_CFG = qk_nope_head_dim + qk_rope_head_dim - V_HEAD_DIM_CFG = v_head_dim - INDEX_HEADS_CFG = index_heads - ATTN_OUT_CFG = num_heads * v_head_dim - INDEX_TOPK_CFG = index_topk - EP_NODES_CFG = ep_nodes - - HIDDEN_BLOCKS = (HIDDEN_CFG + K_CHUNK - 1) // K_CHUNK - QR_BLOCKS = (Q_LORA_RANK_CFG + LORA_CHUNK - 1) // LORA_CHUNK - Q_OUT_BLOCKS = (NUM_HEADS_CFG * QK_HEAD_DIM_CFG + Q_OUT_CHUNK - 1) // Q_OUT_CHUNK - KV_A_OUT = KV_LORA_RANK_CFG + QK_ROPE_HEAD_DIM_CFG - KV_A_BLOCKS = (KV_A_OUT + KV_OUT_CHUNK - 1) // KV_OUT_CHUNK - CACHE_ROWS = BATCH_CFG * MAX_SEQ_CFG - V_OUT_BLOCKS = (V_HEAD_DIM_CFG + V_OUT_CHUNK - 1) // V_OUT_CHUNK - - # Capacity-oriented source tuning guard: - # - stage1_est_bytes models dominant token-tile matmul operands/accumulators. - # - stage2_est_bytes models topk buffers + major per-head vectors in sparse consume. - # The check targets "close to but not exceeding" practical UB usage in source-level tuning. - stage1_est_bytes = ( - TOK_TILE * K_CHUNK * 4 - + TOK_TILE * LORA_CHUNK * 4 - + TOK_TILE * Q_OUT_CHUNK * 4 - + TOK_TILE * KV_OUT_CHUNK * 4 - + TOK_TILE * LOCAL_PAD_WIDTH * 2 - ) - stage2_est_bytes = ( - (1 + 2) * INDEX_TOPK_CFG * 4 # topk vals + blk topk vals - + (1 + 2) * INDEX_TOPK_CFG * 4 # topk idx + blk topk idx - + KV_LORA_RANK_CFG * 4 # oi/ctx_latent dominant row vectors - + QK_ROPE_HEAD_DIM_CFG * 4 # q/pe rope vectors - + V_HEAD_DIM_CFG * 4 # ctx_v - ) - peak_est_bytes = max(stage1_est_bytes, stage2_est_bytes) - if peak_est_bytes > UB_SOFT_LIMIT_BYTES: - raise ValueError( - f"Estimated local working set {peak_est_bytes} bytes exceeds " - f"UB soft limit {UB_SOFT_LIMIT_BYTES} bytes. " - "Reduce TOK_TILE/Q_OUT_CHUNK/K_CHUNK/KV_OUT_CHUNK/LORA_CHUNK." - ) - - @pl.program - class DeepSeekV32PrefillFront: - @pl.function(type=pl.FunctionType.Opaque) - def deepseek_v3_2_prefill_front_layer( - self, - hidden_states: pl.Tensor[[BATCH_CFG, MAX_SEQ_CFG, HIDDEN_CFG], pl.BF16], - seq_lens: pl.Tensor[[BATCH_CFG], pl.INT32], - layer_id_t: pl.Tensor[[1], pl.INT32], - rope_cos: pl.Tensor[[MAX_SEQ_CFG, QK_ROPE_HEAD_DIM_CFG], pl.FP32], - rope_sin: pl.Tensor[[MAX_SEQ_CFG, QK_ROPE_HEAD_DIM_CFG], pl.FP32], - kv_cache: pl.Tensor[[CACHE_ROWS, KV_LORA_RANK_CFG], pl.BF16], - pe_cache: pl.Tensor[[CACHE_ROWS, QK_ROPE_HEAD_DIM_CFG], pl.BF16], - input_rms_weight: pl.Tensor[[1, HIDDEN_CFG], pl.FP32], - wq_a: pl.Tensor[[HIDDEN_CFG, Q_LORA_RANK_CFG], pl.BF16], - q_norm_weight: pl.Tensor[[1, Q_LORA_RANK_CFG], pl.FP32], - wq_b: pl.Tensor[[Q_LORA_RANK_CFG, NUM_HEADS_CFG * QK_HEAD_DIM_CFG], pl.BF16], - wkv_a: pl.Tensor[[HIDDEN_CFG, KV_A_OUT], pl.BF16], - kv_norm_weight: pl.Tensor[[1, KV_LORA_RANK_CFG], pl.FP32], - w_q_nope_to_latent: pl.Tensor[[NUM_HEADS_CFG, QK_NOPE_HEAD_DIM_CFG, KV_LORA_RANK_CFG], pl.BF16], - w_latent_to_v: pl.Tensor[[NUM_HEADS_CFG, KV_LORA_RANK_CFG, V_HEAD_DIM_CFG], pl.BF16], - dispatch_buf: pl.Out[pl.Tensor[[EP_NODES_CFG, BATCH_CFG, MAX_SEQ_CFG, ATTN_OUT_CFG], pl.BF16]], - ) -> pl.Tensor[[EP_NODES_CFG, BATCH_CFG, MAX_SEQ_CFG, ATTN_OUT_CFG], pl.BF16]: - with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.auto_chunk]): - layer_id = pl.tensor.read(layer_id_t, [0]) - - for b in pl.parallel(0, BATCH_CFG, 1, chunk=4): - seq_len_b = pl.tensor.read(seq_lens, [b]) - tok_blocks = (seq_len_b + TOK_TILE - 1) // TOK_TILE - for p0_idx in pl.range(tok_blocks): - p0 = p0_idx * TOK_TILE - valid_tok = pl.min(TOK_TILE, seq_len_b - p0) - - # Scope 1: RMSNorm + Q/K/V projections. - sq_sum = pl.create_tensor([TOK_TILE, 1], dtype=pl.FP32) - sq_sum = pl.mul(sq_sum, 0.0) - # Keep an explicit local Vec pad tensor alive in this - # scope so AllocateMemoryAddr can reflect high occupancy. - usage_pad = pl.create_tensor([TOK_TILE, LOCAL_PAD_WIDTH], dtype=pl.BF16, valid_shape=[valid_tok, LOCAL_PAD_WIDTH]) - usage_pad = pl.mul(usage_pad, 0.0) - usage_pad_fp = pl.cast(usage_pad, target_type=pl.FP32) - usage_pad_sum = pl.row_sum(usage_pad_fp) - for kb in pl.range(HIDDEN_BLOCKS): - k0 = kb * K_CHUNK - x_chunk = pl.cast( - pl.slice(hidden_states, [TOK_TILE, K_CHUNK], [b, p0, k0], valid_shape=[valid_tok, K_CHUNK]), - target_type=pl.FP32, - ) - sq_sum = pl.add(sq_sum, pl.row_sum(pl.mul(x_chunk, x_chunk))) - inv_rms = pl.rsqrt(pl.add(pl.mul(sq_sum, HIDDEN_INV), EPS)) - inv_rms = pl.add(inv_rms, pl.mul(usage_pad_sum, 0.0)) - - q_proj_tile = pl.create_tensor([TOK_TILE, NUM_HEADS_CFG * QK_HEAD_DIM_CFG], dtype=pl.BF16, valid_shape=[valid_tok, NUM_HEADS_CFG * QK_HEAD_DIM_CFG]) - kv_a_tile = pl.create_tensor([TOK_TILE, KV_A_OUT], dtype=pl.BF16, valid_shape=[valid_tok, KV_A_OUT]) - - # Fused Q path (local fusion trial for former incore_0/1): - # directly accumulates q_proj_tile from x -> wq_a -> q_norm -> wq_b - # without materializing full qr_tile. - for ob in pl.parallel(0, Q_OUT_BLOCKS, 1, chunk=8): - q0 = ob * Q_OUT_CHUNK - q_acc = pl.create_tensor([TOK_TILE, Q_OUT_CHUNK], dtype=pl.FP32) - q_acc = pl.mul(q_acc, 0.0) - for kb in pl.range(HIDDEN_BLOCKS): - k0 = kb * K_CHUNK - x_chunk = pl.cast( - pl.slice(hidden_states, [TOK_TILE, K_CHUNK], [b, p0, k0], valid_shape=[valid_tok, K_CHUNK]), - target_type=pl.FP32, - ) - gamma_in = pl.slice(input_rms_weight, [1, K_CHUNK], [0, k0]) - normed = pl.col_expand_mul(pl.row_expand_mul(x_chunk, inv_rms), gamma_in) - for rb in pl.range(QR_BLOCKS): - r0 = rb * LORA_CHUNK - wq_a_chunk = pl.slice(wq_a, [K_CHUNK, LORA_CHUNK], [k0, r0]) - qr_part = pl.matmul(pl.cast(normed, target_type=pl.BF16), wq_a_chunk) - gamma_q = pl.slice(q_norm_weight, [1, LORA_CHUNK], [0, r0]) - qn_part = pl.col_expand_mul(qr_part, gamma_q) - wq_b_chunk = pl.slice(wq_b, [LORA_CHUNK, Q_OUT_CHUNK], [r0, q0]) - q_acc = pl.add(q_acc, pl.matmul(pl.cast(qn_part, target_type=pl.BF16), wq_b_chunk)) - q_proj_tile = pl.assemble(q_proj_tile, pl.cast(q_acc, target_type=pl.BF16), [0, q0]) - - for ob in pl.parallel(0, KV_A_BLOCKS, 1, chunk=8): - kv0 = ob * KV_OUT_CHUNK - kv_acc = pl.create_tensor([TOK_TILE, KV_OUT_CHUNK], dtype=pl.FP32) - kv_acc = pl.mul(kv_acc, 0.0) - for kb in pl.range(HIDDEN_BLOCKS): - k0 = kb * K_CHUNK - x_chunk = pl.cast( - pl.slice(hidden_states, [TOK_TILE, K_CHUNK], [b, p0, k0], valid_shape=[valid_tok, K_CHUNK]), - target_type=pl.FP32, - ) - gamma = pl.slice(input_rms_weight, [1, K_CHUNK], [0, k0]) - normed = pl.col_expand_mul(pl.row_expand_mul(x_chunk, inv_rms), gamma) - wkv_chunk = pl.slice(wkv_a, [K_CHUNK, KV_OUT_CHUNK], [k0, kv0]) - kv_acc = pl.add(kv_acc, pl.matmul(pl.cast(normed, target_type=pl.BF16), wkv_chunk)) - kv_a_tile = pl.assemble(kv_a_tile, pl.cast(kv_acc, target_type=pl.BF16), [0, kv0]) - - # Scope 2: RoPE + cache update + indexer topk + sparse attention. - # Fusion policy (aligned with decode_front): - # - Stage A/B/C all stay in ONE auto_incore scope. - # - A: per-token cache write - # - B1/B2: two-stage topk (block-local then global merge) - # - C: sparse attention consumes merged topk immediately - # This avoids materializing topk intermediates across kernel boundaries. - attn_tile = pl.create_tensor([TOK_TILE, ATTN_OUT_CFG], dtype=pl.FP32, valid_shape=[valid_tok, ATTN_OUT_CFG]) - attn_tile = pl.mul(attn_tile, 0.0) - for ti in pl.range(valid_tok): - pos = p0 + ti - ctx_len = pos + 1 - cos_row = pl.slice(rope_cos, [1, QK_ROPE_HEAD_DIM_CFG], [pos, 0]) - sin_row = pl.slice(rope_sin, [1, QK_ROPE_HEAD_DIM_CFG], [pos, 0]) - cos_lo = pl.slice(cos_row, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, 0]) - cos_hi = pl.slice(cos_row, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, QK_ROPE_HEAD_DIM_CFG // 2]) - sin_lo = pl.slice(sin_row, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, 0]) - sin_hi = pl.slice(sin_row, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, QK_ROPE_HEAD_DIM_CFG // 2]) - - cache_row = b * MAX_SEQ_CFG + pos - kv_row = pl.cast(pl.slice(kv_a_tile, [1, KV_LORA_RANK_CFG], [ti, 0]), target_type=pl.FP32) - kv_gamma = pl.slice(kv_norm_weight, [1, KV_LORA_RANK_CFG], [0, 0]) - kv_normed = pl.col_expand_mul(kv_row, kv_gamma) - pe_row = pl.cast( - pl.slice(kv_a_tile, [1, QK_ROPE_HEAD_DIM_CFG], [ti, KV_LORA_RANK_CFG]), - target_type=pl.FP32, - ) - pe_lo = pl.slice(pe_row, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, 0]) - pe_hi = pl.slice(pe_row, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, QK_ROPE_HEAD_DIM_CFG // 2]) - pe_rot = pl.create_tensor([1, QK_ROPE_HEAD_DIM_CFG], dtype=pl.FP32) - pe_rot = pl.assemble(pe_rot, pl.sub(pl.col_expand_mul(pe_lo, cos_lo), pl.col_expand_mul(pe_hi, sin_lo)), [0, 0]) - pe_rot = pl.assemble(pe_rot, pl.add(pl.col_expand_mul(pe_hi, cos_hi), pl.col_expand_mul(pe_lo, sin_hi)), [0, QK_ROPE_HEAD_DIM_CFG // 2]) - kv_cache = pl.assemble(kv_cache, pl.cast(kv_normed, target_type=pl.BF16), [cache_row, 0]) - pe_cache = pl.assemble(pe_cache, pl.cast(pe_rot, target_type=pl.BF16), [cache_row, 0]) - - # Stage B1: block-local topk (2 blocks, each 2K candidates). - topk_vals = pl.create_tensor([1, INDEX_TOPK_CFG], dtype=pl.FP32) - topk_idx = pl.create_tensor([1, INDEX_TOPK_CFG], dtype=pl.INT32) - blk_topk_vals = pl.create_tensor([2, INDEX_TOPK_CFG], dtype=pl.FP32) - blk_topk_idx = pl.create_tensor([2, INDEX_TOPK_CFG], dtype=pl.INT32) - topk_vals = pl.mul(topk_vals, -3.402823e38) - topk_idx = pl.mul(topk_idx, 0) - blk_topk_vals = pl.mul(blk_topk_vals, -3.402823e38) - blk_topk_idx = pl.mul(blk_topk_idx, 0) - for kk in pl.range(INDEX_TOPK_CFG): - neg_one = pl.create_tensor([1, 1], dtype=pl.INT32) - neg_one = pl.mul(neg_one, 0) - neg_one = pl.add(neg_one, -1) - topk_idx = pl.assemble(topk_idx, neg_one, [0, kk]) - blk_topk_idx = pl.assemble(blk_topk_idx, neg_one, [0, kk]) - blk_topk_idx = pl.assemble(blk_topk_idx, neg_one, [1, kk]) - - q_col0 = 0 - q_nope0 = pl.cast( - pl.slice(q_proj_tile, [1, QK_NOPE_HEAD_DIM_CFG], [ti, q_col0]), - target_type=pl.FP32, - ) - q_pe0 = pl.cast( - pl.slice(q_proj_tile, [1, QK_ROPE_HEAD_DIM_CFG], [ti, q_col0 + QK_NOPE_HEAD_DIM_CFG]), - target_type=pl.FP32, - ) - q0_lo = pl.slice(q_pe0, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, 0]) - q0_hi = pl.slice(q_pe0, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, QK_ROPE_HEAD_DIM_CFG // 2]) - q0_rot = pl.create_tensor([1, QK_ROPE_HEAD_DIM_CFG], dtype=pl.FP32) - q0_rot = pl.assemble(q0_rot, pl.sub(pl.col_expand_mul(q0_lo, cos_lo), pl.col_expand_mul(q0_hi, sin_lo)), [0, 0]) - q0_rot = pl.assemble(q0_rot, pl.add(pl.col_expand_mul(q0_hi, cos_hi), pl.col_expand_mul(q0_lo, sin_hi)), [0, QK_ROPE_HEAD_DIM_CFG // 2]) - q0_nope_latent = pl.matmul( - pl.cast(q_nope0, target_type=pl.BF16), - pl.slice(w_q_nope_to_latent, [QK_NOPE_HEAD_DIM_CFG, KV_LORA_RANK_CFG], [0, 0, 0]), - ) - - sparse_k_gen = pl.min(INDEX_TOPK_CFG, ctx_len) - for blk in pl.range(2): - blk_start = blk * INDEX_TOPK_CFG - blk_end = pl.min(ctx_len, blk_start + INDEX_TOPK_CFG) - for ss in pl.range(INDEX_TOPK_CFG): - s = blk_start + ss - if s < blk_end: - cache_s = b * MAX_SEQ_CFG + s - kv_s = pl.cast(pl.slice(kv_cache, [1, KV_LORA_RANK_CFG], [cache_s, 0]), target_type=pl.FP32) - pe_s = pl.cast(pl.slice(pe_cache, [1, QK_ROPE_HEAD_DIM_CFG], [cache_s, 0]), target_type=pl.FP32) - score_nope = pl.row_sum(pl.mul(q0_nope_latent, kv_s)) - score_pe = pl.row_sum(pl.mul(q0_rot, pe_s)) - score_fp32 = pl.mul(pl.add(score_nope, score_pe), ATTN_SCALE) - score_fp8 = pl.cast(score_fp32, target_type=pl.FP8E4M3FN) - score_a5 = pl.cast(score_fp8, target_type=pl.FP32) - cur_score = pl.tensor.read(score_a5, [0, 0]) - - inserted = pl.create_tensor([1, 1], dtype=pl.INT32) - inserted = pl.mul(inserted, 0) - for kk in pl.range(sparse_k_gen): - ins = pl.tensor.read(inserted, [0, 0]) - kth_val = pl.tensor.read(blk_topk_vals, [blk, kk]) - if ins == 0: - if cur_score > kth_val: - for sh in pl.range(sparse_k_gen - 1, kk, -1): - prev_val = pl.tensor.read(blk_topk_vals, [blk, sh - 1]) - prev_idx = pl.tensor.read(blk_topk_idx, [blk, sh - 1]) - prev_val_t = pl.create_tensor([1, 1], dtype=pl.FP32) - prev_idx_t = pl.create_tensor([1, 1], dtype=pl.INT32) - prev_val_t = pl.mul(prev_val_t, 0.0) - prev_idx_t = pl.mul(prev_idx_t, 0) - prev_val_t = pl.add(prev_val_t, prev_val) - prev_idx_t = pl.add(prev_idx_t, prev_idx) - blk_topk_vals = pl.assemble(blk_topk_vals, prev_val_t, [blk, sh]) - blk_topk_idx = pl.assemble(blk_topk_idx, prev_idx_t, [blk, sh]) - cur_score_t = pl.create_tensor([1, 1], dtype=pl.FP32) - cur_index_t = pl.create_tensor([1, 1], dtype=pl.INT32) - one_t = pl.create_tensor([1, 1], dtype=pl.INT32) - cur_score_t = pl.mul(cur_score_t, 0.0) - cur_index_t = pl.mul(cur_index_t, 0) - one_t = pl.mul(one_t, 0) - cur_score_t = pl.add(cur_score_t, cur_score) - cur_index_t = pl.add(cur_index_t, s) - one_t = pl.add(one_t, 1) - blk_topk_vals = pl.assemble(blk_topk_vals, cur_score_t, [blk, kk]) - blk_topk_idx = pl.assemble(blk_topk_idx, cur_index_t, [blk, kk]) - inserted = pl.assemble(inserted, one_t, [0, 0]) - - # Stage B2: global merge from 2x(local topk) -> final topk. - for blk in pl.range(2): - for kk in pl.range(sparse_k_gen): - cand_idx = pl.tensor.read(blk_topk_idx, [blk, kk]) - if cand_idx >= 0: - cand_val = pl.tensor.read(blk_topk_vals, [blk, kk]) - inserted = pl.create_tensor([1, 1], dtype=pl.INT32) - inserted = pl.mul(inserted, 0) - for tkk in pl.range(sparse_k_gen): - ins = pl.tensor.read(inserted, [0, 0]) - kth_val = pl.tensor.read(topk_vals, [0, tkk]) - if ins == 0: - if cand_val > kth_val: - for sh in pl.range(sparse_k_gen - 1, tkk, -1): - prev_val = pl.tensor.read(topk_vals, [0, sh - 1]) - prev_idx = pl.tensor.read(topk_idx, [0, sh - 1]) - prev_val_t = pl.create_tensor([1, 1], dtype=pl.FP32) - prev_idx_t = pl.create_tensor([1, 1], dtype=pl.INT32) - prev_val_t = pl.mul(prev_val_t, 0.0) - prev_idx_t = pl.mul(prev_idx_t, 0) - prev_val_t = pl.add(prev_val_t, prev_val) - prev_idx_t = pl.add(prev_idx_t, prev_idx) - topk_vals = pl.assemble(topk_vals, prev_val_t, [0, sh]) - topk_idx = pl.assemble(topk_idx, prev_idx_t, [0, sh]) - cand_val_t = pl.create_tensor([1, 1], dtype=pl.FP32) - cand_idx_t = pl.create_tensor([1, 1], dtype=pl.INT32) - one_t = pl.create_tensor([1, 1], dtype=pl.INT32) - cand_val_t = pl.mul(cand_val_t, 0.0) - cand_idx_t = pl.mul(cand_idx_t, 0) - one_t = pl.mul(one_t, 0) - cand_val_t = pl.add(cand_val_t, cand_val) - cand_idx_t = pl.add(cand_idx_t, cand_idx) - one_t = pl.add(one_t, 1) - topk_vals = pl.assemble(topk_vals, cand_val_t, [0, tkk]) - topk_idx = pl.assemble(topk_idx, cand_idx_t, [0, tkk]) - inserted = pl.assemble(inserted, one_t, [0, 0]) - - # Stage C: sparse attention directly consumes merged topk_idx. - attn_row = pl.create_tensor([1, ATTN_OUT_CFG], dtype=pl.FP32) - attn_row = pl.mul(attn_row, 0.0) - for h in pl.parallel(0, NUM_HEADS_CFG, 1, chunk=8): - q_col = h * QK_HEAD_DIM_CFG - q_nope = pl.cast( - pl.slice(q_proj_tile, [1, QK_NOPE_HEAD_DIM_CFG], [ti, q_col]), - target_type=pl.FP32, - ) - q_pe = pl.cast( - pl.slice(q_proj_tile, [1, QK_ROPE_HEAD_DIM_CFG], [ti, q_col + QK_NOPE_HEAD_DIM_CFG]), - target_type=pl.FP32, - ) - q_lo = pl.slice(q_pe, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, 0]) - q_hi = pl.slice(q_pe, [1, QK_ROPE_HEAD_DIM_CFG // 2], [0, QK_ROPE_HEAD_DIM_CFG // 2]) - q_rot = pl.create_tensor([1, QK_ROPE_HEAD_DIM_CFG], dtype=pl.FP32) - q_rot = pl.assemble(q_rot, pl.sub(pl.col_expand_mul(q_lo, cos_lo), pl.col_expand_mul(q_hi, sin_lo)), [0, 0]) - q_rot = pl.assemble(q_rot, pl.add(pl.col_expand_mul(q_hi, cos_hi), pl.col_expand_mul(q_lo, sin_hi)), [0, QK_ROPE_HEAD_DIM_CFG // 2]) - q_nope_latent = pl.matmul( - pl.cast(q_nope, target_type=pl.BF16), - pl.slice(w_q_nope_to_latent, [QK_NOPE_HEAD_DIM_CFG, KV_LORA_RANK_CFG], [h, 0, 0]), - ) - - oi = pl.create_tensor([1, KV_LORA_RANK_CFG], dtype=pl.FP32) - li = pl.create_tensor([1, 1], dtype=pl.FP32) - mi = pl.create_tensor([1, 1], dtype=pl.FP32) - oi = pl.mul(oi, 0.0) - li = pl.mul(li, 0.0) - mi = pl.mul(mi, 0.0) - sparse_k = pl.min(INDEX_TOPK_CFG, ctx_len) - for kk in pl.range(sparse_k): - s = pl.tensor.read(topk_idx, [0, kk]) - if s >= 0: - cache_s = b * MAX_SEQ_CFG + s - kv_s = pl.cast(pl.slice(kv_cache, [1, KV_LORA_RANK_CFG], [cache_s, 0]), target_type=pl.FP32) - pe_s = pl.cast(pl.slice(pe_cache, [1, QK_ROPE_HEAD_DIM_CFG], [cache_s, 0]), target_type=pl.FP32) - score_nope = pl.row_sum(pl.mul(q_nope_latent, kv_s)) - score_pe = pl.row_sum(pl.mul(q_rot, pe_s)) - score = pl.mul(pl.add(score_nope, score_pe), ATTN_SCALE) - cur_mi = score - cur_li = pl.exp(pl.sub(score, cur_mi)) - oi_tmp = pl.row_expand_mul(kv_s, cur_li) - if kk == 0: - oi = oi_tmp - li = cur_li - mi = cur_mi - else: - mi_new = pl.maximum(mi, cur_mi) - alpha = pl.exp(pl.sub(mi, mi_new)) - beta = pl.exp(pl.sub(cur_mi, mi_new)) - li = pl.add(pl.mul(alpha, li), pl.mul(beta, cur_li)) - oi = pl.add(pl.row_expand_mul(oi, alpha), pl.row_expand_mul(oi_tmp, beta)) - mi = mi_new - ctx_latent = pl.row_expand_div(oi, li) - v_col = h * V_HEAD_DIM_CFG - ctx_v = pl.create_tensor([1, V_HEAD_DIM_CFG], dtype=pl.FP32) - ctx_v = pl.mul(ctx_v, 0.0) - for vb in pl.range(V_OUT_BLOCKS): - v0 = vb * V_OUT_CHUNK - wv_tile = pl.slice(w_latent_to_v, [KV_LORA_RANK_CFG, V_OUT_CHUNK], [h, 0, v0]) - v_part = pl.matmul(pl.cast(ctx_latent, target_type=pl.BF16), wv_tile, out_dtype=pl.FP32) - ctx_v = pl.assemble(ctx_v, v_part, [0, v0]) - attn_row = pl.assemble(attn_row, ctx_v, [0, v_col]) - attn_tile = pl.assemble(attn_tile, attn_row, [ti, 0]) - - # Scope 3: dispatch writes and return after dispatch. - for ti in pl.range(valid_tok): - pos = p0 + ti - target_node = (b + pos + layer_id) % EP_NODES_CFG - token_row = pl.cast(pl.slice(attn_tile, [1, ATTN_OUT_CFG], [ti, 0]), target_type=pl.BF16) - dispatch_buf = pl.assemble(dispatch_buf, token_row, [target_node, b, pos, 0]) - - return dispatch_buf - - return DeepSeekV32PrefillFront - - -def build_tensor_specs( - batch: int = BATCH, - max_seq_len: int = MAX_SEQ, - hidden_size: int = HIDDEN, - num_heads: int = NUM_HEADS, - q_lora_rank: int = Q_LORA_RANK, - kv_lora_rank: int = KV_LORA_RANK, - qk_nope_head_dim: int = QK_NOPE_HEAD_DIM, - qk_rope_head_dim: int = QK_ROPE_HEAD_DIM, - v_head_dim: int = V_HEAD_DIM, - index_heads: int = INDEX_HEADS, - index_topk: int = INDEX_TOPK, - ep_nodes: int = EP_NODES, -): - import torch # type: ignore[import] - from pypto.runtime import TensorSpec - - qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - kv_a_out = kv_lora_rank + qk_rope_head_dim - cache_rows = batch * max_seq_len - attn_out = num_heads * v_head_dim - seq_lens_data = torch.randint(1, max_seq_len + 1, (batch,), dtype=torch.int32) - layer_id_data = torch.tensor([0], dtype=torch.int32) - - return [ - TensorSpec("hidden_states", [batch, max_seq_len, hidden_size], torch.bfloat16, init_value=torch.randn), - TensorSpec("seq_lens", [batch], torch.int32, init_value=seq_lens_data), - TensorSpec("layer_id_t", [1], torch.int32, init_value=layer_id_data), - TensorSpec("rope_cos", [max_seq_len, qk_rope_head_dim], torch.float32, init_value=torch.randn), - TensorSpec("rope_sin", [max_seq_len, qk_rope_head_dim], torch.float32, init_value=torch.randn), - TensorSpec("kv_cache", [cache_rows, kv_lora_rank], torch.bfloat16, init_value=torch.randn), - TensorSpec("pe_cache", [cache_rows, qk_rope_head_dim], torch.bfloat16, init_value=torch.randn), - TensorSpec("input_rms_weight", [1, hidden_size], torch.float32, init_value=torch.randn), - TensorSpec("wq_a", [hidden_size, q_lora_rank], torch.bfloat16, init_value=torch.randn), - TensorSpec("q_norm_weight", [1, q_lora_rank], torch.float32, init_value=torch.randn), - TensorSpec("wq_b", [q_lora_rank, num_heads * qk_head_dim], torch.bfloat16, init_value=torch.randn), - TensorSpec("wkv_a", [hidden_size, kv_a_out], torch.bfloat16, init_value=torch.randn), - TensorSpec("kv_norm_weight", [1, kv_lora_rank], torch.float32, init_value=torch.randn), - TensorSpec("w_q_nope_to_latent", [num_heads, qk_nope_head_dim, kv_lora_rank], torch.bfloat16, init_value=torch.randn), - TensorSpec("w_latent_to_v", [num_heads, kv_lora_rank, v_head_dim], torch.bfloat16, init_value=torch.randn), - TensorSpec("dispatch_buf", [ep_nodes, batch, max_seq_len, attn_out], torch.bfloat16, is_output=True), - ] - - -def compile_and_run( - batch: int = BATCH, - max_seq_len: int = MAX_SEQ, - hidden_size: int = HIDDEN, - num_heads: int = NUM_HEADS, - q_lora_rank: int = Q_LORA_RANK, - kv_lora_rank: int = KV_LORA_RANK, - qk_nope_head_dim: int = QK_NOPE_HEAD_DIM, - qk_rope_head_dim: int = QK_ROPE_HEAD_DIM, - v_head_dim: int = V_HEAD_DIM, - index_heads: int = INDEX_HEADS, - index_topk: int = INDEX_TOPK, - ep_nodes: int = EP_NODES, - platform: str = "a2a3", - device_id: int = 0, - work_dir: str | None = None, - dump_passes: bool = True, - enable_l2_swimlane: bool = False, -): - from pypto.backend import BackendType - from pypto.ir.pass_manager import OptimizationStrategy - from pypto.runtime import RunConfig, run - - program = build_deepseek_v3_2_prefill_front_program( - batch=batch, - max_seq_len=max_seq_len, - hidden_size=hidden_size, - num_heads=num_heads, - q_lora_rank=q_lora_rank, - kv_lora_rank=kv_lora_rank, - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - index_heads=index_heads, - index_topk=index_topk, - ep_nodes=ep_nodes, - ) - tensor_specs = build_tensor_specs( - batch=batch, - max_seq_len=max_seq_len, - hidden_size=hidden_size, - num_heads=num_heads, - q_lora_rank=q_lora_rank, - kv_lora_rank=kv_lora_rank, - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - index_heads=index_heads, - index_topk=index_topk, - ep_nodes=ep_nodes, - ) - - if work_dir is None: - work_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "deepseek_v3_2_prefill_front_dump")) - - result = run( - program=program, - specs=tensor_specs, - golden=None, - config=RunConfig( - platform=platform, - device_id=device_id, - rtol=2e-2, - atol=2e-2, - strategy=OptimizationStrategy.Default, - dump_passes=dump_passes, - backend_type=BackendType.CCE, - work_dir=work_dir, - enable_l2_swimlane=enable_l2_swimlane, - ), - ) - return result - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("-p", "--platform", type=str, default="a2a3", - choices=["a2a3", "a2a3sim", "a5", "a5sim"]) - parser.add_argument("-d", "--device", type=int, default=0) - parser.add_argument("--enable-l2-swimlane", action="store_true", default=False) - args = parser.parse_args() - - result = compile_and_run( - platform=args.platform, - device_id=args.device, - enable_l2_swimlane=args.enable_l2_swimlane, - ) - if not result.passed: - if result.error: - print(result.error) - raise SystemExit(1) diff --git a/models/qwen3/32b/qwen3_32b_decode.py b/models/qwen3/32b/qwen3_32b_decode.py index d503893c..3916e0fe 100644 --- a/models/qwen3/32b/qwen3_32b_decode.py +++ b/models/qwen3/32b/qwen3_32b_decode.py @@ -47,13 +47,11 @@ HIDDEN_INV = 1.0 / HIDDEN # Scope 1 tiles -RMSNORM_K_CHUNK = 512 -Q_OUT_CHUNK = 256 -Q_PROJ_K_CHUNK = 128 -KV_OUT_CHUNK = 256 -KV_PROJ_K_CHUNK = 128 -Q_OUT_BLOCKS = HIDDEN // Q_OUT_CHUNK -KV_OUT_BLOCKS = KV_HIDDEN // KV_OUT_CHUNK +RMSNORM_K_TILE = 512 +Q_OUT_TILE = 256 +Q_PROJ_K_TILE = 128 +KV_OUT_TILE = 256 +KV_PROJ_K_TILE = 128 SCOPE2_STAGE_SPMD = 32 # Scope 2 tiles @@ -62,25 +60,19 @@ SEQ_TILE = 256 Q_GROUPS = Q_PER_KV // Q_HEAD_BATCH TOTAL_Q_GROUPS = NUM_KV_HEADS * Q_GROUPS -MAX_CTX_BLOCKS = (MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE # Scope 3 tiles -K_CHUNK = 128 -OUT_PROJ_K_CHUNK = 128 -MLP_OUT_CHUNK = 256 -DOWN_N_CHUNK = 256 -DOWN_K_CHUNK = 128 -OUT_PROJ_BLOCKS = HIDDEN // Q_OUT_CHUNK -MLP_OUT_BLOCKS = INTERMEDIATE // MLP_OUT_CHUNK -DOWN_PROJ_BLOCKS = HIDDEN // DOWN_N_CHUNK +K_TILE = 128 +OUT_PROJ_K_TILE = 128 +MLP_OUT_TILE = 256 +DOWN_N_TILE = 256 +DOWN_K_TILE = 128 - -HIDDEN_BLOCKS = HIDDEN // K_CHUNK BATCH_TILE = BATCH MLP_SPMD_INNER = 2 -MLP_GROUP_CHUNK = MLP_SPMD_INNER * MLP_OUT_CHUNK +MLP_GROUP_TILE = MLP_SPMD_INNER * MLP_OUT_TILE -assert MLP_OUT_BLOCKS % MLP_SPMD_INNER == 0 +assert (INTERMEDIATE // MLP_OUT_TILE) % MLP_SPMD_INNER == 0 def build_qwen3_decode_program(): @@ -114,27 +106,27 @@ def qwen3_decode( normed_states = pl.create_tensor([BATCH, HIDDEN], dtype=pl.BF16) with pl.at(level=pl.Level.CORE_GROUP, name_hint="rmsnorm"): partial_sq = pl.full([1, BATCH], dtype=pl.FP32, value=0.0) - for kb in pl.pipeline(HIDDEN // RMSNORM_K_CHUNK, stage=4): - k0 = kb * RMSNORM_K_CHUNK - x_chunk = pl.cast(hidden_states[:, k0 : k0 + RMSNORM_K_CHUNK], target_type=pl.FP32) + for kb in pl.pipeline(HIDDEN // RMSNORM_K_TILE, stage=4): + k0 = kb * RMSNORM_K_TILE + x_chunk = pl.cast(hidden_states[:, k0 : k0 + RMSNORM_K_TILE], target_type=pl.FP32) partial_sq = pl.add(partial_sq, pl.reshape(pl.row_sum(pl.mul(x_chunk, x_chunk)), [1, BATCH])) variance = pl.reshape(pl.add(pl.mul(partial_sq, HIDDEN_INV), EPS), [BATCH, 1]) inv_rms = pl.recip(pl.sqrt(variance)) - for kb in pl.pipeline(HIDDEN // RMSNORM_K_CHUNK, stage=4): - k0 = kb * RMSNORM_K_CHUNK - x_chunk = pl.cast(hidden_states[:, k0 : k0 + RMSNORM_K_CHUNK], target_type=pl.FP32) - gamma = input_rms_weight[:, k0 : k0 + RMSNORM_K_CHUNK] + for kb in pl.pipeline(HIDDEN // RMSNORM_K_TILE, stage=4): + k0 = kb * RMSNORM_K_TILE + x_chunk = pl.cast(hidden_states[:, k0 : k0 + RMSNORM_K_TILE], target_type=pl.FP32) + gamma = input_rms_weight[:, k0 : k0 + RMSNORM_K_TILE] normed = pl.col_expand_mul(pl.row_expand_mul(x_chunk, inv_rms), gamma) normed_states = pl.assemble(normed_states, pl.cast(normed, target_type=pl.BF16), [0, k0]) # Q projection. - for qi in pl.spmd(Q_OUT_BLOCKS, name_hint="q_proj"): - q0 = qi * Q_OUT_CHUNK - q_acc = pl.create_tensor([BATCH, Q_OUT_CHUNK], dtype=pl.FP32) - for kb in pl.pipeline(0, HIDDEN // Q_PROJ_K_CHUNK, stage=2): - k0 = kb * Q_PROJ_K_CHUNK - tile_a_i = normed_states[:, k0 : k0 + Q_PROJ_K_CHUNK] - tile_b_i = wq[k0 : k0 + Q_PROJ_K_CHUNK, q0 : q0 + Q_OUT_CHUNK] + for qi in pl.spmd(HIDDEN // Q_OUT_TILE, name_hint="q_proj"): + q0 = qi * Q_OUT_TILE + q_acc = pl.create_tensor([BATCH, Q_OUT_TILE], dtype=pl.FP32) + for kb in pl.pipeline(0, HIDDEN // Q_PROJ_K_TILE, stage=2): + k0 = kb * Q_PROJ_K_TILE + tile_a_i = normed_states[:, k0 : k0 + Q_PROJ_K_TILE] + tile_b_i = wq[k0 : k0 + Q_PROJ_K_TILE, q0 : q0 + Q_OUT_TILE] if k0 == 0: q_acc = pl.matmul(tile_a_i, tile_b_i, out_dtype=pl.FP32) else: @@ -142,15 +134,15 @@ def qwen3_decode( q_proj = pl.assemble(q_proj, q_acc, [0, q0]) # K/V projection. - for kvi in pl.spmd(KV_OUT_BLOCKS, name_hint="kv_proj"): - kv0 = kvi * KV_OUT_CHUNK - k_acc = pl.create_tensor([BATCH, KV_OUT_CHUNK], dtype=pl.FP32) - v_acc = pl.create_tensor([BATCH, KV_OUT_CHUNK], dtype=pl.FP32) - for kb in pl.pipeline(0, HIDDEN // KV_PROJ_K_CHUNK, stage=2): - k0 = kb * KV_PROJ_K_CHUNK - tile_a_i = normed_states[:, k0 : k0 + KV_PROJ_K_CHUNK] - tile_wk_i = wk[k0 : k0 + KV_PROJ_K_CHUNK, kv0 : kv0 + KV_OUT_CHUNK] - tile_wv_i = wv[k0 : k0 + KV_PROJ_K_CHUNK, kv0 : kv0 + KV_OUT_CHUNK] + for kvi in pl.spmd(KV_HIDDEN // KV_OUT_TILE, name_hint="kv_proj"): + kv0 = kvi * KV_OUT_TILE + k_acc = pl.create_tensor([BATCH, KV_OUT_TILE], dtype=pl.FP32) + v_acc = pl.create_tensor([BATCH, KV_OUT_TILE], dtype=pl.FP32) + for kb in pl.pipeline(0, HIDDEN // KV_PROJ_K_TILE, stage=2): + k0 = kb * KV_PROJ_K_TILE + tile_a_i = normed_states[:, k0 : k0 + KV_PROJ_K_TILE] + tile_wk_i = wk[k0 : k0 + KV_PROJ_K_TILE, kv0 : kv0 + KV_OUT_TILE] + tile_wv_i = wv[k0 : k0 + KV_PROJ_K_TILE, kv0 : kv0 + KV_OUT_TILE] if k0 == 0: k_acc = pl.matmul(tile_a_i, tile_wk_i, out_dtype=pl.FP32) v_acc = pl.matmul(tile_a_i, tile_wv_i, out_dtype=pl.FP32) @@ -202,7 +194,7 @@ def qwen3_decode( attn_row = pl.create_tensor([1, HIDDEN], dtype=pl.BF16) # Stage 2: QK matmul. - all_raw_scores = pl.create_tensor([TOTAL_Q_GROUPS * MAX_CTX_BLOCKS * Q_HEAD_PAD, SEQ_TILE], dtype=pl.FP32) + all_raw_scores = pl.create_tensor([TOTAL_Q_GROUPS * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD, SEQ_TILE], dtype=pl.FP32) for gi in pl.spmd(TOTAL_Q_GROUPS // 2, name_hint="qk_matmul"): gi0 = gi * 2 gi1 = gi * 2 + 1 @@ -217,17 +209,17 @@ def qwen3_decode( cache_row0_0 = b * NUM_KV_HEADS * MAX_SEQ + kvh0 * MAX_SEQ + s0 k_tile_0 = k_cache[cache_row0_0 : cache_row0_0 + SEQ_TILE, :] raw_scores_0 = pl.matmul(q_padded0, k_tile_0, b_trans=True, out_dtype=pl.FP32) - all_raw_scores = pl.assemble(all_raw_scores, raw_scores_0, [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + all_raw_scores = pl.assemble(all_raw_scores, raw_scores_0, [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) cache_row0_1 = b * NUM_KV_HEADS * MAX_SEQ + kvh1 * MAX_SEQ + s0 k_tile_1 = k_cache[cache_row0_1 : cache_row0_1 + SEQ_TILE, :] raw_scores_1 = pl.matmul(q_padded1, k_tile_1, b_trans=True, out_dtype=pl.FP32) - all_raw_scores = pl.assemble(all_raw_scores, raw_scores_1, [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + all_raw_scores = pl.assemble(all_raw_scores, raw_scores_1, [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) # Stage 3: softmax. - all_exp_padded = pl.create_tensor([TOTAL_Q_GROUPS * MAX_CTX_BLOCKS * Q_HEAD_PAD, SEQ_TILE], dtype=pl.BF16) - all_cur_li = pl.create_tensor([TOTAL_Q_GROUPS * MAX_CTX_BLOCKS * Q_HEAD_BATCH, 1], dtype=pl.FP32) - all_cur_mi = pl.create_tensor([TOTAL_Q_GROUPS * MAX_CTX_BLOCKS * Q_HEAD_BATCH, 1], dtype=pl.FP32) + all_exp_padded = pl.create_tensor([TOTAL_Q_GROUPS * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD, SEQ_TILE], dtype=pl.BF16) + all_cur_li = pl.create_tensor([TOTAL_Q_GROUPS * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH, 1], dtype=pl.FP32) + all_cur_mi = pl.create_tensor([TOTAL_Q_GROUPS * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH, 1], dtype=pl.FP32) for gi in pl.spmd(TOTAL_Q_GROUPS // 2, name_hint="softmax"): gi0 = gi * 2 gi1 = gi * 2 + 1 @@ -235,7 +227,7 @@ def qwen3_decode( s0 = sb * SEQ_TILE valid_len = pl.min(SEQ_TILE, ctx_len - s0) - scores_valid_0 = pl.slice(all_raw_scores, [Q_HEAD_BATCH, SEQ_TILE], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0], valid_shape=[Q_HEAD_BATCH, valid_len]) + scores_valid_0 = pl.slice(all_raw_scores, [Q_HEAD_BATCH, SEQ_TILE], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0], valid_shape=[Q_HEAD_BATCH, valid_len]) scores_padded_0 = pl.fillpad(scores_valid_0, pad_value=pl.PadValue.min) scores_0 = pl.mul(scores_padded_0, ATTN_SCALE) cur_mi_0 = pl.row_max(scores_0) @@ -243,11 +235,11 @@ def qwen3_decode( exp_scores_bf16_0 = pl.cast(exp_scores_0, target_type=pl.BF16) exp_scores_fp32_0 = pl.cast(exp_scores_bf16_0, target_type=pl.FP32) cur_li_0 = pl.row_sum(exp_scores_fp32_0) - all_exp_padded = pl.assemble(all_exp_padded, exp_scores_bf16_0, [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) - all_cur_mi = pl.assemble(all_cur_mi, cur_mi_0, [gi0 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) - all_cur_li = pl.assemble(all_cur_li, cur_li_0, [gi0 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + all_exp_padded = pl.assemble(all_exp_padded, exp_scores_bf16_0, [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + all_cur_mi = pl.assemble(all_cur_mi, cur_mi_0, [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + all_cur_li = pl.assemble(all_cur_li, cur_li_0, [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) - scores_valid_1 = pl.slice(all_raw_scores, [Q_HEAD_BATCH, SEQ_TILE], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0], valid_shape=[Q_HEAD_BATCH, valid_len]) + scores_valid_1 = pl.slice(all_raw_scores, [Q_HEAD_BATCH, SEQ_TILE], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0], valid_shape=[Q_HEAD_BATCH, valid_len]) scores_padded_1 = pl.fillpad(scores_valid_1, pad_value=pl.PadValue.min) scores_1 = pl.mul(scores_padded_1, ATTN_SCALE) cur_mi_1 = pl.row_max(scores_1) @@ -255,12 +247,12 @@ def qwen3_decode( exp_scores_bf16_1 = pl.cast(exp_scores_1, target_type=pl.BF16) exp_scores_fp32_1 = pl.cast(exp_scores_bf16_1, target_type=pl.FP32) cur_li_1 = pl.row_sum(exp_scores_fp32_1) - all_exp_padded = pl.assemble(all_exp_padded, exp_scores_bf16_1, [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) - all_cur_mi = pl.assemble(all_cur_mi, cur_mi_1, [gi1 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) - all_cur_li = pl.assemble(all_cur_li, cur_li_1, [gi1 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + all_exp_padded = pl.assemble(all_exp_padded, exp_scores_bf16_1, [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + all_cur_mi = pl.assemble(all_cur_mi, cur_mi_1, [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + all_cur_li = pl.assemble(all_cur_li, cur_li_1, [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) # Stage 4: SV matmul. - all_oi_tmp = pl.create_tensor([TOTAL_Q_GROUPS * MAX_CTX_BLOCKS * Q_HEAD_PAD, HEAD_DIM], dtype=pl.FP32) + all_oi_tmp = pl.create_tensor([TOTAL_Q_GROUPS * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD, HEAD_DIM], dtype=pl.FP32) for gi in pl.spmd(TOTAL_Q_GROUPS // 2, name_hint="sv_matmul"): gi0 = gi * 2 gi1 = gi * 2 + 1 @@ -269,16 +261,16 @@ def qwen3_decode( for sb in pl.range(ctx_blocks): s0 = sb * SEQ_TILE cache_row0_0 = b * NUM_KV_HEADS * MAX_SEQ + kvh0 * MAX_SEQ + s0 - exp_tile_0 = pl.slice(all_exp_padded, [Q_HEAD_PAD, SEQ_TILE], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + exp_tile_0 = pl.slice(all_exp_padded, [Q_HEAD_PAD, SEQ_TILE], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) v_tile_0 = v_cache[cache_row0_0 : cache_row0_0 + SEQ_TILE, :] oi_tmp_0 = pl.matmul(exp_tile_0, v_tile_0, out_dtype=pl.FP32) - all_oi_tmp = pl.assemble(all_oi_tmp, oi_tmp_0, [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + all_oi_tmp = pl.assemble(all_oi_tmp, oi_tmp_0, [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) cache_row0_1 = b * NUM_KV_HEADS * MAX_SEQ + kvh1 * MAX_SEQ + s0 - exp_tile_1 = pl.slice(all_exp_padded, [Q_HEAD_PAD, SEQ_TILE], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + exp_tile_1 = pl.slice(all_exp_padded, [Q_HEAD_PAD, SEQ_TILE], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) v_tile_1 = v_cache[cache_row0_1 : cache_row0_1 + SEQ_TILE, :] oi_tmp_1 = pl.matmul(exp_tile_1, v_tile_1, out_dtype=pl.FP32) - all_oi_tmp = pl.assemble(all_oi_tmp, oi_tmp_1, [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + all_oi_tmp = pl.assemble(all_oi_tmp, oi_tmp_1, [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) # Stage 5: online softmax accumulation and normalisation. for gi in pl.spmd(TOTAL_Q_GROUPS // 2, name_hint="online_softmax"): @@ -290,16 +282,16 @@ def qwen3_decode( kvh1 = gi1 // Q_GROUPS qg1 = gi1 - kvh1 * Q_GROUPS q_base1 = kvh1 * Q_PER_KV + qg1 * Q_HEAD_BATCH - oi_0 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD, 0]) - mi_0 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_BATCH, 0]) - li_0 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_BATCH, 0]) - oi_1 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD, 0]) - mi_1 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_BATCH, 0]) - li_1 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_BATCH, 0]) + oi_0 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD, 0]) + mi_0 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH, 0]) + li_0 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH, 0]) + oi_1 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD, 0]) + mi_1 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH, 0]) + li_1 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH, 0]) for sb in pl.range(1, ctx_blocks): - oi_tmp_valid_0 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) - cur_mi_0 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) - cur_li_0 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi0 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + oi_tmp_valid_0 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + cur_mi_0 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + cur_li_0 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi0 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) mi_new_0 = pl.maximum(mi_0, cur_mi_0) alpha_0 = pl.exp(pl.sub(mi_0, mi_new_0)) beta_0 = pl.exp(pl.sub(cur_mi_0, mi_new_0)) @@ -307,9 +299,9 @@ def qwen3_decode( oi_0 = pl.add(pl.row_expand_mul(oi_0, alpha_0), pl.row_expand_mul(oi_tmp_valid_0, beta_0)) mi_0 = mi_new_0 - oi_tmp_valid_1 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) - cur_mi_1 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) - cur_li_1 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi1 * MAX_CTX_BLOCKS * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + oi_tmp_valid_1 = pl.slice(all_oi_tmp, [Q_HEAD_BATCH, HEAD_DIM], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD + sb * Q_HEAD_PAD, 0]) + cur_mi_1 = pl.slice(all_cur_mi, [Q_HEAD_BATCH, 1], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) + cur_li_1 = pl.slice(all_cur_li, [Q_HEAD_BATCH, 1], [gi1 * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_BATCH + sb * Q_HEAD_BATCH, 0]) mi_new_1 = pl.maximum(mi_1, cur_mi_1) alpha_1 = pl.exp(pl.sub(mi_1, mi_new_1)) beta_1 = pl.exp(pl.sub(cur_mi_1, mi_new_1)) @@ -330,16 +322,16 @@ def qwen3_decode( resid1_tile = pl.create_tensor([BATCH, HIDDEN], dtype=pl.FP32) # Stage 1 & 2: Output projection + residual addition with hidden_states. - for ob in pl.parallel(0, HIDDEN // Q_OUT_CHUNK, 2): - with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.auto_chunk, pl.split(pl.SplitMode.UP_DOWN)], name_hint="out_proj_residual"): + for ob in pl.parallel(0, HIDDEN // Q_OUT_TILE, 2): + with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.split(pl.SplitMode.UP_DOWN)], name_hint="out_proj_residual"): for oi in pl.range(ob, ob + 2): - o0 = oi * Q_OUT_CHUNK - hidden_chunk = hidden_states[:, o0 : o0 + Q_OUT_CHUNK] - o_acc = pl.create_tensor([BATCH, Q_OUT_CHUNK], dtype=pl.FP32) - for kb in pl.pipeline(0, HIDDEN // OUT_PROJ_K_CHUNK, stage=2): - k0 = kb * OUT_PROJ_K_CHUNK - a_chunk = attn_out[:, k0 : k0 + OUT_PROJ_K_CHUNK] - w_chunk = wo[k0 : k0 + OUT_PROJ_K_CHUNK, o0 : o0 + Q_OUT_CHUNK] + o0 = oi * Q_OUT_TILE + hidden_chunk = hidden_states[:, o0 : o0 + Q_OUT_TILE] + o_acc = pl.create_tensor([BATCH, Q_OUT_TILE], dtype=pl.FP32) + for kb in pl.pipeline(0, HIDDEN // OUT_PROJ_K_TILE, stage=2): + k0 = kb * OUT_PROJ_K_TILE + a_chunk = attn_out[:, k0 : k0 + OUT_PROJ_K_TILE] + w_chunk = wo[k0 : k0 + OUT_PROJ_K_TILE, o0 : o0 + Q_OUT_TILE] if k0 == 0: o_acc = pl.matmul(a_chunk, w_chunk, out_dtype=pl.FP32) else: @@ -352,85 +344,85 @@ def qwen3_decode( post_norm_tile = pl.create_tensor([BATCH, HIDDEN], dtype=pl.BF16) with pl.at(level=pl.Level.CORE_GROUP, name_hint="post_rmsnorm"): sq_sum = pl.full([1, BATCH], dtype=pl.FP32, value=0.0) - for kb in pl.pipeline(HIDDEN // K_CHUNK, stage=2): - k0 = kb * K_CHUNK - resid_chunk = resid1_tile[:, k0 : k0 + K_CHUNK] + for kb in pl.pipeline(HIDDEN // K_TILE, stage=2): + k0 = kb * K_TILE + resid_chunk = resid1_tile[:, k0 : k0 + K_TILE] sq_sum = pl.add(sq_sum, pl.reshape(pl.row_sum(pl.mul(resid_chunk, resid_chunk)), [1, BATCH])) inv_rms_s3 = pl.recip(pl.sqrt(pl.add(pl.mul(sq_sum, HIDDEN_INV), EPS))) inv_rms_s3_col = pl.reshape(inv_rms_s3, [BATCH, 1]) - for kb in pl.pipeline(HIDDEN // K_CHUNK, stage=2): - k0 = kb * K_CHUNK - resid_chunk = resid1_tile[:, k0 : k0 + K_CHUNK] - post_gamma = post_rms_weight[:, k0 : k0 + K_CHUNK] + for kb in pl.pipeline(HIDDEN // K_TILE, stage=2): + k0 = kb * K_TILE + resid_chunk = resid1_tile[:, k0 : k0 + K_TILE] + post_gamma = post_rms_weight[:, k0 : k0 + K_TILE] post_normed = pl.col_expand_mul(pl.row_expand_mul(resid_chunk, inv_rms_s3_col), post_gamma) post_norm_tile = pl.assemble(post_norm_tile, pl.cast(post_normed, target_type=pl.BF16), [0, k0]) # Stage 4~6: keep outer parallel, run smaller SPMD groups and cache per-group only. mlp_tile = pl.create_tensor([BATCH, INTERMEDIATE], dtype=pl.BF16) - for ob_base in pl.parallel(0, MLP_OUT_BLOCKS, MLP_SPMD_INNER): - gate_group = pl.create_tensor([BATCH_TILE, MLP_GROUP_CHUNK], dtype=pl.FP32) - up_group = pl.create_tensor([BATCH_TILE, MLP_GROUP_CHUNK], dtype=pl.FP32) + for ob_base in pl.parallel(0, INTERMEDIATE // MLP_OUT_TILE, MLP_SPMD_INNER): + gate_group = pl.create_tensor([BATCH_TILE, MLP_GROUP_TILE], dtype=pl.FP32) + up_group = pl.create_tensor([BATCH_TILE, MLP_GROUP_TILE], dtype=pl.FP32) # Stage 4: gate projection. for ob in pl.spmd(MLP_SPMD_INNER, name_hint="gate_proj_spmd"): - o0 = (ob_base + ob) * MLP_OUT_CHUNK - g0 = ob * MLP_OUT_CHUNK - post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, 0]) - post_chunk_1 = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, K_CHUNK]) - wg_0 = pl.slice(w_gate, [K_CHUNK, MLP_OUT_CHUNK], [0, o0]) + o0 = (ob_base + ob) * MLP_OUT_TILE + g0 = ob * MLP_OUT_TILE + post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, 0]) + post_chunk_1 = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, K_TILE]) + 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) - wg_1 = pl.slice(w_gate, [K_CHUNK, MLP_OUT_CHUNK], [K_CHUNK, o0]) + wg_1 = pl.slice(w_gate, [K_TILE, MLP_OUT_TILE], [K_TILE, o0]) gate_acc = pl.matmul_acc(gate_acc, post_chunk_1, wg_1) - for kb in pl.pipeline(2, HIDDEN_BLOCKS, stage=2): - 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.pipeline(2, HIDDEN // K_TILE, stage=2): + 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) gate_group = pl.assemble(gate_group, gate_acc, [0, g0]) # Stage 5: up projection. for ob in pl.spmd(MLP_SPMD_INNER, name_hint="up_proj_spmd"): - o0 = (ob_base + ob) * MLP_OUT_CHUNK - g0 = ob * MLP_OUT_CHUNK - post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, 0]) - post_chunk_1 = pl.slice(post_norm_tile, [BATCH_TILE, K_CHUNK], [0, K_CHUNK]) - wu_0 = pl.slice(w_up, [K_CHUNK, MLP_OUT_CHUNK], [0, o0]) + o0 = (ob_base + ob) * MLP_OUT_TILE + g0 = ob * MLP_OUT_TILE + post_chunk_0 = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, 0]) + post_chunk_1 = pl.slice(post_norm_tile, [BATCH_TILE, K_TILE], [0, K_TILE]) + 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) - wu_1 = pl.slice(w_up, [K_CHUNK, MLP_OUT_CHUNK], [K_CHUNK, o0]) + wu_1 = pl.slice(w_up, [K_TILE, MLP_OUT_TILE], [K_TILE, o0]) up_acc = pl.matmul_acc(up_acc, post_chunk_1, wu_1) - for kb in pl.pipeline(2, HIDDEN_BLOCKS, stage=2): - 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.pipeline(2, HIDDEN // K_TILE, stage=2): + 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) up_group = pl.assemble(up_group, up_acc, [0, g0]) # Stage 6: SiLU + gate/up fuse. for ob in pl.spmd(MLP_SPMD_INNER, name_hint="silu_spmd"): - o0 = (ob_base + ob) * MLP_OUT_CHUNK - g0 = ob * MLP_OUT_CHUNK - gate_acc = pl.slice(gate_group, [BATCH_TILE, MLP_OUT_CHUNK], [0, g0]) - up_acc = pl.slice(up_group, [BATCH_TILE, MLP_OUT_CHUNK], [0, g0]) + o0 = (ob_base + ob) * MLP_OUT_TILE + g0 = ob * MLP_OUT_TILE + gate_acc = pl.slice(gate_group, [BATCH_TILE, MLP_OUT_TILE], [0, g0]) + up_acc = pl.slice(up_group, [BATCH_TILE, MLP_OUT_TILE], [0, g0]) sigmoid = pl.recip(pl.add(pl.exp(pl.neg(gate_acc)), 1.0)) mlp_chunk = pl.mul(pl.mul(gate_acc, sigmoid), up_acc) mlp_chunk_bf16 = pl.cast(mlp_chunk, target_type=pl.BF16) mlp_tile = pl.assemble(mlp_tile, mlp_chunk_bf16, [0, o0]) # Stage 7 & 8: Down projection + final residual writeback. - for db in pl.parallel(0, HIDDEN // DOWN_N_CHUNK, 2): - with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.auto_chunk, pl.split(pl.SplitMode.UP_DOWN)], name_hint="down_proj_residual"): + for db in pl.parallel(0, HIDDEN // DOWN_N_TILE, 2): + with pl.at(level=pl.Level.CORE_GROUP, optimizations=[pl.split(pl.SplitMode.UP_DOWN)], name_hint="down_proj_residual"): for di in pl.range(db, db + 2): - d0 = di * DOWN_N_CHUNK - resid1_tile_chunk = resid1_tile[:, d0 : d0 + DOWN_N_CHUNK] - down_acc = pl.create_tensor([BATCH, DOWN_N_CHUNK], dtype=pl.FP32) - for ob in pl.pipeline(0, INTERMEDIATE // DOWN_K_CHUNK, stage=2): - o0 = ob * DOWN_K_CHUNK - down_mlp_chunk = mlp_tile[:, o0 : o0 + DOWN_K_CHUNK] - w_down_chunk = w_down[o0 : o0 + DOWN_K_CHUNK, d0 : d0 + DOWN_N_CHUNK] + d0 = di * DOWN_N_TILE + resid1_tile_chunk = resid1_tile[:, d0 : d0 + DOWN_N_TILE] + down_acc = pl.create_tensor([BATCH, DOWN_N_TILE], dtype=pl.FP32) + for ob in pl.pipeline(0, INTERMEDIATE // DOWN_K_TILE, stage=2): + o0 = ob * DOWN_K_TILE + down_mlp_chunk = mlp_tile[:, o0 : o0 + DOWN_K_TILE] + w_down_chunk = w_down[o0 : o0 + DOWN_K_TILE, d0 : d0 + DOWN_N_TILE] if o0 == 0: down_acc = pl.matmul(down_mlp_chunk, w_down_chunk, out_dtype=pl.FP32) else: