fix(v3_2,qwen3-32b): migrate off retired incore/auto_chunk DSL + tidy naming#671
Conversation
… 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.
📝 WalkthroughWalkthroughThis PR retiles several DeepSeek V3.2 (decode-back, decode-front, prefill-back) and Qwen3-32B decode kernels, replacing chunk/block-based tiling constants and parameterized builders with fixed compile-time TILE constants and adjusted loop/slice indexing. The unused prefill-front draft module is deleted entirely. ChangesDeepSeek V3.2 and Qwen3-32B tiling refactor
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)Skipped — the changes are internal kernel tiling refactors within single files/pipelines rather than multi-component interactions. Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request renames tiling constants from *_CHUNK to *_TILE across several DeepSeek and Qwen3 model files, simplifies program builder signatures by directly referencing global constants, and inlines several block count calculations. Additionally, it optimizes the RMSNorm step in deepseek_v3_2_prefill_back.py using parallel block-chunking and removes the deepseek_v3_2_prefill_front.py file entirely. The reviewer feedback suggests retaining the MAX_CTX_BLOCKS constant in qwen3_32b_decode.py instead of repeating the inline ceiling division expression, which harms code readability and maintainability.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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 | ||
|
|
There was a problem hiding this comment.
Inlining the ceiling division expression ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) over 30 times across the file significantly reduces readability and maintainability. Consider defining MAX_CTX_BLOCKS as a module-level constant (as it was originally) to keep the code clean and avoid repetitive complex expressions.
| MAX_CTX_BLOCKS = (MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE |
|
|
||
| # 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) |
There was a problem hiding this comment.
Use the MAX_CTX_BLOCKS constant here instead of repeating the inline ceiling division expression.
| all_raw_scores = pl.create_tensor([TOTAL_Q_GROUPS * ((MAX_SEQ + SEQ_TILE - 1) // SEQ_TILE) * Q_HEAD_PAD, SEQ_TILE], dtype=pl.FP32) | |
| all_raw_scores = pl.create_tensor([TOTAL_Q_GROUPS * MAX_CTX_BLOCKS * Q_HEAD_PAD, SEQ_TILE], dtype=pl.FP32) |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
models/deepseek/v3_2/deepseek_v3_2_decode_back.py (1)
35-38: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd explicit divisibility guards for the new tile constants.
The new loops use
//bounds and fixed-size slices, so a future dimension/tile change can silently drop tails or slice out of bounds. Add module-level assertions forBATCH,HIDDEN,ATTN_OUT, andINTERMEDIATEtile divisibility.Proposed guard
K_TILE = 128 Q_OUT_TILE = 64 MLP_OUT_TILE = 256 BATCH_TILE = 16 + +assert BATCH % BATCH_TILE == 0 +assert HIDDEN % K_TILE == 0 +assert HIDDEN % Q_OUT_TILE == 0 +assert ATTN_OUT % K_TILE == 0 +assert INTERMEDIATE % MLP_OUT_TILE == 0🤖 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 35 - 38, Add module-level divisibility checks for the tile constants in deepseek_v3_2_decode_back.py, since the fixed-size loops and slices in this module assume exact tiling. Update the top-level constants area around K_TILE, Q_OUT_TILE, MLP_OUT_TILE, and BATCH_TILE to assert that BATCH, HIDDEN, ATTN_OUT, and INTERMEDIATE are evenly divisible by their corresponding tile sizes so future shape changes fail fast instead of silently truncating tails.models/deepseek/v3_2/deepseek_v3_2_decode_front.py (1)
65-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
QREDUCE_BATCH_TILEis defined but unused.This constant is introduced alongside the other retiling constants but doesn't appear anywhere in the Stage 2.8 reduction loop (Lines 374-386), which iterates the batch dimension directly via
pl.range(BATCH)rather than tiling byQREDUCE_BATCH_TILE. If it's leftover from an earlier design, consider removing it; if it was meant to be wired into the batch loop, this is worth revisiting.🤖 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_front.py` at line 65, QREDUCE_BATCH_TILE is currently defined in deepseek_v3_2_decode_front.py but never used, so either remove the unused constant or wire it into the Stage 2.8 reduction loop in the relevant batch-processing logic (for example the loop that currently uses pl.range(BATCH)) if tiling was intended. Check the surrounding reduction/retiling constants and the Stage 2.8 loop so the batch iteration and constants stay consistent.models/qwen3/32b/qwen3_32b_decode.py (1)
65-75: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMake the tile invariants complete before the floor-divided loops.
Line 75 only validates the grouped MLP loop. The retiled kernel now relies on many
//-based bounds; if any dimension is not exactly divisible, the tail is silently skipped or the+2paired loops can overrun. Add explicit assertions for every tile partition, matching the invariant style used by the 14B decode layer context.Suggested invariant coverage
BATCH_TILE = BATCH MLP_SPMD_INNER = 2 MLP_GROUP_TILE = MLP_SPMD_INNER * MLP_OUT_TILE +assert HIDDEN % RMSNORM_K_TILE == 0 +assert HIDDEN % Q_OUT_TILE == 0 +assert (HIDDEN // Q_OUT_TILE) % 2 == 0 +assert HIDDEN % Q_PROJ_K_TILE == 0 +assert KV_HIDDEN % KV_OUT_TILE == 0 +assert HIDDEN % KV_PROJ_K_TILE == 0 +assert HIDDEN % K_TILE == 0 +assert HIDDEN % OUT_PROJ_K_TILE == 0 +assert HIDDEN % DOWN_N_TILE == 0 +assert (HIDDEN // DOWN_N_TILE) % 2 == 0 +assert INTERMEDIATE % MLP_OUT_TILE == 0 assert (INTERMEDIATE // MLP_OUT_TILE) % MLP_SPMD_INNER == 0 +assert INTERMEDIATE % DOWN_K_TILE == 0🤖 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/qwen3/32b/qwen3_32b_decode.py` around lines 65 - 75, The tile invariants in qwen3_32b_decode are incomplete, so add explicit divisibility/assert checks before any floor-divided loop logic that depends on them. Use the existing constants in this module, especially K_TILE, OUT_PROJ_K_TILE, MLP_OUT_TILE, DOWN_N_TILE, DOWN_K_TILE, BATCH_TILE, MLP_SPMD_INNER, and MLP_GROUP_TILE, and mirror the invariant style used in the 14B decode layer context. Make sure every dimension consumed by the retiled kernels and paired loops is proven exactly divisible so the grouped MLP and downstream loops cannot silently drop tails or overrun.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@models/deepseek/v3_2/deepseek_v3_2_decode_back.py`:
- Around line 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.
---
Nitpick comments:
In `@models/deepseek/v3_2/deepseek_v3_2_decode_back.py`:
- Around line 35-38: Add module-level divisibility checks for the tile constants
in deepseek_v3_2_decode_back.py, since the fixed-size loops and slices in this
module assume exact tiling. Update the top-level constants area around K_TILE,
Q_OUT_TILE, MLP_OUT_TILE, and BATCH_TILE to assert that BATCH, HIDDEN, ATTN_OUT,
and INTERMEDIATE are evenly divisible by their corresponding tile sizes so
future shape changes fail fast instead of silently truncating tails.
In `@models/deepseek/v3_2/deepseek_v3_2_decode_front.py`:
- Line 65: QREDUCE_BATCH_TILE is currently defined in
deepseek_v3_2_decode_front.py but never used, so either remove the unused
constant or wire it into the Stage 2.8 reduction loop in the relevant
batch-processing logic (for example the loop that currently uses
pl.range(BATCH)) if tiling was intended. Check the surrounding
reduction/retiling constants and the Stage 2.8 loop so the batch iteration and
constants stay consistent.
In `@models/qwen3/32b/qwen3_32b_decode.py`:
- Around line 65-75: The tile invariants in qwen3_32b_decode are incomplete, so
add explicit divisibility/assert checks before any floor-divided loop logic that
depends on them. Use the existing constants in this module, especially K_TILE,
OUT_PROJ_K_TILE, MLP_OUT_TILE, DOWN_N_TILE, DOWN_K_TILE, BATCH_TILE,
MLP_SPMD_INNER, and MLP_GROUP_TILE, and mirror the invariant style used in the
14B decode layer context. Make sure every dimension consumed by the retiled
kernels and paired loops is proven exactly divisible so the grouped MLP and
downstream loops cannot silently drop tails or overrun.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 4689a4a7-bb43-4c9f-a588-b880d9ce2f84
📒 Files selected for processing (5)
models/deepseek/v3_2/deepseek_v3_2_decode_back.pymodels/deepseek/v3_2/deepseek_v3_2_decode_front.pymodels/deepseek/v3_2/deepseek_v3_2_prefill_back.pymodels/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.pymodels/qwen3/32b/qwen3_32b_decode.py
💤 Files with no reviewable changes (1)
- models/deepseek/v3_2/deepseek_v3_2_prefill_front_draft.py
| 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]: |
There was a problem hiding this comment.
🗄️ 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.
1 participant
Summary
pl.incore/pl.auto_incore/pl.auto_chunkand the loopchunk=kwarg (#1819, #1895), which broke the v3.2decode_frontandprefill_backcases at compile time. Migrated topl.at(level=CORE_GROUP)+ manual loop tiling (parallelstep +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 thedecode_backshape (materializemlp_tile, then a separate down-proj loop merging the final residual).BATCH4 → 1 to keep golden fast.qwen3_32b_decode: dropped the now-redundantpl.auto_chunk, keptpl.split(pl.SplitMode.UP_DOWN).qwen3_32b:_CHUNK→_TILE, flattened the program builders to use module-level constants (dropped_CFG/_SIZEaliases and builder params), inlined_BLOCKS.deepseek_v3_2_prefill_front_draft.py.Validated on a2a3 (
--ptoas 0.46):decode_front,decode_back,prefill_back,qwen3_32b_decodeall PASS.