Skip to content

Commit 2896c25

Browse files
andrejandrej
authored
Merge 181396c into 1f9e5a9 
2 parents 1f9e5a9 + 181396c commit 2896c25

22 files changed

Lines changed: 2247 additions & 91 deletions

File tree

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ The IRON Python API for Ryzen™ AI NPUs is described in the following paper:
4646
| [Element-wise Mul](./aie_kernels/generic/mul.cc) | Element-wise multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/elementwise_mul/](./iron/operators/elementwise_mul/) |
4747
| [GEMM](./aie_kernels/aie2p/mm.cc) | General Matrix Multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gemm/](./iron/operators/gemm/) |
4848
| [GEMV](./aie_kernels/generic/mv.cc) | General Matrix-Vector Multiplication kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/gemv/](./iron/operators/gemv/) |
49-
| [GQA](./aie_kernels/aie2p/mha.cc) | Grouped Query Attention kernel (Single pipeline) | bfloat16 | | ✓ | 🟢 | [iron/operators/mha/](./iron/operators/mha/) |
50-
| [MHA](./aie_kernels/aie2p/mha.cc) | Multi-Head Attention kernel & Grouped Query Attention | bfloat16 | | ✓ | 🟢 | [iron/operators/mha/](./iron/operators/mha/) |
49+
| [GQA](./aie_kernels/aie2p/mha.cc) | Grouped Query Attention kernel (Single pipeline) | bfloat16 | | ✓ | 🟢 | [iron/operators/mha_prefill_df/](./iron/operators/mha_prefill_df/) |
50+
| [MHA](./aie_kernels/aie2p/mha.cc) | Multi-Head Attention kernel & Grouped Query Attention | bfloat16 | | ✓ | 🟢 | [iron/operators/mha_prefill_df/](./iron/operators/mha_prefill_df/) |
5151
| [RMSNorm](./aie_kernels/aie2/rms_norm.cc) | RMSNorm kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rms_norm/](./iron/operators/rms_norm/) |
5252
| [RoPE](./aie_kernels/generic/rope.cc) | Rotary Positional Embedding kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/rope/](./iron/operators/rope/) |
5353
| [SiLU](./aie_kernels/aie2/silu.cc) | Sigmoid Linear Unit activation kernel | bfloat16 | ✓ | ✓ | 🟢 | [iron/operators/silu/](./iron/operators/silu/) |

‎aie_kernels/aie2/softmax.cc‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "lut_based_ops.h"
55

66
#include <aie_api/aie.hpp>
7+
#include <math.h>
78
#include <stdint.h>
89

910
using namespace aie;
@@ -57,13 +58,124 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out
5758
return;
5859
}
5960

61+
// Online (partial / tiled) softmax helpers.
62+
//
63+
// These kernels implement an online softmax that processes a row in sub-tile
64+
// chunks, keeping running max and sum statistics in a small per-core buffer.
65+
// softmax_stats names the two stats slots instead of using hard-coded array
66+
// indices. It occupies the first two bfloat16 elements of the stats buffer.
67+
struct softmax_stats {
68+
bfloat16 max; // running max
69+
bfloat16 sum; // running sum of exp(x - max)
70+
};
71+
72+
void softmax_partial_stats_impl(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size)
73+
{
74+
event0();
75+
76+
const int elem_iters = vector_size / 16;
77+
78+
float running_max = (float)stats->max;
79+
float running_sum = (float)stats->sum;
80+
81+
aie::vector<bfloat16, 16> input_bf16;
82+
aie::accum<accfloat, 16> exp_val_accum = aie::zeros<accfloat, 16>();
83+
84+
auto it_in = aie::cbegin_vector<16>((bfloat16 *)input);
85+
86+
// Single-pass online algorithm: for each vector chunk, check if max
87+
// needs updating, rescale the running sum if so, then accumulate
88+
// exp(x - max).
89+
for (int i = 0; i < elem_iters; i++) {
90+
input_bf16 = *it_in++;
91+
float chunk_max = aie::reduce_max(input_bf16);
92+
93+
if (chunk_max > running_max) {
94+
// Rescale accumulated exp values by exp(old_max - new_max)
95+
aie::vector<bfloat16, 16> correction =
96+
to_v16bfloat16(getExpBf16(aie::broadcast<bfloat16, 16>((bfloat16)(running_max - chunk_max))));
97+
float scale = (float)correction[0];
98+
// Rescale the partial vector accumulator
99+
aie::vector<bfloat16, 16> scale_vec = aie::broadcast<bfloat16, 16>((bfloat16)scale);
100+
exp_val_accum = aie::mul(exp_val_accum.to_vector<bfloat16>(), scale_vec);
101+
// Rescale the running scalar sum from previous chunks
102+
running_sum *= scale;
103+
running_max = chunk_max;
104+
}
105+
106+
aie::vector<bfloat16, 16> shifted = aie::sub(input_bf16, aie::broadcast<bfloat16, 16>((bfloat16)running_max));
107+
aie::vector<bfloat16, 16> exp_val = to_v16bfloat16(getExpBf16(shifted));
108+
exp_val_accum = add(exp_val_accum, exp_val);
109+
}
110+
111+
// Reduce the vector accumulator and add to running sum
112+
aie::vector<float, 16> reduce = exp_val_accum.to_vector<float>();
113+
running_sum += aie::reduce_add(reduce);
114+
115+
stats->max = (bfloat16)running_max;
116+
stats->sum = (bfloat16)running_sum;
117+
118+
event1();
119+
}
120+
121+
void softmax_partial_norm_impl(bfloat16 *restrict input,
122+
bfloat16 *restrict output,
123+
softmax_stats *restrict stats,
124+
const int32_t vector_size)
125+
{
126+
event0();
127+
128+
const int elem_iters = vector_size / 16;
129+
130+
float max_val = (float)stats->max;
131+
float sum_val = (float)stats->sum;
132+
bfloat16 inv_sum = (bfloat16)aie::inv(sum_val);
133+
134+
aie::vector<bfloat16, 16> max_val_vec = aie::broadcast<bfloat16, 16>((bfloat16)max_val);
135+
136+
aie::vector<bfloat16, 16> input_bf16;
137+
aie::accum<accfloat, 16> out_vals;
138+
139+
auto it_in = aie::cbegin_restrict_vector<16>((bfloat16 *)input);
140+
auto it_out = aie::begin_restrict_vector<16>((bfloat16 *)output);
141+
142+
for (int i = 0; i < elem_iters; i++) {
143+
input_bf16 = *it_in++;
144+
aie::vector<bfloat16, 16> shifted = aie::sub(input_bf16, max_val_vec);
145+
aie::vector<bfloat16, 16> exp_val = to_v16bfloat16(getExpBf16(shifted));
146+
out_vals = aie::mul(exp_val, inv_sum);
147+
*it_out++ = out_vals.to_vector<bfloat16>();
148+
}
149+
150+
event1();
151+
}
152+
60153
extern "C" {
61154

62155
void softmax_bf16(bfloat16 *restrict input, bfloat16 *restrict output, const int32_t input_size)
63156
{
64157
softmax_simple_bf16(input, output, input_size);
65158
}
66159

160+
void softmax_partial_init_bf16(softmax_stats *restrict stats)
161+
{
162+
stats->max = (bfloat16)(-INFINITY);
163+
stats->sum = (bfloat16)(0.0f);
164+
}
165+
166+
void softmax_partial_stats_bf16(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size)
167+
{
168+
softmax_partial_stats_impl(input, stats, vector_size);
169+
}
170+
171+
void softmax_partial_norm_bf16(bfloat16 *restrict input,
172+
bfloat16 *restrict output,
173+
softmax_stats *restrict stats,
174+
const int32_t vector_size)
175+
{
176+
softmax_partial_norm_impl(input, output, stats, vector_size);
177+
}
178+
67179
void mask_bf16(bfloat16 *inout, const int32_t unmasked_size, const int32_t total_size)
68180
{
69181
for (int32_t i = unmasked_size; i < total_size; i++) {

‎aie_kernels/aie2p/softmax.cc‎

Lines changed: 130 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// SPDX-License-Identifier: Apache-2.0
33

44
#include <aie_api/aie.hpp>
5+
#include <math.h>
56
#include <stdint.h>
67

78
#define SM_VEC_LEN 64 // 32
@@ -30,7 +31,7 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out
3031
aie::vector<bfloat16, SM_VEC_LEN> in_elems, exp_val, input_bf16, log2e_vec, max_val_vec;
3132
aie::accum<accfloat, SM_VEC_LEN> out_vals, exp_val_accum, scaled_accum, exp_in_accum;
3233

33-
float max_val = 0;
34+
float max_val = -INFINITY;
3435
float accum_exp_val = 0;
3536
float running_max = 0;
3637
bfloat16 col_sum_inv;
@@ -80,6 +81,17 @@ void softmax_simple_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict out
8081
return;
8182
}
8283

84+
// partial_softmax_alias_bf16 is a flash-attention style single-shot softmax
85+
// used by the projected-fused path. It shares the same three-pass structure as
86+
// softmax_simple_bf16 (find max, exp, normalize) but differs in two ways that
87+
// make code sharing awkward: (1) it folds the query scale into the log2e
88+
// multiply instead of using the fixed log2e constant, and (2) it maintains
89+
// per-row running max/sum in an externally-supplied scale_buffer (flash
90+
// accumulation across key tiles) rather than reducing a whole row locally. The
91+
// softmax_partial_stats_impl / softmax_partial_norm_impl pair below implement a
92+
// different (chunked, two-call) online softmax that keeps its running stats in
93+
// a compact softmax_stats buffer; those are used by the standalone softmax
94+
// operator, not by this projected-fused kernel.
8395
void partial_softmax_alias_bf16(bfloat16 *restrict input_vector,
8496
bfloat16 *restrict output_vector,
8597
bfloat16 *restrict scale_buffer,
@@ -159,6 +171,104 @@ void partial_softmax_alias_bf16(bfloat16 *restrict input_vector,
159171
return;
160172
}
161173

174+
// Online (partial / tiled) softmax helpers.
175+
//
176+
// These kernels implement an online softmax that processes a row in sub-tile
177+
// chunks, keeping running max and sum statistics in a small per-core buffer.
178+
// The max is stored scaled by log2e and the sum accumulates exp2(x*log2e -
179+
// max), matching the exp2-based normalization used below. softmax_stats names
180+
// the two stats slots instead of using hard-coded array indices; it occupies
181+
// the first two bfloat16 elements of the stats buffer.
182+
struct softmax_stats {
183+
bfloat16 max; // running max (scaled by log2e)
184+
bfloat16 sum; // running sum of exp2(x*log2e - max)
185+
};
186+
187+
void softmax_partial_stats_impl(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size)
188+
{
189+
event0();
190+
191+
const int elem_iters = vector_size / SM_VEC_LEN;
192+
193+
float running_max = (float)stats->max;
194+
float running_sum = (float)stats->sum;
195+
196+
aie::vector<bfloat16, SM_VEC_LEN> input_bf16;
197+
aie::accum<accfloat, SM_VEC_LEN> scaled_accum, exp_in_accum;
198+
aie::accum<accfloat, SM_VEC_LEN> exp_val_accum = aie::zeros<accfloat, SM_VEC_LEN>();
199+
200+
aie::vector<bfloat16, SM_VEC_LEN> log2e_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)log2e);
201+
202+
auto it_in = aie::cbegin_restrict_vector<SM_VEC_LEN>((bfloat16 *)input);
203+
204+
// Single-pass online algorithm (matches aie_kernels/aie2/softmax.cc): for
205+
// each vector chunk, update the running max if needed -- rescaling the
206+
// partial accumulator and running sum by exp2(old_max - new_max) -- then
207+
// accumulate exp2(x*log2e - max).
208+
for (int i = 0; i < elem_iters; i++) {
209+
input_bf16 = *it_in++;
210+
scaled_accum = aie::mul(input_bf16, log2e_vec);
211+
float chunk_max = aie::reduce_max(scaled_accum.to_vector<bfloat16>());
212+
213+
if (chunk_max > running_max) {
214+
aie::vector<float, SM_VEC_LEN> diff_vec = aie::broadcast<float, SM_VEC_LEN>(running_max - chunk_max);
215+
aie::vector<bfloat16, SM_VEC_LEN> corr = aie::exp2<bfloat16>(diff_vec);
216+
float scale = (float)corr[0];
217+
aie::vector<bfloat16, SM_VEC_LEN> scale_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)scale);
218+
exp_val_accum = aie::mul(exp_val_accum.to_vector<bfloat16>(), scale_vec);
219+
running_sum *= scale;
220+
running_max = chunk_max;
221+
}
222+
223+
aie::vector<bfloat16, SM_VEC_LEN> max_val_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)running_max);
224+
exp_in_accum = aie::sub(scaled_accum, max_val_vec);
225+
aie::vector<bfloat16, SM_VEC_LEN> exp_val = aie::exp2<bfloat16>(exp_in_accum.to_vector<float>());
226+
exp_val_accum = add(exp_val_accum, exp_val);
227+
}
228+
229+
aie::vector<float, SM_VEC_LEN> reduce = exp_val_accum.to_vector<float>();
230+
running_sum += aie::reduce_add(reduce);
231+
232+
stats->max = (bfloat16)running_max;
233+
stats->sum = (bfloat16)running_sum;
234+
235+
event1();
236+
}
237+
238+
void softmax_partial_norm_impl(bfloat16 *restrict input,
239+
bfloat16 *restrict output,
240+
softmax_stats *restrict stats,
241+
const int32_t vector_size)
242+
{
243+
event0();
244+
245+
const int elem_iters = vector_size / SM_VEC_LEN;
246+
247+
float max_val = (float)stats->max;
248+
float sum_val = (float)stats->sum;
249+
bfloat16 inv_sum = (bfloat16)aie::inv(sum_val);
250+
251+
aie::vector<bfloat16, SM_VEC_LEN> log2e_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)log2e);
252+
aie::vector<bfloat16, SM_VEC_LEN> max_val_vec = aie::broadcast<bfloat16, SM_VEC_LEN>((bfloat16)max_val);
253+
254+
aie::vector<bfloat16, SM_VEC_LEN> input_bf16;
255+
aie::accum<accfloat, SM_VEC_LEN> scaled_accum, exp_in_accum, out_vals;
256+
257+
auto it_in = aie::cbegin_restrict_vector<SM_VEC_LEN>((bfloat16 *)input);
258+
auto it_out = aie::begin_restrict_vector<SM_VEC_LEN>((bfloat16 *)output);
259+
260+
for (int i = 0; i < elem_iters; i++) {
261+
input_bf16 = *it_in++;
262+
scaled_accum = aie::mul(input_bf16, log2e_vec);
263+
exp_in_accum = aie::sub(scaled_accum, max_val_vec);
264+
aie::vector<bfloat16, SM_VEC_LEN> exp_val = aie::exp2<bfloat16>(exp_in_accum.to_vector<float>());
265+
out_vals = aie::mul(exp_val, inv_sum);
266+
*it_out++ = out_vals.to_vector<bfloat16>();
267+
}
268+
269+
event1();
270+
}
271+
162272
extern "C" {
163273

164274
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,
177287
partial_softmax_alias_bf16(input, output, scale_buffer, input_size, row_idx, num_rows, scale);
178288
}
179289

290+
void softmax_partial_init_bf16(softmax_stats *restrict stats)
291+
{
292+
stats->max = (bfloat16)(-INFINITY);
293+
stats->sum = (bfloat16)(0.0f);
294+
}
295+
296+
void softmax_partial_stats_bf16(bfloat16 *restrict input, softmax_stats *restrict stats, const int32_t vector_size)
297+
{
298+
softmax_partial_stats_impl(input, stats, vector_size);
299+
}
300+
301+
void softmax_partial_norm_bf16(bfloat16 *restrict input,
302+
bfloat16 *restrict output,
303+
softmax_stats *restrict stats,
304+
const int32_t vector_size)
305+
{
306+
softmax_partial_norm_impl(input, output, stats, vector_size);
307+
}
308+
180309
void mask_bf16(bfloat16 *inout, const int32 unmasked_size, const int32 total_size)
181310
{
182311
// TODO: Optimize this to use vector code

0 commit comments

Comments
 (0)