From d028a7c2d95245880e0e7d36baca563048f31542 Mon Sep 17 00:00:00 2001 From: andrej Date: Thu, 9 Jul 2026 15:59:54 -0600 Subject: [PATCH 1/4] Add single-dispatch layer-by-layer multi-head attention (mha_prefill_lxl) Add a layer-by-layer (LxL) prefill MHA operator built as an OperatorSequence, and rename the existing data-flow (DF) MHA operator for clarity: - Rename iron/operators/mha -> iron/operators/mha_prefill_df (DF = data-flow: single fused kernel). - Add iron/operators/mha_prefill_lxl (LxL = layer-by-layer: chained sub-operators dispatched back-to-back). - Extend axpy with scale / scalar_add / causal-mask modes and softmax with partial/online softmax, used by the LxL design. - Register the benchmark marker in an operator-local conftest.py. Update README and operator registrations to the new names. --- README.md | 4 +- aie_kernels/aie2/softmax.cc | 112 ++++ aie_kernels/aie2p/softmax.cc | 125 +++- aie_kernels/generic/axpy.cc | 103 ++++ iron/operators/__init__.py | 2 +- iron/operators/axpy/design.py | 331 +++++++++-- iron/operators/axpy/op.py | 112 +++- .../{mha => mha_prefill_df}/design.py | 10 + iron/operators/{mha => mha_prefill_df}/op.py | 0 .../{mha => mha_prefill_df}/reference.py | 0 .../operators/{mha => mha_prefill_df}/test.py | 4 +- iron/operators/mha_prefill_lxl/__init__.py | 2 + iron/operators/mha_prefill_lxl/conftest.py | 9 + iron/operators/mha_prefill_lxl/op.py | 554 ++++++++++++++++++ iron/operators/mha_prefill_lxl/reference.py | 283 +++++++++ iron/operators/mha_prefill_lxl/test.py | 305 ++++++++++ iron/operators/softmax/design.py | 207 +++++++ iron/operators/softmax/op.py | 18 +- iron/operators/softmax/test.py | 49 ++ 19 files changed, 2160 insertions(+), 70 deletions(-) rename iron/operators/{mha => mha_prefill_df}/design.py (98%) rename iron/operators/{mha => mha_prefill_df}/op.py (100%) rename iron/operators/{mha => mha_prefill_df}/reference.py (100%) rename iron/operators/{mha => mha_prefill_df}/test.py (95%) create mode 100644 iron/operators/mha_prefill_lxl/__init__.py create mode 100644 iron/operators/mha_prefill_lxl/conftest.py create mode 100644 iron/operators/mha_prefill_lxl/op.py create mode 100644 iron/operators/mha_prefill_lxl/reference.py create mode 100644 iron/operators/mha_prefill_lxl/test.py diff --git a/README.md b/README.md index 8db21ad7..10c5a359 100755 --- a/README.md +++ b/README.md @@ -46,8 +46,8 @@ The IRON Python API for Ryzenβ„’ AI NPUs is described in the following paper: | [Element-wise Mul](./aie_kernels/generic/mul.cc) | Element-wise multiplication kernel | bfloat16 | βœ“ | βœ“ | 🟒 | [iron/operators/elementwise_mul/](./iron/operators/elementwise_mul/) | | [GEMM](./aie_kernels/aie2p/mm.cc) | General Matrix Multiplication kernel | bfloat16 | βœ“ | βœ“ | 🟒 | [iron/operators/gemm/](./iron/operators/gemm/) | | [GEMV](./aie_kernels/generic/mv.cc) | General Matrix-Vector Multiplication kernel | bfloat16 | βœ“ | βœ“ | 🟒 | [iron/operators/gemv/](./iron/operators/gemv/) | -| [GQA](./aie_kernels/aie2p/mha.cc) | Grouped Query Attention kernel (Single pipeline) | bfloat16 | | βœ“ | 🟒 | [iron/operators/mha/](./iron/operators/mha/) | -| [MHA](./aie_kernels/aie2p/mha.cc) | Multi-Head Attention kernel & Grouped Query Attention | bfloat16 | | βœ“ | 🟒 | [iron/operators/mha/](./iron/operators/mha/) | +| [GQA](./aie_kernels/aie2p/mha.cc) | Grouped Query Attention kernel (Single pipeline) | bfloat16 | | βœ“ | 🟒 | [iron/operators/mha_prefill_df/](./iron/operators/mha_prefill_df/) | +| [MHA](./aie_kernels/aie2p/mha.cc) | Multi-Head Attention kernel & Grouped Query Attention | bfloat16 | | βœ“ | 🟒 | [iron/operators/mha_prefill_df/](./iron/operators/mha_prefill_df/) | | [RMSNorm](./aie_kernels/aie2/rms_norm.cc) | RMSNorm kernel | bfloat16 | βœ“ | βœ“ | 🟒 | [iron/operators/rms_norm/](./iron/operators/rms_norm/) | | [RoPE](./aie_kernels/generic/rope.cc) | Rotary Positional Embedding kernel | bfloat16 | βœ“ | βœ“ | 🟒 | [iron/operators/rope/](./iron/operators/rope/) | | [SiLU](./aie_kernels/aie2/silu.cc) | Sigmoid Linear Unit activation kernel | bfloat16 | βœ“ | βœ“ | 🟒 | [iron/operators/silu/](./iron/operators/silu/) | diff --git a/aie_kernels/aie2/softmax.cc b/aie_kernels/aie2/softmax.cc index 919d38f7..07aa4bb1 100644 --- a/aie_kernels/aie2/softmax.cc +++ b/aie_kernels/aie2/softmax.cc @@ -4,6 +4,7 @@ #include "lut_based_ops.h" #include +#include #include using namespace aie; @@ -57,6 +58,98 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out return; } +// --------------------------------------------------------------------------- +// Online (partial / tiled) softmax helpers +// +// These three kernels implement a two-pass online softmax that processes a row +// in sub-tile chunks, keeping running max and sum statistics in a small local +// buffer (`stats`). Layout of the stats buffer (bfloat16[16], only [0..1] +// used): +// stats[0] = running max +// stats[1] = running sum (of exp(x - max)) +// --------------------------------------------------------------------------- + +void softmax_partial_stats_impl(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +{ + event0(); + + const int elem_iters = vector_size / 16; + + float running_max = (float)stats[0]; + float running_sum = (float)stats[1]; + + aie::vector input_bf16; + aie::accum exp_val_accum = aie::zeros(); + + auto it_in = aie::cbegin_vector<16>((bfloat16 *)input); + + // Single-pass online algorithm: for each vector chunk, check if max + // needs updating, rescale the running sum if so, then accumulate + // exp(x - max). + for (int i = 0; i < elem_iters; i++) { + input_bf16 = *it_in++; + float chunk_max = aie::reduce_max(input_bf16); + + if (chunk_max > running_max) { + // Rescale accumulated exp values by exp(old_max - new_max) + aie::vector correction = + to_v16bfloat16(getExpBf16(aie::broadcast((bfloat16)(running_max - chunk_max)))); + float scale = (float)correction[0]; + // Rescale the partial vector accumulator + aie::vector scale_vec = aie::broadcast((bfloat16)scale); + exp_val_accum = aie::mul(exp_val_accum.to_vector(), scale_vec); + // Rescale the running scalar sum from previous chunks + running_sum *= scale; + running_max = chunk_max; + } + + aie::vector shifted = aie::sub(input_bf16, aie::broadcast((bfloat16)running_max)); + aie::vector exp_val = to_v16bfloat16(getExpBf16(shifted)); + exp_val_accum = add(exp_val_accum, exp_val); + } + + // Reduce the vector accumulator and add to running sum + aie::vector reduce = exp_val_accum.to_vector(); + running_sum += aie::reduce_add(reduce); + + stats[0] = (bfloat16)running_max; + stats[1] = (bfloat16)running_sum; + + event1(); +} + +void softmax_partial_norm_impl(bfloat16 *restrict input, + bfloat16 *restrict output, + bfloat16 *stats, + const int32_t vector_size) +{ + event0(); + + const int elem_iters = vector_size / 16; + + float max_val = (float)stats[0]; + float sum_val = (float)stats[1]; + bfloat16 inv_sum = (bfloat16)aie::inv(sum_val); + + aie::vector max_val_vec = aie::broadcast((bfloat16)max_val); + + aie::vector input_bf16; + aie::accum out_vals; + + auto it_in = aie::cbegin_restrict_vector<16>((bfloat16 *)input); + auto it_out = aie::begin_restrict_vector<16>((bfloat16 *)output); + + for (int i = 0; i < elem_iters; i++) { + input_bf16 = *it_in++; + aie::vector shifted = aie::sub(input_bf16, max_val_vec); + aie::vector exp_val = to_v16bfloat16(getExpBf16(shifted)); + out_vals = aie::mul(exp_val, inv_sum); + *it_out++ = out_vals.to_vector(); + } + + event1(); +} + extern "C" { void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int32_t input_size) @@ -64,6 +157,25 @@ void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int softmax_simple_bf16(input, output, input_size); } +void softmax_partial_init_bf16(bfloat16 *stats) +{ + stats[0] = (bfloat16)(-INFINITY); + stats[1] = (bfloat16)(0.0f); +} + +void softmax_partial_stats_bf16(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +{ + softmax_partial_stats_impl(input, stats, vector_size); +} + +void softmax_partial_norm_bf16(bfloat16 *restrict input, + bfloat16 *restrict output, + bfloat16 *stats, + const int32_t vector_size) +{ + softmax_partial_norm_impl(input, output, stats, vector_size); +} + void mask_bf16(bfloat16 *inout, const int32_t unmasked_size, const int32_t total_size) { for (int32_t i = unmasked_size; i < total_size; i++) { diff --git a/aie_kernels/aie2p/softmax.cc b/aie_kernels/aie2p/softmax.cc index 64cca202..eb775266 100644 --- a/aie_kernels/aie2p/softmax.cc +++ b/aie_kernels/aie2p/softmax.cc @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 #include +#include #include #define SM_VEC_LEN 64 // 32 @@ -30,7 +31,7 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out aie::vector in_elems, exp_val, input_bf16, log2e_vec, max_val_vec; aie::accum out_vals, exp_val_accum, scaled_accum, exp_in_accum; - float max_val = 0; + float max_val = -INFINITY; float accum_exp_val = 0; float running_max = 0; bfloat16 col_sum_inv; @@ -159,6 +160,109 @@ void partial_softmax_alias_bf16(bfloat16 *restrict input_vector, return; } +// --------------------------------------------------------------------------- +// Online (partial / tiled) softmax helpers +// +// These three kernels implement a two-pass online softmax that processes a row +// in sub-tile chunks, keeping running max and sum statistics in a small local +// buffer (`stats`). Layout of the stats buffer (bfloat16[16], only [0..1] +// used): +// stats[0] = running max (scaled by log2e) +// stats[1] = running sum (of exp2(x*log2e - max)) +// --------------------------------------------------------------------------- + +void softmax_partial_stats_impl(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +{ + event0(); + + const int elem_iters = vector_size / SM_VEC_LEN; + + aie::vector input_bf16; + aie::accum scaled_accum, exp_in_accum; + aie::accum exp_val_accum = aie::zeros(); + + aie::vector log2e_vec = aie::broadcast((bfloat16)log2e); + + // --- Phase 1: find local max (scaled by log2e) ------------------------- + float local_max = -INFINITY; + auto it_in1 = aie::cbegin_restrict_vector((bfloat16 *)input); + for (int i = 0; i < elem_iters; i++) { + input_bf16 = *it_in1++; + scaled_accum = aie::mul(input_bf16, log2e_vec); + float chunk_max = aie::reduce_max(scaled_accum.to_vector()); + if (chunk_max > local_max) { + local_max = chunk_max; + } + } + + // --- Phase 2: update running max, rescale running sum ------------------ + float old_max = (float)stats[0]; + float old_sum = (float)stats[1]; + + if (local_max > old_max) { + // New max is larger β€” rescale the old sum by exp2(old_max - new_max) + aie::vector diff_vec = aie::broadcast(old_max - local_max); + aie::vector corr = aie::exp2(diff_vec); + old_sum = old_sum * (float)corr[0]; + old_max = local_max; + } + + // --- Phase 3: accumulate exp2(input * log2e - max) for this chunk ------ + aie::vector max_val_vec = aie::broadcast((bfloat16)old_max); + + auto it_in2 = aie::cbegin_restrict_vector((bfloat16 *)input); + for (int i = 0; i < elem_iters; i++) { + input_bf16 = *it_in2++; + scaled_accum = aie::mul(input_bf16, log2e_vec); + exp_in_accum = aie::sub(scaled_accum, max_val_vec); + aie::vector exp_val = aie::exp2(exp_in_accum.to_vector()); + exp_val_accum = add(exp_val_accum, exp_val); + } + + aie::vector reduce = exp_val_accum.to_vector(); + float local_sum = aie::reduce_add(reduce); + + // --- Phase 4: store updated stats -------------------------------------- + stats[0] = (bfloat16)old_max; + stats[1] = (bfloat16)(old_sum + local_sum); + + event1(); +} + +void softmax_partial_norm_impl(bfloat16 *restrict input, + bfloat16 *restrict output, + bfloat16 *stats, + const int32_t vector_size) +{ + event0(); + + const int elem_iters = vector_size / SM_VEC_LEN; + + float max_val = (float)stats[0]; + float sum_val = (float)stats[1]; + bfloat16 inv_sum = (bfloat16)aie::inv(sum_val); + + aie::vector log2e_vec = aie::broadcast((bfloat16)log2e); + aie::vector max_val_vec = aie::broadcast((bfloat16)max_val); + + aie::vector input_bf16; + aie::accum scaled_accum, exp_in_accum, out_vals; + + auto it_in = aie::cbegin_restrict_vector((bfloat16 *)input); + auto it_out = aie::begin_restrict_vector((bfloat16 *)output); + + for (int i = 0; i < elem_iters; i++) { + input_bf16 = *it_in++; + scaled_accum = aie::mul(input_bf16, log2e_vec); + exp_in_accum = aie::sub(scaled_accum, max_val_vec); + aie::vector exp_val = aie::exp2(exp_in_accum.to_vector()); + out_vals = aie::mul(exp_val, inv_sum); + *it_out++ = out_vals.to_vector(); + } + + event1(); +} + extern "C" { void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int32_t input_size) @@ -177,6 +281,25 @@ void partial_softmax_bf16(bfloat16 *restrict input, partial_softmax_alias_bf16(input, output, scale_buffer, input_size, row_idx, num_rows, scale); } +void softmax_partial_init_bf16(bfloat16 *stats) +{ + stats[0] = (bfloat16)(-INFINITY); + stats[1] = (bfloat16)(0.0f); +} + +void softmax_partial_stats_bf16(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +{ + softmax_partial_stats_impl(input, stats, vector_size); +} + +void softmax_partial_norm_bf16(bfloat16 *restrict input, + bfloat16 *restrict output, + bfloat16 *stats, + const int32_t vector_size) +{ + softmax_partial_norm_impl(input, output, stats, vector_size); +} + void mask_bf16(bfloat16 *inout, const int32 unmasked_size, const int32 total_size) { // TODO: Optimize this to use vector code diff --git a/aie_kernels/generic/axpy.cc b/aie_kernels/generic/axpy.cc index 728adb55..45396977 100644 --- a/aie_kernels/generic/axpy.cc +++ b/aie_kernels/generic/axpy.cc @@ -42,4 +42,107 @@ void saxpy_scalar(bfloat16 *x, bfloat16 *y, const bfloat16 a, bfloat16 *z, const } event1(); } + +// z = a * x (scalar-vector multiply; AXPY without the +y term) +void scale_bf16(bfloat16 *restrict x, bfloat16 *restrict z, const float a, const int32_t vector_size) +{ + event0(); + ::aie::vector a_v = ::aie::broadcast(aie::to_float(a, 0)); + for (int i = 0; i < vector_size; i += 64) { + ::aie::vector x_v = ::aie::load_v<64>(x); + x += 64; + ::aie::accum z_v = ::aie::mul(x_v, a_v); + ::aie::vector z_v_converted = z_v.to_vector(); + ::aie::store_v(z, z_v_converted); + z += 64; + } + event1(); +} + +// z = a + y (scalar-vector add; AXPY without the *x term) +void scalar_add_bf16(bfloat16 *restrict y, bfloat16 *restrict z, const float a, const int32_t vector_size) +{ + event0(); + ::aie::vector a_v = ::aie::broadcast(aie::to_float(a, 0)); + for (int i = 0; i < vector_size; i += 64) { + ::aie::vector y_v = ::aie::load_v<64>(y); + y += 64; + ::aie::vector z_v = ::aie::add(y_v, a_v); + ::aie::store_v(z, z_v); + z += 64; + } + event1(); +} + +// z = (col > row) ? a : y applied per-element of one tile, using a tile +// position (chunk_start_col, row_in_head) supplied via the idx_buffer. +// +// The tile is interpreted as a `vector_size`-wide horizontal strip of the +// per-head (S, S) attention-score block; idx[0] is the strip's starting +// column within that block, idx[1] is the strip's row within the block. +// The kernel implements the causal mask in-place by writing `a` to elements +// strictly above the diagonal and copying y -> z everywhere else. This +// avoids materialising an H*S*S mask buffer entirely. +// +// For tiles whose entire range lies at-or-below the diagonal, the kernel +// degenerates to a copy (input still streamed through DMA β€” slightly +// wasteful but simpler than per-tile data-movement skipping). +void scalar_add_causal_bf16(bfloat16 *restrict y, + bfloat16 *restrict z, + int32_t *idx, + const float a, + const int32_t vector_size) +{ + event0(); + + constexpr int VEC = 64; + + int32_t chunk_start_col = idx[0]; + int32_t row_in_head = idx[1]; + + // Index of the first column in the tile that needs to be masked + // (i.e. column index strictly greater than row_in_head). + int32_t mask_start = row_in_head + 1 - chunk_start_col; + if (mask_start < 0) + mask_start = 0; + if (mask_start > vector_size) + mask_start = vector_size; + + bfloat16 s = (bfloat16)a; + ::aie::vector s_v = ::aie::broadcast(s); + int j = 0; + + // ---- Unmasked region [0, mask_start): copy y -> z ---- + // Vectorised body up to the largest VEC-aligned offset <= mask_start. + int mask_start_floor = (mask_start / VEC) * VEC; + for (; j < mask_start_floor; j += VEC) { + ::aie::vector v = ::aie::load_v(y + j); + ::aie::store_v(z + j, v); + } + // Scalar copy for the unmasked remainder (at most VEC - 1 elements). + for (; j < mask_start; j++) { + z[j] = y[j]; + } + + // ---- Masked region [mask_start, vector_size): write scalar ---- + // If mask_start isn't VEC-aligned, scalar-fill up to the next VEC + // boundary (or to vector_size, whichever is smaller). + int next_vec_boundary = ((j + VEC - 1) / VEC) * VEC; + if (next_vec_boundary > vector_size) + next_vec_boundary = vector_size; + for (; j < next_vec_boundary; j++) { + z[j] = s; + } + // Vectorised body of the masked region. + for (; j + VEC <= vector_size; j += VEC) { + ::aie::store_v(z + j, s_v); + } + // Scalar tail when vector_size isn't VEC-aligned (in practice this + // doesn't fire since per_tile_elements is always a multiple of VEC). + for (; j < vector_size; j++) { + z[j] = s; + } + + event1(); +} } \ No newline at end of file diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 4a6c5604..bb83435b 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -5,7 +5,7 @@ from .elementwise_mul.op import ElementwiseMul from .gemm.op import GEMM from .gemv.op import GEMV -from .mha.op import MHA +from .mha_prefill_df.op import MHA from .rms_norm.op import RMSNorm from .rope.op import RoPE from .silu.op import SiLU diff --git a/iron/operators/axpy/design.py b/iron/operators/axpy/design.py index cde40a63..fe85f12c 100644 --- a/iron/operators/axpy/design.py +++ b/iron/operators/axpy/design.py @@ -4,7 +4,7 @@ from ml_dtypes import bfloat16 import numpy as np -from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker, Buffer from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ @@ -16,7 +16,39 @@ def my_axpy( tile_size, trace_size, scalar_factor, + add_y=True, + mul_x=True, + causal_mask=False, + mask_block_dim=0, + rows_per_block=0, + row_offset=0, + func_prefix="", ): + """AXPY-family element-wise design. + + Modes: + mul_x=True, add_y=True β†’ saxpy: Z = a*X + Y + mul_x=True, add_y=False β†’ scale: Z = a*X + mul_x=False, add_y=True β†’ scalar_add: Z = a + Y + mul_x=False, add_y=True, causal_mask=True β†’ + Z[i,j] = a if (j > i within head) else Y[i,j] + (in-place causal mask; supplies row/col-chunk indices to the kernel + via an idx_buffer; data is interpreted as (..., S, S) blocks where + S = mask_block_dim.) + """ + if causal_mask: + return _my_axpy_causal_mask( + dev=dev, + num_elements=num_elements, + num_columns=num_columns, + tile_size=tile_size, + scalar_factor=scalar_factor, + mask_block_dim=mask_block_dim, + rows_per_block=rows_per_block or mask_block_dim, + row_offset=row_offset, + func_prefix=func_prefix, + ) + factor = scalar_factor per_tile_elements = 4096 if tile_size > 4096 else tile_size n = per_tile_elements * num_columns @@ -28,92 +60,269 @@ def my_axpy( chunk = num_elements // num_columns dtype = bfloat16 - # Define tensor types tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]] tile_ty = np.ndarray[(per_tile_elements,), np.dtype[dtype]] - # AIE-array data movement with object fifos (one per column, not per channel) + # Two inputs only when both *X and +Y are kept (saxpy mode). + has_two_inputs = add_y and mul_x + + # AIE-array data movement with object fifos (one per column) of_in1s = [ObjectFifo(tile_ty, name=f"in1_{i}") for i in range(num_columns)] - of_in2s = [ObjectFifo(tile_ty, name=f"in2_{i}") for i in range(num_columns)] + if has_two_inputs: + of_in2s = [ObjectFifo(tile_ty, name=f"in2_{i}") for i in range(num_columns)] of_outs = [ObjectFifo(tile_ty, name=f"out_{i}") for i in range(num_columns)] # AIE Core Function declaration - axpy_bf16_vector = Kernel( - "saxpy", "axpy.o", [tile_ty, tile_ty, np.float32, tile_ty, np.int32] - ) - - # Define a task that will run on a compute tile - def core_body(of_in1, of_in2, of_out, axpy): - # Number of sub-vector "tile" iterations - for _ in range_(N_div_n): - elem_in1 = of_in1.acquire(1) - elem_in2 = of_in2.acquire(1) - elem_out = of_out.acquire(1) - axpy(elem_in1, elem_in2, factor, elem_out, per_tile_elements) - of_in1.release(1) - of_in2.release(1) - of_out.release(1) - - # Create a worker to run the task on a compute tile (one per column) - my_workers = [ - Worker( - core_body, - [ - of_in1s[i].cons(), - of_in2s[i].cons(), - of_outs[i].prod(), - axpy_bf16_vector, - ], + if has_two_inputs: + kernel = Kernel( + f"{func_prefix}saxpy", + f"{func_prefix}axpy.o", + [tile_ty, tile_ty, np.float32, tile_ty, np.int32], ) - for i in range(num_columns) - ] + elif not add_y: + # z = a * x (drop +Y) + kernel = Kernel( + f"{func_prefix}scale_bf16", + f"{func_prefix}axpy.o", + [tile_ty, tile_ty, np.float32, np.int32], + ) + else: + # z = a + y (drop *X) + kernel = Kernel( + f"{func_prefix}scalar_add_bf16", + f"{func_prefix}axpy.o", + [tile_ty, tile_ty, np.float32, np.int32], + ) + + if has_two_inputs: + + def core_body(of_in1, of_in2, of_out, k): + for _ in range_(N_div_n): + e1 = of_in1.acquire(1) + e2 = of_in2.acquire(1) + eo = of_out.acquire(1) + k(e1, e2, factor, eo, per_tile_elements) + of_in1.release(1) + of_in2.release(1) + of_out.release(1) + + else: + + def core_body(of_in1, of_out, k): + for _ in range_(N_div_n): + e1 = of_in1.acquire(1) + eo = of_out.acquire(1) + k(e1, eo, factor, per_tile_elements) + of_in1.release(1) + of_out.release(1) + + if has_two_inputs: + my_workers = [ + Worker( + core_body, + [of_in1s[i].cons(), of_in2s[i].cons(), of_outs[i].prod(), kernel], + ) + for i in range(num_columns) + ] + else: + my_workers = [ + Worker( + core_body, + [of_in1s[i].cons(), of_outs[i].prod(), kernel], + ) + for i in range(num_columns) + ] - # Create a TensorAccessPattern for each column - # to describe the data movement - # The pattern chops the data in equal chunks - # and moves them in parallel across the columns taps = [ TensorAccessPattern( (1, num_elements), - chunk * i, # Start offset for column i + chunk * i, [1, 1, 1, chunk], [0, 0, 0, 1], ) for i in range(num_columns) ] - # Runtime operations to move data to/from the AIE-array rt = Runtime() - with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C): + sequence_types = ( + (tensor_ty, tensor_ty, tensor_ty) if has_two_inputs else (tensor_ty, tensor_ty) + ) + + with rt.sequence(*sequence_types) as bufs: + if has_two_inputs: + A, B, C = bufs + else: + A, C = bufs + rt.start(*my_workers) - # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete. tg = rt.task_group() - - # Fill the input objectFIFOs with data for i in range(num_columns): - rt.fill( - of_in1s[i].prod(), - A, - taps[i], - task_group=tg, + rt.fill(of_in1s[i].prod(), A, taps[i], task_group=tg) + if has_two_inputs: + rt.fill(of_in2s[i].prod(), B, taps[i], task_group=tg) + for i in range(num_columns): + rt.drain(of_outs[i].cons(), C, taps[i], wait=True, task_group=tg) + rt.finish_task_group(tg) + + return Program(dev, rt).resolve_program() + + +def _my_axpy_causal_mask( + dev, + num_elements, + num_columns, + tile_size, + scalar_factor, + mask_block_dim, + rows_per_block, + row_offset, + func_prefix="", +): + """Single-core in-place causal mask via scalar_add_causal_bf16. + + Walks (blocks Γ— rows_per_block Γ— chunks-per-row) with three nested + runtime loops and feeds (chunk_start_col, row_in_head) to the kernel via + an idx buffer. The kernel applies the scalar `a` to elements strictly + above the per-head diagonal and copies y β†’ z elsewhere. Tiles entirely + below the diagonal still get DMA'd (no per-tile data-movement skip) + but the kernel does only a copy in that case. + + Two operating modes (selected by the rows_per_block / row_offset args): + * Multi-block / full-head (rows_per_block = mask_block_dim, row_offset = 0): + walks ``num_blocks`` whole (S, S) blocks; idx[1] resets to 0 at the + start of each block. + * Sub-block (rows_per_block < mask_block_dim or row_offset > 0): + ``num_blocks`` is typically 1 in the MHA caller; processes a contiguous + ``rows_per_block``-tall slice of one block starting at row_offset. + Used at very long S where one (S, S) block exceeds the BD-length cap. + """ + factor = scalar_factor + S = mask_block_dim + per_tile_elements = 4096 if tile_size > 4096 else tile_size + if S % per_tile_elements != 0: + raise ValueError( + f"mask_block_dim ({S}) must be a multiple of per_tile_elements " + f"({per_tile_elements})" + ) + chunks_per_row = S // per_tile_elements + block_elements = rows_per_block * S + if num_elements % block_elements != 0: + raise ValueError( + f"num_elements ({num_elements}) must be a multiple of " + f"rows_per_block * S ({block_elements})" + ) + num_blocks = num_elements // block_elements + + # Two parallelisation modes: + # * block-aligned (num_blocks >= num_columns): each core handles + # blocks_per_core whole (S, S) blocks; idx[1] resets to row_offset at + # every block boundary (same value on every core). + # * within-block (num_blocks == 1, num_columns > 1): a single block is + # too big to split across cores by block, so each core handles a + # contiguous row-range slice of that one block; per-core init_row is + # row_offset + core_idx * rows_per_iter (different per core). The + # kernel logic is unchanged β€” it only cares about (chunk_start_col, + # row_in_block). + if num_blocks >= num_columns: + if num_blocks % num_columns != 0: + raise ValueError( + f"num_blocks ({num_blocks}) must be a multiple of num_columns " + f"({num_columns}); causal_mask multi-core split is block-aligned" ) - rt.fill( - of_in2s[i].prod(), - B, - taps[i], - task_group=tg, + blocks_per_core = num_blocks // num_columns + rows_per_iter = rows_per_block + per_core_init_rows = [row_offset] * num_columns + else: + if num_blocks != 1: + raise ValueError( + f"causal_mask multi-core within-block split requires " + f"num_blocks == 1, got {num_blocks}" ) - # Drain the output objectFIFOs with data - for i in range(num_columns): - rt.drain( - of_outs[i].cons(), - C, - taps[i], - wait=True, # wait for the transfer to complete and data to be available - task_group=tg, + if rows_per_block % num_columns != 0: + raise ValueError( + f"rows_per_block ({rows_per_block}) must be a multiple of " + f"num_columns ({num_columns}) for within-block split" ) + blocks_per_core = 1 + rows_per_iter = rows_per_block // num_columns + per_core_init_rows = [ + row_offset + i * rows_per_iter for i in range(num_columns) + ] + + elements_per_core = num_elements // num_columns + + dtype = bfloat16 + tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]] + tile_ty = np.ndarray[(per_tile_elements,), np.dtype[dtype]] + idx_ty = np.ndarray[(2,), np.dtype[np.int32]] + + of_ins = [ObjectFifo(tile_ty, name=f"in{i}") for i in range(num_columns)] + of_outs = [ObjectFifo(tile_ty, name=f"out{i}") for i in range(num_columns)] + + kernel = Kernel( + f"{func_prefix}scalar_add_causal_bf16", + f"{func_prefix}axpy.o", + [tile_ty, tile_ty, idx_ty, np.float32, np.int32], + ) + + idx_buffers = [ + Buffer( + initial_value=np.zeros((2,), dtype=np.int32), + name=f"causal_mask_idx_{i}", + ) + for i in range(num_columns) + ] + + # Build one core_body per worker so the per-core init_row can be baked + # into the closure (constant within the worker code). + def make_core_body(my_init_row): + def core_body(of_in_, of_out_, k, idx): + # idx[0] = chunk_start_col within the current row of the block + # idx[1] = current row index within the current block + idx[0] = 0 + idx[1] = my_init_row + for _ in range_(blocks_per_core): + for _ in range_(rows_per_iter): + for _ in range_(chunks_per_row): + elem_in = of_in_.acquire(1) + elem_out = of_out_.acquire(1) + k(elem_in, elem_out, idx, factor, per_tile_elements) + of_in_.release(1) + of_out_.release(1) + idx[0] = idx[0] + per_tile_elements + idx[0] = 0 + idx[1] = idx[1] + 1 + idx[1] = my_init_row # reset for next block + + return core_body + + workers = [ + Worker( + make_core_body(per_core_init_rows[i]), + [of_ins[i].cons(), of_outs[i].prod(), kernel, idx_buffers[i]], + ) + for i in range(num_columns) + ] + + taps = [ + TensorAccessPattern( + (1, num_elements), + i * elements_per_core, + [1, 1, 1, elements_per_core], + [0, 0, 0, 1], + ) + for i in range(num_columns) + ] + + rt = Runtime() + with rt.sequence(tensor_ty, tensor_ty) as (A, C): + rt.start(*workers) + tg = rt.task_group() + for i in range(num_columns): + rt.fill(of_ins[i].prod(), A, taps[i], task_group=tg) + for i in range(num_columns): + rt.drain(of_outs[i].cons(), C, taps[i], wait=True, task_group=tg) rt.finish_task_group(tg) - # Place program components (assign them resources on the device) and generate an MLIR module return Program(dev, rt).resolve_program() diff --git a/iron/operators/axpy/op.py b/iron/operators/axpy/op.py index ff4b87a1..f2f543d3 100644 --- a/iron/operators/axpy/op.py +++ b/iron/operators/axpy/op.py @@ -6,6 +6,7 @@ from iron.common import ( BinaryElementwiseOperator, + AIERuntimeArgSpec, KernelObjectArtifact, SourceArtifact, PythonGeneratedMLIRArtifact, @@ -15,14 +16,113 @@ @dataclass class AXPY(BinaryElementwiseOperator): - """AIE-accelerated aX + Y operator""" + """AIE-accelerated aX + Y operator. + + Optional flags select degenerate variants that skip part of the formula: + + * ``add_y=False`` β†’ ``Z = a * X`` (drop the +Y term and the Y buffer) + * ``mul_x=False`` β†’ ``Z = a + Y`` (drop the *X term and the X buffer) + * ``causal_mask=True`` (requires ``mul_x=False``) β†’ + ``Z[i,j] = a if (j > i within head) else Y[i,j]`` + Treats the data as a sequence of (mask_block_dim, mask_block_dim) blocks + and applies a causal (lower-triangular-keep) mask in-place per block. + Tile-position info is supplied to the kernel via an idx_buffer; tiles + that lie entirely below the diagonal degenerate to a kernel-side copy + (no per-tile DMA skipping). Used by MHA to avoid materialising an + H * S * S causal-mask input buffer. + + ``add_y=False`` and ``mul_x=False`` cannot be combined. + """ scalar_factor: float = 3.0 + add_y: bool = True + mul_x: bool = True + causal_mask: bool = False + mask_block_dim: int | None = None + # Sub-block parameters (causal_mask only). Default: process full + # mask_block_dim rows per (S,S) block starting at row 0. Set + # rows_per_block < mask_block_dim and/or row_offset > 0 to process a + # contiguous row-range slice of one block per invocation β€” useful when + # a single full block's element count exceeds the per-invocation BD + # cap (e.g. S>=32K). + rows_per_block: int | None = None + row_offset: int = 0 kernel_name: ClassVar[str] = "axpy" kernel_fn_name: ClassVar[str] = "saxpy" callback_fn: ClassVar[str] = "my_axpy" + def __post_init__(self) -> None: + if not self.add_y and not self.mul_x: + raise ValueError("AXPY requires at least one of add_y or mul_x to be True") + if self.causal_mask: + if self.mul_x: + raise ValueError( + "AXPY causal_mask=True requires mul_x=False (Z = a + Y form)" + ) + if not self.add_y: + raise ValueError( + "AXPY causal_mask=True requires add_y=True (needs the Y buffer)" + ) + if self.mask_block_dim is None: + raise ValueError( + "AXPY causal_mask=True requires mask_block_dim (the (S,S) block dim)" + ) + # Default rows_per_block = mask_block_dim (process full blocks) + if self.rows_per_block is None: + self.rows_per_block = self.mask_block_dim + if self.rows_per_block <= 0 or self.rows_per_block > self.mask_block_dim: + raise ValueError( + f"rows_per_block ({self.rows_per_block}) must be in (0, " + f"mask_block_dim={self.mask_block_dim}]" + ) + if self.row_offset + self.rows_per_block > self.mask_block_dim: + raise ValueError( + f"row_offset ({self.row_offset}) + rows_per_block " + f"({self.rows_per_block}) must be <= mask_block_dim " + f"({self.mask_block_dim})" + ) + block_elements = self.rows_per_block * self.mask_block_dim + if self.size % block_elements != 0: + raise ValueError( + f"size ({self.size}) must be a multiple of " + f"rows_per_block * mask_block_dim ({block_elements})" + ) + # Multi-core split: either block-aligned (each core handles + # whole blocks; num_aie_columns must divide num_blocks) or + # within-block (num_blocks == 1; num_aie_columns must divide + # rows_per_block). + num_blocks = self.size // block_elements + if num_blocks >= self.num_aie_columns: + if num_blocks % self.num_aie_columns != 0: + raise ValueError( + f"AXPY causal_mask block-aligned split: " + f"num_aie_columns ({self.num_aie_columns}) must " + f"divide num_blocks ({num_blocks})" + ) + else: + if num_blocks != 1: + raise ValueError( + f"AXPY causal_mask within-block split requires " + f"num_blocks == 1, got {num_blocks}" + ) + if self.rows_per_block % self.num_aie_columns != 0: + raise ValueError( + f"AXPY causal_mask within-block split: " + f"rows_per_block ({self.rows_per_block}) must be a " + f"multiple of num_aie_columns ({self.num_aie_columns})" + ) + super().__post_init__() + + def get_arg_spec(self) -> list[AIERuntimeArgSpec]: + # When either input is dropped, the design has one input + one output. + if not self.add_y or not self.mul_x: + return [ + AIERuntimeArgSpec("in", (self.size,)), + AIERuntimeArgSpec("out", (self.size,)), + ] + return super().get_arg_spec() + def get_kernel_artifacts(self) -> list[KernelObjectArtifact]: # axpy.cc lives under aie_kernels/generic/ (not device-specific) return [ @@ -37,7 +137,15 @@ def get_kernel_artifacts(self) -> list[KernelObjectArtifact]: ] def _mlir_callback_args(self): - return super()._mlir_callback_args() + [self.scalar_factor] + return super()._mlir_callback_args() + [ + self.scalar_factor, + self.add_y, + self.mul_x, + self.causal_mask, + self.mask_block_dim if self.mask_block_dim is not None else 0, + self.rows_per_block if self.rows_per_block is not None else 0, + self.row_offset, + ] def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact: return PythonGeneratedMLIRArtifact( diff --git a/iron/operators/mha/design.py b/iron/operators/mha_prefill_df/design.py similarity index 98% rename from iron/operators/mha/design.py rename to iron/operators/mha_prefill_df/design.py index d9ee167c..f4386d42 100644 --- a/iron/operators/mha/design.py +++ b/iron/operators/mha_prefill_df/design.py @@ -1,6 +1,16 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +""" +Data-flow (DF) multi-head attention (MHA) prefill design. + +"Data-flow" (DF): the whole attention computation is a single fused data-flow +design on the AIE array (one kernel / one dispatch), with data streaming +between tiles via object FIFOs. Contrast with the layer-by-layer (LxL) design +in iron/operators/mha_prefill_lxl/, which chains one operator per attention +stage into an OperatorSequence. +""" + import argparse import sys import math diff --git a/iron/operators/mha/op.py b/iron/operators/mha_prefill_df/op.py similarity index 100% rename from iron/operators/mha/op.py rename to iron/operators/mha_prefill_df/op.py diff --git a/iron/operators/mha/reference.py b/iron/operators/mha_prefill_df/reference.py similarity index 100% rename from iron/operators/mha/reference.py rename to iron/operators/mha_prefill_df/reference.py diff --git a/iron/operators/mha/test.py b/iron/operators/mha_prefill_df/test.py similarity index 95% rename from iron/operators/mha/test.py rename to iron/operators/mha_prefill_df/test.py index 29c5fd8a..9f481aa1 100755 --- a/iron/operators/mha/test.py +++ b/iron/operators/mha_prefill_df/test.py @@ -4,8 +4,8 @@ import pytest -from iron.operators.mha.op import MHA -from iron.operators.mha.reference import generate_golden_reference +from iron.operators.mha_prefill_df.op import MHA +from iron.operators.mha_prefill_df.reference import generate_golden_reference from iron.common.test_utils import run_test diff --git a/iron/operators/mha_prefill_lxl/__init__.py b/iron/operators/mha_prefill_lxl/__init__.py new file mode 100644 index 00000000..82f09a67 --- /dev/null +++ b/iron/operators/mha_prefill_lxl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/iron/operators/mha_prefill_lxl/conftest.py b/iron/operators/mha_prefill_lxl/conftest.py new file mode 100644 index 00000000..293c4cbb --- /dev/null +++ b/iron/operators/mha_prefill_lxl/conftest.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "benchmark: benchmark-only tests (select with '-m benchmark')", + ) diff --git a/iron/operators/mha_prefill_lxl/op.py b/iron/operators/mha_prefill_lxl/op.py new file mode 100644 index 00000000..2c49e57c --- /dev/null +++ b/iron/operators/mha_prefill_lxl/op.py @@ -0,0 +1,554 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Layer-by-layer (LxL) multi-head attention (MHA) prefill. + +"Layer-by-layer" (LxL): each attention stage (score GEMM, scale, causal mask, +softmax, context GEMM) is a distinct operator, run back-to-back as an +OperatorSequence (a single dispatch on NPU2) rather than fused into one +monolithic kernel. Contrast with the data-flow (DF) design in +iron/operators/mha_prefill_df/. +""" + +import math + +import aie.utils as aie_utils + +from iron.common.context import AIEContext +from iron.common.sequence import OperatorSequence +from iron.operators.axpy.op import AXPY +from iron.operators.gemm.op import GEMM +from iron.operators.rope.op import RoPE +from iron.operators.strided_copy.op import StridedCopy +from iron.operators.repeat.op import Repeat +from iron.operators.softmax.op import Softmax +from iron.operators.transpose.op import Transpose +from iron.operators.elementwise_add.op import ElementwiseAdd + + +def _pick_tile_n(N, num_cols, max_tile_n=64): + tile_n = N // num_cols + while tile_n > max_tile_n: + tile_n //= 2 + assert N % (tile_n * num_cols) == 0 + return tile_n + + +def _build_core_ops(H, G, d, S, elf_ctx, causal_mask=True, num_cols=None): + """Build core attention sub-ops and runlist (no projections/RoPE/GQA). + + Expects pre-processed inputs: + queries: (H, S, d) deinterleaved, contiguous per head + keys: (H, d, S) transposed and GQA-repeated + values: (H, S, d) GQA-repeated + + Produces: + attn_context: (H, S, d) β€” per-head context vectors + + If causal_mask=False, the elementwise-add masking step is omitted. + """ + if num_cols is None: + num_cols = aie_utils.get_current_device().cols + B = 2 # bytes per bf16 element + + # Cap each GEMM invocation's M dimension and split the per-head computation + # into back-to-back GEMM invocations: the per-GEMM runtime sequence grows + # linearly with M and can OOM the compiler at large S. gemm_M_chunk must + # be a multiple of tile_m * n_aie_rows = 64 and divide S evenly. + gemm_M_chunk = min(S, 4096) + assert ( + S % gemm_M_chunk == 0 + ), f"S ({S}) must be a multiple of gemm_M_chunk ({gemm_M_chunk})" + n_m_chunks = S // gemm_M_chunk + + gemm_scores = GEMM( + M=gemm_M_chunk, + K=d, + N=S, + num_aie_columns=num_cols, + tile_m=16, + tile_k=64, + tile_n=_pick_tile_n(S, num_cols), + context=elf_ctx, + ) + # Scale by 1/sqrt(d) in AXPY scale-only mode: the scalar is baked into the + # kernel call instead of a materialised H*S*S broadcast buffer. + scale = AXPY( + size=H * S * S, + tile_size=S * S // num_cols, + num_aie_columns=num_cols, + scalar_factor=1.0 / math.sqrt(d), + add_y=False, + context=elf_ctx, + ) + if causal_mask: + # Apply the causal mask via AXPY scalar-add + causal-mask mode: the + # kernel writes -INF strictly above the per-head diagonal and copies + # elsewhere, avoiding a materialised H*S*S mask buffer. Each + # invocation's transfer must stay under the compiler's int32 byte + # limit (< 2^30 bf16 elements); large S is split into row-range slices + # (one AXPY instance per row_offset). + MASK_MAX_ELEMENTS_PER_INV = (1 << 30) - 1 + if S * S <= MASK_MAX_ELEMENTS_PER_INV: + # Batch as many whole heads per invocation as divide H. + heads_per_mask_inv = max(1, MASK_MAX_ELEMENTS_PER_INV // (S * S)) + while H % heads_per_mask_inv != 0: + heads_per_mask_inv -= 1 + mask_subblocks = 1 + mask_rows_per_block = S + else: + # Split each head into row-range slices under the limit. + heads_per_mask_inv = 1 + mask_subblocks = ( + S * S + MASK_MAX_ELEMENTS_PER_INV - 1 + ) // MASK_MAX_ELEMENTS_PER_INV + while S % mask_subblocks != 0: + mask_subblocks += 1 + mask_rows_per_block = S // mask_subblocks + assert mask_rows_per_block * S <= MASK_MAX_ELEMENTS_PER_INV + + # Multi-core split: block-aligned across heads when batching whole + # heads, otherwise across the row-range of the single block. + if heads_per_mask_inv >= 2: + mask_num_cols = min(num_cols, heads_per_mask_inv) + while heads_per_mask_inv % mask_num_cols != 0: + mask_num_cols -= 1 + else: + mask_num_cols = num_cols + while mask_rows_per_block % mask_num_cols != 0: + mask_num_cols -= 1 + + # One AXPY instance per row-range slice (exactly one when not split). + mask_ops = [ + AXPY( + size=heads_per_mask_inv * mask_rows_per_block * S, + tile_size=min(4096, S), + num_aie_columns=mask_num_cols, + scalar_factor=float("-inf"), + mul_x=False, + add_y=True, + causal_mask=True, + mask_block_dim=S, + rows_per_block=mask_rows_per_block, + row_offset=sub_idx * mask_rows_per_block, + context=elf_ctx, + ) + for sub_idx in range(mask_subblocks) + ] + # Online/partial softmax once full-row FIFO tiles would exhaust the 64 KB + # AIE local memory (each double-buffered FIFO pair uses 4 * tile_size B). + softmax_chunk_size = 1024 if S >= 8192 else None + + # Split the softmax into back-to-back invocations on disjoint row ranges: + # the compiler lowers the BD length through int32 byte arithmetic and + # overflows when a single transfer reaches 2^30 bf16 elements. + SOFTMAX_MAX_ELEMENTS_PER_INV = (1 << 30) - 1 + total_softmax_rows = H * S + if total_softmax_rows * S <= SOFTMAX_MAX_ELEMENTS_PER_INV: + n_softmax_invocations = 1 + else: + # Smallest n that divides total_softmax_rows and keeps each + # invocation's transfer at or below the limit. + n_softmax_invocations = ( + total_softmax_rows * S + SOFTMAX_MAX_ELEMENTS_PER_INV - 1 + ) // SOFTMAX_MAX_ELEMENTS_PER_INV + while ( + total_softmax_rows % n_softmax_invocations != 0 + or (total_softmax_rows // n_softmax_invocations) * S + > SOFTMAX_MAX_ELEMENTS_PER_INV + ): + n_softmax_invocations += 1 + softmax_rows_per_inv = total_softmax_rows // n_softmax_invocations + assert softmax_rows_per_inv % 16 == 0, ( + f"softmax_rows_per_inv ({softmax_rows_per_inv}) must be a multiple of 16; " + f"got total_rows={total_softmax_rows}, n_invocations={n_softmax_invocations}" + ) + + # Largest divisor of softmax_rows_per_inv that is <= num_cols. + softmax_num_cols = num_cols + while softmax_rows_per_inv % softmax_num_cols != 0: + softmax_num_cols -= 1 + + softmax = Softmax( + rows=softmax_rows_per_inv, + cols=S, + num_aie_columns=softmax_num_cols, + num_channels=1, + rtp_vector_size=S, + chunk_size=softmax_chunk_size, + context=elf_ctx, + ) + # Context GEMM capped at 4 cores: with N=d=64 the matmul constraint + # n % (2*tile_n) == 0 forces tile_n*num_aie_columns = 64, so cols <= 4. + gemm_context = GEMM( + M=gemm_M_chunk, + K=S, + N=d, + num_aie_columns=min(4, num_cols), + tile_m=16, + tile_k=64, + tile_n=_pick_tile_n(d, min(4, num_cols)), + context=elf_ctx, + prio_accuracy=True, + ) + + # Per-head byte sizes + qh = S * d * B # queries per head: (S, d) + kdS = d * S * B # keys per head: (d, S) + kSd = S * d * B # values per head: (S, d) + sh = S * S * B # scores/weights per head: (S, S) + ch = S * d * B # context per head: (S, d) + + # Per-M-chunk byte sizes (the M dimension is contiguous in row-major + # storage so M-slices map directly to byte ranges within each head) + q_chunk = gemm_M_chunk * d * B # queries chunk: (M_chunk, d) + s_chunk = gemm_M_chunk * S * B # scores chunk: (M_chunk, S) + w_chunk = gemm_M_chunk * S * B # weights chunk: (M_chunk, S) + c_chunk = gemm_M_chunk * d * B # context chunk: (M_chunk, d) + + # Single (H, S, S) scratch buffer reused in-place throughout the chain + # (score β†’ scale β†’ [mask] β†’ softmax β†’ context) to halve scratch memory. + attn_buf = "attn" + scores_buf = scaled_buf = masked_buf = weights_buf = attn_buf + + score_calls = [ + ( + gemm_scores, + f"queries[{h*qh + i*q_chunk}:{h*qh + (i+1)*q_chunk}]", + f"keys[{h*kdS}:{(h+1)*kdS}]", + f"{scores_buf}[{h*sh + i*s_chunk}:{h*sh + (i+1)*s_chunk}]", + ) + for h in range(H) + for i in range(n_m_chunks) + ] + + context_calls = [ + ( + gemm_context, + f"{weights_buf}[{h*sh + i*w_chunk}:{h*sh + (i+1)*w_chunk}]", + f"values[{h*kSd}:{(h+1)*kSd}]", + f"attn_context[{h*ch + i*c_chunk}:{h*ch + (i+1)*c_chunk}]", + ) + for h in range(H) + for i in range(n_m_chunks) + ] + + # Build the softmax runlist entries (one per invocation when row-split). + softmax_input_buf = masked_buf if causal_mask else scaled_buf + softmax_chunk_bytes = softmax_rows_per_inv * S * B + if n_softmax_invocations == 1: + softmax_calls = [(softmax, softmax_input_buf, weights_buf)] + else: + softmax_calls = [ + ( + softmax, + f"{softmax_input_buf}[{i*softmax_chunk_bytes}:{(i+1)*softmax_chunk_bytes}]", + f"{weights_buf}[{i*softmax_chunk_bytes}:{(i+1)*softmax_chunk_bytes}]", + ) + for i in range(n_softmax_invocations) + ] + + runlist = [ + *score_calls, + (scale, scores_buf, scaled_buf), + ] + + if causal_mask: + # One (input, output) entry per disjoint slice; mask values are baked + # into the kernel call (scalar -INF), so no mask buffer is needed. + n_head_groups = H // heads_per_mask_inv + head_group_bytes = heads_per_mask_inv * S * S * B + sub_chunk_bytes = mask_rows_per_block * S * B + mask_calls = [] + for g in range(n_head_groups): + for sub_idx in range(mask_subblocks): + start = g * head_group_bytes + sub_idx * sub_chunk_bytes + end = start + heads_per_mask_inv * sub_chunk_bytes + mask_calls.append( + ( + mask_ops[sub_idx], + f"{scaled_buf}[{start}:{end}]", + f"{masked_buf}[{start}:{end}]", + ) + ) + if n_head_groups == 1 and mask_subblocks == 1: + # Whole-buffer fast path (avoids slice notation in MLIR). + runlist += [(mask_ops[0], scaled_buf, masked_buf)] + else: + runlist += mask_calls + + runlist += softmax_calls + runlist += context_calls + + buffer_sizes = { + "queries": H * S * d * B, + "keys": H * d * S * B, + "values": H * S * d * B, + "attn": H * S * S * B, + "attn_context": H * S * d * B, + } + + return runlist, buffer_sizes + + +class AttentionPrefillFused(OperatorSequence): + """Fused attention prefill (core, no projections/RoPE). + + Accepts pre-projected Q (S*H,d), K (S*G,d), V (S*G,d) in interleaved layout. + """ + + def __init__( + self, + num_heads, + num_kv_groups, + head_dim, + embedding_dim, + seq_len, + causal_mask=True, + context=None, + dispatch="auto", + ): + assert head_dim == 64 + assert num_heads % num_kv_groups == 0 + assert seq_len % 256 == 0 + assert (num_heads * seq_len) % 16 == 0 + + self.num_heads = num_heads + self.num_kv_groups = num_kv_groups + self.head_dim = head_dim + self.embedding_dim = embedding_dim + self.seq_len = seq_len + + elf_ctx = context or AIEContext() + runlist, buffer_sizes = _build_core_ops( + num_heads, + num_kv_groups, + head_dim, + seq_len, + elf_ctx, + causal_mask=causal_mask, + ) + + mask_suffix = "_causal" if causal_mask else "_nomask" + input_args = ["queries", "keys", "values"] + + super().__init__( + name=f"attention_prefill_fused_{num_heads}h{num_kv_groups}g{head_dim}d{embedding_dim}e{seq_len}s{mask_suffix}", + runlist=runlist, + input_args=input_args, + output_args=["attn_context"], + buffer_sizes=buffer_sizes, + dispatch=dispatch, + context=elf_ctx, + ) + + +class AttentionPrefillProjectedFused(OperatorSequence): + """Fused attention prefill with Q/K/V projections and RoPE. + + Accepts raw input (S, E) and rope_angles (S, d). + """ + + def __init__( + self, + num_heads, + num_kv_groups, + head_dim, + embedding_dim, + seq_len, + causal_mask=True, + context=None, + dispatch="auto", + ): + assert head_dim == 64 + assert num_heads % num_kv_groups == 0 + assert seq_len % 256 == 0 + assert (num_heads * seq_len) % 16 == 0 + + self.num_heads = num_heads + self.num_kv_groups = num_kv_groups + self.head_dim = head_dim + self.embedding_dim = embedding_dim + self.seq_len = seq_len + self._dispatch_arg = dispatch + + H, G, d, E, S = num_heads, num_kv_groups, head_dim, embedding_dim, seq_len + group_size = H // G + B = 2 + num_cols = aie_utils.get_current_device().cols + + elf_ctx = context or AIEContext() + + # ---- Projection + RoPE ---- + gemm_query = GEMM( + M=S, + K=E, + N=H * d, + num_aie_columns=num_cols, + tile_m=16, + tile_k=64, + tile_n=_pick_tile_n(H * d, num_cols), + context=elf_ctx, + ) + gemm_kv = GEMM( + M=S, + K=E, + N=G * d, + num_aie_columns=num_cols, + tile_m=16, + tile_k=64, + tile_n=_pick_tile_n(G * d, num_cols), + context=elf_ctx, + ) + rope_queries = RoPE(rows=S * H, cols=d, angle_rows=S, context=elf_ctx) + rope_keys = RoPE(rows=S * G, cols=d, angle_rows=S, context=elf_ctx) + + # ---- Deinterleave ---- + deinterleave_q = StridedCopy( + input_sizes=(H, S, d), + input_strides=(d, H * d, 1), + input_offset=0, + output_sizes=(H, S, d), + output_strides=(S * d, d, 1), + output_offset=0, + input_buffer_size=S * H * d, + output_buffer_size=H * S * d, + transfer_size=S * d, + num_aie_channels=1, + context=elf_ctx, + ) + deinterleave_kv = StridedCopy( + input_sizes=(G, S, d), + input_strides=(d, G * d, 1), + input_offset=0, + output_sizes=(G, S, d), + output_strides=(S * d, d, 1), + output_offset=0, + input_buffer_size=S * G * d, + output_buffer_size=G * S * d, + transfer_size=S * d, + num_aie_channels=1, + context=elf_ctx, + ) + + # ---- Transpose keys + GQA repeat ---- + transpose_keys = Transpose( + M=S, + N=d, + num_aie_columns=2, + num_channels=1, + m=256, + n=32, + s=8, + context=elf_ctx, + ) + repeat_kv = Repeat( + rows=G, + cols=d * S, + repeat=group_size, + transfer_size=d, + context=elf_ctx, + ) + + kSd = S * d * B + kdS = d * S * B + + prefix_runlist = [ + (gemm_query, "input", "W_query", "queries_projected"), + (gemm_kv, "input", "W_key", "keys_projected"), + (gemm_kv, "input", "W_value", "values_projected"), + (rope_queries, "queries_projected", "rope_angles", "queries_roped"), + (rope_keys, "keys_projected", "rope_angles", "keys_roped"), + (deinterleave_q, "queries_roped", "queries"), + (deinterleave_kv, "keys_roped", "keys_deint"), + (deinterleave_kv, "values_projected", "values_deint"), + *[ + ( + transpose_keys, + f"keys_deint[{g*kSd}:{(g+1)*kSd}]", + f"keys_transposed[{g*kdS}:{(g+1)*kdS}]", + ) + for g in range(G) + ], + (repeat_kv, "keys_transposed", "keys"), + (repeat_kv, "values_deint", "values"), + ] + prefix_buffer_sizes = { + "queries_projected": S * H * d * B, + "keys_projected": S * G * d * B, + "values_projected": S * G * d * B, + "queries_roped": S * H * d * B, + "keys_roped": S * G * d * B, + "keys_deint": G * S * d * B, + "values_deint": G * S * d * B, + "keys_transposed": G * d * S * B, + } + + core_runlist, core_buffer_sizes = _build_core_ops( + H, + G, + d, + S, + elf_ctx, + causal_mask=causal_mask, + num_cols=num_cols, + ) + + # ---- Reinterleave + output projection ---- + reinterleave = StridedCopy( + input_sizes=(1, 1, 1, H * S * d), + input_strides=(0, 0, 0, 1), + input_offset=0, + output_sizes=(H, 256, S // 256, d), + output_strides=(d, 256 * H * d, H * d, 1), + output_offset=0, + input_buffer_size=H * S * d, + output_buffer_size=S * H * d, + transfer_size=S * d, + num_aie_channels=1, + context=elf_ctx, + ) + gemm_output = GEMM( + M=S, + K=H * d, + N=E, + num_aie_columns=num_cols, + tile_m=16, + tile_k=64, + tile_n=_pick_tile_n(E, num_cols), + context=elf_ctx, + prio_accuracy=True, + ) + + suffix_runlist = [ + (reinterleave, "attn_context", "context_interleaved"), + (gemm_output, "context_interleaved", "W_output", "attn_output"), + ] + suffix_buffer_sizes = { + "context_interleaved": S * H * d * B, + } + + mask_suffix = "_causal" if causal_mask else "_nomask" + input_args = [ + "input", + "rope_angles", + "W_query", + "W_key", + "W_value", + "W_output", + ] + + super().__init__( + name=f"attention_prefill_projected_fused_{H}h{G}g{d}d{E}e{S}s{mask_suffix}", + runlist=prefix_runlist + core_runlist + suffix_runlist, + input_args=input_args, + output_args=["attn_output"], + buffer_sizes={ + **prefix_buffer_sizes, + **core_buffer_sizes, + **suffix_buffer_sizes, + }, + dispatch=dispatch, + context=elf_ctx, + ) diff --git a/iron/operators/mha_prefill_lxl/reference.py b/iron/operators/mha_prefill_lxl/reference.py new file mode 100644 index 00000000..df8a49ae --- /dev/null +++ b/iron/operators/mha_prefill_lxl/reference.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def generate_random_inputs(H, G, d, E, S, causal=True, seed=42): + """Generate just the *inputs* needed for the core MHA attention test, with + no expensive PyTorch reference computation. + + Suitable for the benchmark test (which doesn't verify the full output) and + for very large sequence lengths where the full golden reference is + impractical. The causal mask is built with a single ``torch.triu`` instead + of the H*SΒ² nested-loop construction in :func:`generate_golden_reference`. + + Returned dict matches the input keys consumed by the benchmark test + (``queries_deinterleaved``, ``keys_for_scores``, ``values_for_context``), + plus ``_scale`` (the scalar 1/sqrt(d)) and ``_causal`` (the bool flag) for + use by sample verification. + + Note: the scale factor and causal mask are baked into the operator's + kernels (no host-side buffers), so this function deliberately does *not* + materialise the H*S*S scale/mask tensors β€” at large S those would be tens + of GB each and OOM the machine (e.g. ~24 GB apiece at S=32768, H=12). + """ + torch.manual_seed(seed) + val_range = 0.5 + + # Pre-deinterleaved/transposed/repeated Q, K, V (the layout the + # AttentionPrefillFused operator consumes directly). + queries_deinterleaved = (torch.randn(H, S, d) * val_range).to(torch.bfloat16) + keys_for_scores = (torch.randn(H, d, S) * val_range).to(torch.bfloat16) + values_for_context = (torch.randn(H, S, d) * val_range).to(torch.bfloat16) + + scale = 1.0 / (d**0.5) + + out = { + "queries_deinterleaved": queries_deinterleaved, + "keys_for_scores": keys_for_scores, + "values_for_context": values_for_context, + "_scale": scale, + "_causal": causal, + } + + return out + + +def compute_attn_context_at_rows( + queries_deinterleaved, + keys_for_scores, + values_for_context, + scale, + causal, + sample_hms, +): + """Compute the expected ``attn_context[h, m, :]`` row for each (h, m) in + ``sample_hms``. + + Cheap even at very large S: per sample is O(S * d) β€” a single (1, d) @ (d, S) + matmul plus an O(S) softmax plus an (S,) @ (S, d) reduction. + + Args: + queries_deinterleaved: (H, S, d) bfloat16 tensor + keys_for_scores: (H, d, S) bfloat16 tensor + values_for_context: (H, S, d) bfloat16 tensor + scale: 1 / sqrt(d) (Python float) + causal: bool β€” apply causal mask (zero out k > m) + sample_hms: iterable of (h, m) tuples to compute + + Returns: + dict mapping (h, m) -> torch.Tensor of shape (d,) in bfloat16. + """ + out = {} + for h, m in sample_hms: + q = queries_deinterleaved[h, m, :].float() # (d,) + k = keys_for_scores[h, :, :].float() # (d, S) + scores = q @ k # (S,) + scaled = scores * scale # (S,) + if causal: + # Match the operator's behaviour: positions strictly greater than m + # receive -inf and contribute zero after softmax. + scaled = scaled.clone() + scaled[m + 1 :] = float("-inf") + weights = torch.softmax(scaled, dim=-1) # (S,) + v = values_for_context[h, :, :].float() # (S, d) + out[(h, m)] = (weights @ v).to(torch.bfloat16) # (d,) + return out + + +def _apply_rope_4d(x, angles): + """Apply RoPE to a 4D tensor using interleaved cos/sin angles. + + x: (batch, heads, seq_len, head_dim) + angles: (seq_len, head_dim) with interleaved [cos_0, sin_0, cos_1, sin_1, ...] + Returns: same shape as x with RoPE applied (two-halves method). + """ + half = x.shape[-1] // 2 + cos = angles[:, ::2].unsqueeze(0).unsqueeze(0) # (1, 1, S, half) + sin = angles[:, 1::2].unsqueeze(0).unsqueeze(0) # (1, 1, S, half) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +def _bf16_matmul(a, b): + """(float32 matmul) β†’ bfloat16, matching NPU accumulation.""" + return (a.float() @ b.float()).to(torch.bfloat16) + + +def generate_golden_reference( + num_heads, + num_kv_groups, + head_dim, + embedding_dim, + seq_len, + seed=42, +): + """Generate golden reference for fused attention prefill. + + Parameters: + num_heads (H): number of query attention heads + num_kv_groups (G): number of KV heads (G=H for MHA, G 1: + keys_for_scores = ( + keys_transposed.reshape(G, d * S) + .repeat_interleave(group_size, dim=0) + .reshape(H, d, S) + ) + values_for_context = ( + values_deinterleaved.reshape(G, S * d) + .repeat_interleave(group_size, dim=0) + .reshape(H, S, d) + ) + else: + keys_for_scores = keys_transposed # (H, d, S) + values_for_context = values_deinterleaved # (H, S, d) + + # ---- Score GEMM per head ---- + attn_scores = torch.stack( + [_bf16_matmul(queries_deinterleaved[h], keys_for_scores[h]) for h in range(H)] + ) # (H, S, S) + + # ---- Scale ---- + attn_scores_scaled = (attn_scores.float() * scale).to(torch.bfloat16) + + # ---- Causal mask ---- + attn_scores_masked = ( + attn_scores_scaled.reshape(H * S, S).float() + causal_mask.float() + ).to(torch.bfloat16) + + # ---- Softmax ---- + attn_weights = torch.nn.functional.softmax( + attn_scores_masked.float().reshape(H, S, S), dim=-1 + ).to( + torch.bfloat16 + ) # (H, S, S) + + # ---- Context GEMM per head ---- + attn_context = torch.stack( + [_bf16_matmul(attn_weights[h], values_for_context[h]) for h in range(H)] + ) # (H, S, d) + + # ---- Re-interleave context: (H, S, d) β†’ (S, H*d) ---- + context_interleaved = attn_context.transpose(0, 1).contiguous().reshape(S, H * d) + + # ---- Output projection ---- + attn_output = _bf16_matmul(context_interleaved, W_output) + + return { + "input": x, + "rope_angles": rope_angles, + "W_query": W_query, + "W_key": W_key, + "W_value": W_value, + "W_output": W_output, + "attn_scale_factor": attn_scale_factor, + "causal_mask": causal_mask, + "queries_raw": queries_raw, + "keys_raw": keys_raw, + "values_raw": values_raw, + "queries_roped": queries_roped, + "keys_roped": keys_roped, + "queries_deinterleaved": queries_deinterleaved, + "keys_deinterleaved": keys_deinterleaved, + "keys_transposed": keys_transposed, + "values_deinterleaved": values_deinterleaved, + "keys_for_scores": keys_for_scores, + "values_for_context": values_for_context, + "attn_scores": attn_scores, + "attn_scores_scaled": attn_scores_scaled, + "attn_scores_masked": attn_scores_masked, + "attn_weights": attn_weights, + "attn_context": attn_context, + "context_interleaved": context_interleaved, + "attn_output": attn_output, + } diff --git a/iron/operators/mha_prefill_lxl/test.py b/iron/operators/mha_prefill_lxl/test.py new file mode 100644 index 00000000..5fb73aff --- /dev/null +++ b/iron/operators/mha_prefill_lxl/test.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +import torch +from ml_dtypes import bfloat16 + +from iron.common.test_utils import verify_buffer + +from iron.operators.mha_prefill_lxl.op import ( + AttentionPrefillFused, + AttentionPrefillProjectedFused, +) +from iron.operators.mha_prefill_lxl.reference import ( + generate_golden_reference, + generate_random_inputs, + compute_attn_context_at_rows, +) + +REL_TOL = 0.08 +ABS_TOL = 2.0 +MAX_ERROR_RATE = 0.03 + + +def get_params(): + return [ + pytest.param(2, 2, 64, 256, 256, id="H2"), + pytest.param(32, 8, 64, 2048, 256, id="Llama3.2-256seq"), + pytest.param(12, 12, 64, 768, 256, id="GPT2-Small-256seq"), + ] + + +def get_benchmark_params(): + """GPT-2 Small across sequence lengths 256..32768, with/without causal mask.""" + params = [] + S = 256 + while S <= 32768: + for mask in [True, False]: + for dispatch in ["auto", "separate"]: + tag = "causal" if mask else "nomask" + suffix = f"-{dispatch}" if dispatch != "auto" else "" + params.append( + pytest.param( + 12, + 12, + 64, + 768, + S, + mask, + dispatch, + id=f"GPT2-S{S}-{tag}{suffix}", + ) + ) + S *= 2 + return params + + +def _load_input(fc, name, tensor): + """Load a tensor into a named sub-buffer of the fused callable.""" + np_buf = tensor.contiguous().view(torch.uint16).numpy().view(bfloat16) + fc.get_buffer(name).data[:] = np_buf.flatten() + + +def _get_scratch_tensor(fc, name, shape): + """Read a named buffer from the fused callable's scratch space.""" + fc.scratch_buffer._sync_from_device() + sub = fc.get_buffer(name) + return sub.data[: int(np.prod(shape))].reshape(shape).astype(np.float32) + + +def _get_output_tensor(fc, name, shape): + """Read a named buffer from the fused callable's output space.""" + fc.output_buffer._sync_from_device() + sub = fc.get_buffer(name) + return sub.data[: int(np.prod(shape))].reshape(shape).astype(np.float32) + + +def _verify_output(fc, golden, H, d, S, E): + """Chain-consistent output verification shared by both test variants.""" + npu_context = torch.from_numpy( + _get_scratch_tensor(fc, "context_interleaved", (S, H * d)) + ).bfloat16() + chain_ref = (npu_context.float() @ golden["W_output"].float()).to(torch.bfloat16) + + fc.output_buffer._sync_from_device() + output_np = fc.get_buffer("attn_output").data + output = torch.from_numpy(output_np.reshape(S, E).astype(np.float32)).bfloat16() + + errors = verify_buffer( + output, + "attn_output", + chain_ref.reshape(S, E), + rel_tol=REL_TOL, + abs_tol=ABS_TOL, + max_error_rate=MAX_ERROR_RATE, + ) + assert not errors, f"Output verification failed with {len(errors)} errors" + + +def _core_gemm_flops(H, G, d, E, S): + """Count GEMM FLOPs for the core attention operator.""" + score_flops = H * 2 * S * d * S # H x (S,d)@(d,S) + context_flops = H * 2 * S * S * d # H x (S,S)@(S,d) + return score_flops + context_flops + + +def _projected_gemm_flops(H, G, d, E, S): + """Count GEMM FLOPs for the projected attention operator.""" + query_proj = 2 * S * E * (H * d) # (S,E)@(E,H*d) + kv_proj = 2 * (2 * S * E * (G * d)) # key + value: (S,E)@(E,G*d) each + output_proj = 2 * S * (H * d) * E # (S,H*d)@(H*d,E) + return query_proj + kv_proj + _core_gemm_flops(H, G, d, E, S) + output_proj + + +# --------------------------------------------------------------------------- +# Core attention tests (pre-projected Q, K, V) +# --------------------------------------------------------------------------- + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Throughput=r"Throughput: (?P[\d\.e\+-]+) GFLOP/s", +) +@pytest.mark.parametrize("H,G,d,E,S", get_params()) +def test_mha_prefill_lxl(H, G, d, E, S): + """Core attention: score GEMM -> scale -> mask -> softmax -> context GEMM.""" + golden = generate_golden_reference(H, G, d, E, S) + + op = AttentionPrefillFused(H, G, d, E, S) + op.compile() + fc = op.get_callable() + + _load_input(fc, "queries", golden["queries_deinterleaved"]) + _load_input(fc, "keys", golden["keys_for_scores"]) + _load_input(fc, "values", golden["values_for_context"]) + + fc() + + latency_us = fc.last_elapsed * 1e6 + gflops = _core_gemm_flops(H, G, d, E, S) / (fc.last_elapsed) / 1e9 + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Throughput: {gflops:.6e} GFLOP/s") + + actual = _get_output_tensor(fc, "attn_context", (H, S, d)) + expected = golden["attn_context"].float().numpy().reshape(H, S, d) + errors = verify_buffer( + torch.from_numpy(actual).bfloat16(), + "attn_context", + torch.from_numpy(expected).bfloat16().reshape(H, S, d), + rel_tol=REL_TOL, + abs_tol=ABS_TOL, + max_error_rate=MAX_ERROR_RATE, + ) + assert not errors, f"Output verification failed with {len(errors)} errors" + + +# --------------------------------------------------------------------------- +# Projected attention tests (with Q/K/V projections + RoPE) +# --------------------------------------------------------------------------- + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Throughput=r"Throughput: (?P[\d\.e\+-]+) GFLOP/s", +) +@pytest.mark.parametrize("H,G,d,E,S", get_params()) +def test_attention_prefill_projected_fused(H, G, d, E, S): + """Projected attention: Q/K/V proj -> RoPE -> GQA -> attention -> output proj.""" + golden = generate_golden_reference(H, G, d, E, S) + + op = AttentionPrefillProjectedFused(H, G, d, E, S) + op.compile() + fc = op.get_callable() + + _load_input(fc, "input", golden["input"]) + _load_input(fc, "rope_angles", golden["rope_angles"]) + _load_input(fc, "W_query", golden["W_query"]) + _load_input(fc, "W_key", golden["W_key"]) + _load_input(fc, "W_value", golden["W_value"]) + _load_input(fc, "W_output", golden["W_output"]) + + fc() + + latency_us = fc.last_elapsed * 1e6 + gflops = _projected_gemm_flops(H, G, d, E, S) / (fc.last_elapsed) / 1e9 + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Throughput: {gflops:.6e} GFLOP/s") + + _verify_output(fc, golden, H, d, S, E) + + +# --------------------------------------------------------------------------- +# Benchmark: GPT-2 Small core MHA across sequence lengths, +/- causal mask +# --------------------------------------------------------------------------- + + +@pytest.mark.benchmark +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Throughput=r"Throughput: (?P[\d\.e\+-]+) GFLOP/s", +) +@pytest.mark.parametrize("H,G,d,E,S,causal,dispatch", get_benchmark_params()) +def test_mha_prefill_benchmark(H, G, d, E, S, causal, dispatch): + """Benchmark core MHA for GPT-2 Small across sequence lengths. + + Uses cheap random inputs (no full PyTorch reference) and verifies + correctness by recomputing the expected ``attn_context`` row for a small + number of randomly-chosen (head, row) positions β€” feasible at any S. + """ + inputs = generate_random_inputs(H, G, d, E, S, causal=causal) + + op = AttentionPrefillFused(H, G, d, E, S, causal_mask=causal, dispatch=dispatch) + op.compile() + fc = op.get_callable() + + _load_input(fc, "queries", inputs["queries_deinterleaved"]) + _load_input(fc, "keys", inputs["keys_for_scores"]) + _load_input(fc, "values", inputs["values_for_context"]) + + fc() + + latency_us = fc.last_elapsed * 1e6 + gflops = _core_gemm_flops(H, G, d, E, S) / (fc.last_elapsed) / 1e9 + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Throughput: {gflops:.6e} GFLOP/s") + + # ---- Sample-based correctness check ---- + # Pick a handful of random (head, row) pairs and recompute the expected + # attn_context row for each (cheap: O(S*d) per sample). + actual_context = _get_output_tensor(fc, "attn_context", (H, S, d)) + rng = np.random.default_rng(seed=0) + n_samples = 16 + sample_hms = [ + (int(rng.integers(0, H)), int(rng.integers(0, S))) for _ in range(n_samples) + ] + expected_rows = compute_attn_context_at_rows( + inputs["queries_deinterleaved"], + inputs["keys_for_scores"], + inputs["values_for_context"], + inputs["_scale"], + causal, + sample_hms, + ) + failures = [] + for (h, m), exp in expected_rows.items(): + act = torch.from_numpy(actual_context[h, m, :]).bfloat16() + diff = (act.float() - exp.float()).abs() + rel = diff / (exp.float().abs() + 1e-6) + # An element fails only if it exceeds BOTH abs_tol and rel_tol + bad = (diff > ABS_TOL) & (rel > REL_TOL) + if bad.any(): + failures.append( + f"(h={h}, m={m}): {int(bad.sum())}/{d} bad, " + f"max_abs={diff.max().item():.4f}, max_rel={rel.max().item():.4f}" + ) + assert not failures, "Sample verification failed:\n " + "\n ".join(failures) + + +# --------------------------------------------------------------------------- +# Intermediate checks (extensive, not run by default) +# --------------------------------------------------------------------------- + +INTERMEDIATE_CHECKS = [ + ("attn_scores", "attn_scores", lambda H, G, S, d: (H, S, S), "scratch"), + ( + "attn_scores_masked", + "attn_scores_masked", + lambda H, G, S, d: (H, S, S), + "scratch", + ), + ("attn_weights", "attn_weights", lambda H, G, S, d: (H, S, S), "scratch"), + ("attn_context", "attn_context", lambda H, G, S, d: (H, S, d), "output"), +] + + +@pytest.mark.extensive +@pytest.mark.parametrize("H,G,d,E,S", get_params()) +def test_mha_prefill_lxl_intermediates(H, G, d, E, S): + """Check intermediate buffers of core attention (for debugging).""" + golden = generate_golden_reference(H, G, d, E, S) + + op = AttentionPrefillFused(H, G, d, E, S) + op.compile() + fc = op.get_callable() + + _load_input(fc, "queries", golden["queries_deinterleaved"]) + _load_input(fc, "keys", golden["keys_for_scores"]) + _load_input(fc, "values", golden["values_for_context"]) + + fc() + + for buf_name, golden_key, shape_fn, buf_type in INTERMEDIATE_CHECKS: + shape = shape_fn(H, G, S, d) + if buf_type == "output": + actual = _get_output_tensor(fc, buf_name, shape) + else: + actual = _get_scratch_tensor(fc, buf_name, shape) + expected = golden[golden_key].float().numpy().reshape(shape) + diff = np.abs(actual - expected) + print( + f" [{buf_name}] shape={shape} " + f"nan={int(np.isnan(actual).sum())} " + f"max_abs_err={diff.max():.4f} mean_abs_err={diff.mean():.6f}" + ) diff --git a/iron/operators/softmax/design.py b/iron/operators/softmax/design.py index 6492b9c7..1004227d 100644 --- a/iron/operators/softmax/design.py +++ b/iron/operators/softmax/design.py @@ -20,6 +20,198 @@ from ml_dtypes import bfloat16 +def _softmax_partial( + dev, + num_elements, + num_aie_columns, + num_channels, + tile_size, + chunk_size, + func_prefix="", + kernel_obj_file="softmax.o", +): + """Online / tiled softmax that processes each row in sub-tile chunks. + + Each row of *tile_size* elements is processed in two passes: + 1. Stats pass – reads chunks, accumulates running max and sum(exp). + 2. Norm pass – reads the same chunks again, writes exp(x-max)/sum. + + Two separate input ObjectFifos are used so that the DMA can feed each pass + independently from the same DDR source buffer. + """ + total_cores = num_aie_columns * num_channels + per_core_elements = num_elements // total_cores + if num_elements % total_cores != 0: + raise ValueError( + f"Number of elements ({num_elements}) must be a multiple of {total_cores}." + ) + + rows_per_core = per_core_elements // tile_size + chunks_per_row = tile_size // chunk_size + dtype = bfloat16 + + # Tensor / tile types + tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]] + chunk_ty = np.ndarray[(chunk_size,), np.dtype[dtype]] + stats_ty = np.ndarray[(16,), np.dtype[dtype]] # only [0..1] used + + chunk = num_elements // num_aie_columns // num_channels + + # --- Object FIFOs ------------------------------------------------------- + of_in_stats = [ + ObjectFifo(chunk_ty, name=f"in_stats_{i}_{j}") + for i in range(num_aie_columns) + for j in range(num_channels) + ] + of_in_norm = [ + ObjectFifo(chunk_ty, name=f"in_norm_{i}_{j}") + for i in range(num_aie_columns) + for j in range(num_channels) + ] + of_outs = [ + ObjectFifo(chunk_ty, name=f"out_{i}_{j}") + for i in range(num_aie_columns) + for j in range(num_channels) + ] + + # --- Kernel declarations ------------------------------------------------ + init_kernel = Kernel( + f"{func_prefix}softmax_partial_init_bf16", + f"{func_prefix}{kernel_obj_file}", + [stats_ty], + ) + stats_kernel = Kernel( + f"{func_prefix}softmax_partial_stats_bf16", + f"{func_prefix}{kernel_obj_file}", + [chunk_ty, stats_ty, np.int32], + ) + norm_kernel = Kernel( + f"{func_prefix}softmax_partial_norm_bf16", + f"{func_prefix}{kernel_obj_file}", + [chunk_ty, chunk_ty, stats_ty, np.int32], + ) + + # --- Local stats buffers (one per core) --------------------------------- + stats_buffers = [ + Buffer( + initial_value=np.zeros(16, dtype=dtype), + name=f"stats_{i}_{j}", + ) + for i in range(num_aie_columns) + for j in range(num_channels) + ] + + barriers = [ + WorkerRuntimeBarrier() + for i in range(num_aie_columns) + for j in range(num_channels) + ] + + # --- Worker body -------------------------------------------------------- + def core_body( + of_s, + of_n, + of_out, + init_k, + stats_k, + norm_k, + stats_buf, + barrier, + ): + barrier.wait_for_value(1) + for _ in range_(rows_per_core): + # Reset running max / sum for the new row + init_k(stats_buf) + + # Pass 1 – accumulate max and sum(exp) + for _ in range_(chunks_per_row): + elem = of_s.acquire(1) + stats_k(elem, stats_buf, chunk_size) + of_s.release(1) + + # Pass 2 – normalise: exp(x - max) / sum + for _ in range_(chunks_per_row): + elem_in = of_n.acquire(1) + elem_out = of_out.acquire(1) + norm_k(elem_in, elem_out, stats_buf, chunk_size) + of_n.release(1) + of_out.release(1) + + # --- Workers ------------------------------------------------------------ + def _worker_args(k): + return [ + of_in_stats[k].cons(), + of_in_norm[k].cons(), + of_outs[k].prod(), + init_kernel, + stats_kernel, + norm_kernel, + stats_buffers[k], + barriers[k], + ] + + workers = [ + Worker(core_body, _worker_args(i * num_channels + j), stack_size=0xD00) + for i in range(num_aie_columns) + for j in range(num_channels) + ] + + # --- Tensor access patterns (identical for both input FIFOs) ------------ + taps = [ + TensorAccessPattern( + (1, num_elements), + chunk * i * num_channels + chunk * j, + [1, 1, 1, chunk], + [0, 0, 0, 1], + ) + for i in range(num_aie_columns) + for j in range(num_channels) + ] + + # --- Runtime sequence --------------------------------------------------- + rt = Runtime() + with rt.sequence(tensor_ty, tensor_ty) as (A, C): + rt.start(*workers) + + for k in range(num_aie_columns * num_channels): + rt.set_barrier(barriers[k], 1) + + tg = rt.task_group() + + for i in range(num_aie_columns): + for j in range(num_channels): + k = i * num_channels + j + # Feed the stats-pass FIFO + rt.fill( + of_in_stats[k].prod(), + A, + taps[k], + task_group=tg, + ) + # Feed the norm-pass FIFO (same source data) + rt.fill( + of_in_norm[k].prod(), + A, + taps[k], + task_group=tg, + ) + + for i in range(num_aie_columns): + for j in range(num_channels): + k = i * num_channels + j + rt.drain( + of_outs[k].cons(), + C, + taps[k], + wait=True, + task_group=tg, + ) + + rt.finish_task_group(tg) + + return Program(dev, rt).resolve_program() + + def softmax( dev, num_elements, @@ -31,7 +223,22 @@ def softmax( vector_size_parameter=None, func_prefix="", kernel_obj_file="softmax.o", + chunk_size=None, ): + # ---- Partial (online) softmax path ---- + if chunk_size is not None: + return _softmax_partial( + dev, + num_elements, + num_aie_columns, + num_channels, + tile_size, + chunk_size, + func_prefix, + kernel_obj_file, + ) + + # ---- Full-row softmax path (original) ---- per_tile_elements = tile_size if rtp_vector_size is None: rtp_vector_size = per_tile_elements diff --git a/iron/operators/softmax/op.py b/iron/operators/softmax/op.py index 9da4564f..d9eb2476 100644 --- a/iron/operators/softmax/op.py +++ b/iron/operators/softmax/op.py @@ -20,7 +20,12 @@ @dataclass class Softmax(MLIROperator): - """AIE-accelerated Softmax operation""" + """AIE-accelerated Softmax operation + + When *chunk_size* is set (and < cols), uses an online / tiled softmax + that processes each row in two passes with sub-tile chunks, avoiding the + local-memory exhaustion that occurs with very long rows (e.g. S >= 8192). + """ rows: int cols: int @@ -28,6 +33,7 @@ class Softmax(MLIROperator): num_channels: int = 1 rtp_vector_size: int | None = None vector_size_parameter: str | None = None + chunk_size: int | None = None context: object = field(default=None, repr=False) @property @@ -43,6 +49,15 @@ def __post_init__(self): raise ValueError( f"rows ({self.rows}) must be a multiple of num_aie_columns ({self.num_aie_columns})" ) + if self.chunk_size is not None: + if self.cols % self.chunk_size != 0: + raise ValueError( + f"cols ({self.cols}) must be a multiple of chunk_size ({self.chunk_size})" + ) + if self.chunk_size % 64 != 0: + raise ValueError( + f"chunk_size ({self.chunk_size}) must be a multiple of 64" + ) MLIROperator.__init__(self, context=self.context) @property @@ -69,6 +84,7 @@ def get_mlir_artifact(self): "rtp_vector_size": self.rtp_vector_size, "vector_size_parameter": self.vector_size_parameter, "kernel_obj_file": self._kernel_link_file, + "chunk_size": self.chunk_size, }, ), ) diff --git a/iron/operators/softmax/test.py b/iron/operators/softmax/test.py index 066d2309..8cca9f69 100755 --- a/iron/operators/softmax/test.py +++ b/iron/operators/softmax/test.py @@ -85,3 +85,52 @@ def test_softmax(input_length, num_aie_columns, num_channels, tile_size, aie_con print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") assert not errors, f"Test failed with errors: {errors}" + + +# --------------------------------------------------------------------------- +# Partial (online / tiled) softmax tests β€” enables long rows (S >= 8192) +# --------------------------------------------------------------------------- + + +def get_partial_params(): + """GPT-2 style: 12 heads Γ— S rows of S columns, tested via partial softmax.""" + params = [] + for S in [8192]: + H = 12 # GPT-2 Small heads + rows = H * S + cols = S + chunk_size = 1024 + params.append(pytest.param(rows, cols, chunk_size, id=f"GPT2-S{S}")) + return params + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", +) +@pytest.mark.parametrize("rows,cols,chunk_size", get_partial_params()) +def test_softmax_partial(rows, cols, chunk_size, aie_context): + """Test partial / online softmax with sub-tile chunks for long rows.""" + + golden_ref = generate_golden_reference(rows=rows, cols=cols) + + operator = Softmax( + rows=rows, + cols=cols, + num_aie_columns=1, + num_channels=1, + chunk_size=chunk_size, + context=aie_context, + ) + + input_buffers = {"in": golden_ref["input"]} + output_buffers = {"output": golden_ref["output"]} + + errors, latency_us, bandwidth_gbps = run_test( + operator, input_buffers, output_buffers, rel_tol=0.08, abs_tol=1e-6 + ) + + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") + + assert not errors, f"Test failed with errors: {errors}" From 3a476a428fa7419906637b528117246270d0bb70 Mon Sep 17 00:00:00 2001 From: andrej Date: Mon, 13 Jul 2026 12:33:44 -0600 Subject: [PATCH 2/4] Fix CI benchmark marker + address PR #91 review comments CI: register 'benchmark' pytest marker and fold benchmark tests into the 'extensive' suite (conftest adds extensive marker before -m deselection), so they no longer run in the small suite where large-S benchmarks OOM'd the process and left tests_latest.csv unwritten. Kernels: - softmax (aie2/aie2p): stats buffer -> softmax_stats struct with *restrict; drop '// ---' comment banners to match file style. - softmax (aie2p): rewrite partial-stats to the single-pass online algorithm used in aie2; document how partial_softmax_alias differs from the online pair. - axpy: collapse saxpy/scale/scalar_add into one templated saxpy_family with extern-C specializations; generalize the causal mask into a templated triangular_fill_impl and update comments (general triangular fill). sequence.py: clarify 'single dispatch' terminology in dispatch docstrings. --- aie_kernels/aie2/softmax.cc | 46 ++++----- aie_kernels/aie2p/softmax.cc | 104 ++++++++++--------- aie_kernels/generic/axpy.cc | 187 ++++++++++++++++++++--------------- conftest.py | 6 ++ iron/common/sequence.py | 14 ++- pytest.ini | 1 + 6 files changed, 203 insertions(+), 155 deletions(-) diff --git a/aie_kernels/aie2/softmax.cc b/aie_kernels/aie2/softmax.cc index 07aa4bb1..a331777e 100644 --- a/aie_kernels/aie2/softmax.cc +++ b/aie_kernels/aie2/softmax.cc @@ -58,25 +58,25 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out return; } -// --------------------------------------------------------------------------- -// Online (partial / tiled) softmax helpers +// Online (partial / tiled) softmax helpers. // -// These three kernels implement a two-pass online softmax that processes a row -// in sub-tile chunks, keeping running max and sum statistics in a small local -// buffer (`stats`). Layout of the stats buffer (bfloat16[16], only [0..1] -// used): -// stats[0] = running max -// stats[1] = running sum (of exp(x - max)) -// --------------------------------------------------------------------------- - -void softmax_partial_stats_impl(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +// These kernels implement an online softmax that processes a row in sub-tile +// chunks, keeping running max and sum statistics in a small per-core buffer. +// softmax_stats names the two stats slots instead of using hard-coded array +// indices. It occupies the first two bfloat16 elements of the stats buffer. +struct softmax_stats { + bfloat16 max; // running max + bfloat16 sum; // running sum of exp(x - max) +}; + +void softmax_partial_stats_impl(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size) { event0(); const int elem_iters = vector_size / 16; - float running_max = (float)stats[0]; - float running_sum = (float)stats[1]; + float running_max = (float)stats->max; + float running_sum = (float)stats->sum; aie::vector input_bf16; aie::accum exp_val_accum = aie::zeros(); @@ -112,23 +112,23 @@ void softmax_partial_stats_impl(bfloat16 *restrict input, bfloat16 *stats, const aie::vector reduce = exp_val_accum.to_vector(); running_sum += aie::reduce_add(reduce); - stats[0] = (bfloat16)running_max; - stats[1] = (bfloat16)running_sum; + stats->max = (bfloat16)running_max; + stats->sum = (bfloat16)running_sum; event1(); } void softmax_partial_norm_impl(bfloat16 *restrict input, bfloat16 *restrict output, - bfloat16 *stats, + softmax_stats *restrict stats, const int32_t vector_size) { event0(); const int elem_iters = vector_size / 16; - float max_val = (float)stats[0]; - float sum_val = (float)stats[1]; + float max_val = (float)stats->max; + float sum_val = (float)stats->sum; bfloat16 inv_sum = (bfloat16)aie::inv(sum_val); aie::vector max_val_vec = aie::broadcast((bfloat16)max_val); @@ -157,20 +157,20 @@ void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int softmax_simple_bf16(input, output, input_size); } -void softmax_partial_init_bf16(bfloat16 *stats) +void softmax_partial_init_bf16(softmax_stats *restrict stats) { - stats[0] = (bfloat16)(-INFINITY); - stats[1] = (bfloat16)(0.0f); + stats->max = (bfloat16)(-INFINITY); + stats->sum = (bfloat16)(0.0f); } -void softmax_partial_stats_bf16(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +void softmax_partial_stats_bf16(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size) { softmax_partial_stats_impl(input, stats, vector_size); } void softmax_partial_norm_bf16(bfloat16 *restrict input, bfloat16 *restrict output, - bfloat16 *stats, + softmax_stats *restrict stats, const int32_t vector_size) { softmax_partial_norm_impl(input, output, stats, vector_size); diff --git a/aie_kernels/aie2p/softmax.cc b/aie_kernels/aie2p/softmax.cc index eb775266..d9eecbde 100644 --- a/aie_kernels/aie2p/softmax.cc +++ b/aie_kernels/aie2p/softmax.cc @@ -81,6 +81,17 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out return; } +// partial_softmax_alias_bf16 is a flash-attention style single-shot softmax +// used by the projected-fused path. It shares the same three-pass structure as +// softmax_simple_bf16 (find max, exp, normalize) but differs in two ways that +// make code sharing awkward: (1) it folds the query scale into the log2e +// multiply instead of using the fixed log2e constant, and (2) it maintains +// per-row running max/sum in an externally-supplied scale_buffer (flash +// accumulation across key tiles) rather than reducing a whole row locally. The +// softmax_partial_stats_impl / softmax_partial_norm_impl pair below implement a +// different (chunked, two-call) online softmax that keeps its running stats in +// a compact softmax_stats buffer; those are used by the standalone softmax +// operator, not by this projected-fused kernel. void partial_softmax_alias_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict output_vector, bfloat16 *restrict scale_buffer, @@ -160,86 +171,81 @@ void partial_softmax_alias_bf16(bfloat16 *restrict input_vector, return; } -// --------------------------------------------------------------------------- -// Online (partial / tiled) softmax helpers +// Online (partial / tiled) softmax helpers. // -// These three kernels implement a two-pass online softmax that processes a row -// in sub-tile chunks, keeping running max and sum statistics in a small local -// buffer (`stats`). Layout of the stats buffer (bfloat16[16], only [0..1] -// used): -// stats[0] = running max (scaled by log2e) -// stats[1] = running sum (of exp2(x*log2e - max)) -// --------------------------------------------------------------------------- - -void softmax_partial_stats_impl(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +// These kernels implement an online softmax that processes a row in sub-tile +// chunks, keeping running max and sum statistics in a small per-core buffer. +// The max is stored scaled by log2e and the sum accumulates exp2(x*log2e - +// max), matching the exp2-based normalization used below. softmax_stats names +// the two stats slots instead of using hard-coded array indices; it occupies +// the first two bfloat16 elements of the stats buffer. +struct softmax_stats { + bfloat16 max; // running max (scaled by log2e) + bfloat16 sum; // running sum of exp2(x*log2e - max) +}; + +void softmax_partial_stats_impl(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size) { event0(); const int elem_iters = vector_size / SM_VEC_LEN; + float running_max = (float)stats->max; + float running_sum = (float)stats->sum; + aie::vector input_bf16; aie::accum scaled_accum, exp_in_accum; aie::accum exp_val_accum = aie::zeros(); aie::vector log2e_vec = aie::broadcast((bfloat16)log2e); - // --- Phase 1: find local max (scaled by log2e) ------------------------- - float local_max = -INFINITY; - auto it_in1 = aie::cbegin_restrict_vector((bfloat16 *)input); + auto it_in = aie::cbegin_restrict_vector((bfloat16 *)input); + + // Single-pass online algorithm (matches aie_kernels/aie2/softmax.cc): for + // each vector chunk, update the running max if needed -- rescaling the + // partial accumulator and running sum by exp2(old_max - new_max) -- then + // accumulate exp2(x*log2e - max). for (int i = 0; i < elem_iters; i++) { - input_bf16 = *it_in1++; + input_bf16 = *it_in++; scaled_accum = aie::mul(input_bf16, log2e_vec); float chunk_max = aie::reduce_max(scaled_accum.to_vector()); - if (chunk_max > local_max) { - local_max = chunk_max; - } - } - // --- Phase 2: update running max, rescale running sum ------------------ - float old_max = (float)stats[0]; - float old_sum = (float)stats[1]; - - if (local_max > old_max) { - // New max is larger β€” rescale the old sum by exp2(old_max - new_max) - aie::vector diff_vec = aie::broadcast(old_max - local_max); - aie::vector corr = aie::exp2(diff_vec); - old_sum = old_sum * (float)corr[0]; - old_max = local_max; - } - - // --- Phase 3: accumulate exp2(input * log2e - max) for this chunk ------ - aie::vector max_val_vec = aie::broadcast((bfloat16)old_max); + if (chunk_max > running_max) { + aie::vector diff_vec = aie::broadcast(running_max - chunk_max); + aie::vector corr = aie::exp2(diff_vec); + float scale = (float)corr[0]; + aie::vector scale_vec = aie::broadcast((bfloat16)scale); + exp_val_accum = aie::mul(exp_val_accum.to_vector(), scale_vec); + running_sum *= scale; + running_max = chunk_max; + } - auto it_in2 = aie::cbegin_restrict_vector((bfloat16 *)input); - for (int i = 0; i < elem_iters; i++) { - input_bf16 = *it_in2++; - scaled_accum = aie::mul(input_bf16, log2e_vec); + aie::vector max_val_vec = aie::broadcast((bfloat16)running_max); exp_in_accum = aie::sub(scaled_accum, max_val_vec); aie::vector exp_val = aie::exp2(exp_in_accum.to_vector()); exp_val_accum = add(exp_val_accum, exp_val); } aie::vector reduce = exp_val_accum.to_vector(); - float local_sum = aie::reduce_add(reduce); + running_sum += aie::reduce_add(reduce); - // --- Phase 4: store updated stats -------------------------------------- - stats[0] = (bfloat16)old_max; - stats[1] = (bfloat16)(old_sum + local_sum); + stats->max = (bfloat16)running_max; + stats->sum = (bfloat16)running_sum; event1(); } void softmax_partial_norm_impl(bfloat16 *restrict input, bfloat16 *restrict output, - bfloat16 *stats, + softmax_stats *restrict stats, const int32_t vector_size) { event0(); const int elem_iters = vector_size / SM_VEC_LEN; - float max_val = (float)stats[0]; - float sum_val = (float)stats[1]; + float max_val = (float)stats->max; + float sum_val = (float)stats->sum; bfloat16 inv_sum = (bfloat16)aie::inv(sum_val); aie::vector log2e_vec = aie::broadcast((bfloat16)log2e); @@ -281,20 +287,20 @@ void partial_softmax_bf16(bfloat16 *restrict input, partial_softmax_alias_bf16(input, output, scale_buffer, input_size, row_idx, num_rows, scale); } -void softmax_partial_init_bf16(bfloat16 *stats) +void softmax_partial_init_bf16(softmax_stats *restrict stats) { - stats[0] = (bfloat16)(-INFINITY); - stats[1] = (bfloat16)(0.0f); + stats->max = (bfloat16)(-INFINITY); + stats->sum = (bfloat16)(0.0f); } -void softmax_partial_stats_bf16(bfloat16 *restrict input, bfloat16 *stats, const int32_t vector_size) +void softmax_partial_stats_bf16(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size) { softmax_partial_stats_impl(input, stats, vector_size); } void softmax_partial_norm_bf16(bfloat16 *restrict input, bfloat16 *restrict output, - bfloat16 *stats, + softmax_stats *restrict stats, const int32_t vector_size) { softmax_partial_norm_impl(input, output, stats, vector_size); diff --git a/aie_kernels/generic/axpy.cc b/aie_kernels/generic/axpy.cc index 45396977..fc16c576 100644 --- a/aie_kernels/generic/axpy.cc +++ b/aie_kernels/generic/axpy.cc @@ -12,86 +12,61 @@ #include -extern "C" { -void saxpy(bfloat16 *restrict x, bfloat16 *restrict y, const float a, bfloat16 *restrict z, const int32_t vector_size) -{ - event0(); - ::aie::vector a_v = - ::aie::broadcast(aie::to_float(a, 0)); // Convert to bfloat16 - // #pragma clang loop min_iteration_count(4) - for (int i = 0; i < vector_size; i += 64) { - ::aie::vector x_v = ::aie::load_v<64>(x); - x += 64; - ::aie::vector y_v = ::aie::load_v<64>(y); - y += 64; - ::aie::accum ax_v = ::aie::mul(x_v, a_v); - ::aie::accum z_v = ::aie::add(ax_v, y_v); - ::aie::vector z_v_converted = z_v.to_vector(); - ::aie::store_v(z, z_v_converted); - z += 64; - } - event1(); -} - -void saxpy_scalar(bfloat16 *x, bfloat16 *y, const bfloat16 a, bfloat16 *z, const int32_t vector_size) -{ - event0(); - float a_f = a; - for (int i = 0; i < vector_size; ++i) { - z[i] = a_f * x[i] + y[i]; - } - event1(); -} - -// z = a * x (scalar-vector multiply; AXPY without the +y term) -void scale_bf16(bfloat16 *restrict x, bfloat16 *restrict z, const float a, const int32_t vector_size) -{ - event0(); - ::aie::vector a_v = ::aie::broadcast(aie::to_float(a, 0)); - for (int i = 0; i < vector_size; i += 64) { - ::aie::vector x_v = ::aie::load_v<64>(x); - x += 64; - ::aie::accum z_v = ::aie::mul(x_v, a_v); - ::aie::vector z_v_converted = z_v.to_vector(); - ::aie::store_v(z, z_v_converted); - z += 64; - } - event1(); -} - -// z = a + y (scalar-vector add; AXPY without the *x term) -void scalar_add_bf16(bfloat16 *restrict y, bfloat16 *restrict z, const float a, const int32_t vector_size) +// One vectorised AXPY-family kernel. The compile-time flags select which terms +// are emitted so the public entry points share a single loop body: +// Mul && Add : z = a * x + y (saxpy) +// Mul : z = a * x (scale_bf16) +// Add : z = a + y (scalar_add_bf16) +template +static inline void saxpy_family(const bfloat16 *restrict x, + const bfloat16 *restrict y, + const float a, + bfloat16 *restrict z, + const int32_t vector_size) { event0(); ::aie::vector a_v = ::aie::broadcast(aie::to_float(a, 0)); for (int i = 0; i < vector_size; i += 64) { - ::aie::vector y_v = ::aie::load_v<64>(y); - y += 64; - ::aie::vector z_v = ::aie::add(y_v, a_v); - ::aie::store_v(z, z_v); + ::aie::vector result; + if constexpr (Mul && Add) { + ::aie::vector x_v = ::aie::load_v<64>(x); + x += 64; + ::aie::vector y_v = ::aie::load_v<64>(y); + y += 64; + ::aie::accum ax_v = ::aie::mul(x_v, a_v); + result = ::aie::add(ax_v, y_v).to_vector(); + } else if constexpr (Mul) { + ::aie::vector x_v = ::aie::load_v<64>(x); + x += 64; + result = ::aie::mul(x_v, a_v).to_vector(); + } else { + ::aie::vector y_v = ::aie::load_v<64>(y); + y += 64; + result = ::aie::add(y_v, a_v); + } + ::aie::store_v(z, result); z += 64; } event1(); } -// z = (col > row) ? a : y applied per-element of one tile, using a tile -// position (chunk_start_col, row_in_head) supplied via the idx_buffer. -// -// The tile is interpreted as a `vector_size`-wide horizontal strip of the -// per-head (S, S) attention-score block; idx[0] is the strip's starting -// column within that block, idx[1] is the strip's row within the block. -// The kernel implements the causal mask in-place by writing `a` to elements -// strictly above the diagonal and copying y -> z everywhere else. This -// avoids materialising an H*S*S mask buffer entirely. +// General in-place triangular fill over a horizontal strip of a square (N, N) +// block. The strip is one `vector_size`-wide row segment; idx[0] is the +// segment's starting column within the block row and idx[1] is the row index. +// Elements strictly above the diagonal (column > row) are updated from the +// scalar `a`; elements on or below the diagonal copy the input y -> z. This +// realises e.g. a causal attention mask (a = -inf) without materialising a +// full (N, N) mask buffer. Segments lying entirely on/below the diagonal +// degenerate to a copy. // -// For tiles whose entire range lies at-or-below the diagonal, the kernel -// degenerates to a copy (input still streamed through DMA β€” slightly -// wasteful but simpler than per-tile data-movement skipping). -void scalar_add_causal_bf16(bfloat16 *restrict y, - bfloat16 *restrict z, - int32_t *idx, - const float a, - const int32_t vector_size) +// The masked update is templated: FillConst writes the raw scalar `a`, while +// !FillConst adds `a` to the input (an AXPY-like masked update). +template +static inline void triangular_fill_impl(const bfloat16 *restrict y, + bfloat16 *restrict z, + const int32_t *restrict idx, + const float a, + const int32_t vector_size) { event0(); @@ -100,7 +75,7 @@ void scalar_add_causal_bf16(bfloat16 *restrict y, int32_t chunk_start_col = idx[0]; int32_t row_in_head = idx[1]; - // Index of the first column in the tile that needs to be masked + // Index of the first column in the strip that needs to be masked // (i.e. column index strictly greater than row_in_head). int32_t mask_start = row_in_head + 1 - chunk_start_col; if (mask_start < 0) @@ -112,8 +87,8 @@ void scalar_add_causal_bf16(bfloat16 *restrict y, ::aie::vector s_v = ::aie::broadcast(s); int j = 0; - // ---- Unmasked region [0, mask_start): copy y -> z ---- - // Vectorised body up to the largest VEC-aligned offset <= mask_start. + // Unmasked region [0, mask_start): copy y -> z. Vectorised body up to the + // largest VEC-aligned offset <= mask_start. int mask_start_floor = (mask_start / VEC) * VEC; for (; j < mask_start_floor; j += VEC) { ::aie::vector v = ::aie::load_v(y + j); @@ -124,25 +99,75 @@ void scalar_add_causal_bf16(bfloat16 *restrict y, z[j] = y[j]; } - // ---- Masked region [mask_start, vector_size): write scalar ---- - // If mask_start isn't VEC-aligned, scalar-fill up to the next VEC - // boundary (or to vector_size, whichever is smaller). + // Masked region [mask_start, vector_size). If mask_start isn't VEC-aligned, + // scalar-update up to the next VEC boundary to keep the vector body aligned. int next_vec_boundary = ((j + VEC - 1) / VEC) * VEC; if (next_vec_boundary > vector_size) next_vec_boundary = vector_size; for (; j < next_vec_boundary; j++) { - z[j] = s; + if constexpr (FillConst) + z[j] = s; + else + z[j] = (bfloat16)(a + (float)y[j]); } // Vectorised body of the masked region. for (; j + VEC <= vector_size; j += VEC) { - ::aie::store_v(z + j, s_v); + if constexpr (FillConst) { + ::aie::store_v(z + j, s_v); + } else { + ::aie::vector v = ::aie::load_v(y + j); + ::aie::store_v(z + j, ::aie::add(v, s_v)); + } } - // Scalar tail when vector_size isn't VEC-aligned (in practice this - // doesn't fire since per_tile_elements is always a multiple of VEC). + // Scalar tail when vector_size isn't VEC-aligned (in practice this doesn't + // fire since per_tile_elements is always a multiple of VEC). for (; j < vector_size; j++) { - z[j] = s; + if constexpr (FillConst) + z[j] = s; + else + z[j] = (bfloat16)(a + (float)y[j]); } event1(); } + +extern "C" { +void saxpy(bfloat16 *restrict x, bfloat16 *restrict y, const float a, bfloat16 *restrict z, const int32_t vector_size) +{ + saxpy_family(x, y, a, z, vector_size); +} + +void saxpy_scalar(bfloat16 *x, bfloat16 *y, const bfloat16 a, bfloat16 *z, const int32_t vector_size) +{ + event0(); + float a_f = a; + for (int i = 0; i < vector_size; ++i) { + z[i] = a_f * x[i] + y[i]; + } + event1(); +} + +// z = a * x (scalar-vector multiply; AXPY without the +y term) +void scale_bf16(bfloat16 *restrict x, bfloat16 *restrict z, const float a, const int32_t vector_size) +{ + saxpy_family(x, nullptr, a, z, vector_size); +} + +// z = a + y (scalar-vector add; AXPY without the *x term) +void scalar_add_bf16(bfloat16 *restrict y, bfloat16 *restrict z, const float a, const int32_t vector_size) +{ + saxpy_family(nullptr, y, a, z, vector_size); +} + +// Causal attention mask: write `a` (typically -inf) strictly above the +// per-row diagonal and copy the input elsewhere. Thin wrapper over the general +// triangular fill (constant-fill specialization). +void scalar_add_causal_bf16(bfloat16 *restrict y, + bfloat16 *restrict z, + int32_t *idx, + const float a, + const int32_t vector_size) +{ + triangular_fill_impl(y, z, idx, a, vector_size); +} } \ No newline at end of file diff --git a/conftest.py b/conftest.py index 7ef67529..faa4a67c 100644 --- a/conftest.py +++ b/conftest.py @@ -173,6 +173,12 @@ def pytest_configure(config): def pytest_collection_modifyitems(config, items): device = aie_utils.DefaultNPURuntime.device().resolve().name for item in items: + # Benchmark tests are part of the extensive suite. Adding the marker here + # (before pytest's built-in '-m' deselection runs) ensures they are + # excluded by '-m "not extensive"' just like any other extensive test. + if item.get_closest_marker("benchmark"): + item.add_marker(pytest.mark.extensive) + marker = item.get_closest_marker("supported_devices") if marker and device not in marker.args: item.add_marker( diff --git a/iron/common/sequence.py b/iron/common/sequence.py index f29ea6f7..6474e2eb 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -56,7 +56,9 @@ def make_callable(self, seq): class AutoDispatch(SequenceDispatch): - """Selects the platform default: full-ELF on NPU2, chained-xclbin elsewhere.""" + """Selects the platform default: a single dispatch (full-ELF) on NPU2, and + a chained-xclbin dispatch elsewhere (NPU1 has no full-ELF support, so the + sequence degenerates into one dispatch per operator).""" name = "auto" @@ -67,7 +69,11 @@ def resolve(self, device): class FusedDispatch(SequenceDispatch): - """Single-ELF dispatch (NPU2 only): all operators fused into one ELF.""" + """Single-dispatch execution (NPU2 only): every operator in the sequence is + inlined into one ELF and launched with a single dispatch. The array is + still reconfigured between operators; "fused" refers to collapsing the + per-operator dispatches into one, not to merging the compute. NPU1 has no + full-ELF support, so it cannot use this mode (see SeparateDispatch).""" name = "fused" @@ -138,6 +144,10 @@ class SeparateDispatch(SequenceDispatch): """Chained-xclbin dispatch: one xclbin+insts per unique operator, linked via ``--xclbin-input`` and invoked sequentially. Owns the compiled per-operator xclbin/insts maps consumed by the runtime callable. + + This is the default (and only) mode on NPU1. On NPU2 it can be selected + explicitly to debug or to benchmark the impact of single-dispatch + (``fused``) execution against per-operator dispatches. """ name = "separate" diff --git a/pytest.ini b/pytest.ini index 6f9f5106..03209f83 100644 --- a/pytest.ini +++ b/pytest.ini @@ -8,5 +8,6 @@ python_classes = Test* python_functions = test_* markers = extensive: extensive test suite (deselect with '-m "not extensive"') + benchmark: benchmark test (implies 'extensive'; only run as part of the extensive suite) supported_devices(*devices): mark test as only supported on the given devices (e.g. "npu1", "npu2"). All devices supported by default. addopts = -v --tb=short --import-mode=importlib From 55b50d5456a161acf3df489f511121e464b7af8a Mon Sep 17 00:00:00 2001 From: andrej Date: Mon, 13 Jul 2026 14:15:56 -0600 Subject: [PATCH 3/4] fix sync regression --- aie_kernels/generic/axpy.cc | 2 +- iron/common/sequence.py | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/aie_kernels/generic/axpy.cc b/aie_kernels/generic/axpy.cc index fc16c576..7dba157c 100644 --- a/aie_kernels/generic/axpy.cc +++ b/aie_kernels/generic/axpy.cc @@ -25,7 +25,7 @@ static inline void saxpy_family(const bfloat16 *restrict x, const int32_t vector_size) { event0(); - ::aie::vector a_v = ::aie::broadcast(aie::to_float(a, 0)); + ::aie::vector a_v = ::aie::broadcast((bfloat16)a); for (int i = 0; i < vector_size; i += 64) { ::aie::vector result; if constexpr (Mul && Add) { diff --git a/iron/common/sequence.py b/iron/common/sequence.py index 6474e2eb..a8c2ec36 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -591,12 +591,10 @@ def get_buffer(self, buffer_name): return sub def _sync_inputs(self): - # Sub-views handed out by get_buffer() propagate their host-dirty state - # to this parent, so the parent syncs to the device here. - self.input_buffer.to("npu") + self.input_buffer._sync_to_device() def _sync_outputs(self): - self.output_buffer.to("cpu") + self.output_buffer._sync_from_device() def _run(self): self.run_handle.start() @@ -640,13 +638,19 @@ def get_buffer(self, buffer_name): return self._buffer_cache[buffer_name] def _sync_inputs(self): + # Direct host writes via get_buffer(name).data[:] don't flip the buffer's + # device flag, so Tensor.to("npu") would no-op (the flag reads "npu" + # from construction) and strand the writes on the host. Force the + # host->device upload for every input buffer. for name in self.op.input_args: - self._buffers[name].to("npu") + self._buffers[name]._sync_to_device() def _sync_outputs(self): + # Force device->host for every non-input buffer so callers read fresh + # device results regardless of the (unchanged) device flag. for name in self.op.subbuffer_layout: if name not in self.op.input_args: - self._buffers[name].to("cpu") + self._buffers[name]._sync_from_device() class SequenceXclbinCallable(_PerBufferCallable): From 181396cd54afd3372179efa5a540b59ee87177ac Mon Sep 17 00:00:00 2001 From: andrej Date: Mon, 13 Jul 2026 14:56:11 -0600 Subject: [PATCH 4/4] make test reporting more resilient --- conftest.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/conftest.py b/conftest.py index faa4a67c..148d18db 100644 --- a/conftest.py +++ b/conftest.py @@ -82,6 +82,9 @@ def add_result( def finalize_results(self): """Compute statistics for all collected metrics""" + # Rebuild from scratch so this is idempotent and safe to call after + # every test (incremental flushing) as well as at session finish. + self.results = [] for (test_path, test_name), data in self.test_metrics.items(): row = { "Commit": self.commit, @@ -161,10 +164,21 @@ def pytest_runtest_makereport(item, call): test_path, test_name, passed, captured, metric_patterns ) + # Flush the CSV incrementally after every test. If a later test + # crashes the interpreter (e.g. a native segfault), pytest_sessionfinish + # never runs, but the results gathered so far are already on disk so + # downstream CI steps (commit_results/pretty.py) still find the file. + csv_reporter.finalize_results() + csv_reporter.write_csv() + def pytest_configure(config): csv_path = config.getoption("--csv-output") config._csv_reporter = CSVReporter(csv_path) + # Create the (empty) CSV up front so the file always exists for downstream + # CI steps, even if zero tests are collected or the very first test crashes + # the interpreter before any results are recorded. + config._csv_reporter.write_csv() config.addinivalue_line( "markers", "metrics(**patterns): specify metric patterns for this test" )