From be67b66fc9e5c626d46b4c056b91f7dfd3dd4211 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Thu, 9 Jul 2026 07:23:05 -0700 Subject: [PATCH] Re write of x86 kernel to better utilize memory It should also simplify adding more kernels for qk sV It also aligns logic with arm and is a step in direction to merge two + Small cleanups PiperOrigin-RevId: 945103083 --- BUILD.bazel | 1 - gemma/flash_attention.cc | 970 +++++++++++++++++++++++++++------------ ops/ops-inl.h | 614 +++++++++---------------- 3 files changed, 889 insertions(+), 696 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index 7a506341..8f06a282 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -705,7 +705,6 @@ cc_library( ":configs", ":flash_structs", ":kv_cache", - ":kv_transcoding", ":mat", ":matmul", ":matmul_env", diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index c8969523..1147b8f3 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -333,7 +333,7 @@ HWY_INLINE void QDotKTile148BF16NotNative( const size_t kNF = hn::Lanes(df); const float* HWY_RESTRICT q_base[kVTileSize]; for (size_t i = 0; i < kVTileSize; ++i) { - q_base[i] = reinterpret_cast(q + q_offsets[i]); + q_base[i] = HWY_RCAST_ALIGNED(const float*, q + q_offsets[i]); } const BF16* HWY_RESTRICT k_base = k.Row(pos / (2 * kNF)); for (size_t i = 0; i < half_cols; ++i, k_base += kNF * 4) { @@ -435,7 +435,7 @@ HWY_INLINE void QDotKTile148BF16Native( const size_t kNF = hn::Lanes(df); const float* HWY_RESTRICT q_base[kVTileSize]; for (size_t i = 0; i < kVTileSize; ++i) { - q_base[i] = reinterpret_cast(q + q_offsets[i]); + q_base[i] = HWY_RCAST_ALIGNED(const float*, q + q_offsets[i]); } const BF16* HWY_RESTRICT k_base = k.Row(pos / (2 * kNF)); for (size_t i = 0; i < half_cols; ++i, k_base += kNF * 4) { @@ -907,36 +907,18 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidth( HWY_DASSERT(kNumQueries <= 8); HWY_DASSERT(gcpp::KVCache::kTileSize >= hn::Lanes(df) * 2); // So we can decompress 2 lanes at a time. - sum0_p0 = hn::Zero(df); - sum0_p1 = hn::Zero(df); - if constexpr (kNumQueries >= 2) { - sum1_p0 = hn::Zero(df); - sum1_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 3) { - sum2_p0 = hn::Zero(df); - sum2_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 4) { - sum3_p0 = hn::Zero(df); - sum3_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 5) { - sum4_p0 = hn::Zero(df); - sum4_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 6) { - sum5_p0 = hn::Zero(df); - sum5_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 7) { - sum6_p0 = hn::Zero(df); - sum6_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 8) { - sum7_p0 = hn::Zero(df); - sum7_p1 = hn::Zero(df); - } + auto zero_init = [&](VQ_T& s0, VQ_T& s1) HWY_ATTR { + s0 = hn::Zero(df); + s1 = hn::Zero(df); + }; + zero_init(sum0_p0, sum0_p1); + if constexpr (kNumQueries >= 2) zero_init(sum1_p0, sum1_p1); + if constexpr (kNumQueries >= 3) zero_init(sum2_p0, sum2_p1); + if constexpr (kNumQueries >= 4) zero_init(sum3_p0, sum3_p1); + if constexpr (kNumQueries >= 5) zero_init(sum4_p0, sum4_p1); + if constexpr (kNumQueries >= 6) zero_init(sum5_p0, sum5_p1); + if constexpr (kNumQueries >= 7) zero_init(sum6_p0, sum6_p1); + if constexpr (kNumQueries >= 8) zero_init(sum7_p0, sum7_p1); for (size_t i = 0; i < qkv_dim; ++i) { VQ_T k_vec1, k_vec2; @@ -1124,36 +1106,18 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidthBF16( HWY_DASSERT(kNumQueries <= 8); HWY_DASSERT(gcpp::KVCache::kTileSize >= hn::Lanes(df) * 2); // So we can decompress 2 lanes at a time. - sum0_p0 = hn::Zero(df); - sum0_p1 = hn::Zero(df); - if constexpr (kNumQueries >= 2) { - sum1_p0 = hn::Zero(df); - sum1_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 3) { - sum2_p0 = hn::Zero(df); - sum2_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 4) { - sum3_p0 = hn::Zero(df); - sum3_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 5) { - sum4_p0 = hn::Zero(df); - sum4_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 6) { - sum5_p0 = hn::Zero(df); - sum5_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 7) { - sum6_p0 = hn::Zero(df); - sum6_p1 = hn::Zero(df); - } - if constexpr (kNumQueries >= 8) { - sum7_p0 = hn::Zero(df); - sum7_p1 = hn::Zero(df); - } + auto zero_init = [&](VF& s0, VF& s1) HWY_ATTR { + s0 = hn::Zero(df); + s1 = hn::Zero(df); + }; + zero_init(sum0_p0, sum0_p1); + if constexpr (kNumQueries >= 2) zero_init(sum1_p0, sum1_p1); + if constexpr (kNumQueries >= 3) zero_init(sum2_p0, sum2_p1); + if constexpr (kNumQueries >= 4) zero_init(sum3_p0, sum3_p1); + if constexpr (kNumQueries >= 5) zero_init(sum4_p0, sum4_p1); + if constexpr (kNumQueries >= 6) zero_init(sum5_p0, sum5_p1); + if constexpr (kNumQueries >= 7) zero_init(sum6_p0, sum6_p1); + if constexpr (kNumQueries >= 8) zero_init(sum7_p0, sum7_p1); VF helper_sum0_p0 = hn::Zero(df), helper_sum0_p1 = hn::Zero(df); VF helper_sum1_p0 = hn::Zero(df), helper_sum1_p1 = hn::Zero(df); VF helper_sum2_p0 = hn::Zero(df), helper_sum2_p1 = hn::Zero(df); @@ -1205,36 +1169,25 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidthBF16( } } #if HWY_NATIVE_DOT_BF16 == 0 - sum0_p0 = hn::Add(sum0_p0, helper_sum0_p0); - sum0_p1 = hn::Add(sum0_p1, helper_sum0_p1); - if constexpr (kNumQueries >= 2) { - sum1_p0 = hn::Add(sum1_p0, helper_sum1_p0); - sum1_p1 = hn::Add(sum1_p1, helper_sum1_p1); - } - if constexpr (kNumQueries >= 3) { - sum2_p0 = hn::Add(sum2_p0, helper_sum2_p0); - sum2_p1 = hn::Add(sum2_p1, helper_sum2_p1); - } - if constexpr (kNumQueries >= 4) { - sum3_p0 = hn::Add(sum3_p0, helper_sum3_p0); - sum3_p1 = hn::Add(sum3_p1, helper_sum3_p1); - } - if constexpr (kNumQueries >= 5) { - sum4_p0 = hn::Add(sum4_p0, helper_sum4_p0); - sum4_p1 = hn::Add(sum4_p1, helper_sum4_p1); - } - if constexpr (kNumQueries >= 6) { - sum5_p0 = hn::Add(sum5_p0, helper_sum5_p0); - sum5_p1 = hn::Add(sum5_p1, helper_sum5_p1); - } - if constexpr (kNumQueries >= 7) { - sum6_p0 = hn::Add(sum6_p0, helper_sum6_p0); - sum6_p1 = hn::Add(sum6_p1, helper_sum6_p1); - } - if constexpr (kNumQueries >= 8) { - sum7_p0 = hn::Add(sum7_p0, helper_sum7_p0); - sum7_p1 = hn::Add(sum7_p1, helper_sum7_p1); - } + auto add_helper = [&](VF& s0, VF& s1, const VF& h0, const VF& h1) HWY_ATTR { + s0 = hn::Add(s0, h0); + s1 = hn::Add(s1, h1); + }; + add_helper(sum0_p0, sum0_p1, helper_sum0_p0, helper_sum0_p1); + if constexpr (kNumQueries >= 2) + add_helper(sum1_p0, sum1_p1, helper_sum1_p0, helper_sum1_p1); + if constexpr (kNumQueries >= 3) + add_helper(sum2_p0, sum2_p1, helper_sum2_p0, helper_sum2_p1); + if constexpr (kNumQueries >= 4) + add_helper(sum3_p0, sum3_p1, helper_sum3_p0, helper_sum3_p1); + if constexpr (kNumQueries >= 5) + add_helper(sum4_p0, sum4_p1, helper_sum4_p0, helper_sum4_p1); + if constexpr (kNumQueries >= 6) + add_helper(sum5_p0, sum5_p1, helper_sum5_p0, helper_sum5_p1); + if constexpr (kNumQueries >= 7) + add_helper(sum6_p0, sum6_p1, helper_sum6_p0, helper_sum6_p1); + if constexpr (kNumQueries >= 8) + add_helper(sum7_p0, sum7_p1, helper_sum7_p0, helper_sum7_p1); #endif } @@ -1262,8 +1215,509 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidthBF16( // Need to be have multiple of 4 elements alocated and // be initizalized If you need to compute over multiple chunks of kv's you can // keep values between calls to this function and avoid explicit merge. -template -HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogits( +#ifndef BENCHMARK_BLOCK_SIZE_BF16 +#define BENCHMARK_BLOCK_SIZE_BF16 128 +#endif + +#ifndef BENCHMARK_BLOCK_SIZE_INT8 +#define BENCHMARK_BLOCK_SIZE_INT8 128 +#endif + +template +static HWY_INLINE void ComputeQKMacroTile( + DF df, DU du, const hwy::Span>& kvs, + size_t current_kv_idx, size_t current_kv_start_offset, + const Q_T* HWY_RESTRICT q_base, size_t qkv_dim, size_t q_count, + size_t position, size_t actual_block_size, size_t step_size, + const size_t* HWY_RESTRICT min_start_pos_per_group, + const size_t* HWY_RESTRICT max_last_pos_per_group, + float* HWY_RESTRICT softmax_buf_ptr, size_t kBlockSize, + size_t macro_block_start_pos) { + using VF = hn::Vec; + constexpr int kTileSize = gcpp::KVCache::kTileSize; + size_t actual_steps = actual_block_size / step_size; + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + size_t step_pos = position + step_idx * step_size; + + size_t kv_idx = current_kv_idx; + size_t kv_offset = current_kv_start_offset; + // Check and shift KV chunks continuously + while (kv_idx + 1 < kvs.size() && + step_pos - kv_offset >= kvs[kv_idx].Rows() * kTileSize) { + kv_offset += kvs[kv_idx].Rows() * kTileSize; + kv_idx++; + } + + size_t tile_idx = (step_pos - kv_offset) / kTileSize; + // Cap to valid bounds in case step_pos exceeds the total allocated KV + // caches + if (tile_idx >= kvs[kv_idx].Rows()) { + tile_idx = kvs[kv_idx].Rows() - 1; + } + + const size_t pos_in_tile = step_pos % kTileSize; + const KV_T* tile_base = + reinterpret_cast(kvs[kv_idx].RowBytes(tile_idx)); + + auto process_q_bundle = [&](size_t query_idx) HWY_ATTR { + size_t loop_idx = query_idx / kNumQueriesPerLoop; + if (step_pos + step_size <= min_start_pos_per_group[loop_idx] || + step_pos > max_last_pos_per_group[loop_idx]) { + return; + } + + VF x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1; + VF x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1; + const Q_T* HWY_RESTRICT q_group = q_base + query_idx * qkv_dim; + + if constexpr (IsF32()) { + const KV_T* k_transposed_tile = tile_base + pos_in_tile; + QDotKTilexUpTo8TransposedKDoubleWidth( + df, q_group, k_transposed_tile, qkv_dim, x_0_p_0, x_0_p_1, x_1_p_0, + x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, + x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); + } else if constexpr (IsBF16()) { + const KV_T* k_transposed_tile = tile_base + pos_in_tile * 2; + QDotKTilexUpTo8TransposedKDoubleWidthBF16( + df, q_group, k_transposed_tile, qkv_dim, x_0_p_0, x_0_p_1, x_1_p_0, + x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, + x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); + } else if constexpr (IsInt16()) { + const KV_T* k_transposed_tile = tile_base + pos_in_tile * 2; + QDotKTilexUpTo8TransposedKDoubleWidthInt16( + df, q_group, k_transposed_tile, qkv_dim, x_0_p_0, x_0_p_1, x_1_p_0, + x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, + x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); + } + + float* softmax_buf_step_ptr = softmax_buf_ptr + query_idx * kBlockSize + + (step_pos - macro_block_start_pos); + auto store_logits = [&](const VF& x_p0, const VF& x_p1, + size_t q) HWY_ATTR { + hn::StoreU(x_p0, df, softmax_buf_step_ptr + q * kBlockSize + 0); + hn::StoreU(x_p1, df, + softmax_buf_step_ptr + q * kBlockSize + hn::Lanes(df)); + }; + + if constexpr (kNumQueries >= 1) store_logits(x_0_p_0, x_0_p_1, 0); + if constexpr (kNumQueries >= 2) store_logits(x_1_p_0, x_1_p_1, 1); + if constexpr (kNumQueries >= 3) store_logits(x_2_p_0, x_2_p_1, 2); + if constexpr (kNumQueries >= 4) store_logits(x_3_p_0, x_3_p_1, 3); + if constexpr (kNumQueries >= 5) store_logits(x_4_p_0, x_4_p_1, 4); + if constexpr (kNumQueries >= 6) store_logits(x_5_p_0, x_5_p_1, 5); + if constexpr (kNumQueries >= 7) store_logits(x_6_p_0, x_6_p_1, 6); + if constexpr (kNumQueries >= 8) store_logits(x_7_p_0, x_7_p_1, 7); + }; + + size_t q_idx = 0; + while (q_idx + kNumQueriesPerLoop <= q_count) { + process_q_bundle.template operator()(q_idx); + q_idx += kNumQueriesPerLoop; + } + if (q_idx < q_count) { + size_t rem_q = q_count - q_idx; + if (rem_q >= 8) { + process_q_bundle.template operator()<8>(q_idx); + q_idx += 8; + rem_q -= 8; + } + if (rem_q >= 4) { + process_q_bundle.template operator()<4>(q_idx); + q_idx += 4; + rem_q -= 4; + } + switch (rem_q) { + case 1: + process_q_bundle.template operator()<1>(q_idx); + break; + case 2: + process_q_bundle.template operator()<2>(q_idx); + break; + case 3: + process_q_bundle.template operator()<3>(q_idx); + break; + default: + break; + } + } + } +} + +template > +static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( + DF df, DU du, float* HWY_RESTRICT q_logits, size_t actual_block_size, + size_t q, // global query index + float* HWY_RESTRICT max_logits, float* HWY_RESTRICT exp_denominator_sums, + size_t q_offset, // local query index in the group + float* HWY_RESTRICT scales_old, float* HWY_RESTRICT q_scales_new, + size_t step_size, void* HWY_RESTRICT step_consts_buf, + float* HWY_RESTRICT step_q_scales_s_buf_ptr, size_t actual_steps, + const BF16* const* HWY_RESTRICT step_microscaling_v_ptrs, + const BF16* const* HWY_RESTRICT step_microscaling_k_ptrs, float q_scale_val, + float att_cap, float one_over_cap, size_t position, size_t first_pos, + size_t last_pos) { + constexpr size_t kMaxLanes = hn::MaxLanes(df); + const size_t L_f = hn::Lanes(df); + const size_t unroll_step = 4 * L_f; + using DI16 = hn::Repartition; + const DI16 di16; + using DBF = hn::Repartition; + const DBF dbf; + using D1 = hn::CappedTag; + const D1 d1; + + VF v_max = hn::Set(df, kMaskedLogitVal); + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + size_t step_pos = position + step_idx * step_size; + float* ptr = q_logits + step_idx * step_size; + VF p0 = hn::LoadU(df, ptr); + VF p1 = hn::LoadU(df, ptr + L_f); + + if constexpr (IsInt8()) { + const PackedSpan scales_span = + MakeConstSpan(step_microscaling_k_ptrs[step_idx], 2 * L_f); + VF scales_p0, scales_p1; + Decompress2(df, scales_span, 0, scales_p0, scales_p1); + p0 = hn::Mul(p0, scales_p0); + p1 = hn::Mul(p1, scales_p1); + } + if constexpr (IsInt16()) { + VF s = hn::Set(df, q_scale_val); + p0 = hn::Mul(p0, s); + p1 = hn::Mul(p1, s); + } + if (att_cap > 0.0f) { + VF cap = hn::Set(df, att_cap); + VF one_over_cap_vec = hn::Set(df, one_over_cap); + p0 = hn::Mul(cap, hn::CallFastTanh(df, hn::Mul(p0, one_over_cap_vec))); + p1 = hn::Mul(cap, hn::CallFastTanh(df, hn::Mul(p1, one_over_cap_vec))); + } + if (step_pos < first_pos || step_pos + step_size - 1 > last_pos) { + ApplyMasking<1>(df, du, step_pos, &first_pos, &last_pos, p0, p1, p0, p1, + p0, p1, p0, p1, p0, p1, p0, p1, p0, p1, p0, p1); + } + + hn::StoreU(p0, df, ptr); + hn::StoreU(p1, df, ptr + L_f); + v_max = hn::Max(v_max, hn::Max(p0, p1)); + } + float known_block_max = hn::ReduceMax(df, v_max); + + float old_m = max_logits[q]; + float old_sum = exp_denominator_sums[q]; + float new_m = std::max(old_m, known_block_max); + + float block_sum = 0.0f; + VF v_sum0 = hn::Zero(df); + VF v_sum1 = hn::Zero(df); + VF v_sum2 = hn::Zero(df); + VF v_sum3 = hn::Zero(df); + VF v_new_m = hn::Set(df, new_m); + + size_t t = 0; + for (; t + unroll_step <= actual_block_size; t += unroll_step) { + VF v_logits0 = hn::LoadU(df, q_logits + t); + VF v_logits1 = hn::LoadU(df, q_logits + t + L_f); + VF v_logits2 = hn::LoadU(df, q_logits + t + 2 * L_f); + VF v_logits3 = hn::LoadU(df, q_logits + t + 3 * L_f); + + VF v_exp0 = hn::FastExpMinusOrZero(df, hn::Sub(v_logits0, v_new_m)); + VF v_exp1 = hn::FastExpMinusOrZero(df, hn::Sub(v_logits1, v_new_m)); + VF v_exp2 = hn::FastExpMinusOrZero(df, hn::Sub(v_logits2, v_new_m)); + VF v_exp3 = hn::FastExpMinusOrZero(df, hn::Sub(v_logits3, v_new_m)); + + hn::StoreU(v_exp0, df, q_logits + t); + hn::StoreU(v_exp1, df, q_logits + t + L_f); + hn::StoreU(v_exp2, df, q_logits + t + 2 * L_f); + hn::StoreU(v_exp3, df, q_logits + t + 3 * L_f); + + v_sum0 = hn::Add(v_sum0, v_exp0); + v_sum1 = hn::Add(v_sum1, v_exp1); + v_sum2 = hn::Add(v_sum2, v_exp2); + v_sum3 = hn::Add(v_sum3, v_exp3); + } + for (; t + L_f <= actual_block_size; t += L_f) { + VF v_logits = hn::LoadU(df, q_logits + t); + VF v_exp = hn::FastExpMinusOrZero(df, hn::Sub(v_logits, v_new_m)); + hn::StoreU(v_exp, df, q_logits + t); + v_sum0 = hn::Add(v_sum0, v_exp); + } + v_sum0 = hn::Add(hn::Add(v_sum0, v_sum1), hn::Add(v_sum2, v_sum3)); + if (t < actual_block_size) { + const size_t remaining = actual_block_size - t; + auto mask = hn::FirstN(df, remaining); + VF v_logits = hn::LoadN(df, q_logits + t, remaining); + VF v_exp = hn::FastExpMinusOrZero(df, hn::Sub(v_logits, v_new_m)); + hn::StoreN(v_exp, df, q_logits + t, remaining); + v_sum0 = hn::Add(v_sum0, hn::IfThenElseZero(mask, v_exp)); + } + block_sum = hn::ReduceSum(df, v_sum0); + + float exp_diff = 1.0f; + if (old_m != new_m) { + auto v_diff = hn::Set(d1, old_m - new_m); + auto v_exp = hn::FastExpMinusOrZero(d1, v_diff); + exp_diff = hn::GetLane(v_exp); + } + float new_sum = old_sum * exp_diff + block_sum; + + float scale_old = (new_sum > 0.0f) ? (old_sum * exp_diff) / new_sum : 1.0f; + float q_scale = (new_sum > 0.0f) ? (1.0f / new_sum) : 0.0f; + + scales_old[q_offset] = scale_old; + q_scales_new[q_offset] = q_scale; + max_logits[q] = new_m; + exp_denominator_sums[q] = new_sum; + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + const float* ptr = q_logits + step_idx * step_size; + VF p0 = hn::LoadU(df, ptr); + VF p1 = hn::LoadU(df, ptr + L_f); + + if constexpr (IsF32() || IsBF16()) { + VF s = hn::Set(df, q_scale); + p0 = hn::Mul(p0, s); + p1 = hn::Mul(p1, s); + } else if constexpr (IsInt16()) { + float step_max_v_scale = 1.0f; + if constexpr (IsInt8()) { + const PackedSpan scales_span = + MakeConstSpan(step_microscaling_v_ptrs[step_idx], 2 * L_f); + VF v_scales_p0, v_scales_p1; + Decompress2(df, scales_span, 0, v_scales_p0, v_scales_p1); + step_max_v_scale = hn::ReduceMax(df, hn::Max(v_scales_p0, v_scales_p1)); + } + float step_max_val = std::max( + std::max(hn::ReduceMax(df, p0), hn::ReduceMax(df, p1)), 0.0f); + if (step_max_val > 1e-10f && q_scale > 0.0f) { + float eff_v_scale = std::max(step_max_v_scale, 1e-10f); + float scale_to_int16 = 32767.0f / (step_max_val * eff_v_scale); + VF s_int16 = hn::Set(df, scale_to_int16); + p0 = hn::Mul(p0, s_int16); + p1 = hn::Mul(p1, s_int16); + step_q_scales_s_buf_ptr[step_idx * kNumQueries + q_offset] = + q_scale * step_max_val * (eff_v_scale / 32767.0f); + } else { + p0 = hn::Zero(df); + p1 = hn::Zero(df); + step_q_scales_s_buf_ptr[step_idx * kNumQueries + q_offset] = 0.0f; + } + } + + if constexpr (IsInt8()) { + const PackedSpan scales_span = + MakeConstSpan(step_microscaling_v_ptrs[step_idx], 2 * L_f); + VF v_scales_p0, v_scales_p1; + Decompress2(df, scales_span, 0, v_scales_p0, v_scales_p1); + p0 = hn::Mul(p0, v_scales_p0); + p1 = hn::Mul(p1, v_scales_p1); + } + + if constexpr (IsF32()) { + float* dst = ((float*)step_consts_buf) + + step_idx * (kNumQueries * 2 * kMaxLanes) + + q_offset * 2 * kMaxLanes; + hn::Store(p0, df, dst); + hn::Store(p1, df, dst + kMaxLanes); + } else if constexpr (IsInt16()) { + int16_t* dst = ((int16_t*)step_consts_buf) + + step_idx * (kNumQueries * 2 * kMaxLanes) + + q_offset * 2 * kMaxLanes; + auto i0 = + hn::OrderedDemote2To(di16, hn::NearestInt(p0), hn::NearestInt(p1)); + hn::Store(i0, di16, dst); + } else { + BF16* dst = ((BF16*)step_consts_buf) + + step_idx * (kNumQueries * 2 * kMaxLanes) + + q_offset * 2 * kMaxLanes; + auto bf0 = hn::OrderedDemote2To(dbf, p0, p1); + hn::Store(bf0, dbf, dst); + } + } +} + +template +using FlashStepBufT = + std::conditional_t(), float, + std::conditional_t(), int16_t, BF16>>; + +template +static HWY_INLINE void ComputeSoftmaxAndSVBundle( + DF df, DU du, const hwy::Span>& kvs, + size_t current_kv_idx, size_t current_kv_start_offset, size_t q_base_idx, + size_t actual_q_count, size_t qkv_dim, const float* HWY_RESTRICT q_scales, + size_t position, size_t actual_block_size, size_t step_size, float att_cap, + float one_over_cap, const size_t* HWY_RESTRICT start_pos_per_query, + const size_t* HWY_RESTRICT last_pos_per_query, + const size_t* HWY_RESTRICT min_start_pos_per_group, + const size_t* HWY_RESTRICT max_last_pos_per_group, + const float* HWY_RESTRICT softmax_buf_ptr, float* HWY_RESTRICT max_logits, + float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT C_accumulators_ptr, size_t kBlockSize, + FlashStepBufT* HWY_RESTRICT step_consts_buf, + const KV_T** HWY_RESTRICT step_v_tiles, + const BF16** HWY_RESTRICT step_microscaling_v_ptrs, + const BF16** HWY_RESTRICT step_microscaling_k_ptrs, + const float** HWY_RESTRICT step_q_scales_s_ptrs, + float* HWY_RESTRICT step_q_scales_s_buf) { + constexpr int kTileSize = gcpp::KVCache::kTileSize; + size_t actual_steps = actual_block_size / step_size; + size_t loop_idx = q_base_idx / kNumQueriesPerLoop; + if (position + actual_block_size <= min_start_pos_per_group[loop_idx] || + position > max_last_pos_per_group[loop_idx]) { + return; + } + + HWY_ALIGN float scales_old[kNumQueriesPerLoop]; + HWY_ALIGN float q_scales_new[kNumQueriesPerLoop]; + bool any_query_active = false; + for (size_t i = 0; i < kNumQueriesPerLoop; ++i) { + scales_old[i] = 1.0f; + q_scales_new[i] = 0.0f; + } + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + size_t step_pos = position + step_idx * step_size; + size_t kv_idx = current_kv_idx; + size_t kv_offset = current_kv_start_offset; + // Check and shift KV chunks continuously + while (kv_idx + 1 < kvs.size() && + step_pos - kv_offset >= kvs[kv_idx].Rows() * kTileSize) { + kv_offset += kvs[kv_idx].Rows() * kTileSize; + kv_idx++; + } + + size_t tile_idx = (step_pos - kv_offset) / kTileSize; + if (tile_idx >= kvs[kv_idx].Rows()) { + tile_idx = kvs[kv_idx].Rows() - 1; + } + + size_t pos_in_tile = step_pos % kTileSize; + const KV_T* tile_base = + HWY_RCAST_ALIGNED(const KV_T*, kvs[kv_idx].RowBytes(tile_idx)); + step_v_tiles[step_idx] = + tile_base + qkv_dim * kTileSize + pos_in_tile * qkv_dim; + + if constexpr (IsInt8()) { + const BF16* microscaling_scales_k = + HWY_RCAST_ALIGNED(const BF16*, tile_base + qkv_dim * 2 * kTileSize) + + pos_in_tile; + const BF16* microscaling_scales_v = microscaling_scales_k + kTileSize; + step_microscaling_k_ptrs[step_idx] = microscaling_scales_k; + step_microscaling_v_ptrs[step_idx] = microscaling_scales_v; + } + if constexpr (IsInt16()) { + step_q_scales_s_ptrs[step_idx] = + step_q_scales_s_buf + step_idx * kNumQueriesPerLoop; + } + } + + for (size_t q_offset = 0; q_offset < actual_q_count; ++q_offset) { + size_t q = q_base_idx + q_offset; + if (position + actual_block_size <= start_pos_per_query[q] || + position > last_pos_per_query[q]) { + continue; + } + any_query_active = true; + float* q_logits = ((float*)softmax_buf_ptr) + q * kBlockSize; + float q_scale_val = q_scales ? q_scales[q] : 1.0f; + size_t first_pos = start_pos_per_query[q]; + size_t last_pos = last_pos_per_query[q]; + + UpdateOnlineSoftmaxAndPackSingleQuery( + df, du, q_logits, actual_block_size, q, max_logits, + exp_denominator_sums, q_offset, scales_old, q_scales_new, step_size, + step_consts_buf, IsInt16() ? step_q_scales_s_buf : nullptr, + actual_steps, step_microscaling_v_ptrs, step_microscaling_k_ptrs, + q_scale_val, att_cap, one_over_cap, position, first_pos, last_pos); + } + if (!any_query_active) { + return; + } + + MatPtrT group_out("group_out", Extents2D(kNumQueriesPerLoop, qkv_dim)); + group_out.SetPtr(C_accumulators_ptr + q_base_idx * qkv_dim, qkv_dim); + + if constexpr (IsF32()) { + MulByConstAndAddTileUpTo8( + df, scales_old, actual_steps, step_consts_buf, step_v_tiles, group_out); + } else if constexpr (IsInt16()) { + MulByConstAndAddTileUpTo8_BF16_Int16( + df, scales_old, actual_steps, step_consts_buf, step_v_tiles, group_out, + step_q_scales_s_ptrs); + } else { + MulByConstAndAddTileUpTo8_BF16( + df, scales_old, actual_steps, step_consts_buf, step_v_tiles, group_out); + } +} + +template +struct TileFlashAttentionWorkspaceLayout { + using StepBufT = FlashStepBufT; + + size_t padded_q_count; + size_t c_accum_offset, c_accum_bytes; + size_t softmax_offset, softmax_bytes; + size_t pos_data_offset, pos_data_bytes; + size_t step_consts_offset, step_consts_bytes; + size_t step_v_tiles_offset, step_v_tiles_bytes; + size_t step_mv_offset, step_mv_bytes; + size_t step_mk_offset, step_mk_bytes; + size_t step_qp_offset, step_qp_bytes; + size_t step_qb_offset, step_qb_bytes; + size_t total_bytes; + + constexpr TileFlashAttentionWorkspaceLayout(size_t q_count, size_t qkv_dim, + size_t kBlockSize, + size_t num_loops) + : padded_q_count(hwy::RoundUpTo( + q_count, HWY_MAX(size_t{16}, size_t(kNumQueriesPerLoop)))), + c_accum_offset(0), + c_accum_bytes( + hwy::RoundUpTo(padded_q_count * qkv_dim * sizeof(float), 64)), + softmax_offset(c_accum_offset + c_accum_bytes), + softmax_bytes( + hwy::RoundUpTo(padded_q_count * kBlockSize * sizeof(float), 64)), + pos_data_offset(softmax_offset + softmax_bytes), + pos_data_bytes(hwy::RoundUpTo(2 * num_loops * sizeof(size_t), 64)), + step_consts_offset(pos_data_offset + pos_data_bytes), + step_consts_bytes(hwy::RoundUpTo( + 256 * kNumQueriesPerLoop * 2 * 16 * sizeof(StepBufT), 64)), + step_v_tiles_offset(step_consts_offset + step_consts_bytes), + step_v_tiles_bytes(hwy::RoundUpTo(256 * sizeof(const KV_T*), 64)), + step_mv_offset(step_v_tiles_offset + step_v_tiles_bytes), + step_mv_bytes(hwy::RoundUpTo( + (IsInt8() ? 256 : 1) * sizeof(const BF16*), 64)), + step_mk_offset(step_mv_offset + step_mv_bytes), + step_mk_bytes(hwy::RoundUpTo( + (IsInt8() ? 256 : 1) * sizeof(const BF16*), 64)), + step_qp_offset(step_mk_offset + step_mk_bytes), + step_qp_bytes(hwy::RoundUpTo( + (IsInt16() ? 256 : 1) * sizeof(const float*), 64)), + step_qb_offset(step_qp_offset + step_qp_bytes), + step_qb_bytes(hwy::RoundUpTo( + (IsInt16() ? (256 * kNumQueriesPerLoop) : 1) * sizeof(float), + 64)), + total_bytes(step_qb_offset + step_qb_bytes) {} +}; + +template +static size_t ComputeTileFlashAttentionWorkspaceBytes(size_t q_count, + size_t qkv_dim, + size_t kBlockSize, + size_t num_loops) { + return TileFlashAttentionWorkspaceLayout( + q_count, qkv_dim, kBlockSize, num_loops) + .total_bytes; +} + +template +HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( const hwy::Span> kvs, size_t q_count, const Q_T* HWY_RESTRICT q_base, const hwy::Span q_scales, hwy::Span start_pos_per_query, @@ -1272,260 +1726,181 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogits( float* HWY_RESTRICT max_logits) { using DF = hn::ScalableTag; const DF df; - using VF = hn::Vec; using DU = hn::ScalableTag; - [[maybe_unused]] const DU du; + const DU du; constexpr int kTileSize = gcpp::KVCache::kTileSize; HWY_LANES_CONSTEXPR size_t kHTileSize = hn::Lanes(df); - constexpr int kNumQueriesPerLoop = - (!HWY_ARCH_X86 || (HWY_TARGET <= HWY_AVX3)) ? 8 : 4; - const size_t num_loops = hwy::DivCeil(q_count, kNumQueriesPerLoop); const size_t qkv_dim = att_out.Cols(); HWY_DASSERT(kHTileSize <= hn::MaxLanes(df)); + constexpr size_t kBlockSize = + IsInt8() ? BENCHMARK_BLOCK_SIZE_INT8 : BENCHMARK_BLOCK_SIZE_BF16; HWY_LANES_CONSTEXPR size_t step_size = 2 * kHTileSize; + size_t smallest_start_pos = std::numeric_limits::max(); size_t largest_last_pos = std::numeric_limits::min(); for (size_t i = 0; i < start_pos_per_query.size(); ++i) { smallest_start_pos = std::min(smallest_start_pos, start_pos_per_query[i]); largest_last_pos = std::max(largest_last_pos, last_pos_per_query[i]); } - // start / end positions per group of 4 queries. - std::vector> pos_data(num_loops * 4); - hwy::Span min_start_pos_per_group(pos_data.data(), num_loops); - hwy::Span max_start_pos_per_group(pos_data.data() + num_loops, - num_loops); - hwy::Span min_last_pos_per_group(pos_data.data() + 2 * num_loops, - num_loops); - hwy::Span max_last_pos_per_group(pos_data.data() + 3 * num_loops, - num_loops); + + TileFlashAttentionWorkspaceLayout layout( + q_count, qkv_dim, kBlockSize, num_loops); + const size_t padded_q_count = layout.padded_q_count; + + // Use AllocateAlignedBytes to avoid the cost of zeroing out + // (value-initialization). + auto workspace = hwy::AlignedUniquePtr( + static_cast(hwy::AllocateAlignedBytes(layout.total_bytes)), + hwy::AlignedDeleter()); + uint8_t* raw_ptr = workspace.get(); + + float* C_accumulators_ptr = + HWY_RCAST_ALIGNED(float*, raw_ptr + layout.c_accum_offset); + float* softmax_buf_ptr = + HWY_RCAST_ALIGNED(float*, raw_ptr + layout.softmax_offset); + size_t* pos_data_ptr = + HWY_RCAST_ALIGNED(size_t*, raw_ptr + layout.pos_data_offset); + FlashStepBufT* step_consts_buf = HWY_RCAST_ALIGNED( + FlashStepBufT*, raw_ptr + layout.step_consts_offset); + const KV_T** step_v_tiles = + HWY_RCAST_ALIGNED(const KV_T**, raw_ptr + layout.step_v_tiles_offset); + const BF16** step_microscaling_v_ptrs = + HWY_RCAST_ALIGNED(const BF16**, raw_ptr + layout.step_mv_offset); + const BF16** step_microscaling_k_ptrs = + HWY_RCAST_ALIGNED(const BF16**, raw_ptr + layout.step_mk_offset); + const float** step_q_scales_s_ptrs = + HWY_RCAST_ALIGNED(const float**, raw_ptr + layout.step_qp_offset); + float* step_q_scales_s_buf = + HWY_RCAST_ALIGNED(float*, raw_ptr + layout.step_qb_offset); + + hwy::Span c_accumulators_span(C_accumulators_ptr, + padded_q_count * qkv_dim); + hwy::Span softmax_buf_span(softmax_buf_ptr, + padded_q_count * kBlockSize); + hwy::Span min_start_pos_per_group(pos_data_ptr, num_loops); + hwy::Span max_last_pos_per_group(pos_data_ptr + num_loops, num_loops); for (size_t i = 0; i < num_loops; ++i) { size_t min_start = std::numeric_limits::max(); - size_t max_start = 0; - size_t min_last = std::numeric_limits::max(); size_t max_last = 0; for (int j = 0; j < kNumQueriesPerLoop; ++j) { if (i * kNumQueriesPerLoop + j < q_count) { min_start = std::min(min_start, start_pos_per_query[i * kNumQueriesPerLoop + j]); - max_start = std::max(max_start, - start_pos_per_query[i * kNumQueriesPerLoop + j]); - min_last = - std::min(min_last, last_pos_per_query[i * kNumQueriesPerLoop + j]); max_last = std::max(max_last, last_pos_per_query[i * kNumQueriesPerLoop + j]); } } min_start_pos_per_group[i] = min_start; - max_start_pos_per_group[i] = max_start; - min_last_pos_per_group[i] = min_last; max_last_pos_per_group[i] = max_last; } + const size_t base_pos = smallest_start_pos - (smallest_start_pos % kTileSize); const size_t rem = smallest_start_pos % kTileSize; const size_t num_skipped_sub_tiles = rem / step_size; size_t position = base_pos + num_skipped_sub_tiles * step_size; [[maybe_unused]] float one_over_cap = 1.0f / att_cap; - std::vector> att_out_per_query; - att_out_per_query.reserve(num_loops); - for (size_t i = 0; i < num_loops; ++i) { - att_out_per_query.emplace_back("att_out", - Extents2D(kNumQueriesPerLoop, qkv_dim)); - att_out_per_query.back().SetPtr(att_out.Row(i * kNumQueriesPerLoop), - att_out.Stride()); - } + size_t current_kv_start_offset = 0; size_t current_kv_idx = 0; - HWY_ALIGN float q_scales_s[8]; - float* q_scales_s_ptr = nullptr; - if constexpr (IsInt16()) { - q_scales_s_ptr = q_scales_s; - } - float max_v_scale = 1.0f; - auto inner_loop = [&](size_t query_idx) HWY_ATTR { - size_t loop_idx = query_idx / kNumQueriesPerLoop; - if (position + step_size <= min_start_pos_per_group[loop_idx] || - position > max_last_pos_per_group[loop_idx]) { - return; - } - VF x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1; - VF x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1; - const size_t pos_in_tile = position % kTileSize; - // tile base can point to same tile as previous loop iteration, hence no - // HWY_RESTRICT - // KVs are unaligned and we only use unaligned loads in this implementation. - const KV_T* tile_base = - reinterpret_cast(kvs[current_kv_idx].RowBytes( - (position - current_kv_start_offset) / kTileSize)); - - const KV_T* v_tile = - tile_base + qkv_dim * kTileSize + (pos_in_tile)*qkv_dim; - - const Q_T* HWY_RESTRICT q_group = q_base + query_idx * qkv_dim; - - if constexpr (IsF32()) { - const KV_T* k_transposed_tile = tile_base + pos_in_tile; - QDotKTilexUpTo8TransposedKDoubleWidth( - df, q_group, k_transposed_tile, qkv_dim, x_0_p_0, x_0_p_1, x_1_p_0, - x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, - x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } else if constexpr (IsBF16()) { - const KV_T* k_transposed_tile = tile_base + pos_in_tile * 2; - QDotKTilexUpTo8TransposedKDoubleWidthBF16( - df, q_group, k_transposed_tile, qkv_dim, x_0_p_0, x_0_p_1, x_1_p_0, - x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, - x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } else if constexpr (IsInt16()) { - const KV_T* k_transposed_tile = tile_base + pos_in_tile * 2; - QDotKTilexUpTo8TransposedKDoubleWidthInt16( - df, q_group, k_transposed_tile, qkv_dim, x_0_p_0, x_0_p_1, x_1_p_0, - x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, - x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } else { - static_assert(false, - "Query type not supported, only float, BF16, and " - "Int16 are supported"); - } - // microscaling - // TODO: Change to more generic function to inform if we should use - // microscaling or not. - constexpr bool kUseMicroScaling = IsInt8(); - if constexpr (kUseMicroScaling) { - // After end of the tile, we have kTileSize * 2 bfloat16 for the - // microscaling scales for K and V. - const BF16* microscaling_scales_k = - reinterpret_cast(tile_base + qkv_dim * 2 * kTileSize) + - pos_in_tile; - MultiplyByScale(df, microscaling_scales_k, x_0_p_0, x_0_p_1, - x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, - x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, - x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } - if constexpr (IsInt16()) { - ApplyQuantizationScale( - df, q_scales.data(), query_idx, x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, - x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, - x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } - - constexpr int kFirstHalfAmountOfQueries = std::min(kNumQueries, 4); - constexpr int kSecondHalfAmountOfQueries = - kNumQueries - kFirstHalfAmountOfQueries; - ApplySoftCap( - df, att_cap, one_over_cap, x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, - x_2_p_1, x_3_p_0, x_3_p_1); - if constexpr (kNumQueries > 4) { - ApplySoftCap( - df, att_cap, one_over_cap, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, - x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } - - if (position < max_start_pos_per_group[loop_idx] || - position + step_size - 1 > min_last_pos_per_group[loop_idx]) { - ApplyMasking( - df, du, position, start_pos_per_query.data() + query_idx, - last_pos_per_query.data() + query_idx, x_0_p_0, x_0_p_1, x_1_p_0, - x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, - x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } - HWY_ALIGN float scales[kNumQueriesPerLoop]; - - for (size_t i = 0; i < kNumQueriesPerLoop; ++i) { - scales[i] = 1.0f; - } - - if constexpr (IsInt16() && kUseMicroScaling) { - if (query_idx == 0) { // update only when needed - const BF16* microscaling_scales_v = - reinterpret_cast(tile_base + qkv_dim * 2 * kTileSize) + - kTileSize + pos_in_tile; - const PackedSpan scales_span = - MakeConstSpan(microscaling_scales_v, 2 * hn::Lanes(df)); - VF v_scales_p0, v_scales_p1; - Decompress2(df, scales_span, 0, v_scales_p0, v_scales_p1); - max_v_scale = hn::ReduceMax(df, hn::Max(v_scales_p0, v_scales_p1)); - } - } - - FlashAttentionTileStepAndApplySoftCap( - df, 0.0f, 1.0f, x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, - x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, x_6_p_0, x_6_p_1, - x_7_p_0, x_7_p_1, max_logits, exp_denominator_sums, scales, query_idx, - q_scales_s_ptr, max_v_scale); - if constexpr (kUseMicroScaling) { - const BF16* microscaling_scales_v = - reinterpret_cast(tile_base + qkv_dim * 2 * kTileSize) + - kTileSize + pos_in_tile; - MultiplyByScale(df, microscaling_scales_v, x_0_p_0, x_0_p_1, - x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, x_3_p_0, - x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, - x_6_p_0, x_6_p_1, x_7_p_0, x_7_p_1); - } - MatPtrT offset_out = att_out_per_query[loop_idx]; - const size_t group_offset = query_idx % kNumQueriesPerLoop; - if (group_offset > 0) { - offset_out.SetPtr(offset_out.Row(group_offset), offset_out.Stride()); - } - - if constexpr (IsF32()) { - MulByConstAndAddTileUpTo8( - df, scales, x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, - x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, x_6_p_0, - x_6_p_1, x_7_p_0, x_7_p_1, v_tile, offset_out); - } else if constexpr (IsInt16()) { - MulByConstAndAddTileUpTo8_BF16_Int16( - df, scales, x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, - x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, x_6_p_0, - x_6_p_1, x_7_p_0, x_7_p_1, v_tile, offset_out, q_scales_s); - } else { - MulByConstAndAddTileUpTo8_BF16( - df, scales, x_0_p_0, x_0_p_1, x_1_p_0, x_1_p_1, x_2_p_0, x_2_p_1, - x_3_p_0, x_3_p_1, x_4_p_0, x_4_p_1, x_5_p_0, x_5_p_1, x_6_p_0, - x_6_p_1, x_7_p_0, x_7_p_1, v_tile, offset_out); - } - }; - while (position <= largest_last_pos) { - while (position - current_kv_start_offset >= - kvs[current_kv_idx].Rows() * kTileSize) { - current_kv_start_offset += kvs[current_kv_idx].Rows() * kTileSize; - current_kv_idx++; - } - size_t query_idx = 0; - for (; query_idx + kNumQueriesPerLoop <= q_count; - query_idx += kNumQueriesPerLoop) { - inner_loop.template operator()(query_idx); - } - if (query_idx < q_count) { - size_t rem = q_count - query_idx; - if (rem >= 4) { - inner_loop.template operator()<4>(query_idx); - query_idx += 4; - rem -= 4; - } - switch (rem) { - case 1: - inner_loop.template operator()<1>(query_idx); + std::fill(softmax_buf_ptr, softmax_buf_ptr + padded_q_count * kBlockSize, + kMaskedLogitVal); + + // Call precisely with kBlockSize and allow bounded logic inside macros + // instead of manually chunking + ComputeQKMacroTile( + df, du, kvs, current_kv_idx, current_kv_start_offset, q_base, qkv_dim, + q_count, position, kBlockSize, step_size, + min_start_pos_per_group.data(), max_last_pos_per_group.data(), + softmax_buf_ptr, kBlockSize, position); + + auto dispatch_sv = [&](size_t q_base_idx, + size_t actual_q_count) HWY_ATTR { + ComputeSoftmaxAndSVBundle( + df, du, kvs, current_kv_idx, current_kv_start_offset, q_base_idx, + actual_q_count, qkv_dim, q_scales.data(), position, kBlockSize, + step_size, att_cap, one_over_cap, start_pos_per_query.data(), + last_pos_per_query.data(), min_start_pos_per_group.data(), + max_last_pos_per_group.data(), softmax_buf_ptr, max_logits, + exp_denominator_sums, C_accumulators_ptr, kBlockSize, step_consts_buf, + step_v_tiles, step_microscaling_v_ptrs, step_microscaling_k_ptrs, + step_q_scales_s_ptrs, step_q_scales_s_buf); + }; + + for (size_t q_base_idx = 0; q_base_idx < q_count; + q_base_idx += kNumQueriesPerLoop) { + size_t actual_q_count = + std::min(size_t(kNumQueriesPerLoop), q_count - q_base_idx); + switch (actual_q_count) { + case 8: + dispatch_sv.template operator()<8>(q_base_idx, actual_q_count); break; - case 2: - inner_loop.template operator()<2>(query_idx); + case 7: + dispatch_sv.template operator()<7>(q_base_idx, actual_q_count); + break; + case 6: + dispatch_sv.template operator()<6>(q_base_idx, actual_q_count); + break; + case 5: + dispatch_sv.template operator()<5>(q_base_idx, actual_q_count); + break; + case 4: + dispatch_sv.template operator()<4>(q_base_idx, actual_q_count); break; case 3: - inner_loop.template operator()<3>(query_idx); + dispatch_sv.template operator()<3>(q_base_idx, actual_q_count); + break; + case 2: + dispatch_sv.template operator()<2>(q_base_idx, actual_q_count); + break; + case 1: + dispatch_sv.template operator()<1>(q_base_idx, actual_q_count); break; default: break; } } - position += step_size; + position += kBlockSize; + while (position <= largest_last_pos && current_kv_idx + 1 < kvs.size() && + position - current_kv_start_offset >= + kvs[current_kv_idx].Rows() * kTileSize) { + current_kv_start_offset += kvs[current_kv_idx].Rows() * kTileSize; + current_kv_idx++; + } + } + + using VF = hn::Vec; + for (size_t qi = 0; qi < q_count; ++qi) { + float* out = att_out.Row(qi); + const float* accum = C_accumulators_ptr + qi * qkv_dim; + for (size_t d = 0; d < qkv_dim; d += hn::Lanes(df)) { + VF v = hn::LoadU(df, accum + d); + hn::StoreU(v, df, out + d); + } } } +template +HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogits( + const hwy::Span> kvs, size_t q_count, + const Q_T* HWY_RESTRICT q_base, const hwy::Span q_scales, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, const float att_cap, + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT max_logits) { + constexpr int kDefaultChunkSize = + (!HWY_ARCH_X86 || (HWY_TARGET <= HWY_AVX3)) ? 8 : 4; + TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( + kvs, q_count, q_base, q_scales, start_pos_per_query, last_pos_per_query, + att_cap, att_out, exp_denominator_sums, max_logits); +} + void DispatchTileFlashAttentionReturnExpSumsAndMaxLogits( hwy::Span kvs, size_t q_count, const float* HWY_RESTRICT q_base, @@ -2289,7 +2664,6 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, } } -// NOLINTNEXTLINE(google-readability-namespace-comments) } // namespace HWY_NAMESPACE } // namespace gcpp HWY_AFTER_NAMESPACE(); diff --git a/ops/ops-inl.h b/ops/ops-inl.h index 9d0f2a7a..5dd97185 100644 --- a/ops/ops-inl.h +++ b/ops/ops-inl.h @@ -1041,158 +1041,48 @@ static HWY_INLINE void LoadAndMulUpTo8Times2( template , typename VType> HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8( - DF df, const float* HWY_RESTRICT scales, const VF& c0_p0, const VF& c0_p1, - const VF& c1_p0, const VF& c1_p1, const VF& c2_p0, const VF& c2_p1, - const VF& c3_p0, const VF& c3_p1, const VF& c4_p0, const VF& c4_p1, - const VF& c5_p0, const VF& c5_p1, const VF& c6_p0, const VF& c6_p1, - const VF& c7_p0, const VF& c7_p1, const VType* HWY_RESTRICT v_tile, - MatPtrT& out) { + DF df, const float* HWY_RESTRICT scales_old, size_t actual_steps, + const float* HWY_RESTRICT step_consts, + const VType* const* HWY_RESTRICT step_v_tiles, MatPtrT& out) { static_assert(N <= 8); namespace hn = hwy::HWY_NAMESPACE; const size_t qkv_dim = out.Cols(); constexpr size_t kMaxLanes = hn::MaxLanes(df); HWY_LANES_CONSTEXPR size_t NF = hn::Lanes(df); - PackedSpan v_span = MakeConstSpan(v_tile, qkv_dim * 2 * NF); - size_t i = 0; HWY_DASSERT(qkv_dim % (NF * 2) == 0); - HWY_ALIGN float consts_buffer[kMaxLanes * N * 2]; - hn::Store(c0_p0, df, consts_buffer); - hn::Store(c0_p1, df, consts_buffer + kMaxLanes); - if constexpr (N >= 2) { - hn::Store(c1_p0, df, consts_buffer + 2 * kMaxLanes); - hn::Store(c1_p1, df, consts_buffer + 3 * kMaxLanes); - } - if constexpr (N >= 3) { - hn::Store(c2_p0, df, consts_buffer + 4 * kMaxLanes); - hn::Store(c2_p1, df, consts_buffer + 5 * kMaxLanes); - } - if constexpr (N >= 4) { - hn::Store(c3_p0, df, consts_buffer + 6 * kMaxLanes); - hn::Store(c3_p1, df, consts_buffer + 7 * kMaxLanes); - } - if constexpr (N >= 5) { - hn::Store(c4_p0, df, consts_buffer + 8 * kMaxLanes); - hn::Store(c4_p1, df, consts_buffer + 9 * kMaxLanes); - } - if constexpr (N >= 6) { - hn::Store(c5_p0, df, consts_buffer + 10 * kMaxLanes); - hn::Store(c5_p1, df, consts_buffer + 11 * kMaxLanes); - } - if constexpr (N >= 7) { - hn::Store(c6_p0, df, consts_buffer + 12 * kMaxLanes); - hn::Store(c6_p1, df, consts_buffer + 13 * kMaxLanes); - } - if constexpr (N >= 8) { - hn::Store(c7_p0, df, consts_buffer + 14 * kMaxLanes); - hn::Store(c7_p1, df, consts_buffer + 15 * kMaxLanes); - } - HWY_DASSERT(qkv_dim % (NF * 2) == 0); while (i + NF * 2 <= qkv_dim) { VF out0_0, out1_0, out2_0, out3_0, out4_0, out5_0, out6_0, out7_0; VF out0_1, out1_1, out2_1, out3_1, out4_1, out5_1, out6_1, out7_1; - LoadAndMulUpTo8Times2(df, out, i, scales, out0_0, out0_1, out1_0, out1_1, - out2_0, out2_1, out3_0, out3_1, out4_0, out4_1, - out5_0, out5_1, out6_0, out6_1, out7_0, out7_1); - for (int lane = 0; lane < NF; ++lane) { - VF xI1, xI2; - Decompress2(df, v_span, qkv_dim * lane + i, xI1, xI2); - - out0_0 = hn::MulAdd(xI1, hn::Set(df, consts_buffer[lane + 0 * kMaxLanes]), - out0_0); - out0_1 = hn::MulAdd(xI2, hn::Set(df, consts_buffer[lane + 0 * kMaxLanes]), - out0_1); - if constexpr (N >= 2) { - out1_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 2 * kMaxLanes]), out1_0); - out1_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 2 * kMaxLanes]), out1_1); - } - if constexpr (N >= 3) { - out2_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 4 * kMaxLanes]), out2_0); - out2_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 4 * kMaxLanes]), out2_1); - } - if constexpr (N >= 4) { - out3_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 6 * kMaxLanes]), out3_0); - out3_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 6 * kMaxLanes]), out3_1); - } - if constexpr (N >= 5) { - out4_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 8 * kMaxLanes]), out4_0); - out4_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 8 * kMaxLanes]), out4_1); - } - if constexpr (N >= 6) { - out5_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 10 * kMaxLanes]), out5_0); - out5_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 10 * kMaxLanes]), out5_1); - } - if constexpr (N >= 7) { - out6_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 12 * kMaxLanes]), out6_0); - out6_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 12 * kMaxLanes]), out6_1); - } - if constexpr (N >= 8) { - out7_0 = hn::MulAdd( - xI1, hn::Set(df, consts_buffer[lane + 14 * kMaxLanes]), out7_0); - out7_1 = hn::MulAdd( - xI2, hn::Set(df, consts_buffer[lane + 14 * kMaxLanes]), out7_1); - } - VF xI3, xI4; - Decompress2(df, v_span, qkv_dim * (NF + lane) + i, xI3, xI4); - - out0_0 = hn::MulAdd(xI3, hn::Set(df, consts_buffer[lane + 1 * kMaxLanes]), - out0_0); - out0_1 = hn::MulAdd(xI4, hn::Set(df, consts_buffer[lane + 1 * kMaxLanes]), - out0_1); - if constexpr (N >= 2) { - out1_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 3 * kMaxLanes]), out1_0); - out1_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 3 * kMaxLanes]), out1_1); - } - if constexpr (N >= 3) { - out2_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 5 * kMaxLanes]), out2_0); - out2_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 5 * kMaxLanes]), out2_1); - } - if constexpr (N >= 4) { - out3_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 7 * kMaxLanes]), out3_0); - out3_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 7 * kMaxLanes]), out3_1); - } - if constexpr (N >= 5) { - out4_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 9 * kMaxLanes]), out4_0); - out4_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 9 * kMaxLanes]), out4_1); - } - if constexpr (N >= 6) { - out5_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 11 * kMaxLanes]), out5_0); - out5_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 11 * kMaxLanes]), out5_1); - } - if constexpr (N >= 7) { - out6_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 13 * kMaxLanes]), out6_0); - out6_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 13 * kMaxLanes]), out6_1); - } - if constexpr (N >= 8) { - out7_0 = hn::MulAdd( - xI3, hn::Set(df, consts_buffer[lane + 15 * kMaxLanes]), out7_0); - out7_1 = hn::MulAdd( - xI4, hn::Set(df, consts_buffer[lane + 15 * kMaxLanes]), out7_1); + LoadAndMulUpTo8Times2(df, out, i, scales_old, out0_0, out0_1, out1_0, + out1_1, out2_0, out2_1, out3_0, out3_1, out4_0, + out4_1, out5_0, out5_1, out6_0, out6_1, out7_0, + out7_1); + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + const float* consts_buffer = step_consts + step_idx * (N * 2 * kMaxLanes); + const VType* v_tile = step_v_tiles[step_idx]; + PackedSpan v_span = MakeConstSpan(v_tile, qkv_dim * 2 * NF); + + for (size_t t = 0; t < 2 * NF; ++t) { + VF xI1, xI2; + Decompress2(df, v_span, qkv_dim * t + i, xI1, xI2); + + auto mul_add = [&](int j, VF& o0, VF& o1) HWY_ATTR { + const auto c = hn::Set(df, consts_buffer[t + 2 * j * kMaxLanes]); + o0 = hn::MulAdd(xI1, c, o0); + o1 = hn::MulAdd(xI2, c, o1); + }; + + mul_add(0, out0_0, out0_1); + if constexpr (N >= 2) mul_add(1, out1_0, out1_1); + if constexpr (N >= 3) mul_add(2, out2_0, out2_1); + if constexpr (N >= 4) mul_add(3, out3_0, out3_1); + if constexpr (N >= 5) mul_add(4, out4_0, out4_1); + if constexpr (N >= 6) mul_add(5, out5_0, out5_1); + if constexpr (N >= 7) mul_add(6, out6_0, out6_1); + if constexpr (N >= 8) mul_add(7, out7_0, out7_1); } } StoreUpTo8Times2(df, out, i, out0_0, out0_1, out1_0, out1_1, out2_0, @@ -1204,15 +1094,13 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8( HWY_DASSERT(qkv_dim == i); } -// Specialized version for BF16 models that uses int16 quantization for V. template > HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16_Int16( - DF df, const float* HWY_RESTRICT scales, const VF& c0_p0, const VF& c0_p1, - const VF& c1_p0, const VF& c1_p1, const VF& c2_p0, const VF& c2_p1, - const VF& c3_p0, const VF& c3_p1, const VF& c4_p0, const VF& c4_p1, - const VF& c5_p0, const VF& c5_p1, const VF& c6_p0, const VF& c6_p1, - const VF& c7_p0, const VF& c7_p1, const int8_t* HWY_RESTRICT v_tile, - MatPtrT& out, const float* HWY_RESTRICT q_scales_s) { + DF df, const float* HWY_RESTRICT scales_old, size_t actual_steps, + const int16_t* HWY_RESTRICT step_cs_i16, + const int8_t* const* HWY_RESTRICT step_v_tiles, + MatPtrT& out, + const float* const* HWY_RESTRICT step_q_scales_s) { static_assert(N <= 8); namespace hn = hwy::HWY_NAMESPACE; const size_t qkv_dim = out.Cols(); @@ -1230,117 +1118,102 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16_Int16( const hn::Half di8_half; HWY_LANES_CONSTEXPR size_t kInt16Lanes = hn::Lanes(di16); - HWY_ALIGN int16_t cs_i16[N * kMaxLanes * 2]; - - auto quantize_s_and_store = [&](int j, const VF& p0, const VF& p1) HWY_ATTR { - auto i0 = - hn::OrderedDemote2To(di16, hn::NearestInt(p0), hn::NearestInt(p1)); - hn::Store(i0, di16, cs_i16 + j * kMaxLanes * 2); - }; - - quantize_s_and_store(0, c0_p0, c0_p1); - if constexpr (N >= 2) quantize_s_and_store(1, c1_p0, c1_p1); - if constexpr (N >= 3) quantize_s_and_store(2, c2_p0, c2_p1); - if constexpr (N >= 4) quantize_s_and_store(3, c3_p0, c3_p1); - if constexpr (N >= 5) quantize_s_and_store(4, c4_p0, c4_p1); - if constexpr (N >= 6) quantize_s_and_store(5, c5_p0, c5_p1); - if constexpr (N >= 7) quantize_s_and_store(6, c6_p0, c6_p1); - if constexpr (N >= 8) quantize_s_and_store(7, c7_p0, c7_p1); - size_t i = 0; HWY_DASSERT(qkv_dim % (NF * 2) == 0); while (i + 2 * NF <= qkv_dim) { - VI32 acc0_0 = hn::Zero(di32), acc0_1 = hn::Zero(di32); - VI32 acc1_0 = hn::Zero(di32), acc1_1 = hn::Zero(di32); - VI32 acc2_0 = hn::Zero(di32), acc2_1 = hn::Zero(di32); - VI32 acc3_0 = hn::Zero(di32), acc3_1 = hn::Zero(di32); - VI32 acc4_0 = hn::Zero(di32), acc4_1 = hn::Zero(di32); - VI32 acc5_0 = hn::Zero(di32), acc5_1 = hn::Zero(di32); - VI32 acc6_0 = hn::Zero(di32), acc6_1 = hn::Zero(di32); - VI32 acc7_0 = hn::Zero(di32), acc7_1 = hn::Zero(di32); - - VI32 acc0_o_0 = hn::Zero(di32), acc0_o_1 = hn::Zero(di32); - VI32 acc1_o_0 = hn::Zero(di32), acc1_o_1 = hn::Zero(di32); - VI32 acc2_o_0 = hn::Zero(di32), acc2_o_1 = hn::Zero(di32); - VI32 acc3_o_0 = hn::Zero(di32), acc3_o_1 = hn::Zero(di32); - VI32 acc4_o_0 = hn::Zero(di32), acc4_o_1 = hn::Zero(di32); - VI32 acc5_o_0 = hn::Zero(di32), acc5_o_1 = hn::Zero(di32); - VI32 acc6_o_0 = hn::Zero(di32), acc6_o_1 = hn::Zero(di32); - VI32 acc7_o_0 = hn::Zero(di32), acc7_o_1 = hn::Zero(di32); - - for (int lane = 0; lane < NF; ++lane) { - VI16 vi_first8, vi_next8; - - const int8_t* v_ptr = v_tile + 2 * qkv_dim * lane + i * 2; - - auto v8_t0 = hn::LoadU(di8_half, v_ptr); - auto v16_t0 = hn::PromoteTo(di16, v8_t0); - - auto v8_t1 = hn::LoadU(di8_half, v_ptr + kInt16Lanes); - auto v16_t1 = hn::PromoteTo(di16, v8_t1); - - vi_first8 = v16_t0; - vi_next8 = v16_t1; - - auto mul_acc = [&](int j, VI32& a0, VI32& a_o0, VI32& a1, - VI32& a_o1) HWY_ATTR { - int16_t s0 = cs_i16[2 * lane + j * kMaxLanes * 2]; - int16_t s1 = cs_i16[2 * lane + 1 + j * kMaxLanes * 2]; - - int32_t s01; - hwy::CopySameSize(&s0, reinterpret_cast(&s01)); - hwy::CopySameSize(&s1, reinterpret_cast(&s01) + 1); - VI16 sj = hn::BitCast(di16, hn::Set(di32, s01)); - - a0 = hn::ReorderWidenMulAccumulate(di32, vi_first8, sj, a0, a_o0); - a1 = hn::ReorderWidenMulAccumulate(di32, vi_next8, sj, a1, a_o1); - }; - - mul_acc(0, acc0_0, acc0_o_0, acc0_1, acc0_o_1); - if constexpr (N >= 2) mul_acc(1, acc1_0, acc1_o_0, acc1_1, acc1_o_1); - if constexpr (N >= 3) mul_acc(2, acc2_0, acc2_o_0, acc2_1, acc2_o_1); - if constexpr (N >= 4) mul_acc(3, acc3_0, acc3_o_0, acc3_1, acc3_o_1); - if constexpr (N >= 5) mul_acc(4, acc4_0, acc4_o_0, acc4_1, acc4_o_1); - if constexpr (N >= 6) mul_acc(5, acc5_0, acc5_o_0, acc5_1, acc5_o_1); - if constexpr (N >= 7) mul_acc(6, acc6_0, acc6_o_0, acc6_1, acc6_o_1); - if constexpr (N >= 8) mul_acc(7, acc7_0, acc7_o_0, acc7_1, acc7_o_1); - } - - auto convert_and_add = [&](int j, VI32& a0, VI32& a_o0, VI32& a1, - VI32& a_o1) HWY_ATTR { - VF f0 = hn::ConvertTo(df, hn::RearrangeToOddPlusEven(a0, a_o0)); - VF f1 = hn::ConvertTo(df, hn::RearrangeToOddPlusEven(a1, a_o1)); - - VF o0 = hn::Load(df, out.Row(j) + i); - VF o1 = hn::Load(df, out.Row(j) + i + NF); - - VF scale_old = hn::Set(df, scales[j]); - o0 = hn::Mul(o0, scale_old); - o1 = hn::Mul(o1, scale_old); + VF out0_0, out1_0, out2_0, out3_0, out4_0, out5_0, out6_0, out7_0; + VF out0_1, out1_1, out2_1, out3_1, out4_1, out5_1, out6_1, out7_1; + LoadAndMulUpTo8Times2(df, out, i, scales_old, out0_0, out0_1, out1_0, + out1_1, out2_0, out2_1, out3_0, out3_1, out4_0, + out4_1, out5_0, out5_1, out6_0, out6_1, out7_0, + out7_1); + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + const int16_t* cs_i16 = step_cs_i16 + step_idx * (N * kMaxLanes * 2); + const int8_t* v_tile = step_v_tiles[step_idx]; + const float* q_scales_s = step_q_scales_s[step_idx]; + + VI32 acc0_0 = hn::Zero(di32), acc0_1 = hn::Zero(di32); + VI32 acc1_0 = hn::Zero(di32), acc1_1 = hn::Zero(di32); + VI32 acc2_0 = hn::Zero(di32), acc2_1 = hn::Zero(di32); + VI32 acc3_0 = hn::Zero(di32), acc3_1 = hn::Zero(di32); + VI32 acc4_0 = hn::Zero(di32), acc4_1 = hn::Zero(di32); + VI32 acc5_0 = hn::Zero(di32), acc5_1 = hn::Zero(di32); + VI32 acc6_0 = hn::Zero(di32), acc6_1 = hn::Zero(di32); + VI32 acc7_0 = hn::Zero(di32), acc7_1 = hn::Zero(di32); + + VI32 acc0_o_0 = hn::Zero(di32), acc0_o_1 = hn::Zero(di32); + VI32 acc1_o_0 = hn::Zero(di32), acc1_o_1 = hn::Zero(di32); + VI32 acc2_o_0 = hn::Zero(di32), acc2_o_1 = hn::Zero(di32); + VI32 acc3_o_0 = hn::Zero(di32), acc3_o_1 = hn::Zero(di32); + VI32 acc4_o_0 = hn::Zero(di32), acc4_o_1 = hn::Zero(di32); + VI32 acc5_o_0 = hn::Zero(di32), acc5_o_1 = hn::Zero(di32); + VI32 acc6_o_0 = hn::Zero(di32), acc6_o_1 = hn::Zero(di32); + VI32 acc7_o_0 = hn::Zero(di32), acc7_o_1 = hn::Zero(di32); + + for (size_t token_pair_idx = 0; token_pair_idx < NF; ++token_pair_idx) { + VI16 vi_first8, vi_next8; + + const int8_t* v_ptr = v_tile + 2 * qkv_dim * token_pair_idx + i * 2; + + auto v8_t0 = hn::LoadU(di8_half, v_ptr); + auto v16_t0 = hn::PromoteTo(di16, v8_t0); + + auto v8_t1 = hn::LoadU(di8_half, v_ptr + kInt16Lanes); + auto v16_t1 = hn::PromoteTo(di16, v8_t1); + + vi_first8 = v16_t0; + vi_next8 = v16_t1; + + auto mul_acc = [&](int j, VI32& a0, VI32& a_o0, VI32& a1, + VI32& a_o1) HWY_ATTR { + const int32_t* s_ptr = reinterpret_cast( + &cs_i16[2 * token_pair_idx + j * kMaxLanes * 2]); + VI16 sj = hn::BitCast(di16, hn::Set(di32, *s_ptr)); + + a0 = hn::ReorderWidenMulAccumulate(di32, vi_first8, sj, a0, a_o0); + a1 = hn::ReorderWidenMulAccumulate(di32, vi_next8, sj, a1, a_o1); + }; + + mul_acc(0, acc0_0, acc0_o_0, acc0_1, acc0_o_1); + if constexpr (N >= 2) mul_acc(1, acc1_0, acc1_o_0, acc1_1, acc1_o_1); + if constexpr (N >= 3) mul_acc(2, acc2_0, acc2_o_0, acc2_1, acc2_o_1); + if constexpr (N >= 4) mul_acc(3, acc3_0, acc3_o_0, acc3_1, acc3_o_1); + if constexpr (N >= 5) mul_acc(4, acc4_0, acc4_o_0, acc4_1, acc4_o_1); + if constexpr (N >= 6) mul_acc(5, acc5_0, acc5_o_0, acc5_1, acc5_o_1); + if constexpr (N >= 7) mul_acc(6, acc6_0, acc6_o_0, acc6_1, acc6_o_1); + if constexpr (N >= 8) mul_acc(7, acc7_0, acc7_o_0, acc7_1, acc7_o_1); + } - VF scale_new = hn::Set(df, q_scales_s[j]); - o0 = hn::MulAdd(f0, scale_new, o0); - o1 = hn::MulAdd(f1, scale_new, o1); + auto convert_and_add = [&](int j, VI32& a0, VI32& a_o0, VI32& a1, + VI32& a_o1, VF& o0, VF& o1) HWY_ATTR { + VF f0 = hn::ConvertTo(df, hn::RearrangeToOddPlusEven(a0, a_o0)); + VF f1 = hn::ConvertTo(df, hn::RearrangeToOddPlusEven(a1, a_o1)); - hn::Store(o0, df, out.Row(j) + i); - hn::Store(o1, df, out.Row(j) + i + NF); - }; + VF scale_new = hn::Set(df, q_scales_s[j]); + o0 = hn::MulAdd(f0, scale_new, o0); + o1 = hn::MulAdd(f1, scale_new, o1); + }; - convert_and_add(0, acc0_0, acc0_o_0, acc0_1, acc0_o_1); - if constexpr (N >= 2) - convert_and_add(1, acc1_0, acc1_o_0, acc1_1, acc1_o_1); - if constexpr (N >= 3) - convert_and_add(2, acc2_0, acc2_o_0, acc2_1, acc2_o_1); - if constexpr (N >= 4) - convert_and_add(3, acc3_0, acc3_o_0, acc3_1, acc3_o_1); - if constexpr (N >= 5) - convert_and_add(4, acc4_0, acc4_o_0, acc4_1, acc4_o_1); - if constexpr (N >= 6) - convert_and_add(5, acc5_0, acc5_o_0, acc5_1, acc5_o_1); - if constexpr (N >= 7) - convert_and_add(6, acc6_0, acc6_o_0, acc6_1, acc6_o_1); - if constexpr (N >= 8) - convert_and_add(7, acc7_0, acc7_o_0, acc7_1, acc7_o_1); + convert_and_add(0, acc0_0, acc0_o_0, acc0_1, acc0_o_1, out0_0, out0_1); + if constexpr (N >= 2) + convert_and_add(1, acc1_0, acc1_o_0, acc1_1, acc1_o_1, out1_0, out1_1); + if constexpr (N >= 3) + convert_and_add(2, acc2_0, acc2_o_0, acc2_1, acc2_o_1, out2_0, out2_1); + if constexpr (N >= 4) + convert_and_add(3, acc3_0, acc3_o_0, acc3_1, acc3_o_1, out3_0, out3_1); + if constexpr (N >= 5) + convert_and_add(4, acc4_0, acc4_o_0, acc4_1, acc4_o_1, out4_0, out4_1); + if constexpr (N >= 6) + convert_and_add(5, acc5_0, acc5_o_0, acc5_1, acc5_o_1, out5_0, out5_1); + if constexpr (N >= 7) + convert_and_add(6, acc6_0, acc6_o_0, acc6_1, acc6_o_1, out6_0, out6_1); + if constexpr (N >= 8) + convert_and_add(7, acc7_0, acc7_o_0, acc7_1, acc7_o_1, out7_0, out7_1); + } + StoreUpTo8Times2(df, out, i, out0_0, out0_1, out1_0, out1_1, out2_0, + out2_1, out3_0, out3_1, out4_0, out4_1, out5_0, out5_1, + out6_0, out6_1, out7_0, out7_1); i += 2 * NF; } @@ -1348,10 +1221,9 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16_Int16( template , typename VType> HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16( - DF df, const float* HWY_RESTRICT scales, VF c0_p0, VF c0_p1, VF c1_p0, - VF c1_p1, VF c2_p0, VF c2_p1, VF c3_p0, VF c3_p1, VF c4_p0, VF c4_p1, - VF c5_p0, VF c5_p1, VF c6_p0, VF c6_p1, VF c7_p0, VF c7_p1, - VType* HWY_RESTRICT v_tile, MatPtrT& out) { + DF df, const float* HWY_RESTRICT scales_old, size_t actual_steps, + const BF16* HWY_RESTRICT step_cs, + const VType* const* HWY_RESTRICT step_v_tiles, MatPtrT& out) { static_assert(N <= 8); namespace hn = hwy::HWY_NAMESPACE; const size_t qkv_dim = out.Cols(); @@ -1360,32 +1232,6 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16( using DBF = hn::ScalableTag; const DBF dbf; using VBF = hn::Vec; - PackedSpan v_span = MakeConstSpan(v_tile, qkv_dim * 2 * NF); - HWY_ALIGN BF16 cs[N * kMaxLanes * 2]; - PackedSpan cs_span = MakeSpan(cs, N * kMaxLanes * 2); - float* cs_as_float = HWY_RCAST_ALIGNED(float*, cs); - Compress2(df, c0_p0, c0_p1, cs_span, 0); - if constexpr (N >= 2) { - Compress2(df, c1_p0, c1_p1, cs_span, kMaxLanes * 2); - } - if constexpr (N >= 3) { - Compress2(df, c2_p0, c2_p1, cs_span, 2 * kMaxLanes * 2); - } - if constexpr (N >= 4) { - Compress2(df, c3_p0, c3_p1, cs_span, 3 * kMaxLanes * 2); - } - if constexpr (N >= 5) { - Compress2(df, c4_p0, c4_p1, cs_span, 4 * kMaxLanes * 2); - } - if constexpr (N >= 6) { - Compress2(df, c5_p0, c5_p1, cs_span, 5 * kMaxLanes * 2); - } - if constexpr (N >= 7) { - Compress2(df, c6_p0, c6_p1, cs_span, 6 * kMaxLanes * 2); - } - if constexpr (N >= 8) { - Compress2(df, c7_p0, c7_p1, cs_span, 7 * kMaxLanes * 2); - } size_t i = 0; HWY_DASSERT(qkv_dim % (NF * 2) == 0); while (i + NF * 2 <= qkv_dim) { @@ -1393,6 +1239,10 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16( VF out0_1, out1_1, out2_1, out3_1; VF out4_0, out5_0, out6_0, out7_0; VF out4_1, out5_1, out6_1, out7_1; + LoadAndMulUpTo8Times2(df, out, i, scales_old, out0_0, out0_1, out1_0, + out1_1, out2_0, out2_1, out3_0, out3_1, out4_0, + out4_1, out5_0, out5_1, out6_0, out6_1, out7_0, + out7_1); VF helper_out0_0 = hn::Zero(df), helper_out0_1 = hn::Zero(df), helper_out1_0 = hn::Zero(df), helper_out1_1 = hn::Zero(df), helper_out2_0 = hn::Zero(df), helper_out2_1 = hn::Zero(df), @@ -1401,122 +1251,61 @@ HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileUpTo8_BF16( helper_out5_0 = hn::Zero(df), helper_out5_1 = hn::Zero(df), helper_out6_0 = hn::Zero(df), helper_out6_1 = hn::Zero(df), helper_out7_0 = hn::Zero(df), helper_out7_1 = hn::Zero(df); - LoadAndMulUpTo8Times2(df, out, i, scales, out0_0, out0_1, out1_0, out1_1, - out2_0, out2_1, out3_0, out3_1, out4_0, out4_1, - out5_0, out5_1, out6_0, out6_1, out7_0, out7_1); - for (int lane = 0; lane < NF; ++lane) { - VBF xI, xI2; - Decompress2(dbf, v_span, 2 * qkv_dim * lane + i * 2, xI, xI2); - - // Set pair of c scales for 2 value vectors - out0_0 = hn::ReorderWidenMulAccumulate( - df, xI, hn::BitCast(dbf, hn::Set(df, cs_as_float[lane])), out0_0, - helper_out0_0); - out0_1 = hn::ReorderWidenMulAccumulate( - df, xI2, hn::BitCast(dbf, hn::Set(df, cs_as_float[lane])), out0_1, - helper_out0_1); - if constexpr (N >= 2) { - out1_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + kMaxLanes])), - out1_0, helper_out1_0); - out1_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + kMaxLanes])), - out1_1, helper_out1_1); - } - if constexpr (N >= 3) { - out2_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 2 * kMaxLanes])), - out2_0, helper_out2_0); - out2_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 2 * kMaxLanes])), - out2_1, helper_out2_1); - } - if constexpr (N >= 4) { - out3_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 3 * kMaxLanes])), - out3_0, helper_out3_0); - out3_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 3 * kMaxLanes])), - out3_1, helper_out3_1); - } - if constexpr (N >= 5) { - out4_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 4 * kMaxLanes])), - out4_0, helper_out4_0); - out4_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 4 * kMaxLanes])), - out4_1, helper_out4_1); - } - if constexpr (N >= 6) { - out5_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 5 * kMaxLanes])), - out5_0, helper_out5_0); - out5_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 5 * kMaxLanes])), - out5_1, helper_out5_1); - } - if constexpr (N >= 7) { - out6_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 6 * kMaxLanes])), - out6_0, helper_out6_0); - out6_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 6 * kMaxLanes])), - out6_1, helper_out6_1); - } - if constexpr (N >= 8) { - out7_0 = hn::ReorderWidenMulAccumulate( - df, xI, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 7 * kMaxLanes])), - out7_0, helper_out7_0); - out7_1 = hn::ReorderWidenMulAccumulate( - df, xI2, - hn::BitCast(dbf, hn::Set(df, cs_as_float[lane + 7 * kMaxLanes])), - out7_1, helper_out7_1); + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + const BF16* cs = step_cs + step_idx * (N * kMaxLanes * 2); + const float* cs_as_float = HWY_RCAST_ALIGNED(const float*, cs); + const VType* v_tile = step_v_tiles[step_idx]; + PackedSpan v_span = MakeConstSpan(v_tile, qkv_dim * 2 * NF); + + for (size_t token_pair_idx = 0; token_pair_idx < NF; ++token_pair_idx) { + VBF xI, xI2; + Decompress2(dbf, v_span, 2 * qkv_dim * token_pair_idx + i * 2, xI, xI2); + + // Set pair of c scales for 2 value vectors + auto mul_acc = [&](int j, VF& o0, VF& h0, VF& o1, VF& h1) HWY_ATTR { + VBF cs_pair = hn::BitCast( + dbf, hn::Set(df, cs_as_float[token_pair_idx + j * kMaxLanes])); + o0 = hn::ReorderWidenMulAccumulate(df, xI, cs_pair, o0, h0); + o1 = hn::ReorderWidenMulAccumulate(df, xI2, cs_pair, o1, h1); + }; + + mul_acc(0, out0_0, helper_out0_0, out0_1, helper_out0_1); + if constexpr (N >= 2) + mul_acc(1, out1_0, helper_out1_0, out1_1, helper_out1_1); + if constexpr (N >= 3) + mul_acc(2, out2_0, helper_out2_0, out2_1, helper_out2_1); + if constexpr (N >= 4) + mul_acc(3, out3_0, helper_out3_0, out3_1, helper_out3_1); + if constexpr (N >= 5) + mul_acc(4, out4_0, helper_out4_0, out4_1, helper_out4_1); + if constexpr (N >= 6) + mul_acc(5, out5_0, helper_out5_0, out5_1, helper_out5_1); + if constexpr (N >= 7) + mul_acc(6, out6_0, helper_out6_0, out6_1, helper_out6_1); + if constexpr (N >= 8) + mul_acc(7, out7_0, helper_out7_0, out7_1, helper_out7_1); } } #if HWY_NATIVE_DOT_BF16 == 0 - out0_0 = hn::Add(out0_0, helper_out0_0); - out0_1 = hn::Add(out0_1, helper_out0_1); - if constexpr (N >= 2) { - out1_0 = hn::Add(out1_0, helper_out1_0); - out1_1 = hn::Add(out1_1, helper_out1_1); - } - if constexpr (N >= 3) { - out2_0 = hn::Add(out2_0, helper_out2_0); - out2_1 = hn::Add(out2_1, helper_out2_1); - } - if constexpr (N >= 4) { - out3_0 = hn::Add(out3_0, helper_out3_0); - out3_1 = hn::Add(out3_1, helper_out3_1); - } - if constexpr (N >= 5) { - out4_0 = hn::Add(out4_0, helper_out4_0); - out4_1 = hn::Add(out4_1, helper_out4_1); - } - if constexpr (N >= 6) { - out5_0 = hn::Add(out5_0, helper_out5_0); - out5_1 = hn::Add(out5_1, helper_out5_1); - } - if constexpr (N >= 7) { - out6_0 = hn::Add(out6_0, helper_out6_0); - out6_1 = hn::Add(out6_1, helper_out6_1); - } - if constexpr (N >= 8) { - out7_0 = hn::Add(out7_0, helper_out7_0); - out7_1 = hn::Add(out7_1, helper_out7_1); - } + auto add_helper = [&](VF& o0, VF& o1, const VF& h0, const VF& h1) HWY_ATTR { + o0 = hn::Add(o0, h0); + o1 = hn::Add(o1, h1); + }; + add_helper(out0_0, out0_1, helper_out0_0, helper_out0_1); + if constexpr (N >= 2) + add_helper(out1_0, out1_1, helper_out1_0, helper_out1_1); + if constexpr (N >= 3) + add_helper(out2_0, out2_1, helper_out2_0, helper_out2_1); + if constexpr (N >= 4) + add_helper(out3_0, out3_1, helper_out3_0, helper_out3_1); + if constexpr (N >= 5) + add_helper(out4_0, out4_1, helper_out4_0, helper_out4_1); + if constexpr (N >= 6) + add_helper(out5_0, out5_1, helper_out5_0, helper_out5_1); + if constexpr (N >= 7) + add_helper(out6_0, out6_1, helper_out6_0, helper_out6_1); + if constexpr (N >= 8) + add_helper(out7_0, out7_1, helper_out7_0, helper_out7_1); #endif StoreUpTo8Times2(df, out, i, out0_0, out0_1, out1_0, out1_1, out2_0, out2_1, out3_0, out3_1, out4_0, out4_1, out5_0, out5_1, @@ -2006,7 +1795,7 @@ HWY_INLINE void ApplySoftCap(DF df, float att_cap, float one_over_cap, VF& x0, template , typename DU = hn::ScalableTag, class VU = hn::Vec> -HWY_NOINLINE void ApplyMasking(DF df, DU du, size_t position, +HWY_INLINE void ApplyMasking(DF df, DU du, size_t position, const size_t* HWY_RESTRICT first_pos_per_query, const size_t* HWY_RESTRICT last_pos_per_query, VF& x0_p0, VF& x0_p1, VF& x1_p0, VF& x1_p1, @@ -2148,6 +1937,37 @@ HWY_INLINE void ApplyQuantizationScale( } } +template +HWY_INLINE void MultiplyByScaleAndQuantScale( + DF df, const BF16* scales, const float* HWY_RESTRICT q_scales, + size_t query_idx, VF& x0_p0, VF& x0_p1, VF& x1_p0, VF& x1_p1, VF& x2_p0, + VF& x2_p1, VF& x3_p0, VF& x3_p1, VF& x4_p0, VF& x4_p1, VF& x5_p0, VF& x5_p1, + VF& x6_p0, VF& x6_p1, VF& x7_p0, VF& x7_p1) { + const size_t kTileSize = hn::Lanes(df); + const PackedSpan scales_span = + MakeConstSpan(scales, 2 * kTileSize); + VF scales_p0, scales_p1; + Decompress2(df, scales_span, 0, scales_p0, scales_p1); + + auto apply_both = [&](size_t i, VF& x_p0, VF& x_p1) HWY_ATTR { + size_t scale_idx = query_idx + i; + VF s = hn::Set(df, q_scales[scale_idx]); + x_p0 = hn::Mul(x_p0, scales_p0); + x_p1 = hn::Mul(x_p1, scales_p1); + x_p0 = hn::Mul(x_p0, s); + x_p1 = hn::Mul(x_p1, s); + }; + + if constexpr (kNumQueries >= 1) apply_both(0, x0_p0, x0_p1); + if constexpr (kNumQueries >= 2) apply_both(1, x1_p0, x1_p1); + if constexpr (kNumQueries >= 3) apply_both(2, x2_p0, x2_p1); + if constexpr (kNumQueries >= 4) apply_both(3, x3_p0, x3_p1); + if constexpr (kNumQueries >= 5) apply_both(4, x4_p0, x4_p1); + if constexpr (kNumQueries >= 6) apply_both(5, x5_p0, x5_p1); + if constexpr (kNumQueries >= 7) apply_both(6, x6_p0, x6_p1); + if constexpr (kNumQueries >= 8) apply_both(7, x7_p0, x7_p1); +} + template HWY_INLINE V SumReduceSegments(D d, V v) { constexpr size_t L = HWY_MAX_LANES_D(D);