Normalize TPUSH/TPOP usage and add interim onboard lane-offset bridge

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.
 
 ---
 

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)

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)).

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<uint32_t>(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),

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<SijPipeT, AccTile_QK, TileSplitAxis::TILE_UP_DOWN>(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<OiPipeT, AccTile_PV, TileSplitAxis::TILE_UP_DOWN>(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 <typename Cfg, int K, int N>
@@ -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<int>(block_size) * static_cast<int>(sizeof(float));
-    int pij_sub_offset =
-        sub_block_id * Cfg::SUB_QT * static_cast<int>(block_size) * static_cast<int>(sizeof(bfloat16_t));
-    int oi_sub_offset = sub_block_id * Cfg::SUB_QT * HEAD_DIM * static_cast<int>(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<int>(block_size) * static_cast<int>(sizeof(float))
+    );
+    pij_pipe.prod.setEntryOffset(
+        sub_block_id * Cfg::SUB_QT * static_cast<int>(block_size) * static_cast<int>(sizeof(bfloat16_t))
+    );
+    oi_pipe.cons.setEntryOffset(sub_block_id * Cfg::SUB_QT * HEAD_DIM * static_cast<int>(sizeof(float)));
 
     __gm__ float *out_base = reinterpret_cast<__gm__ float *>(out_t->buffer.addr) + out_t->start_offset;
 

tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py

@@ -9,19 +9,13 @@
 # -----------------------------------------------------------------------------------------------------------
 """Paged attention unroll with TPUSH/TPOP: MIX kernel AIC+AIV cooperative pipeline."""
 
-import pytest
 import torch
 from simpler.task_interface import ArgDirection as D
 
 from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test
 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):