From 6accc66d1ec28f6dce9987a2c3f474e11b40a1c5 Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:19:21 -0700 Subject: [PATCH] Add: READMEs for the four a2a3 paged_attention variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four sibling directories share a name prefix and nothing tells you what separates them. Diffing establishes it, and each README now leads with the answer: - `paged_attention` is the baseline: four tasks per KV block (QK / SF / PV / UP) inside a plain scope, with TensorMap deriving the ordering from the tensors the tasks share. - `paged_attention_manual_scope` has **byte-identical kernels** — all four — and differs only in orchestration, 45 lines of 292. It is the cheapest way to see what MANUAL scope costs, and it deliberately uses both dependency APIs in one file: the primitive `set_dependencies(deps, n)` for tasks with a fixed predecessor set, and `L0TaskArgsWithDeps<>::add_dep` for `UP`, whose edges are conditional. - `paged_attention_unroll_manual_scope` is **not** the manual-scope variant plus an unroll. Every file differs; `aiv_softmax_prepare.cpp` changes 312 lines against a 156-line original. It batches KV blocks into groups of `N_UNROLL` and submits four tasks per group rather than per block. - `paged_attention_ringbuffer` contains **no kernels at all** — it points at the `tests/st` batch-paged-attention sources and exists to stress ring rotation at `ring_task_window: 64`, 256x below the 16384 default. The directory index carried a wrong one-liner for the unroll variant ("manual scope plus loop unrolling"); it now describes a second implementation, and the manual-scope row states the byte-identical kernels. Co-Authored-By: Claude Opus 5 (1M context) --- .../a2a3/tensormap_and_ringbuffer/README.md | 4 +- .../paged_attention/README.md | 44 ++++++++++++++ .../paged_attention_manual_scope/README.md | 56 ++++++++++++++++++ .../paged_attention_ringbuffer/README.md | 56 ++++++++++++++++++ .../README.md | 58 +++++++++++++++++++ 5 files changed, 216 insertions(+), 2 deletions(-) create mode 100644 examples/a2a3/tensormap_and_ringbuffer/paged_attention/README.md create mode 100644 examples/a2a3/tensormap_and_ringbuffer/paged_attention_manual_scope/README.md create mode 100644 examples/a2a3/tensormap_and_ringbuffer/paged_attention_ringbuffer/README.md create mode 100644 examples/a2a3/tensormap_and_ringbuffer/paged_attention_unroll_manual_scope/README.md diff --git a/examples/a2a3/tensormap_and_ringbuffer/README.md b/examples/a2a3/tensormap_and_ringbuffer/README.md index ad725573b..0e800bebb 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/README.md +++ b/examples/a2a3/tensormap_and_ringbuffer/README.md @@ -26,8 +26,8 @@ For the `Worker` API underneath the framework, see | ------- | --------------- | | [`benchmark_bgemm/`](benchmark_bgemm/) | Runtime-configurable tiled matmul `C = sum(k) A[k] @ B[k]`, shaped for measurement. Runs on sim. | | [`paged_attention/`](paged_attention/) | Online softmax with AIC/AIV subgraph splitting, bfloat16. The baseline the three variants below are compared against. | -| [`paged_attention_manual_scope/`](paged_attention_manual_scope/) | The same computation with explicit scope control — see [`docs/manual-scope.md`](../../../docs/manual-scope.md). | -| [`paged_attention_unroll_manual_scope/`](paged_attention_unroll_manual_scope/) | Manual scope plus loop unrolling. | +| [`paged_attention_manual_scope/`](paged_attention_manual_scope/) | The same computation with explicit scope control — kernels byte-identical to the baseline's, only the orchestration differs. See [`docs/manual-scope.md`](../../../docs/manual-scope.md). | +| [`paged_attention_unroll_manual_scope/`](paged_attention_unroll_manual_scope/) | A second implementation, not a patch on the baseline: KV blocks batched into groups of `N_UNROLL`, four tasks per group instead of per block, with the kernels rewritten to match. | | [`paged_attention_ringbuffer/`](paged_attention_ringbuffer/) | Deliberately undersized rings, driven per case through `config.runtime_env` (`ring_task_window` / `ring_heap` / `ring_dep_pool`) rather than process-global env. A stress test for rotation and reclamation. | | [`merge_pipeline_barrier/`](merge_pipeline_barrier/) | Three pipeline stages merged into **one** `block_num=8` SPMD task, ordered by an intra-task cross-core barrier instead of by three scheduled tasks. | | [`qwen3_14b_decode/`](qwen3_14b_decode/README.md) | The whole Qwen3-14B 40-layer decode stack as a single fused dispatch, using CANN fused attention. The largest example in the repo. | diff --git a/examples/a2a3/tensormap_and_ringbuffer/paged_attention/README.md b/examples/a2a3/tensormap_and_ringbuffer/paged_attention/README.md new file mode 100644 index 000000000..cd63f1552 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/paged_attention/README.md @@ -0,0 +1,44 @@ +# paged_attention — the baseline + +Online-softmax paged attention in bfloat16, split across AIC and AIV. This is +the reference the three sibling variants are read against; start here. + +## The four-task loop + +Per (batch, head) and per KV block, the orchestration submits four tasks: + +| Task | Core | Computation | +| ---- | ---- | ----------- | +| `QK` | AIC | `qi @ K^T` for the block | +| `SF` | AIV | softmax prepare — running `mi`, `li` | +| `PV` | AIC | `P @ V` | +| `UP` | AIV | online-softmax accumulation into the running output | + +Ordering is **not written down**. The tasks sit inside a plain `PTO2_SCOPE()` +with a `PTO2_SCOPE_GUARD()`, and TensorMap derives `QK → SF → PV → UP` from the +tensors they share. That is the whole point of the baseline: the dependency +graph is a consequence of the data, not of the code. + +## Cases + +Seven, all `platforms=["a2a3"]` with `aicpu_thread_num: 4`: `Case1`–`Case3` at +production scale (up to batch 256, 16 heads, 8192 context), `CaseSmall1` / +`CaseSmall2`, and `CaseVarSeq2` / `CaseVarSeq4` for ragged sequence lengths. +The four kernels are registered as sub-callables named `QK`, `SF`, `PV`, `UP`. + +## Run + +```bash +pytest examples/a2a3/tensormap_and_ringbuffer/paged_attention --platform a2a3 --device 0 +pytest examples/a2a3/tensormap_and_ringbuffer/paged_attention --platform a2a3 --device 0 -k Case1 +``` + +Onboard only. Wrap in `task-submit` on a shared box. + +## The variants + +| Directory | Differs how | +| --------- | ----------- | +| [`../paged_attention_manual_scope/`](../paged_attention_manual_scope/) | **Orchestration only.** All four kernels are byte-identical to this directory's; 45 of 292 orchestration lines change to declare dependencies explicitly. | +| [`../paged_attention_unroll_manual_scope/`](../paged_attention_unroll_manual_scope/) | **A different implementation.** Kernels and orchestration are both substantially rewritten — it batches KV blocks into groups instead of one task set per block. | +| [`../paged_attention_ringbuffer/`](../paged_attention_ringbuffer/) | **No kernels of its own.** A ring-sizing stress harness that borrows the `tests/st` batch-paged-attention kernels. | diff --git a/examples/a2a3/tensormap_and_ringbuffer/paged_attention_manual_scope/README.md b/examples/a2a3/tensormap_and_ringbuffer/paged_attention_manual_scope/README.md new file mode 100644 index 000000000..7e0044946 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/paged_attention_manual_scope/README.md @@ -0,0 +1,56 @@ +# paged_attention_manual_scope — the same computation, dependencies declared + +Identical arithmetic to [`../paged_attention/`](../paged_attention/), with +automatic dependency wiring replaced by explicit task-to-task edges. + +**All four kernels are byte-identical to the baseline's.** The only file that +differs is `kernels/orchestration/paged_attention_orch.cpp`, by 45 of 292 +lines. If you want to know what "manual scope" costs you in practice, that diff +is the answer — nothing about the computation changes, only who states the +ordering. + +```bash +diff -r ../paged_attention/kernels ./kernels # only the orch file differs +``` + +## What changes + +`PTO2_SCOPE()` + `PTO2_SCOPE_GUARD()` becomes `PTO2_SCOPE(PTO2ScopeMode::MANUAL)`, +and each submit names its predecessors. See +[`docs/manual-scope.md`](../../../../docs/manual-scope.md) for the mode itself. + +The variant deliberately uses **both dependency APIs**, so one file shows the +choice: + +| API | Shape | Suited to | +| --- | ----- | --------- | +| Primitive | `PTO2TaskId deps[] = {...}; params.set_dependencies(deps, n);` | Codegen, and any task whose dependency set is fixed and known at the call site. The caller owns the buffer; the `Arg` stores `(ptr, count)`. | +| Convenience | `L0TaskArgsWithDeps<> params; params.add_dep(id);` | Hand-written orchestration where the set is assembled conditionally across branches. The wrapper owns a stack-sized buffer and binds it on submit. | + +`SF` and `PV` use the primitive form (one predecessor each). `UP` uses the +convenience form, because its edges are conditional: always the `PV` output, +plus the previous group's `UP` when there is one, plus the allocation task on +the last iteration — that last edge exists only to keep the scratch buffers +alive until their final consumer, not to order anything. + +Note what is *not* declared: `UP` reads `SF`'s `mi` / `li`, but `SF → PV → UP` +already orders that transitively, so only the `PV` edge is stated. Manual mode +means you own the edges, including the ones you can leave out. + +## Cases + +The same seven as the baseline, with one addition: `Case1` sets +`runtime_env: {ring_task_window: 32768}`. Long-context cases submit more than +16384 in-flight tasks into a single MANUAL scope, and the default per-ring task +window is 16384 — it can fill before the oldest task retires and wedge the +orchestrator with `FLOW_CONTROL_DEADLOCK` (code 3). Doubling the window buys +the headroom. Worth knowing before you put a long loop in one manual scope: the +scope's in-flight task count becomes a sizing constraint you have to check. + +## Run + +```bash +pytest examples/a2a3/tensormap_and_ringbuffer/paged_attention_manual_scope --platform a2a3 --device 0 +``` + +Onboard only. Wrap in `task-submit` on a shared box. diff --git a/examples/a2a3/tensormap_and_ringbuffer/paged_attention_ringbuffer/README.md b/examples/a2a3/tensormap_and_ringbuffer/paged_attention_ringbuffer/README.md new file mode 100644 index 000000000..236f3718a --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/paged_attention_ringbuffer/README.md @@ -0,0 +1,56 @@ +# paged_attention_ringbuffer — deliberately undersized rings + +Not really a paged-attention example. **This directory contains no kernels at +all** — a single test file, which points its `CALLABLE` at the batch-paged- +attention kernels under `tests/st`: + +```python +PA_KERNELS = "../../../../tests/st/a2a3/tensormap_and_ringbuffer/batch_paged_attention/kernels" +``` + +The attention math is borrowed, and correct output is only the pass criterion. +What is under test is the **ring buffers**: whether rotation and reclamation +survive a real workload when the rings are far too small for it. + +## The knobs + +One case, `ringbuffer_stress`, sized to hurt: + +| `runtime_env` key | Value here | Note | +| ----------------- | ---------- | ---- | +| `ring_task_window` | 64 | Default is 16384 (`PTO2_TASK_WINDOW_SIZE`, `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h`) — **256× smaller** | +| `ring_heap` | 4 MiB | bytes per ring | +| `ring_dep_pool` | 256 | | + +At a 64-task window, a workload that submits thousands of tasks must rotate the +window continuously — every task record is overwritten many times over during +one run. If reclamation is off by one, or a record is reused while still +referenced, the golden check fails or the orchestrator wedges. + +Two things this example is the reference for: + +- **Per-case ring sizing through `config.runtime_env`**, not the process-global + `PTO2_RING_*` environment variables. The rings are sized per task, so one + suite can mix a stress case with normally-sized ones. +- **Non-power-of-2 sizes are accepted.** 4 MiB is chosen to keep the stress + intent compact, not because the runtime requires a round number. + +It also exercises `INOUT` tensors, bfloat16, and mixed AIC+AIV execution +incidentally — the `UP` kernel takes four `INOUT` arguments. + +## Run + +```bash +pytest examples/a2a3/tensormap_and_ringbuffer/paged_attention_ringbuffer --platform a2a3 --device 0 +``` + +Onboard only. Wrap in `task-submit` on a shared box. + +A failure here is usually a ring bug, not an attention bug — the same kernels +pass in `tests/st/a2a3/tensormap_and_ringbuffer/batch_paged_attention/` at +default sizes, which is the first thing to check. + +## See also + +- [`../paged_attention/`](../paged_attention/) — the baseline implementation and its variants +- [`docs/troubleshooting/device-error-codes/capacity.md`](../../../../docs/troubleshooting/device-error-codes/capacity.md) — the sub-class codes each of these three rings raises when it overflows, and what to do about it diff --git a/examples/a2a3/tensormap_and_ringbuffer/paged_attention_unroll_manual_scope/README.md b/examples/a2a3/tensormap_and_ringbuffer/paged_attention_unroll_manual_scope/README.md new file mode 100644 index 000000000..4f896aa6b --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/paged_attention_unroll_manual_scope/README.md @@ -0,0 +1,58 @@ +# paged_attention_unroll_manual_scope — fewer, larger tasks + +Not "the manual-scope variant plus a loop unroll". **Every file here differs +from the baseline**, kernels included — `aiv_softmax_prepare.cpp` is a full +rewrite (312 changed lines against a 156-line original). Read it as a second +implementation of paged attention, not as a patch on the first. + +Changed lines against the baseline's copy of the same file, over that file's +original length: + +| File | Changed / original | +| ---- | ------------------ | +| `kernels/aic/aic_qk_matmul.cpp` | 145 / 115 | +| `kernels/aic/aic_pv_matmul.cpp` | 164 / 114 | +| `kernels/aiv/aiv_softmax_prepare.cpp` | 312 / 156 | +| `kernels/aiv/aiv_online_update.cpp` | 27 / 256 | +| `kernels/orchestration/paged_attention_orch.cpp` | 226 / 292 | + +## What it does differently + +The baseline submits four tasks **per KV block**. This one batches up to +`N_UNROLL` blocks into a group and submits four tasks **per group**: + +| Task | Computation over the whole group | +| ---- | -------------------------------- | +| `QK` | `qi @ K^T` for `n_blocks` → `sij_buf` shaped `(q_tile, n_blocks * block_size)` | +| `SF` | two-pass softmax over the full `sij_buf` → `pij_buf`, `mi`, `li` | +| `PV` | SplitK-accumulated `P @ V` → `oi_new` | +| `UP` | online-softmax accumulation with group-level `mi`, `li`, `oi_new` | + +So the task count drops by roughly `N_UNROLL`, and each task carries +proportionally more work — the trade the whole directory exists to demonstrate. +Softmax becomes a two-pass over a wider buffer, which is why that kernel is +rewritten rather than adjusted. + +`N_UNROLL` is a `#define` at the top of the orchestration source, currently +**64**. (The file's header comment still says `N_UNROLL=8`; the `#define` is +what takes effect.) + +Dependencies are declared explicitly inside +`PTO2_SCOPE(PTO2ScopeMode::MANUAL)`, using the primitive +`set_dependencies(deps, n)` form throughout — +[`../paged_attention_manual_scope/`](../paged_attention_manual_scope/) is the +directory to read for that API on its own, against unchanged kernels. + +## Cases + +Three — `Case1`–`Case3`, all `platforms=["a2a3"]` with +`aicpu_thread_num: 4`. The baseline's small and variable-sequence cases have no +counterpart here. + +## Run + +```bash +pytest examples/a2a3/tensormap_and_ringbuffer/paged_attention_unroll_manual_scope --platform a2a3 --device 0 +``` + +Onboard only. Wrap in `task-submit` on a shared box.