From bf4a58d978da450c8d568befe78b19f9b39969eb Mon Sep 17 00:00:00 2001 From: Noah Baumann Date: Mon, 3 Aug 2026 16:00:26 +0200 Subject: [PATCH] Add: host_build_graph copies of the paged_attention examples paged_attention and paged_attention_manual_scope existed only for tensormap_and_ringbuffer. pto_orchestration_api.h is byte-identical between the two runtimes and compile_orchestration() selects the include dirs of whichever runtime @scene_test names, so the kernels and orchestration C++ are unchanged copies; the delta is the runtime= decorator, the class name, and ring sizing. host_build_graph builds the whole graph before the device schedules anything, so nothing is reclaimed mid-orchestration and the entire graph must be resident. Case1 (65792 tasks) and Case2 (32832) therefore carry ring_task_window/ring_heap in their own CASES[*]["config"]["runtime_env"], needing no PTO2_RING_* env var. Case3 (head_dim 256) is omitted: it fails on both runtimes, giving max_diff=nan here and max_diff=0.1255 on the unmodified tensormap_and_ringbuffer example, with a normal 16 ms device_wall and no allocator fatal. That is a pre-existing kernel numerics defect at head_dim > 128, not a property of this runtime. Validated on a2a3 with golden checking: 12/12 cases pass across both examples. --- .../paged_attention/README.md | 65 ++++ .../kernels/aic/aic_pv_matmul.cpp | 114 +++++++ .../kernels/aic/aic_qk_matmul.cpp | 115 +++++++ .../kernels/aiv/aiv_online_update.cpp | 256 ++++++++++++++ .../kernels/aiv/aiv_softmax_prepare.cpp | 156 +++++++++ .../orchestration/paged_attention_orch.cpp | 292 ++++++++++++++++ .../paged_attention/test_paged_attention.py | 202 ++++++++++++ .../paged_attention_manual_scope/README.md | 49 +++ .../kernels/aic/aic_pv_matmul.cpp | 114 +++++++ .../kernels/aic/aic_qk_matmul.cpp | 115 +++++++ .../kernels/aiv/aiv_online_update.cpp | 256 ++++++++++++++ .../kernels/aiv/aiv_softmax_prepare.cpp | 156 +++++++++ .../orchestration/paged_attention_orch.cpp | 311 ++++++++++++++++++ .../test_paged_attention_manual_scope.py | 200 +++++++++++ 14 files changed, 2401 insertions(+) create mode 100644 examples/a2a3/host_build_graph/paged_attention/README.md create mode 100644 examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_pv_matmul.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_qk_matmul.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp create mode 100644 examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py diff --git a/examples/a2a3/host_build_graph/paged_attention/README.md b/examples/a2a3/host_build_graph/paged_attention/README.md new file mode 100644 index 0000000000..ef624e642f --- /dev/null +++ b/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 +``` diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..0220a6bbb3 --- /dev/null +++ b/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 +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +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, Stride>; + using GlobalB = GlobalTensor, Stride>; + using GlobalOut = GlobalTensor, Stride>; + + 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; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + 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(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); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..efd423bd6e --- /dev/null +++ b/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 +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +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, Stride>; + // kj stored as (N, K) row-major = (K, N) column-major -> DN layout + using GlobalB = GlobalTensor, Stride, Layout::DN>; + using GlobalOut = GlobalTensor, Stride>; + + 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; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + 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(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); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..ded4dcad87 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,256 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Online Softmax Update + Normalize Kernel (AIV) +// +// Operates on full tiles where M=q_tile_size, N=head_dim (128): +// Case1: oi/oi_new are (16, 128), mij/lij/mi/li are 16-element vectors +// Case2: oi/oi_new are (64, 128), mij/lij/mi/li are 64-element vectors +// +// Scalar layout strategy using TRESHAPE (zero-copy UB reshape): +// Scalars loaded as DN ColMajor (M, 1) for TROWEXPANDMUL/TROWEXPANDDIV. +// For element-wise ops (TMAX, TSUB, TEXP, etc.), TRESHAPE to RowMajor (1, M). +// After arithmetic, TRESHAPE back to ColMajor (M, 1) for row-broadcast ops. +// This eliminates the GM round-trip (TSTORE ND → TLOAD DN) used in the original. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_impl( + __gm__ Tensor *mij, __gm__ Tensor *lij, __gm__ Tensor *oi_new, __gm__ Tensor *mi, __gm__ Tensor *li, + __gm__ Tensor *oi, uint64_t is_first, uint64_t is_last, __gm__ Tensor *dst +) { + __gm__ float *mij_ptr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_ptr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + __gm__ float *oi_new_ptr = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr); + __gm__ float *mi_ptr = reinterpret_cast<__gm__ float *>(mi->buffer.addr); + __gm__ float *li_ptr = reinterpret_cast<__gm__ float *>(li->buffer.addr); + __gm__ float *oi_ptr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + __gm__ float *dst_ptr = reinterpret_cast<__gm__ float *>(dst->buffer.addr); + + // Aligned rows for ColMajor DN tiles (32-byte alignment) + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + // --- GlobalTensor types --- + + // Data (M, N) RowMajor + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + + // Scalar DN: M contiguous floats as (kAlignedRows, 1) ColMajor for TROWEXPAND ops and loading + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + // Scalar ND: for storing mi_new and li_new back to GM + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + using GlobalScalarND = + GlobalTensor, Stride<1, 1, 1, kScalarCols, 1>>; + + // --- GlobalTensor instances --- + + GlobalDataMxN oiNewGlobal(oi_new_ptr + oi_new->start_offset); + GlobalDataMxN oiGlobal(oi_ptr + oi->start_offset); + GlobalDataMxN dstGlobal(dst_ptr + dst->start_offset); + + // DN globals for loading scalars as ColMajor + GlobalScalarDN mijGlobalDN(mij_ptr + mij->start_offset); + GlobalScalarDN lijGlobalDN(lij_ptr + lij->start_offset); + GlobalScalarDN miGlobalDN(mi_ptr + mi->start_offset); + GlobalScalarDN liGlobalDN(li_ptr + li->start_offset); + + // ND globals for storing scalar results + GlobalScalarND miGlobalND(mi_ptr + mi->start_offset); + GlobalScalarND liGlobalND(li_ptr + li->start_offset); + + // --- Tile types --- + + using TileDataMxN = Tile; + using TileScalarDN = Tile; + + // RowMajor (1, M) tiles for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // ND tile for storing back to GM + using TileScalarND = + Tile; + + // --- UB memory layout --- + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Data tiles + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + // Scalar DN tiles loaded from GM (ColMajor) + TileScalarDN mijDN, lijDN, miDN, liDN; + + // Temporary DN tiles for results + TileScalarDN miNewDN, alphaDN, betaDN, liNewDN, tmpDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijDN, 2 * kDataBytes); + TASSIGN(lijDN, 2 * kDataBytes + kScalarDNBytes); + TASSIGN(miDN, 2 * kDataBytes + 2 * kScalarDNBytes); + TASSIGN(liDN, 2 * kDataBytes + 3 * kScalarDNBytes); + TASSIGN(miNewDN, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaDN, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaDN, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewDN, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpDN, 2 * kDataBytes + 8 * kScalarDNBytes); + + if (is_first) { + // --- First block: copy inputs to accumulators --- + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Store mi = mij, li = lij, oi = oi_new + // Alias ND tiles to the same UB as DN tiles for storing as ND format + TileScalarND mijND, lijND; + TASSIGN(mijND, 2 * kDataBytes); // alias same UB as mijDN + TASSIGN(lijND, 2 * kDataBytes + kScalarDNBytes); // alias same UB as lijDN + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); // mi = mij + TSTORE(liGlobalND, lijND); // li = lij + TSTORE(oiGlobal, oiNewTile); // oi = oi_new + + if (is_last) { + // Single block: normalize dst = oi_new / lij + // lijDN already in ColMajor DN format, use directly for TROWEXPANDDIV + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TROWEXPANDDIV(oiNewTile, oiNewTile, lijDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + // --- Subsequent blocks: accumulate --- + + // Load all inputs + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + TLOAD(miDN, miGlobalDN); + TLOAD(liDN, liGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise arithmetic + TileScalarRow miRow, mijRow, liRow, lijRow; + TRESHAPE(miRow, miDN); + TRESHAPE(mijRow, mijDN); + TRESHAPE(liRow, liDN); + TRESHAPE(lijRow, lijDN); + + // Scalar arithmetic in RowMajor (1, M) layout + TileScalarRow miNewRow, alphaRow, betaRow, liNewRow, tmpRow; + TASSIGN(miNewRow, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaRow, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaRow, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewRow, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpRow, 2 * kDataBytes + 8 * kScalarDNBytes); + + TMAX(miNewRow, miRow, mijRow); // mi_new = max(mi, mij) + pipe_barrier(PIPE_V); + TSUB(alphaRow, miRow, miNewRow); // alpha_exp = mi - mi_new + pipe_barrier(PIPE_V); + TEXP(alphaRow, alphaRow); // alpha = exp(mi - mi_new) + pipe_barrier(PIPE_V); + TSUB(betaRow, mijRow, miNewRow); // beta_exp = mij - mi_new + pipe_barrier(PIPE_V); + TEXP(betaRow, betaRow); // beta = exp(mij - mi_new) + pipe_barrier(PIPE_V); + TMUL(tmpRow, alphaRow, liRow); // alpha * li + pipe_barrier(PIPE_V); + TMUL(liNewRow, betaRow, lijRow); // beta * lij + pipe_barrier(PIPE_V); + TADD(liNewRow, tmpRow, liNewRow); // li_new = alpha*li + beta*lij + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for TROWEXPANDMUL + TRESHAPE(alphaDN, alphaRow); + TRESHAPE(betaDN, betaRow); + + // Scale data tiles using row-broadcast multiply + TROWEXPANDMUL(oiTile, oiTile, alphaDN); // oi *= alpha + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); // oi_new *= beta + pipe_barrier(PIPE_V); + TADD(oiTile, oiTile, oiNewTile); // oi = alpha*oi + beta*oi_new + + // Store mi_new and li_new to GM (ND format) + // Alias ND tiles to the same UB locations as miNewRow and liNewRow + TileScalarND miNewND, liNewND; + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(liNewND, 2 * kDataBytes + 7 * kScalarDNBytes); + + if (is_last) { + // Normalize and output: dst = oi / li_new + TRESHAPE(liNewDN, liNewRow); + pipe_barrier(PIPE_V); + TROWEXPANDDIV(oiTile, oiTile, liNewDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(dstGlobal, oiTile); + } else { + // Store updated accumulators + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(oiGlobal, oiTile); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *dst = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t q_tile_size = static_cast(mij->shapes[0]); + + if (q_tile_size == 16 && oi_new->shapes[1] <= 16) { + online_update_impl<16, 16>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else if (q_tile_size == 16) { + online_update_impl<16, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else { + online_update_impl<64, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..8f0c41775d --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,156 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Softmax Preparation Kernel (AIV) with partial block masking +// +// Operates on (M, N) tile where M=q_tile_size, N=block_size: +// Case1: sij is (16, 128) +// Case2: sij is (64, 64) +// +// For partial blocks (valid_len < N), positions [valid_len, N) in sij are +// filled with -inf via TFILLPAD_INPLACE before softmax, ensuring exp(-inf)=0 +// so that invalid key positions contribute zero attention weight. +// +// Computes: +// sij_masked = TFILLPAD(sij, valid_len, pad=-inf) +// sij_scale = sij_masked * scale +// mij = row_max(sij_scale) -> (M, 1) +// pij = exp(sij_scale - mij) -> (M, N) +// lij = row_sum(pij) -> (M, 1) + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void softmax_prepare_impl( + __gm__ Tensor *sij, float scale_value, __gm__ Tensor *pij, __gm__ Tensor *mij, __gm__ Tensor *lij +) { + uint64_t valid_len = static_cast(sij->shapes[1]); + __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr); + __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr); + __gm__ float *mij_addr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_addr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + GlobalDataMxN sijGlobal(sij_addr + sij->start_offset); + GlobalDataMxN_bf16 pijGlobal(pij_addr + pij->start_offset); + GlobalScalarDN mijGlobal(mij_addr + mij->start_offset); + GlobalScalarDN lijGlobal(lij_addr + lij->start_offset); + + // Dynamic-cols tile: marks which columns are valid for TFILLPAD boundary + using TileSijDyn = Tile; + // Padded tile: TFILLPAD_INPLACE fills positions [valid_len, N) with -inf + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijDyn sijDynTile(static_cast(valid_len)); + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + // All sij tiles share UB address 0x0 (in-place masking) + TASSIGN(sijTile, 0x0); + TASSIGN(sijDynTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + // Load full sij (M, N) tile from GM - all N columns including garbage for partial blocks + // printf("sij addr incore %x\n", sij->buffer.addr); + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Mask columns [valid_len, N) with -inf. sijDynTile provides the valid boundary, + // sijPadTile provides PadValue::Min as the fill value. No-op when valid_len == N. + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + pipe_barrier(PIPE_V); + + TMULS(sijTile, sijTile, scale_value); + pipe_barrier(PIPE_V); + TROWMAX(maxTile, sijTile, tmpTile); + pipe_barrier(PIPE_V); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + pipe_barrier(PIPE_V); + TEXP(pijTile, pijTile); + // Truncate pij to bf16 first + pipe_barrier(PIPE_V); + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // pij bf16 ready, can store early + + // Continue computing: bf16 → f32 and rowsum while pij store proceeds in parallel + pipe_barrier(PIPE_V); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + TROWSUM(sumTile, pijTile, tmpTile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); // sum ready + + // Store pij (overlaps with TCVT + TROWSUM above) + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + + // Store max and sum + TSTORE(mijGlobal, maxTile); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(lijGlobal, sumTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[3]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[4]); + float scale_value = scale_conv.f; + uint64_t q_tile_size = static_cast(sij->shapes[0]); + + if (q_tile_size == 16 && pij->shapes[1] <= 16) { + softmax_prepare_impl<16, 16>(sij, scale_value, pij, mij, lij); + } else if (q_tile_size == 16) { + softmax_prepare_impl<16, 128>(sij, scale_value, pij, mij, lij); + } else { + softmax_prepare_impl<64, 64>(sij, scale_value, pij, mij, lij); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..e87df0bf14 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,292 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Paged Attention Orchestration Function - 16x16 Version + * + * Simplified for 16x16 framework-generated matmul kernels. + * Each block processes a single 16x16 matmul operation. + * + * Memory Layout: + * Query: (batch, 16, 16) - one 16x16 tile per batch + * Key: (total_blocks, 16, 16) - stored as K^T for direct matmul + * Value: (total_blocks, 16, 16) - direct format + */ + +#include +#include +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 +constexpr uint64_t PLATFORM_PROF_SYS_CNT_FREQ = 50000000; // 50 MHz + +inline double cycles_to_us(uint64_t cycles) { + return (static_cast(cycles) / PLATFORM_PROF_SYS_CNT_FREQ) * 1000000.0; +} + +inline uint64_t get_sys_cnt_aicpu() { +#if defined(__aarch64__) + uint64_t ticks; + asm volatile("mrs %0, cntvct_el0" : "=r"(ticks)); + return ticks; +#elif defined(__x86_64__) + return 0; +#else + return 0; +#endif +} + +#ifdef ENABLE_PROFILING +#define CYCLE_COUNT_START() uint64_t _t0 = get_sys_cnt_aicpu(), _t1 +#define CYCLE_COUNT_LAP(acc) \ + do { \ + _t1 = get_sys_cnt_aicpu(); \ + acc += (_t1 - _t0); \ + _t0 = _t1; \ + } while (0) +#define PROF_INC(counter, n) (counter) += (n) +#else +#define CYCLE_COUNT_START() (void)0 +#define CYCLE_COUNT_LAP(acc) (void)0 +#define PROF_INC(counter, n) (void)0 +#endif + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { +#ifdef ENABLE_PROFILING + uint64_t prof_param_extract = 0; + uint64_t prof_ext_tensor = 0; + uint64_t prof_scope = 0; + uint64_t prof_make_tensor = 0; + uint64_t prof_tensor_view = 0; + uint64_t prof_param_setup = 0; + uint64_t prof_submit_task = 0; + int prof_submit_count = 0; + int prof_make_count = 0; + int prof_view_count = 0; +#endif + + CYCLE_COUNT_START(); + + // Read dimensions from tensor metadata + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + uint64_t scale_value = orch_args.scalar(0); + + uint64_t q_head_num = num_heads; + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (q_head_num + q_tile - 1) / q_tile; + CYCLE_COUNT_LAP(prof_param_extract); + + LOG_INFO(">>>>>> batch = %" PRIu64, batch); + + // Reshape tensors for kernel consumption (2D flattened) + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint32_t key_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t value_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32); + CYCLE_COUNT_LAP(prof_ext_tensor); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + // Create infos are loop-invariant — shapes depend only on q_tile/head_dim/block_size + uint32_t tile2d_shapes[2] = {static_cast(q_tile), static_cast(head_dim)}; + uint32_t scalar_shapes[1] = {static_cast(q_tile)}; + uint32_t sij_shapes[2] = {static_cast(q_tile), static_cast(block_size)}; + TensorCreateInfo tile2d_ci(tile2d_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_ci(scalar_shapes, 1, DataType::FLOAT32); + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_f16_ci(sij_shapes, 2, data_type); + + PROF_INC(prof_make_count, 4); + CYCLE_COUNT_LAP(prof_make_tensor); + + for (uint64_t b_idx = 0; b_idx < batch; b_idx++) { + uint32_t cl_idx[1] = {static_cast(b_idx)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size; + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + PTO2_SCOPE() { + CYCLE_COUNT_LAP(prof_scope); + uint64_t cur_offset = b_idx * q_head_num + q_idx * q_tile; + + uint32_t qi_offsets[2] = {static_cast(cur_offset), 0}; + Tensor qi = query.view(tile2d_shapes, qi_offsets); + uint32_t out_view_offsets[2] = {static_cast(cur_offset), 0}; + Tensor out_view = out.view(tile2d_shapes, out_view_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors alloc_outs = alloc_tensors(tile2d_ci, scalar_ci, scalar_ci); + const Tensor &oi = alloc_outs.get_ref(0); + const Tensor &li_update = alloc_outs.get_ref(1); + const Tensor &mi_update = alloc_outs.get_ref(2); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + for (uint64_t bn = 0; bn < bn_this_batch; bn++) { + PTO2_SCOPE_GUARD(); + + uint32_t bt_idx[2] = {static_cast(b_idx), static_cast(bn)}; + uint64_t cur_block_idx = static_cast(get_tensor_data(block_table, 2, bt_idx)); + uint64_t valid_len = std::min(block_size, cur_seq - bn * block_size); + CYCLE_COUNT_LAP(prof_param_extract); + + uint32_t kv_shapes[2] = {static_cast(block_size), static_cast(head_dim)}; + uint32_t kv_offsets[2] = {static_cast(cur_block_idx * block_size), 0}; + Tensor kj = key_cache.view(kv_shapes, kv_offsets); + Tensor vj = value_cache.view(kv_shapes, kv_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + L0TaskArgs params_qk; + params_qk.add_input(qi); + params_qk.add_input(kj); + params_qk.add_output(sij_ci); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij = qk_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint32_t sij_valid_shapes[2] = {static_cast(q_tile), static_cast(valid_len)}; + uint32_t sij_valid_offsets[2] = {0, 0}; + Tensor sij_valid = sij.view(sij_valid_shapes, sij_valid_offsets); + PROF_INC(prof_view_count, 1); + CYCLE_COUNT_LAP(prof_tensor_view); + + L0TaskArgs params_sf; + params_sf.add_input(sij_valid); + params_sf.add_output(pij_f16_ci); + params_sf.add_output(scalar_ci); + params_sf.add_output(scalar_ci); + params_sf.add_scalar(scale_value); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_f16 = sf_outs.get_ref(0); + const Tensor &mi = sf_outs.get_ref(1); + const Tensor &li = sf_outs.get_ref(2); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + L0TaskArgs params_pv; + params_pv.add_input(pij_f16); + params_pv.add_input(vj); + params_pv.add_output(tile2d_ci); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_tmp = pv_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn == bn_this_batch - 1) ? 1 : 0; + CYCLE_COUNT_LAP(prof_param_extract); + + L0TaskArgs params_up; + params_up.add_input(mi); + params_up.add_input(li); + params_up.add_input(oi_tmp); + params_up.add_inout(mi_update); + params_up.add_inout(li_update); + params_up.add_inout(oi); + params_up.add_inout(out_view); + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + CYCLE_COUNT_LAP(prof_param_setup); + rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + } + } + CYCLE_COUNT_LAP(prof_scope); + } + } + +#ifdef ENABLE_PROFILING + uint64_t total = prof_param_extract + prof_ext_tensor + prof_make_tensor + prof_tensor_view + prof_param_setup + + prof_submit_task + prof_scope; + LOG_INFO( + "=== PagedAttn Orch Profiling: %d submits, %d makes, %d views, total=%.3fus ===", prof_submit_count, + prof_make_count, prof_view_count, cycles_to_us(total) + ); + if (total > 0) { + LOG_INFO( + " param_extract : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_extract), + prof_param_extract * 100.0 / total + ); + LOG_INFO( + " ext_tensor(x4) : %7.3fus (%5.1f%%)", cycles_to_us(prof_ext_tensor), prof_ext_tensor * 100.0 / total + ); + LOG_INFO( + " create_info(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_make_count, cycles_to_us(prof_make_tensor), + prof_make_tensor * 100.0 / total, + prof_make_count > 0 ? cycles_to_us(prof_make_tensor) / prof_make_count : 0.0 + ); + LOG_INFO( + " tensor_view(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_view_count, cycles_to_us(prof_tensor_view), + prof_tensor_view * 100.0 / total, + prof_view_count > 0 ? cycles_to_us(prof_tensor_view) / prof_view_count : 0.0 + ); + LOG_INFO( + " param_setup : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_setup), prof_param_setup * 100.0 / total + ); + LOG_INFO(" scope : %7.3fus (%5.1f%%)", cycles_to_us(prof_scope), prof_scope * 100.0 / total); + LOG_INFO( + " submit_task(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_submit_count, cycles_to_us(prof_submit_task), + prof_submit_task * 100.0 / total, + prof_submit_count > 0 ? cycles_to_us(prof_submit_task) / prof_submit_count : 0.0 + ); + } +#endif +} + +} // extern "C" diff --git a/examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py b/examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py new file mode 100644 index 0000000000..490f6f18b2 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention: online softmax with AIC/AIV subgraph splitting (bfloat16), host_build_graph runtime. + +The orchestration and kernel sources are the same as the tensormap_and_ringbuffer +variant: `pto_orchestration_api.h` is identical between the two runtimes, and the +framework compiles the orchestration against the include dirs of whichever +runtime `@scene_test` names. + +host_build_graph populates the whole task graph on the host before the device +begins scheduling, so a ring slot cannot be reclaimed mid-orchestration: the +ring window and GM heap must hold every task of a case at once. Cases above the +default window carry explicit `runtime_env` sizing below. +""" + +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 + +# tasks = batch * (4 * ceil(context_len / block_size) + 1). ring_task_window must +# be a power of two in [4, INT32_MAX] and is rejected otherwise by the runtime. +_RING_65K = {"ring_task_window": 131072, "ring_heap": 2 * 1024 * 1024 * 1024} +_RING_33K = {"ring_task_window": 65536, "ring_heap": 1024 * 1024 * 1024} + + +@scene_test(level=2, runtime="host_build_graph") +class TestPagedAttentionHostBuildGraph(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + # 65 792 tasks. + "name": "Case1", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_65K}, + "manual": True, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + # 32 832 tasks. + "name": "Case2", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_33K}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall1", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + # context_lens_list makes the per-batch block counts differ, so the + # graph is ragged rather than a uniform batch * blocks grid. + "name": "CaseVarSeq2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md b/examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md new file mode 100644 index 0000000000..781c03c47f --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md @@ -0,0 +1,49 @@ +# paged_attention_manual_scope — host_build_graph + +Same computation as `../paged_attention/`, with the automatic same-scope +dependency wiring replaced by explicit task-to-task edges inside +`PTO2_SCOPE(PTO2ScopeMode::MANUAL)`. Read it against the baseline to see what +the automatic mode was deriving. + +## The two dependency APIs it demonstrates + +| API | Shape | Suited to | +| --- | ----- | --------- | +| `Arg::set_dependencies(buf, n)` | caller owns the buffer, `Arg` stores `(ptr, count)` | codegen, fixed dep sets | +| `L0TaskArgsWithDeps<>::add_dep(id)` | wrapper owns a stack-sized buffer, incremental | hand-written orch, deps assembled across branches | + +`SF` and `PV` use the primitive form; `UP` uses the convenience form because its +dep set is conditional — it always takes the `PV` edge, adds the previous +`UP` task when there is one, and adds the alloc task on the last block so the +scratch buffers outlive their final consumer. + +Both APIs exist unchanged on `host_build_graph`: the orchestration header is +byte-identical between the two runtimes, so this variant is the same C++ as its +`tensormap_and_ringbuffer` sibling. + +## Ring sizing + +`host_build_graph` submits the whole graph before the device schedules, so every +task of a case is live at once and the ring must hold all of them — `Case1` +65 792 and `Case2` 32 832, both above the default 16384 window. Each case +carries its own `runtime_env` sizing in `CASES[*]["config"]`. + +This is a stronger constraint than the `tensormap_and_ringbuffer` variant faces, +where slots retire as orchestration proceeds and only the *in-flight* depth of a +single MANUAL scope has to fit. + +## Cases + +`Case1`, `Case2`, `CaseSmall1`, `CaseSmall2`, `CaseVarSeq2`, `CaseVarSeq4`. All +but `CaseSmall1` are `manual`. The upstream `Case3` (`head_dim: 256`) is absent +because it does not produce correct results on either runtime — see +`KNOWN_ISSUES.md`. + +## Run + +```bash +python examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py \ + -p a2a3 -d 0 # CaseSmall1, golden checked +python examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py \ + -p a2a3 -d 0 --manual include --case Case1 --rounds 2 +``` diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_pv_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..0220a6bbb3 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/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 +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +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, Stride>; + using GlobalB = GlobalTensor, Stride>; + using GlobalOut = GlobalTensor, Stride>; + + 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; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + 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(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); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_qk_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..efd423bd6e --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/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 +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +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, Stride>; + // kj stored as (N, K) row-major = (K, N) column-major -> DN layout + using GlobalB = GlobalTensor, Stride, Layout::DN>; + using GlobalOut = GlobalTensor, Stride>; + + 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; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + 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(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); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..ded4dcad87 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,256 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Online Softmax Update + Normalize Kernel (AIV) +// +// Operates on full tiles where M=q_tile_size, N=head_dim (128): +// Case1: oi/oi_new are (16, 128), mij/lij/mi/li are 16-element vectors +// Case2: oi/oi_new are (64, 128), mij/lij/mi/li are 64-element vectors +// +// Scalar layout strategy using TRESHAPE (zero-copy UB reshape): +// Scalars loaded as DN ColMajor (M, 1) for TROWEXPANDMUL/TROWEXPANDDIV. +// For element-wise ops (TMAX, TSUB, TEXP, etc.), TRESHAPE to RowMajor (1, M). +// After arithmetic, TRESHAPE back to ColMajor (M, 1) for row-broadcast ops. +// This eliminates the GM round-trip (TSTORE ND → TLOAD DN) used in the original. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_impl( + __gm__ Tensor *mij, __gm__ Tensor *lij, __gm__ Tensor *oi_new, __gm__ Tensor *mi, __gm__ Tensor *li, + __gm__ Tensor *oi, uint64_t is_first, uint64_t is_last, __gm__ Tensor *dst +) { + __gm__ float *mij_ptr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_ptr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + __gm__ float *oi_new_ptr = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr); + __gm__ float *mi_ptr = reinterpret_cast<__gm__ float *>(mi->buffer.addr); + __gm__ float *li_ptr = reinterpret_cast<__gm__ float *>(li->buffer.addr); + __gm__ float *oi_ptr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + __gm__ float *dst_ptr = reinterpret_cast<__gm__ float *>(dst->buffer.addr); + + // Aligned rows for ColMajor DN tiles (32-byte alignment) + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + // --- GlobalTensor types --- + + // Data (M, N) RowMajor + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + + // Scalar DN: M contiguous floats as (kAlignedRows, 1) ColMajor for TROWEXPAND ops and loading + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + // Scalar ND: for storing mi_new and li_new back to GM + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + using GlobalScalarND = + GlobalTensor, Stride<1, 1, 1, kScalarCols, 1>>; + + // --- GlobalTensor instances --- + + GlobalDataMxN oiNewGlobal(oi_new_ptr + oi_new->start_offset); + GlobalDataMxN oiGlobal(oi_ptr + oi->start_offset); + GlobalDataMxN dstGlobal(dst_ptr + dst->start_offset); + + // DN globals for loading scalars as ColMajor + GlobalScalarDN mijGlobalDN(mij_ptr + mij->start_offset); + GlobalScalarDN lijGlobalDN(lij_ptr + lij->start_offset); + GlobalScalarDN miGlobalDN(mi_ptr + mi->start_offset); + GlobalScalarDN liGlobalDN(li_ptr + li->start_offset); + + // ND globals for storing scalar results + GlobalScalarND miGlobalND(mi_ptr + mi->start_offset); + GlobalScalarND liGlobalND(li_ptr + li->start_offset); + + // --- Tile types --- + + using TileDataMxN = Tile; + using TileScalarDN = Tile; + + // RowMajor (1, M) tiles for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // ND tile for storing back to GM + using TileScalarND = + Tile; + + // --- UB memory layout --- + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Data tiles + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + // Scalar DN tiles loaded from GM (ColMajor) + TileScalarDN mijDN, lijDN, miDN, liDN; + + // Temporary DN tiles for results + TileScalarDN miNewDN, alphaDN, betaDN, liNewDN, tmpDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijDN, 2 * kDataBytes); + TASSIGN(lijDN, 2 * kDataBytes + kScalarDNBytes); + TASSIGN(miDN, 2 * kDataBytes + 2 * kScalarDNBytes); + TASSIGN(liDN, 2 * kDataBytes + 3 * kScalarDNBytes); + TASSIGN(miNewDN, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaDN, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaDN, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewDN, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpDN, 2 * kDataBytes + 8 * kScalarDNBytes); + + if (is_first) { + // --- First block: copy inputs to accumulators --- + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Store mi = mij, li = lij, oi = oi_new + // Alias ND tiles to the same UB as DN tiles for storing as ND format + TileScalarND mijND, lijND; + TASSIGN(mijND, 2 * kDataBytes); // alias same UB as mijDN + TASSIGN(lijND, 2 * kDataBytes + kScalarDNBytes); // alias same UB as lijDN + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); // mi = mij + TSTORE(liGlobalND, lijND); // li = lij + TSTORE(oiGlobal, oiNewTile); // oi = oi_new + + if (is_last) { + // Single block: normalize dst = oi_new / lij + // lijDN already in ColMajor DN format, use directly for TROWEXPANDDIV + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TROWEXPANDDIV(oiNewTile, oiNewTile, lijDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + // --- Subsequent blocks: accumulate --- + + // Load all inputs + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + TLOAD(miDN, miGlobalDN); + TLOAD(liDN, liGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise arithmetic + TileScalarRow miRow, mijRow, liRow, lijRow; + TRESHAPE(miRow, miDN); + TRESHAPE(mijRow, mijDN); + TRESHAPE(liRow, liDN); + TRESHAPE(lijRow, lijDN); + + // Scalar arithmetic in RowMajor (1, M) layout + TileScalarRow miNewRow, alphaRow, betaRow, liNewRow, tmpRow; + TASSIGN(miNewRow, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaRow, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaRow, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewRow, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpRow, 2 * kDataBytes + 8 * kScalarDNBytes); + + TMAX(miNewRow, miRow, mijRow); // mi_new = max(mi, mij) + pipe_barrier(PIPE_V); + TSUB(alphaRow, miRow, miNewRow); // alpha_exp = mi - mi_new + pipe_barrier(PIPE_V); + TEXP(alphaRow, alphaRow); // alpha = exp(mi - mi_new) + pipe_barrier(PIPE_V); + TSUB(betaRow, mijRow, miNewRow); // beta_exp = mij - mi_new + pipe_barrier(PIPE_V); + TEXP(betaRow, betaRow); // beta = exp(mij - mi_new) + pipe_barrier(PIPE_V); + TMUL(tmpRow, alphaRow, liRow); // alpha * li + pipe_barrier(PIPE_V); + TMUL(liNewRow, betaRow, lijRow); // beta * lij + pipe_barrier(PIPE_V); + TADD(liNewRow, tmpRow, liNewRow); // li_new = alpha*li + beta*lij + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for TROWEXPANDMUL + TRESHAPE(alphaDN, alphaRow); + TRESHAPE(betaDN, betaRow); + + // Scale data tiles using row-broadcast multiply + TROWEXPANDMUL(oiTile, oiTile, alphaDN); // oi *= alpha + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); // oi_new *= beta + pipe_barrier(PIPE_V); + TADD(oiTile, oiTile, oiNewTile); // oi = alpha*oi + beta*oi_new + + // Store mi_new and li_new to GM (ND format) + // Alias ND tiles to the same UB locations as miNewRow and liNewRow + TileScalarND miNewND, liNewND; + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(liNewND, 2 * kDataBytes + 7 * kScalarDNBytes); + + if (is_last) { + // Normalize and output: dst = oi / li_new + TRESHAPE(liNewDN, liNewRow); + pipe_barrier(PIPE_V); + TROWEXPANDDIV(oiTile, oiTile, liNewDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(dstGlobal, oiTile); + } else { + // Store updated accumulators + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(oiGlobal, oiTile); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *dst = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t q_tile_size = static_cast(mij->shapes[0]); + + if (q_tile_size == 16 && oi_new->shapes[1] <= 16) { + online_update_impl<16, 16>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else if (q_tile_size == 16) { + online_update_impl<16, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else { + online_update_impl<64, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..8f0c41775d --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,156 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Softmax Preparation Kernel (AIV) with partial block masking +// +// Operates on (M, N) tile where M=q_tile_size, N=block_size: +// Case1: sij is (16, 128) +// Case2: sij is (64, 64) +// +// For partial blocks (valid_len < N), positions [valid_len, N) in sij are +// filled with -inf via TFILLPAD_INPLACE before softmax, ensuring exp(-inf)=0 +// so that invalid key positions contribute zero attention weight. +// +// Computes: +// sij_masked = TFILLPAD(sij, valid_len, pad=-inf) +// sij_scale = sij_masked * scale +// mij = row_max(sij_scale) -> (M, 1) +// pij = exp(sij_scale - mij) -> (M, N) +// lij = row_sum(pij) -> (M, 1) + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void softmax_prepare_impl( + __gm__ Tensor *sij, float scale_value, __gm__ Tensor *pij, __gm__ Tensor *mij, __gm__ Tensor *lij +) { + uint64_t valid_len = static_cast(sij->shapes[1]); + __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr); + __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr); + __gm__ float *mij_addr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_addr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + GlobalDataMxN sijGlobal(sij_addr + sij->start_offset); + GlobalDataMxN_bf16 pijGlobal(pij_addr + pij->start_offset); + GlobalScalarDN mijGlobal(mij_addr + mij->start_offset); + GlobalScalarDN lijGlobal(lij_addr + lij->start_offset); + + // Dynamic-cols tile: marks which columns are valid for TFILLPAD boundary + using TileSijDyn = Tile; + // Padded tile: TFILLPAD_INPLACE fills positions [valid_len, N) with -inf + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijDyn sijDynTile(static_cast(valid_len)); + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + // All sij tiles share UB address 0x0 (in-place masking) + TASSIGN(sijTile, 0x0); + TASSIGN(sijDynTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + // Load full sij (M, N) tile from GM - all N columns including garbage for partial blocks + // printf("sij addr incore %x\n", sij->buffer.addr); + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Mask columns [valid_len, N) with -inf. sijDynTile provides the valid boundary, + // sijPadTile provides PadValue::Min as the fill value. No-op when valid_len == N. + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + pipe_barrier(PIPE_V); + + TMULS(sijTile, sijTile, scale_value); + pipe_barrier(PIPE_V); + TROWMAX(maxTile, sijTile, tmpTile); + pipe_barrier(PIPE_V); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + pipe_barrier(PIPE_V); + TEXP(pijTile, pijTile); + // Truncate pij to bf16 first + pipe_barrier(PIPE_V); + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // pij bf16 ready, can store early + + // Continue computing: bf16 → f32 and rowsum while pij store proceeds in parallel + pipe_barrier(PIPE_V); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + TROWSUM(sumTile, pijTile, tmpTile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); // sum ready + + // Store pij (overlaps with TCVT + TROWSUM above) + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + + // Store max and sum + TSTORE(mijGlobal, maxTile); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(lijGlobal, sumTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[3]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[4]); + float scale_value = scale_conv.f; + uint64_t q_tile_size = static_cast(sij->shapes[0]); + + if (q_tile_size == 16 && pij->shapes[1] <= 16) { + softmax_prepare_impl<16, 16>(sij, scale_value, pij, mij, lij); + } else if (q_tile_size == 16) { + softmax_prepare_impl<16, 128>(sij, scale_value, pij, mij, lij); + } else { + softmax_prepare_impl<64, 64>(sij, scale_value, pij, mij, lij); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..37d545dd17 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,311 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Paged Attention Orchestration Function - manual-scope variant + * + * Matches the small-case paged_attention orchestration shape while replacing + * the automatic same-scope dependency wiring with explicit task-to-task deps + * inside PTO2_SCOPE(PTO2ScopeMode::MANUAL). + */ + +#include +#include +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 +constexpr uint64_t PLATFORM_PROF_SYS_CNT_FREQ = 50000000; // 50 MHz + +inline double cycles_to_us(uint64_t cycles) { + return (static_cast(cycles) / PLATFORM_PROF_SYS_CNT_FREQ) * 1000000.0; +} + +inline uint64_t get_sys_cnt_aicpu() { +#if defined(__aarch64__) + uint64_t ticks; + asm volatile("mrs %0, cntvct_el0" : "=r"(ticks)); + return ticks; +#elif defined(__x86_64__) + return 0; +#else + return 0; +#endif +} + +#ifdef ENABLE_PROFILING +#define CYCLE_COUNT_START() uint64_t _t0 = get_sys_cnt_aicpu(), _t1 +#define CYCLE_COUNT_LAP(acc) \ + do { \ + _t1 = get_sys_cnt_aicpu(); \ + acc += (_t1 - _t0); \ + _t0 = _t1; \ + } while (0) +#define PROF_INC(counter, n) (counter) += (n) +#else +#define CYCLE_COUNT_START() (void)0 +#define CYCLE_COUNT_LAP(acc) (void)0 +#define PROF_INC(counter, n) (void)0 +#endif + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { +#ifdef ENABLE_PROFILING + uint64_t prof_param_extract = 0; + uint64_t prof_ext_tensor = 0; + uint64_t prof_scope = 0; + uint64_t prof_make_tensor = 0; + uint64_t prof_tensor_view = 0; + uint64_t prof_param_setup = 0; + uint64_t prof_submit_task = 0; + int prof_submit_count = 0; + int prof_make_count = 0; + int prof_view_count = 0; +#endif + + CYCLE_COUNT_START(); + + // Read dimensions from tensor metadata + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + uint64_t scale_value = orch_args.scalar(0); + + uint64_t q_head_num = num_heads; + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (q_head_num + q_tile - 1) / q_tile; + CYCLE_COUNT_LAP(prof_param_extract); + + LOG_INFO(">>>>>> batch = %" PRIu64, batch); + + // Reshape tensors for kernel consumption (2D flattened) + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint32_t key_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t value_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32); + CYCLE_COUNT_LAP(prof_ext_tensor); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + // Create infos are loop-invariant — shapes depend only on q_tile/head_dim/block_size + uint32_t tile2d_shapes[2] = {static_cast(q_tile), static_cast(head_dim)}; + uint32_t scalar_shapes[1] = {static_cast(q_tile)}; + uint32_t sij_shapes[2] = {static_cast(q_tile), static_cast(block_size)}; + TensorCreateInfo tile2d_ci(tile2d_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_ci(scalar_shapes, 1, DataType::FLOAT32); + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_f16_ci(sij_shapes, 2, data_type); + + PROF_INC(prof_make_count, 4); + CYCLE_COUNT_LAP(prof_make_tensor); + + for (uint64_t b_idx = 0; b_idx < batch; b_idx++) { + uint32_t cl_idx[1] = {static_cast(b_idx)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size; + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + PTO2_SCOPE(PTO2ScopeMode::MANUAL) { + CYCLE_COUNT_LAP(prof_scope); + uint64_t cur_offset = b_idx * q_head_num + q_idx * q_tile; + + uint32_t qi_offsets[2] = {static_cast(cur_offset), 0}; + Tensor qi = query.view(tile2d_shapes, qi_offsets); + uint32_t out_view_offsets[2] = {static_cast(cur_offset), 0}; + Tensor out_view = out.view(tile2d_shapes, out_view_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors alloc_outs = alloc_tensors(tile2d_ci, scalar_ci, scalar_ci); + const Tensor &oi = alloc_outs.get_ref(0); + const Tensor &li_update = alloc_outs.get_ref(1); + const Tensor &mi_update = alloc_outs.get_ref(2); + PTO2TaskId alloc_task = alloc_outs.task_id(); + PTO2TaskId prev_update_task = PTO2TaskId::invalid(); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + for (uint64_t bn = 0; bn < bn_this_batch; bn++) { + uint32_t bt_idx[2] = {static_cast(b_idx), static_cast(bn)}; + uint64_t cur_block_idx = static_cast(get_tensor_data(block_table, 2, bt_idx)); + uint64_t valid_len = std::min(block_size, cur_seq - bn * block_size); + CYCLE_COUNT_LAP(prof_param_extract); + + uint32_t kv_shapes[2] = {static_cast(block_size), static_cast(head_dim)}; + uint32_t kv_offsets[2] = {static_cast(cur_block_idx * block_size), 0}; + Tensor kj = key_cache.view(kv_shapes, kv_offsets); + Tensor vj = value_cache.view(kv_shapes, kv_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + L0TaskArgs params_qk; + params_qk.add_input(qi); + params_qk.add_input(kj); + params_qk.add_output(sij_ci); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij = qk_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint32_t sij_valid_shapes[2] = {static_cast(q_tile), static_cast(valid_len)}; + uint32_t sij_valid_offsets[2] = {0, 0}; + Tensor sij_valid = sij.view(sij_valid_shapes, sij_valid_offsets); + PROF_INC(prof_view_count, 1); + CYCLE_COUNT_LAP(prof_tensor_view); + + // --- Primitive dep API (Arg + set_dependencies) --- + // Caller owns the deps buffer; Arg stores (ptr, count). + // Suited for codegen and for cases with a fixed dep set. + L0TaskArgs params_sf; + params_sf.add_input(sij_valid); + params_sf.add_output(pij_f16_ci); + params_sf.add_output(scalar_ci); + params_sf.add_output(scalar_ci); + PTO2TaskId sf_deps[] = {qk_outs.task_id()}; + params_sf.set_dependencies(sf_deps, 1); + params_sf.add_scalar(scale_value); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_f16 = sf_outs.get_ref(0); + const Tensor &mi = sf_outs.get_ref(1); + const Tensor &li = sf_outs.get_ref(2); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + L0TaskArgs params_pv; + params_pv.add_input(pij_f16); + params_pv.add_input(vj); + params_pv.add_output(tile2d_ci); + PTO2TaskId pv_deps[] = {sf_outs.task_id()}; + params_pv.set_dependencies(pv_deps, 1); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_tmp = pv_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn == bn_this_batch - 1) ? 1 : 0; + CYCLE_COUNT_LAP(prof_param_extract); + + // --- Convenience dep API (L0TaskArgsWithDeps + add_dep) --- + // Wrapper owns a stack-sized deps buffer and accepts + // incremental add_dep() calls; the submit overload binds + // them to the underlying Arg via set_dependencies(...). + // Suited for hand-written orch where the dep set is + // assembled conditionally across branches. + L0TaskArgsWithDeps<> params_up; + params_up.add_input(mi); + params_up.add_input(li); + params_up.add_input(oi_tmp); + params_up.add_inout(mi_update); + params_up.add_inout(li_update); + params_up.add_inout(oi); + params_up.add_inout(out_view); + // UP reads SF's mi/li, but SF -> PV -> UP already orders it; only the PV edge is explicit. + params_up.add_dep(pv_outs.task_id()); + if (prev_update_task.is_valid()) { + params_up.add_dep(prev_update_task); + } + // alloc completes inline; this dep only keeps the scratch buffers alive until the last consumer. + if (is_last) { + params_up.add_dep(alloc_task); + } + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors up_outs = rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + prev_update_task = up_outs.task_id(); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + } + } + CYCLE_COUNT_LAP(prof_scope); + } + } + +#ifdef ENABLE_PROFILING + uint64_t total = prof_param_extract + prof_ext_tensor + prof_make_tensor + prof_tensor_view + prof_param_setup + + prof_submit_task + prof_scope; + LOG_INFO( + "=== PagedAttn Orch Profiling: %d submits, %d makes, %d views, total=%.3fus ===", prof_submit_count, + prof_make_count, prof_view_count, cycles_to_us(total) + ); + if (total > 0) { + LOG_INFO( + " param_extract : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_extract), + prof_param_extract * 100.0 / total + ); + LOG_INFO( + " ext_tensor(x4) : %7.3fus (%5.1f%%)", cycles_to_us(prof_ext_tensor), prof_ext_tensor * 100.0 / total + ); + LOG_INFO( + " create_info(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_make_count, cycles_to_us(prof_make_tensor), + prof_make_tensor * 100.0 / total, + prof_make_count > 0 ? cycles_to_us(prof_make_tensor) / prof_make_count : 0.0 + ); + LOG_INFO( + " tensor_view(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_view_count, cycles_to_us(prof_tensor_view), + prof_tensor_view * 100.0 / total, + prof_view_count > 0 ? cycles_to_us(prof_tensor_view) / prof_view_count : 0.0 + ); + LOG_INFO( + " param_setup : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_setup), prof_param_setup * 100.0 / total + ); + LOG_INFO(" scope : %7.3fus (%5.1f%%)", cycles_to_us(prof_scope), prof_scope * 100.0 / total); + LOG_INFO( + " submit_task(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_submit_count, cycles_to_us(prof_submit_task), + prof_submit_task * 100.0 / total, + prof_submit_count > 0 ? cycles_to_us(prof_submit_task) / prof_submit_count : 0.0 + ); + } +#endif +} + +} // extern "C" diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py b/examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py new file mode 100644 index 0000000000..7c4cdb8ca8 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention manual-scope wrapper for A2A3 host_build_graph. + +Replaces the automatic same-scope dependency wiring with explicit task-to-task +deps inside PTO2_SCOPE(PTO2ScopeMode::MANUAL). The orchestration source is the +same one the tensormap_and_ringbuffer variant compiles; only the runtime named +by @scene_test differs. + +host_build_graph submits the whole graph before the device schedules, so the +ring window must hold every task of a case at once — larger than the +tensormap_and_ringbuffer variant needs, where slots retire as orchestration +proceeds. +""" + +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 + +# tasks = batch * (4 * ceil(context_len / block_size) + 1). ring_task_window must +# be a power of two in [4, INT32_MAX] and is rejected otherwise by the runtime. +_RING_65K = {"ring_task_window": 131072, "ring_heap": 2 * 1024 * 1024 * 1024} +_RING_33K = {"ring_task_window": 65536, "ring_heap": 1024 * 1024 * 1024} + + +@scene_test(level=2, runtime="host_build_graph") +class TestPagedAttentionManualScopeHostBuildGraph(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + # 65 792 tasks, all live at once, in a single MANUAL scope. + "name": "Case1", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_65K}, + "manual": True, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + # 32 832 tasks. + "name": "Case2", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_33K}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall1", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__)