Add: host_build_graph copies of the paged_attention examples

examples/a2a3/host_build_graph/paged_attention/README.md

@@ -0,0 +1,65 @@
+# paged_attention — host_build_graph
+
+Online-softmax paged attention in bfloat16, split across AIC and AIV, on the
+`host_build_graph` runtime. The kernels and the orchestration source are the
+same ones the `tensormap_and_ringbuffer` sibling compiles.
+
+## Why the same orchestration compiles on both runtimes
+
+`orchestration/pto_orchestration_api.h` is identical under
+`src/a2a3/runtime/host_build_graph/` and
+`src/a2a3/runtime/tensormap_and_ringbuffer/` — same entry points, same macros.
+`KernelCompiler.compile_orchestration(runtime, source)` picks the include dirs
+of whichever runtime `@scene_test(runtime=...)` names, so switching runtime is a
+one-line change in the test class and nothing in the C++.
+
+## 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 in a plain `PTO2_SCOPE()` with a
+`PTO2_SCOPE_GUARD()` and the dependency graph is derived from the tensors they
+share. See `../paged_attention_manual_scope/` for the variant that wires the
+same edges by hand.
+
+## Ring sizing is the one place host_build_graph differs
+
+`host_build_graph` builds the entire task graph on the host before the device
+starts scheduling, so no ring slot can be reclaimed mid-orchestration — the ring
+window and GM heap must hold every task of a case at once. Task count is
+
+```text
+tasks = batch * (4 * ceil(context_len / block_size) + 1)
+```
+
+so `Case1` needs 65 792 slots and `Case2` needs 32 832, both above the default
+window of 16384. Each carries its own `runtime_env` in `CASES[*]["config"]`
+(`ring_task_window` must be a power of two in `[4, INT32_MAX]`; `ring_heap` at
+least 1024), so no `PTO2_RING_*` environment variable is needed.
+
+## Cases
+
+`Case1` (65 792 tasks) and `Case2` (32 832) at production scale, `CaseSmall1` /
+`CaseSmall2`, and `CaseVarSeq2` / `CaseVarSeq4` for ragged sequence lengths.
+All but `CaseSmall1` are `manual`, so the default run executes `CaseSmall1`
+only; add `--manual include` to reach the rest.
+
+The upstream `tensormap_and_ringbuffer` example also carries a `Case3`
+(`head_dim: 256`). It is absent here because it does not produce correct
+results on either runtime — see `KNOWN_ISSUES.md`.
+
+## Run
+
+```bash
+python examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py \
+    -p a2a3 -d 0                                   # CaseSmall1, golden checked
+python examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py \
+    -p a2a3 -d 0 --manual include --case Case2 --rounds 2
+```

examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) PyPTO Contributors.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ * -----------------------------------------------------------------------------------------------------------
+ */
+// PV Matmul Kernel: pij(M, K) @ vj(K, N) -> oi_new(M, N)
+//
+// Supports two tile configurations via runtime dispatch:
+//   Case1: (16, 128) @ (128, 128) -> (16, 128)
+//   Case2: (64,  64) @ ( 64, 128) -> (64, 128)
+//
+// pij is bfloat16 (converted from fp32 in softmax_prepare via TCVT).
+// vj is stored as (K, N) = (block_size, head_dim) in row-major (ND) layout.
+// Standard non-transposed B pattern: ND GlobalB + ColMajor/RowMajor TileMatB.
+
+#include <cstdint>
+#include <pto/pto-inst.hpp>
+
+#include "tensor.h"
+
+using namespace pto;
+
+#include "pipe_sync.h"
+
+#ifndef __gm__
+#define __gm__
+#endif
+
+#ifndef __aicore__
+#define __aicore__ [aicore]
+#endif
+
+template <int M, int K, int N>
+static __aicore__ void pv_matmul_impl(__gm__ Tensor *pij, __gm__ Tensor *vj, __gm__ Tensor *oi) {
+    __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr);
+    __gm__ bfloat16_t *vj_addr = reinterpret_cast<__gm__ bfloat16_t *>(vj->buffer.addr);
+    __gm__ float *oi_addr = reinterpret_cast<__gm__ float *>(oi->buffer.addr);
+
+    // pij (M, K) bf16, vj (K, N) bf16 in ND (row-major), oi_new (M, N) fp32
+    using GlobalA = GlobalTensor<bfloat16_t, Shape<1, 1, 1, M, K>, Stride<M * K, M * K, M * K, K, 1>>;
+    using GlobalB = GlobalTensor<bfloat16_t, Shape<1, 1, 1, K, N>, Stride<K * N, K * N, K * N, N, 1>>;
+    using GlobalOut = GlobalTensor<float, Shape<1, 1, 1, M, N>, Stride<M * N, M * N, M * N, N, 1>>;
+
+    GlobalA pijGlobal(pij_addr + pij->start_offset);
+    GlobalB vjGlobal(vj_addr + vj->start_offset);
+    GlobalOut oiGlobal(oi_addr + oi->start_offset);
+
+    // L1 Mat tiles: standard ND pattern for both A and B
+    using TileMatA = Tile<TileType::Mat, bfloat16_t, M, K, BLayout::ColMajor, M, K, SLayout::RowMajor, 512>;
+    using TileMatB = Tile<TileType::Mat, bfloat16_t, K, N, BLayout::ColMajor, K, N, SLayout::RowMajor, 512>;
+
+    // L0 tiles
+    using LeftTile = TileLeft<bfloat16_t, M, K, M, K>;
+    using RightTile = TileRight<bfloat16_t, K, N, K, N>;
+    using AccTile = TileAcc<float, M, N, M, N>;
+
+    TileMatA aMatTile;
+    TileMatB bMatTile;
+    TASSIGN(aMatTile, 0x0);
+    TASSIGN(bMatTile, 0x20000);
+
+    LeftTile aTile;
+    RightTile bTile;
+    AccTile cTile;
+    TASSIGN(aTile, 0x0);
+    TASSIGN(bTile, 0x0);
+    TASSIGN(cTile, 0x0);
+
+    // Load pij and vj to L1 with separate events for pipeline overlap
+    TLOAD(aMatTile, pijGlobal);
+    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);  // A load done
+    TLOAD(bMatTile, vjGlobal);
+    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);  // B load done
+
+    // Move A to L0A as soon as A load completes (B may still be loading)
+    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
+    TMOV(aTile, aMatTile);
+    // Move B to L0B after B load completes
+    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
+    TMOV(bTile, bMatTile);
+
+    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
+    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
+
+    // Single matmul: (M,K) x (K,N) -> (M,N)
+    TMATMUL(cTile, aTile, bTile);
+
+    set_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
+    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
+
+    TSTORE(oiGlobal, cTile);
+
+    pipe_sync();
+}
+
+extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
+    __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[0]);
+    __gm__ Tensor *vj = reinterpret_cast<__gm__ Tensor *>(args[1]);
+    __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]);
+    uint64_t q_tile_size = static_cast<uint64_t>(pij->shapes[0]);
+
+    if (q_tile_size == 16 && pij->shapes[1] <= 16) {
+        pv_matmul_impl<16, 16, 16>(pij, vj, oi_new);
+    } else if (q_tile_size == 16) {
+        pv_matmul_impl<16, 128, 128>(pij, vj, oi_new);
+    } else {
+        pv_matmul_impl<64, 64, 128>(pij, vj, oi_new);
+    }
+}

examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp

@@ -0,0 +1,115 @@
+/*
+ * Copyright (c) PyPTO Contributors.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ * -----------------------------------------------------------------------------------------------------------
+ */
+// QK Matmul Kernel: qi(M, K) @ kj.T(K, N) -> sij(M, N)
+//
+// Supports two tile configurations via runtime dispatch:
+//   Case1: (16, 128) @ (128, 128).T -> (16, 128)
+//   Case2: (64, 128) @ (128,  64).T -> (64,  64)
+//
+// kj is stored as (N, K) = (block_size, head_dim) in row-major memory.
+// This is equivalent to (K, N) in column-major (DN) layout.
+// Using DN GlobalB + RowMajor/ColMajor TileMatB to handle the transposed B pattern.
+
+#include <cstdint>
+#include <pto/pto-inst.hpp>
+
+#include "tensor.h"
+
+using namespace pto;
+
+#include "pipe_sync.h"
+
+#ifndef __gm__
+#define __gm__
+#endif
+
+#ifndef __aicore__
+#define __aicore__ [aicore]
+#endif
+
+template <int M, int K, int N>
+static __aicore__ void qk_matmul_impl(__gm__ Tensor *qi, __gm__ Tensor *kj, __gm__ Tensor *sij) {
+    __gm__ bfloat16_t *qi_addr = reinterpret_cast<__gm__ bfloat16_t *>(qi->buffer.addr);
+    __gm__ bfloat16_t *kj_addr = reinterpret_cast<__gm__ bfloat16_t *>(kj->buffer.addr);
+    __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr);
+
+    // qi (M, K) bf16 in ND (row-major) layout
+    using GlobalA = GlobalTensor<bfloat16_t, Shape<1, 1, 1, M, K>, Stride<M * K, M * K, M * K, K, 1>>;
+    // kj stored as (N, K) row-major = (K, N) column-major -> DN layout
+    using GlobalB = GlobalTensor<bfloat16_t, Shape<1, 1, 1, K, N>, Stride<K * N, K * N, K * N, 1, K>, Layout::DN>;
+    using GlobalOut = GlobalTensor<float, Shape<1, 1, 1, M, N>, Stride<M * N, M * N, M * N, N, 1>>;
+
+    GlobalA qiGlobal(qi_addr + qi->start_offset);
+    GlobalB kjGlobal(kj_addr + kj->start_offset);
+    GlobalOut sijGlobal(sij_addr + sij->start_offset);
+
+    // L1 Mat tiles: A is standard ND, B uses transposed-B pattern (RowMajor/ColMajor)
+    using TileMatA = Tile<TileType::Mat, bfloat16_t, M, K, BLayout::ColMajor, M, K, SLayout::RowMajor, 512>;
+    using TileMatB = Tile<TileType::Mat, bfloat16_t, K, N, BLayout::RowMajor, K, N, SLayout::ColMajor, 512>;
+
+    // L0 tiles
+    using LeftTile = TileLeft<bfloat16_t, M, K, M, K>;
+    using RightTile = TileRight<bfloat16_t, K, N, K, N>;
+    using AccTile = TileAcc<float, M, N, M, N>;
+
+    TileMatA aMatTile;
+    TileMatB bMatTile;
+    TASSIGN(aMatTile, 0x0);
+    TASSIGN(bMatTile, 0x20000);
+
+    LeftTile aTile;
+    RightTile bTile;
+    AccTile cTile;
+    TASSIGN(aTile, 0x0);
+    TASSIGN(bTile, 0x0);
+    TASSIGN(cTile, 0x0);
+
+    // Load A and B to L1 with separate events for pipeline overlap
+    TLOAD(aMatTile, qiGlobal);
+    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);  // A load done
+    TLOAD(bMatTile, kjGlobal);
+    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);  // B load done
+
+    // Move A to L0A as soon as A load completes (B may still be loading)
+    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
+    TMOV(aTile, aMatTile);
+    // Move B to L0B after B load completes
+    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
+    TMOV(bTile, bMatTile);
+
+    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
+    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
+
+    // Matmul
+    TMATMUL(cTile, aTile, bTile);
+
+    set_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
+    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
+
+    TSTORE(sijGlobal, cTile);
+
+    pipe_sync();
+}
+
+extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
+    __gm__ Tensor *qi = reinterpret_cast<__gm__ Tensor *>(args[0]);
+    __gm__ Tensor *kj = reinterpret_cast<__gm__ Tensor *>(args[1]);
+    __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[2]);
+    uint64_t q_tile_size = static_cast<uint64_t>(qi->shapes[0]);
+
+    if (q_tile_size == 16 && qi->shapes[1] <= 16) {
+        qk_matmul_impl<16, 16, 16>(qi, kj, sij);
+    } else if (q_tile_size == 16) {
+        qk_matmul_impl<16, 128, 128>(qi, kj, sij);
+    } else {
+        qk_matmul_impl<64, 128, 64>(qi, kj, sij);
+    }
+}