Add single-dispatch layer-by-layer multi-head attention

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

aie_kernels/aie2/softmax.cc

@@ -4,6 +4,7 @@
 #include "lut_based_ops.h"
 
 #include <aie_api/aie.hpp>
+#include <math.h>
 #include <stdint.h>
 
 using namespace aie;
@@ -57,13 +58,124 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out
     return;
 }
 
+// Online (partial / tiled) softmax helpers.
+//
+// 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->max;
+    float running_sum = (float)stats->sum;
+
+    aie::vector<bfloat16, 16> input_bf16;
+    aie::accum<accfloat, 16> exp_val_accum = aie::zeros<accfloat, 16>();
+
+    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<bfloat16, 16> correction =
+                to_v16bfloat16(getExpBf16(aie::broadcast<bfloat16, 16>((bfloat16)(running_max - chunk_max))));
+            float scale = (float)correction[0];
+            // Rescale the partial vector accumulator
+            aie::vector<bfloat16, 16> scale_vec = aie::broadcast<bfloat16, 16>((bfloat16)scale);
+            exp_val_accum = aie::mul(exp_val_accum.to_vector<bfloat16>(), scale_vec);
+            // Rescale the running scalar sum from previous chunks
+            running_sum *= scale;
+            running_max = chunk_max;
+        }
+
+        aie::vector<bfloat16, 16> shifted = aie::sub(input_bf16, aie::broadcast<bfloat16, 16>((bfloat16)running_max));
+        aie::vector<bfloat16, 16> 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<float, 16> reduce = exp_val_accum.to_vector<float>();
+    running_sum += aie::reduce_add(reduce);
+
+    stats->max = (bfloat16)running_max;
+    stats->sum = (bfloat16)running_sum;
+
+    event1();
+}
+
+void softmax_partial_norm_impl(bfloat16 *restrict input,
+                               bfloat16 *restrict output,
+                               softmax_stats *restrict stats,
+                               const int32_t vector_size)
+{
+    event0();
+
+    const int elem_iters = vector_size / 16;
+
+    float max_val = (float)stats->max;
+    float sum_val = (float)stats->sum;
+    bfloat16 inv_sum = (bfloat16)aie::inv(sum_val);
+
+    aie::vector<bfloat16, 16> max_val_vec = aie::broadcast<bfloat16, 16>((bfloat16)max_val);
+
+    aie::vector<bfloat16, 16> input_bf16;
+    aie::accum<accfloat, 16> 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<bfloat16, 16> shifted = aie::sub(input_bf16, max_val_vec);
+        aie::vector<bfloat16, 16> exp_val = to_v16bfloat16(getExpBf16(shifted));
+        out_vals = aie::mul(exp_val, inv_sum);
+        *it_out++ = out_vals.to_vector<bfloat16>();
+    }
+
+    event1();
+}
+
 extern "C" {
 
 void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int32_t input_size)
 {
     softmax_simple_bf16(input, output, input_size);
 }
 
+void softmax_partial_init_bf16(softmax_stats *restrict stats)
+{
+    stats->max = (bfloat16)(-INFINITY);
+    stats->sum = (bfloat16)(0.0f);
+}
+
+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,
+                               softmax_stats *restrict 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++) {

aie_kernels/aie2p/softmax.cc

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include <aie_api/aie.hpp>
+#include <math.h>
 #include <stdint.h>
 
 #define SM_VEC_LEN 64   // 32
@@ -30,7 +31,7 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out
     aie::vector<bfloat16, SM_VEC_LEN> in_elems, exp_val, input_bf16, log2e_vec, max_val_vec;
     aie::accum<accfloat, SM_VEC_LEN> 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;
@@ -80,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,
@@ -159,6 +171,104 @@ void partial_softmax_alias_bf16(bfloat16 *restrict input_vector,
     return;
 }
 
+// Online (partial / tiled) softmax helpers.
+//
+// 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<bfloat16, SM_VEC_LEN> input_bf16;
+    aie::accum<accfloat, SM_VEC_LEN> scaled_accum, exp_in_accum;
+    aie::accum<accfloat, SM_VEC_LEN> exp_val_accum = aie::zeros<accfloat, SM_VEC_LEN>();
+
+    aie::vector<bfloat16, SM_VEC_LEN> log2e_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)log2e);
+
+    auto it_in = aie::cbegin_restrict_vector<SM_VEC_LEN>((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_in++;
+        scaled_accum = aie::mul(input_bf16, log2e_vec);
+        float chunk_max = aie::reduce_max(scaled_accum.to_vector<bfloat16>());
+
+        if (chunk_max > running_max) {
+            aie::vector<float, SM_VEC_LEN> diff_vec = aie::broadcast<float, SM_VEC_LEN>(running_max - chunk_max);
+            aie::vector<bfloat16, SM_VEC_LEN> corr = aie::exp2<bfloat16>(diff_vec);
+            float scale = (float)corr[0];
+            aie::vector<bfloat16, SM_VEC_LEN> scale_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)scale);
+            exp_val_accum = aie::mul(exp_val_accum.to_vector<bfloat16>(), scale_vec);
+            running_sum *= scale;
+            running_max = chunk_max;
+        }
+
+        aie::vector<bfloat16, SM_VEC_LEN> max_val_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)running_max);
+        exp_in_accum = aie::sub(scaled_accum, max_val_vec);
+        aie::vector<bfloat16, SM_VEC_LEN> exp_val = aie::exp2<bfloat16>(exp_in_accum.to_vector<float>());
+        exp_val_accum = add(exp_val_accum, exp_val);
+    }
+
+    aie::vector<float, SM_VEC_LEN> reduce = exp_val_accum.to_vector<float>();
+    running_sum += aie::reduce_add(reduce);
+
+    stats->max = (bfloat16)running_max;
+    stats->sum = (bfloat16)running_sum;
+
+    event1();
+}
+
+void softmax_partial_norm_impl(bfloat16 *restrict input,
+                               bfloat16 *restrict output,
+                               softmax_stats *restrict stats,
+                               const int32_t vector_size)
+{
+    event0();
+
+    const int elem_iters = vector_size / SM_VEC_LEN;
+
+    float max_val = (float)stats->max;
+    float sum_val = (float)stats->sum;
+    bfloat16 inv_sum = (bfloat16)aie::inv(sum_val);
+
+    aie::vector<bfloat16, SM_VEC_LEN> log2e_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)log2e);
+    aie::vector<bfloat16, SM_VEC_LEN> max_val_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)max_val);
+
+    aie::vector<bfloat16, SM_VEC_LEN> input_bf16;
+    aie::accum<accfloat, SM_VEC_LEN> scaled_accum, exp_in_accum, out_vals;
+
+    auto it_in = aie::cbegin_restrict_vector<SM_VEC_LEN>((bfloat16 *)input);
+    auto it_out = aie::begin_restrict_vector<SM_VEC_LEN>((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<bfloat16, SM_VEC_LEN> exp_val = aie::exp2<bfloat16>(exp_in_accum.to_vector<float>());
+        out_vals = aie::mul(exp_val, inv_sum);
+        *it_out++ = out_vals.to_vector<bfloat16>();
+    }
+
+    event1();
+}
+
 extern "C" {
 
 void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int32_t input_size)
@@ -177,6 +287,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(softmax_stats *restrict stats)
+{
+    stats->max = (bfloat16)(-INFINITY);
+    stats->sum = (bfloat16)(0.0f);
+}
+
+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,
+                               softmax_stats *restrict 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