From 22f9c804831a74b2fb11ff10ddcdb9e6b38e5e76 Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Fri, 3 Jul 2026 17:27:54 +0800 Subject: [PATCH] Normalize TPUSH/TPOP usage and add interim onboard lane-offset bridge Converge MIX GM FIFO tile-pipe kernels to pto-isa default semantics: - Remove redundant manual prod.record() and setRecordStatus(false) / setAllocateStatus(false) / setFreeStatus(false) from spmd_paged_attention - Rely on built-in TPUSH/TPOP/TFREE record and back-pressure; keep existing set_flag/wait_flag pipeline sync unchanged - Fix bgemm AIV GM accumulator row addressing via get_sub_block_id(args) instead of stale CCE get_subblockid() Work around a known pinned pto-isa bug on 1C2V MIX where get_subblockid() returns 0 for both AIV lanes, breaking TILE_UP_DOWN lane offsets inside the library. Until pto-isa/dispatch fixes sub-block identity, inject GlobalContext.sub_block_id through intrinsic.h tpipe_lane_byte_offset() + setEntryOffset as an interim bridge. Docs and validation: - Add docs/tpush-tpop.md; update aicore-kernel-programming.md - Drop spmd_paged_attention pytest skip - Onboard SmallCase1 / Case2 / Case1 golden checks pass --- docs/aicore-kernel-programming.md | 53 +++++------- docs/tpush-tpop.md | 83 +++++++++++++++++++ .../qwen3_14b_decode/README.md | 14 ++-- .../bgemm/kernels/mix/kernel_bgemm.cpp | 7 +- .../kernels/mix/paged_attention_parallel.cpp | 55 +++--------- .../test_spmd_paged_attention.py | 6 -- 6 files changed, 129 insertions(+), 89 deletions(-) create mode 100644 docs/tpush-tpop.md diff --git a/docs/aicore-kernel-programming.md b/docs/aicore-kernel-programming.md index 5e26fb3ee..9a4c7c794 100644 --- a/docs/aicore-kernel-programming.md +++ b/docs/aicore-kernel-programming.md @@ -184,18 +184,25 @@ down from `kernel_entry` into `UnpadAttentionDecoderAic::SetArgs` and The porting checklist above assumes you own every `get_subblockid()` call site. You do **not** when the kernel drives the pto-isa tile-pipe library (`TPUSH` / `TPOP` / `TFREE` with `TileSplitAxis::TILE_UP_DOWN` or -`TILE_LEFT_RIGHT`): those templates compute their per-AIV FIFO offset from the +`TILE_LEFT_RIGHT`): those templates compute per-AIV FIFO offset from the **no-arg** `get_subblockid()` internally -(`TPush.hpp::pushVec2GMFiFo` / `popVecTileFromGMFiFo` — -`subAIVOffset = get_subblockid() * rows * cols * sizeof(T)`), and you cannot +(`TPush.hpp::pushVec2GMFiFo` / `popVecTileFromGMFiFo`), and you cannot thread `args` into a third-party library template. -Under simpler's MIX dispatch that library call returns 0 for both lanes, so its -auto-split contributes nothing and **both AIV lanes read/write the same FIFO -half**. The other half is never produced; attention then reads uninitialised GM -and the output is wrong/partial (observed as `NaN` in the qwen3 case here). +**Recommended usage** (see [`docs/tpush-tpop.md`](tpush-tpop.md)): -**Do not** bridge it with a file-scope cache: +- Call `TPUSH` / `TPOP` / `TFREE` directly — record, back-pressure, and + `TILE_UP_DOWN` lane offset are already implemented inside those templates. +- Do **not** add manual `pipe.prod.record()` or batch `setRecordStatus(false)` / + `setAllocateStatus(false)` / `setFreeStatus(false)` unless a reviewed pipeline + analysis requires it. +- Do **not** `setEntryOffset(get_sub_block_id(args) * …)` for lane split when + `get_subblockid()` is correct — the library already adds + `get_subblockid() * tile_bytes_per_lane`. +- For **non-tile-pipe** GM addressing (output rows, head partitioning), keep + using `get_sub_block_id(args)` from this header. + +**Do not** bridge `get_subblockid()` with a file-scope cache: ```cpp // WRONG — the .o will not load. See §4. @@ -205,25 +212,11 @@ and the output is wrong/partial (observed as `NaN` in the qwen3 case here). ``` A `[[block_local]]` (or any non-const) static is read via a `.text` relocation -that the AICore loader rejects (§4), and AICore forbids ordinary global/static -data, so there is no relocation-free global to redirect into either. - -The working pattern (as in `spmd_paged_attention`) leaves the library's -`get_subblockid()` at its native 0 and adds the per-lane offset **explicitly** -on the tile-pipe, computed from `get_sub_block_id(args)`: - -```cpp -int32_t lane = get_sub_block_id(args); // 0 or 1, real lane -// TILE_UP_DOWN splits an M×N tile by rows; AIV1 starts one sub-tile in. -pipe.cons.setEntryOffset(lane * Rows * Cols * sizeof(ConsT)); // C2V pop side -pipe.prod.setEntryOffset(lane * Rows * Cols * sizeof(ProdT)); // V2C push side -``` - -The library adds `get_subblockid() * bytes` (= 0) **plus** your `entryOffset`, -so the explicit offset now carries the entire lane split. Set it once after the -pipe is constructed; `entryOffset` persists across `TPUSH` / `TPOP`. See -`examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_fused_aiv.cpp` -for a worked example. +that the AICore loader rejects (§4). If onboard `get_subblockid()` does not +match `get_sub_block_id(args)`, prefer fixing platform/launch identity; until +then add the lane split explicitly with `setEntryOffset` computed inline from +`get_sub_block_id(args)` (see the `run_aiv` `setEntryOffset` call sites in +[`spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp`](../tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp)). --- @@ -258,9 +251,9 @@ readelf -SW kernel.o | grep -E '\.text' # want only ".text"; ".text._Z*" or ".r readelf -r kernel.o # want: no relocation entries ``` -This is exactly why §3's sub-block-id fix uses an explicit `setEntryOffset` -argument rather than a cached `[[block_local]]` static: the static would -reintroduce a `.rela.text` the loader rejects. +This is exactly why §3 rejects a cached `[[block_local]]` static to redirect +`get_subblockid()`: the static would reintroduce a `.rela.text` the loader +rejects. --- diff --git a/docs/tpush-tpop.md b/docs/tpush-tpop.md new file mode 100644 index 000000000..87af744d4 --- /dev/null +++ b/docs/tpush-tpop.md @@ -0,0 +1,83 @@ +# TPUSH/TPOP Usage Guidelines (Advisory) + +Recommended patterns for MIX kernels that use GM FIFO tile-pipe (`TPipe` + +`TPUSH` / `TPOP` / `TFREE`, commonly with `TileSplitAxis::TILE_UP_DOWN`). +New kernels and refactors should follow these by default; special pipelines +may deviate after review. This is not a hard lint. + +Implementation reference: `pto-isa/include/pto/npu/*/TPush.hpp`, `TPop.hpp`. + +## Why (overview) + +pto-isa's `TPUSH` / `TPOP` / `TFREE` already embed record handling, +allocate/back-pressure, and lane offset computation under +`TILE_UP_DOWN`. Re-implementing the same logic in kernel code is usually +redundant and can diverge from the built-in path. Prefer calling TPUSH/TPOP +directly and leave the cross-core FIFO protocol to the library. + +## 1. Prefer TPUSH/TPOP directly — do not manual-record + +**Why**: `TPUSH_IMPL` / `TPOP_IMPL` call `prod.record()` / consumer wait +after push/pop when `getRecordStatus()` is true (the default). Cross-core +handshake is already inside TPUSH/TPOP. + +- Default `TPUSH_IMPL` flow: `allocate` → `push` (TSTORE to GM FIFO) → + `tileIndex++` → `record()`. +- **Prefer**: call `TPUSH(...)` / `TPOP(...)` only; do not also call + `pipe.prod.record()`. +- **Prefer**: do not call `setRecordStatus(false)` to disable built-in + record; combining with manual record is usually redundant and easy to + miss a signal. +- Pipeline `set_flag` / `wait_flag` and TPUSH record are **independent**: + the former synchronizes MTE/M/V pipeline stages; the latter handles + cross-core FIFO handshake. Keep each where needed — this rule does not + constrain flag usage. +- If C2V AccTile still has record timing issues on a specific pto-isa + version, prefer upgrading pto-isa or the platform bridge over keeping + manual record long-term. + +## 2. Prefer default back-pressure — do not bulk-disable status flags + +**Why**: TPUSH already calls `prod.allocate()` when needed via +`getAllocateStatus()` (default true). The consumer path with `TFREE` uses +consumer free sync. That back-pressure lives in TPUSH/TPOP/TFREE; turning +it off requires the kernel to guarantee slots never overflow. + +- Producer default `isAllocate=true`: when the FIFO is full, `allocate()` + waits for consumer `TFREE`. +- Consumer default enables free sync: after consume, the slot is released + to the producer. +- **Prefer**: do not batch `setAllocateStatus(false)` / `setFreeStatus(false)` + at kernel entry unless pipeline analysis proves `FIFO_DEPTH` and scheduling + order already guarantee safe slot reuse — and document the rationale in a + comment. + +## 3. Rely on TILE_UP_DOWN auto lane offset — do not manual setEntryOffset + +**Why**: under `TileSplitAxis::TILE_UP_DOWN`, `TPUSH` / `TPOP` / `TFREE` +compute `subAIVOffset` via `get_subblockid()` inside +`pushVec2GMFiFo` / `popVecTileFromGMFiFo` and related paths. Calling +`setEntryOffset(get_sub_block_id(args) * …)` on top duplicates or stacks +with the library offset. + +- Library path: `subAIVOffset = get_subblockid() * tile_bytes_per_lane` + (expanded by split axis and tile shape). +- **Prefer**: do not use + `pipe.cons/prod.setEntryOffset(get_sub_block_id(args) * …)` for lane + splitting; ensure the platform provides a correct `get_subblockid()`. +- **Known issue (current pinned pto-isa)**: on 1C2V (1 Cube + 2 Vector) + MIX, CCE `get_subblockid()` returns 0 for both AIVs, so library + `subAIVOffset` cannot distinguish lanes — a pto-isa / launch identity + bug, not kernel misuse. +- **Onboard bridge (interim)**: when that applies (or simpler dispatch has + not yet programmed sub-block registers), add the lane split explicitly on + the tile-pipe with `setEntryOffset(get_sub_block_id(args) * sub_rows * cols * + elem_bytes)` (`GlobalContext.sub_block_id`) — see the `run_aiv` + `setEntryOffset` call sites in + [`paged_attention_parallel.cpp`](../tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp). + Remove once pto-isa / dispatch fixes `get_subblockid()`. + +## Reference examples + +- [`tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp`](../tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp) +- [`examples/a5/tensormap_and_ringbuffer/bgemm/kernels/mix/kernel_bgemm.cpp`](../examples/a5/tensormap_and_ringbuffer/bgemm/kernels/mix/kernel_bgemm.cpp) diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/README.md b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/README.md index 6478b9123..7a5c2cc09 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/README.md +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/README.md @@ -44,9 +44,11 @@ per-layer KV pools. The kernels are used essentially as generated, with **one hand-edit** to `fa_fused_aiv`: the codegen emitted the AIV sub-block id as a `[[block_local]] static` (`pypto_runtime_subblock_id`) whose **non-branch** `.text` relocation -simpler's strict `.text`-only loader rejects. The fix swaps it for the pto-isa -explicit-offset pattern (`setEntryOffset(get_sub_block_id(args) * …)`) — no -per-core static, no relocation. See Status below. +simpler's strict `.text`-only loader rejects. The fix swaps it for +`setEntryOffset(get_sub_block_id(args) * …)` (legacy workaround). New MIX +kernels should prefer direct `TPUSH`/`TPOP` per +[docs/tpush-tpop.md](../../../docs/tpush-tpop.md). See +Status below. To regenerate: in pypto-lib, set `_CHUNK_NLAYERS=2`, `PTO2_MANUAL_MAX_SEQ=5500`, build the ×2-stacked inputs, and `decode_fwd_layers.compile_for_test(...)`; then @@ -99,6 +101,6 @@ static and macro, and apply the per-lane split explicitly on the tile-pipe in nothing) and the explicit offset carries the lane separation. No per-core static, no relocation — loads under the strict loader and computes correctly. -The cleaner long-term fix is upstream codegen: have pto-isa emit the explicit -`sub_block_id` offset directly, or have the runtime program the FFTS sub-block -register so native `get_subblockid()` returns 0/1. +The cleaner long-term fix is direct `TPUSH`/`TPOP` with platform-correct +`get_subblockid()` (see [docs/tpush-tpop.md](../../../docs/tpush-tpop.md) +and [`spmd_paged_attention`](../../../../tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp)). diff --git a/examples/a5/tensormap_and_ringbuffer/bgemm/kernels/mix/kernel_bgemm.cpp b/examples/a5/tensormap_and_ringbuffer/bgemm/kernels/mix/kernel_bgemm.cpp index 66321dc45..4a8f2decc 100644 --- a/examples/a5/tensormap_and_ringbuffer/bgemm/kernels/mix/kernel_bgemm.cpp +++ b/examples/a5/tensormap_and_ringbuffer/bgemm/kernels/mix/kernel_bgemm.cpp @@ -50,9 +50,11 @@ using pto::TPipe; #endif #ifndef __aicore__ -#define __aicore__ +#define __aicore__ [aicore] #endif +#include "intrinsic.h" + #ifdef __DAV_CUBE__ constexpr bool DAV_CUBE = true; #else @@ -159,7 +161,8 @@ extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { // Vector side: TPOP result from cube → TLOAD C from GM → TADD → TSTORE // ========================================================================= if constexpr (DAV_VEC) { - uint32_t subBlockIdx = get_subblockid(); + int32_t sub_id = get_sub_block_id(args); + uint32_t subBlockIdx = (sub_id >= 0) ? static_cast(sub_id) : 0U; __gm__ float *c_ptr = reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; // Each vector sub-core handles its half: sub-core 0 → rows [0, VEC_M), diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp index 32e9e3302..c922bd423 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp @@ -198,12 +198,9 @@ static __aicore__ void aic_qk_step( set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - // TPUSH sij (C2V): AccTile L0C -> GM. Ensure prior MTE3 is done, - // then push, then wait for MTE3 DMA to complete before signaling consumer. TPUSH(sij_pipe, cTile_QK); set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - sij_pipe.prod.record(); } // Helper: PV matmul for block i — TPOP pij, load value, move to L0, matmul, TPUSH oi @@ -255,11 +252,9 @@ static __aicore__ void aic_pv_step( set_flag(PIPE_M, PIPE_FIX, EVENT_ID1); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID1); - // TPUSH oi (C2V): AccTile L0C -> GM. Same manual record pattern as sij. TPUSH(oi_pipe, cTile_PV); set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - oi_pipe.prod.record(); } template @@ -642,30 +637,6 @@ static __aicore__ void run_aic( typename Cfg::PijPipeT pij_pipe(pij_fifo_base, 0U, Cfg::PIJ_L1_BASE); typename Cfg::OiPipeT oi_pipe(oi_fifo_base, Cfg::OI_UB_BASE, 0U); - // Disable auto-record on C2V pipes: AccTile TSTORE goes through FIX → MTE3, - // but auto-record fires on PIPE_FIX which may complete before MTE3 DMA writes - // to GM. Manual pipe_barrier(PIPE_MTE3) + record() in each step ensures the - // cross-core signal fires only after the GM write is visible. - sij_pipe.prod.setRecordStatus(false); - oi_pipe.prod.setRecordStatus(false); - - // Disable reverse-dependency sync (back-pressure). Forward dependency chain - // (AIC: QK-first; AIV: SF-first; FIFO_DEPTH=2) guarantees producer is at - // most SLOT_NUM=2 tiles ahead of consumer: - // sij: AIC pushes sij[i+1] only after TPOP(pij[i-1]), which requires - // AIV TPOP(sij[i-1]) — slot reuse safe. - // oi : AIC pushes oi[i+1] only after TPOP(pij[i+1]), which requires - // AIV's iter i+1 SF, by which time AIV iter i finished UP[i-1] - // i.e. TPOP(oi[i-1]) — slot reuse safe. - // pij: AIV pushes pij[i+1] only after TPOP(sij[i+1]), which fires after - // AIC iter i+1 starts and AIC iter i has finished TPOP(pij[i-1]) - // — slot reuse safe. - // If the QK-first/SF-first interleaving or FIFO_DEPTH changes, restore - // these flags. - sij_pipe.prod.setAllocateStatus(false); - oi_pipe.prod.setAllocateStatus(false); - pij_pipe.cons.setFreeStatus(false); - __gm__ Tensor *query_t = reinterpret_cast<__gm__ Tensor *>(args[0]); __gm__ Tensor *key_cache_t = reinterpret_cast<__gm__ Tensor *>(args[1]); __gm__ Tensor *value_cache_t = reinterpret_cast<__gm__ Tensor *>(args[2]); @@ -718,22 +689,16 @@ static __aicore__ void run_aiv( int32_t sub_block_id = get_sub_block_id(args); int64_t row_offset = sub_block_id * Cfg::SUB_QT; - // Entry offsets depend on the actual tile width (block_size for sij/pij, HEAD_DIM for oi). - // TILE_UP_DOWN splits Q_TILE rows into two SUB_QT halves; AIV1's data starts at - // SUB_QT * tile_width * sizeof(element) within the contiguous TPUSH'd tile. - int sij_sub_offset = sub_block_id * Cfg::SUB_QT * static_cast(block_size) * static_cast(sizeof(float)); - int pij_sub_offset = - sub_block_id * Cfg::SUB_QT * static_cast(block_size) * static_cast(sizeof(bfloat16_t)); - int oi_sub_offset = sub_block_id * Cfg::SUB_QT * HEAD_DIM * static_cast(sizeof(float)); - sij_pipe.cons.setEntryOffset(sij_sub_offset); - pij_pipe.prod.setEntryOffset(pij_sub_offset); - oi_pipe.cons.setEntryOffset(oi_sub_offset); - - // Mirror reverse-dependency disable on the AIV side (see run_aic for - // the full forward-chain argument). - pij_pipe.prod.setAllocateStatus(false); - sij_pipe.cons.setFreeStatus(false); - oi_pipe.cons.setFreeStatus(false); + // pto-isa TPUSH/TPOP add `get_subblockid() * sub_rows * cols * elem_bytes` internally, but the CCE + // `get_subblockid()` register is 0 for both lanes under simpler onboard MIX dispatch; add the lane split + // explicitly from GlobalContext.sub_block_id (block_size wide for sij/pij, HEAD_DIM for oi). + sij_pipe.cons.setEntryOffset( + sub_block_id * Cfg::SUB_QT * static_cast(block_size) * static_cast(sizeof(float)) + ); + pij_pipe.prod.setEntryOffset( + sub_block_id * Cfg::SUB_QT * static_cast(block_size) * static_cast(sizeof(bfloat16_t)) + ); + oi_pipe.cons.setEntryOffset(sub_block_id * Cfg::SUB_QT * HEAD_DIM * static_cast(sizeof(float))); __gm__ float *out_base = reinterpret_cast<__gm__ float *>(out_t->buffer.addr) + out_t->start_offset; diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py index a3bd99790..56c0e2bb6 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py @@ -9,7 +9,6 @@ # ----------------------------------------------------------------------------------------------------------- """Paged attention unroll with TPUSH/TPOP: MIX kernel AIC+AIV cooperative pipeline.""" -import pytest import torch from simpler.task_interface import ArgDirection as D @@ -17,11 +16,6 @@ from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs -# Case1 flakily aborts with AICore error 507018 on a2a3 (the st-onboard-a2a3 job -# also flakes on main, independent of the pinned pto-isa). Skipped to keep the -# pto-isa pin bump green; re-enable once the 507018 paged-attention flake is fixed. -pytestmark = pytest.mark.skip(reason="paged-attention flakily aborts with AICore 507018 on a2a3 (known flake)") - @scene_test(level=2, runtime="tensormap_and_ringbuffer") class TestPagedAttentionUnrollTpushPop(SceneTestCase):