feat(step3p5): initial in-tree drop of step3p5 decoder + Phase A split-scope refactor#510
feat(step3p5): initial in-tree drop of step3p5 decoder + Phase A split-scope refactor#510csy0225 wants to merge 2 commits into
Conversation
…t-scope refactor
Adds the step3p5 model implementation (35 files, ~27k lines) under
models/step3p5/. This is the first commit of code that has lived as
untracked working-tree WIP for several weeks (TASK-22).
The 5 files below carry this session's Phase A refactor on top of the
in-tree drop. Phase A splits mixed AIC+AIV scopes into pure-cube +
pure-vec spmds (FP32 GM scratch handoff), mirroring qwen3/32b's idiom,
in an attempt to unblock the Phase 15 single-rank 507018 VEC UB align
crash. The refactor compiles clean (smoke probe rc=0) but does NOT
unblock the runtime crash — see upstream issue simpler#1036 for the
full diagnosis.
Phase A files (with diff intent):
- attention_full.py : full_fa_fused -> 4 spmds + full_out_proj
cube/cast split + full_rope_kv_cache qwen3
full-row-cast idiom + full_rmsnorm_zc to
CORE_GROUP+pipeline form
- attention_swa.py : mirror of split + out_proj split
- decode_layer.py : dense_gate_up_silu_tp + dense_down_proj_tp
each split into cube/cast
- prefill_attention_full.py : prefill_full_out_proj cube/cast split
+ TP=1 tp_all_reduce guard
- prefill_attention_swa.py : same mirrors
All other files (config.py, _ops.py, collectives.py, moe.py, etc.) are
the existing WIP at the state they were before Phase A started.
Status:
- Phase 14 frontend smoke probe: rc=0 (see _smoke_program_build.py)
- Phase 15 single-rank NPU: blocked at simpler#1036 (507018 VEC UB
not aligned). Phase A refactors did not move the fault but are
net-positive cleanup and align with qwen3/32b patterns.
- Phase 16 multi-rank: blocked at simpler#1037 (driver
support_shmem_map_exbus=0).
Related issues:
- hw-native-sys/simpler#1036 (Phase 15 single-rank blocker)
- hw-native-sys/simpler#1037 (Phase 16 multi-rank blocker)
- hw-native-sys/simpler#1023 (earlier 507018 zero-shape view fix
— landed but not sufficient for Phase 15 unblock)
- hw-native-sys/simpler#1018 (comm_init segfault — verified fixed by
--no-as-needed link patch this session)
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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 implements the Step3p5 language model on the PyPTO framework with TP=8 and EP=8, introducing kernel modules for attention, MoE, MTP, and final LM head projections. The code review identified several critical issues that must be addressed before deployment: an undefined variable current_hidden in the orchestrator, multiple SRAM memory limit violations (exceeding 192 KB) in the routed expert and prefill attention modules, and potential compiler failures in the shared expert due to auto-K-split. Additionally, performance can be optimized by reducing redundant global memory reads in the dispatch, combine, and MoE modules, avoiding redundant writes in the sliding-window attention path, and replacing the deprecated pld.tile.remote_store API with the supported pl.store and pld.tensor.put pattern.
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.
| # outside the @pl.program (the per-layer programs each | ||
| # build their own host_orch — see decode_layer.py). |
There was a problem hiding this comment.
The variable current_hidden is undefined in the host_orch scope. The parameter is named hidden_states in the function signature. This will cause a compilation or runtime NameError when executing the orchestrator.
self.chip_orch(
hidden_states[r],| pair lands on it (degenerate hot-routing): ``EP_WORLD_SIZE * BATCH * | ||
| TOPK`` rows. We size ``LOCAL_RECV_MAX`` to that worst case. With | ||
| ``BATCH=16, TOPK=8, EP_WORLD_SIZE=8`` -> 1024 rows, which at |
There was a problem hiding this comment.
The tensor h_tile is allocated with shape [MAX_TILE, INTER] in FP32, which requires 1024 * 1280 * 4 = 5.24 MB of memory. This exceeds the 192 KB Vector UB (SRAM) limit on the Ascend NPU, leading to compilation or runtime failures. Please implement row-tiling (e.g., using RECV_TILE = 32 as done in moe.py) to keep the SRAM footprint within the hardware limits.
| # pyright: reportUndefinedVariable=false | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import pypto.language as pl | ||
|
|
||
| from .config import ( |
There was a problem hiding this comment.
Using a single K-tile of size DOWN_K_CHUNK = INTER = 160 directly will trigger the backend auto-K-split, which emits an invalid tmov acc→acc instruction and causes compilation or runtime failures. Please implement explicit K-chunking at the user level (e.g., using _SH_DOWN_K_INNER = 32 as done in moe.py) to ensure the compiler generates a structured matmul + matmul_acc chain.
| def build_inverse_map( | ||
| indices, # pl.Tensor[[T, TOPK], pl.INT32] | ||
| pub_counts, # cross-rank read-only — [N_RANKS*N_RANKS, N_LOCAL_EXPERTS] | ||
| inverse_map, # pl.Tensor[[T, TOPK], pl.INT32] — packed (dst_rank, dst_row) | ||
| my_rank, # pl.Scalar[pl.INT32] | ||
| ): | ||
| """For each (t, k), encode (dst_rank, dst_row_in_recv_buf) into one INT32. | ||
|
|
||
| The receiver's recv buffer layout is (src_rank-major within each | ||
| local_eid). My rank's contribution to dst's expert ``loc_e`` lands | ||
| at slot ``Σ_{s<my_rank} pub_counts[s*N_RANKS + dst, loc_e] + my_cursor`` | ||
| inside the (loc_e) block. The (loc_e) block itself starts at | ||
| ``Σ_{prev_e} (Σ_s pub_counts[s*N_RANKS + dst, prev_e])`` rows from | ||
| the recv buffer base. | ||
|
|
||
| For combine.py we only need the per-(t, k) destination row inside | ||
| the receiver's recv buffer; combine.py walks the same arithmetic in | ||
| reverse using the same pub_counts window. | ||
|
|
||
| The packed encoding ``dst_rank * LOCAL_RECV_MAX + dst_row`` fits in |
There was a problem hiding this comment.
The nested loops inside build_inverse_map perform up to T * TOPK * N_LOCAL_EXPERTS * N_RANKS (e.g., 16 * 8 * 36 * 8 = 36,864) scalar reads from the pub_counts DistributedTensor in Global Memory (GM). Scalar reads from GM are extremely slow on the Ascend NPU due to high latency. Precomputing the cumulative offsets for all 288 buckets once in a flat local tensor of size 288 reduces the GM reads to O(1) lookups inside the main loop, significantly improving performance.
| for qkn_idx in pl.spmd( | ||
| (PREFILL_T // TOK_TILE) * 1, | ||
| name_hint="prefill_swa_qk_norm_zc", | ||
| ): | ||
| tg_idx2 = qkn_idx // 1 | ||
| kh = qkn_idx % 1 | ||
| tg = tg_idx2 * TOK_TILE | ||
| q_col = kh * 12 * HEAD_DIM | ||
| q_chunk = pl.reshape( | ||
| pl.slice( |
There was a problem hiding this comment.
The prefill_swa_qk_norm_zc loop uses TOK_TILE (32) directly, resulting in a q_chunk of shape [384, 128] (since q_per_kv = 12). This requires 196.6 KB for a single FP32 tensor, which exceeds the 192 KB Vector UB limit and will cause compilation or runtime failures. Please use a smaller sub-tile size (e.g., QK_NORM_T_TILE = TOK_TILE // 4 as done in prefill_attention_full.py) to keep the SRAM footprint within hardware limits.
| # Advance e_cursor by total rows for this local expert e. | ||
| total_e = pl.cast(0, pl.INT32) | ||
| for src2 in pl.range(N_RANKS): | ||
| total_e = total_e + pl.read( | ||
| pub_counts, [src2 * N_RANKS + my_rank, e], | ||
| ) | ||
| e_cursor = e_cursor + total_e |
There was a problem hiding this comment.
The loop over src2 to compute total_e is completely redundant because src_off already holds the exact sum of n over all src ranks from the previous loop. You can directly use src_off to advance e_cursor, avoiding 36 * 8 = 288 redundant GM reads.
| # Advance e_cursor by total rows for this local expert e. | |
| total_e = pl.cast(0, pl.INT32) | |
| for src2 in pl.range(N_RANKS): | |
| total_e = total_e + pl.read( | |
| pub_counts, [src2 * N_RANKS + my_rank, e], | |
| ) | |
| e_cursor = e_cursor + total_e | |
| e_cursor = e_cursor + src_off |
| if src == my_rank: | ||
| pl.store(tile, [r_route, 0], routed_y_buf) | ||
| else: | ||
| pld.tile.remote_store( | ||
| tile, | ||
| target=routed_y_buf, |
There was a problem hiding this comment.
The loop over src2 to compute total_e is completely redundant because src_off already holds the exact sum of n over all src ranks from the previous loop. You can directly use src_off to advance e_cursor, avoiding 36 * 8 = 288 redundant GM reads.
e_cursor = e_cursor + src_off| [local_recv_max, HIDDEN], dtype=pl.BF16, | ||
| ) | ||
| local_expert_offset = pl.create_tensor( | ||
| [n_local_experts], dtype=pl.INT32, | ||
| ) | ||
| local_expert_count = pl.create_tensor( |
There was a problem hiding this comment.
The loop over src2 to compute total_e is completely redundant because src_off already holds the exact sum of n over all src ranks from the previous loop. You can directly use src_off to advance e_cursor, avoiding 36 * 8 = 288 redundant GM reads.
| [local_recv_max, HIDDEN], dtype=pl.BF16, | |
| ) | |
| local_expert_offset = pl.create_tensor( | |
| [n_local_experts], dtype=pl.INT32, | |
| ) | |
| local_expert_count = pl.create_tensor( | |
| e_cursor = e_cursor + src_off |
| pld.tile.remote_store( | ||
| tile, | ||
| target=routed_y_buf, | ||
| peer=src, | ||
| offsets=[r_route, 0], | ||
| ) |
There was a problem hiding this comment.
The pld.tile.remote_store API is deprecated/unsupported. Please use the canonical pl.store followed by pld.tensor.put pattern as done in prefill_moe.py to ensure compatibility and correct execution.
| local_row = ( | ||
| pl.cast(e_cursor, pl.INDEX) | ||
| + pl.cast(src_off, pl.INDEX) + row | ||
| ) | ||
| tile = pl.load( | ||
| local_routed_y, |
There was a problem hiding this comment.
The pld.tile.remote_store API is deprecated/unsupported. Please use the canonical pl.store followed by pld.tensor.put pattern as done in prefill_moe.py to ensure compatibility and correct execution.
Eliminate intermediate cos_row/sin_row temporaries; slice each lo/hi half directly from rope_cos / rope_sin in one pl.slice. 6 slices -> 4, no behavior change.
1 participant
Summary
Initial in-tree drop of
models/step3p5/(35 files, ~27k lines).This is the first commit of code that has lived as untracked working-tree WIP
for several weeks (TASK-22 in our local backlog). The 5 attention/MLP files
listed below also carry this session's Phase A refactor on top of the
initial drop — splitting mixed AIC+AIV scopes into pure-cube + pure-vec
spmds (FP32 GM scratch handoff), mirroring qwen3/32b's working idiom.
What works
python -m models.step3p5._smoke_program_build) returns rc=0readystate on real NPU (a2a3, single rank)simpler / pto-isa / PTOAS / Python — so the runtime stack is healthy
What does NOT work (upstream-blocked)
dispatch with
aclrtSynchronizeStreamWithTimeout (AICPU) failed: 507018,plog
errcode 0x800 (VEC UB not aligned) tslot:6. Filed assimpler#1036. The
Phase A refactor in this PR does NOT unblock this (8 different model-side
mitigations all leave the fault hash unchanged — see issue for full matrix).
support_shmem_map_exbus=0→
aclrtIpcMemImportByKey 507899. Filed assimpler#1037.
Why open the PR now
Even though the code can't run end-to-end on NPU yet (waiting on upstream
fixes for simpler#1036), getting it into review NOW makes sense because:
when investigating the VEC UB align fault
eventually unblocks 507018 — eliminates mixed AIC+AIV dispatch and
aligns with the working qwen3/32b patterns
Phase A files (with diff intent)
attention_full.py:full_fa_fused-> 4 spmds +full_out_projcube/cast split +
full_rope_kv_cacheqwen3 full-row-cast idiom +full_rmsnorm_zcto CORE_GROUP+pipeline formattention_swa.py: mirror of split + out_proj splitdecode_layer.py:dense_gate_up_silu_tp+dense_down_proj_tpeachsplit into cube/cast
prefill_attention_full.py:prefill_full_out_projcube/cast splittp_all_reduceguardprefill_attention_swa.py: same mirrorsAll other 30 files (config.py, _ops.py, collectives.py, moe.py, etc.) are
the WIP state immediately before Phase A.
Test plan
python -m models.step3p5._smoke_program_buildreturns=== probe rc=0 ===python -m models.step3p5.step3p5_decode -p a2a3 -d 0 --no-smoke --dummy-weightsreturns clean
next_hidden_out— blocked on simpler#1036Related issues
comm_initfails inos.fork()'d chip worker — rank 0 dies during HCCL bring-up, version-invariant simpler#1018 — comm_init segfault (verified fixed via--no-as-needed)Status
Draft — opened for early review and so upstream maintainers can
reference the exact code from simpler#1036. Not ready to merge until
the runtime crash is resolved upstream.