diff --git a/gemma/configs.cc b/gemma/configs.cc index f06fcc39..d60b3a77 100644 --- a/gemma/configs.cc +++ b/gemma/configs.cc @@ -863,6 +863,7 @@ constexpr std::pair kAttentionImplNameToEnum[] = { {"flash_transposed_qs", AttentionImpl::kFlashTransposedQs}, {"flash_transposed_qs_bf16", AttentionImpl::kFlashTransposedQsBF16}, {"flash_transposed_qs_int16", AttentionImpl::kFlashTransposedQsInt16}, + {"flash_transposed_qs_int8", AttentionImpl::kFlashTransposedQsInt8}, {"flash_matrix_accumulation", AttentionImpl::kFlashMatrixAccumulation}, {"int8_matrix_accumulation", AttentionImpl::kInt8MatrixAccumulation}, }; @@ -897,6 +898,8 @@ std::string KVEncodingToString(KVEncoding encoding) { return "Int8"; case KVEncoding::kInt8TwoTranspositions: return "Int8TwoTranspositions"; + case KVEncoding::kInt8VNNITwoTranspositions: + return "Int8VNNITwoTranspositions"; case KVEncoding::kBF16MatrixAccumulation: return "BF16MatrixAccumulation"; case KVEncoding::kInt8MatrixAccumulation: diff --git a/gemma/configs.h b/gemma/configs.h index 38809b6e..dd9370dd 100644 --- a/gemma/configs.h +++ b/gemma/configs.h @@ -94,6 +94,7 @@ enum class KVEncoding { kInt8TwoTranspositions = 6, kBF16MatrixAccumulation = 7, kInt8MatrixAccumulation = 8, + kInt8VNNITwoTranspositions = 9, }; // Returns a string representation of the KVEncoding. @@ -110,6 +111,7 @@ enum class AttentionImpl { kFlashTransposedQsInt16, kFlashMatrixAccumulation, kInt8MatrixAccumulation, + kFlashTransposedQsInt8, kSentinel, }; diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index bd17edbf..a3ea26a8 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -15,10 +15,10 @@ #include #include -#include #include #include +#include #include #include #include @@ -63,7 +63,6 @@ HWY_BEFORE_NAMESPACE(); namespace gcpp { namespace HWY_NAMESPACE { - // Updates q in place for RMSNorm and positional encoding. void RMSNormAndPositionalEncoding(const size_t num_tokens, const QBatch& qbatch, MatPtrT& q, @@ -562,336 +561,221 @@ HWY_INLINE float DoubleFlashAttentionRowVector(DF df, size_t start_pos, return scale; } -// Handles Up to 4 Q rows by NF*2 timesteps of flash attention. template > -static HWY_INLINE void FlashAttentionTileStepAndApplySoftCap4( - DF df, float att_cap, float one_over_att_cap, VF& x_0_p0, VF& x_0_p1, - VF& x_1_p0, VF& x_1_p1, VF& x_2_p0, VF& x_2_p1, VF& x_3_p0, VF& x_3_p1, - float* HWY_RESTRICT old_max, float* HWY_RESTRICT old_d, - float* HWY_RESTRICT scales, float* HWY_RESTRICT q_scales_s = nullptr, - float max_v_scale = 1.0f) { - using DF4 = hn::CappedTag; - const DF4 df4; - using VF4 = hn::Vec; - static_assert(kNumQueries >= 1 && kNumQueries <= 4); - VF4 new_max = hn::Set(df4, kMaskedLogitVal); - VF max_0 = hn::Zero(df), max_1 = hn::Zero(df), max_2 = hn::Zero(df), - max_3 = hn::Zero(df); - max_0 = hn::Max(x_0_p0, x_0_p1); - if constexpr (kNumQueries >= 2) { - max_1 = hn::Max(x_1_p0, x_1_p1); - } - if constexpr (kNumQueries >= 3) { - max_2 = hn::Max(x_2_p0, x_2_p1); - } - if constexpr (kNumQueries >= 4) { - max_3 = hn::Max(x_3_p0, x_3_p1); - } - if constexpr (kNumQueries == 1) { - new_max = hn::InsertLane(new_max, 0, hn::ReduceMax(df, max_0)); - } else { - new_max = Reduce4(df, max_0, max_1, max_2, max_3, - [](auto a, auto b) HWY_ATTR { return hn::Max(a, b); }); - } - if (att_cap > 0.0f) { - VF4 cap = hn::Set(df4, att_cap); - VF4 one_over_cap = hn::Set(df4, one_over_att_cap); - new_max = hn::Mul(cap, hn::FastTanh(df4, hn::Mul(new_max, one_over_cap))); - } - VF4 local_max = new_max; - VF4 old_max_vf = hn::Set(df4, kMaskedLogitVal); - old_max_vf = hn::LoadU(df4, old_max); - new_max = hn::Max(new_max, old_max_vf); - auto changed_max = hn::Gt(new_max, hn::Set(df4, kMaskedLogitVal)); - hn::StoreU(new_max, df4, old_max); - auto apply_exp = [&](int i, VF& x_p0, VF& x_p1) HWY_ATTR { - const VF new_max_i = hn::Set(df, old_max[i]); - x_p0 = hn::FastExpMinusOrZero(df, hn::Sub(x_p0, new_max_i)); - x_p1 = hn::FastExpMinusOrZero(df, hn::Sub(x_p1, new_max_i)); - }; - if constexpr (kNumQueries >= 1) { - apply_exp(0, x_0_p0, x_0_p1); - } - if constexpr (kNumQueries >= 2) { - apply_exp(1, x_1_p0, x_1_p1); - } - if constexpr (kNumQueries >= 3) { - apply_exp(2, x_2_p0, x_2_p1); - } - if constexpr (kNumQueries >= 4) { - apply_exp(3, x_3_p0, x_3_p1); - } - VF4 old_d_vf = hn::Set(df4, 0.0f); - old_d_vf = hn::LoadU(df4, old_d); +static HWY_INLINE void QDotKTilexVNNI( + DF df, const int8_t* HWY_RESTRICT q_base, + const int8_t* HWY_RESTRICT k_transposed_tile, size_t qkv_dim, + const int32_t* HWY_RESTRICT k_sums, VF& sum0_p0, VF& sum0_p1, VF& sum1_p0, + VF& sum1_p1, VF& sum2_p0, VF& sum2_p1, VF& sum3_p0, VF& sum3_p1, + VF& sum4_p0, VF& sum4_p1, VF& sum5_p0, VF& sum5_p1, VF& sum6_p0, + VF& sum6_p1, VF& sum7_p0, VF& sum7_p1) { + using DI32 = hn::Repartition; + const DI32 di32; + using VI32 = hn::Vec; + using DI8 = hn::Repartition; + const DI8 di8; + using DU8 = hn::Repartition; + const DU8 du8; + constexpr size_t kTileSize = gcpp::KVCache::kTileSize; - VF4 x_sum = hn::Zero(df4); - if constexpr (kNumQueries == 1) { - x_sum = hn::Set(df4, hn::ReduceSum(df, x_0_p0) + hn::ReduceSum(df, x_0_p1)); - } else { - VF x_0_sum = hn::Add(x_0_p0, x_0_p1); - VF x_1_sum = hn::Add(x_1_p0, x_1_p1); - VF x_2_sum = hn::Add(x_2_p0, x_2_p1); - VF x_3_sum = hn::Add(x_3_p0, x_3_p1); - x_sum = Reduce4(df, x_0_sum, x_1_sum, x_2_sum, x_3_sum, - [](auto a, auto b) HWY_ATTR { return hn::Add(a, b); }); - } - VF4 scale = hn::Mul( - old_d_vf, hn::FastExpMinusOrZero(df4, hn::Sub(old_max_vf, new_max))); - old_d_vf = hn::Add(scale, x_sum); - auto non_zero_mask = hn::Gt(old_d_vf, hn::Set(df4, 0.0f)); - const VF zero = hn::Zero(df); - const VF4 zero4 = hn::Zero(df4); - const VF4 one_over_d = - hn::MaskedDivOr(zero4, non_zero_mask, hn::Set(df4, 1.0f), old_d_vf); - VF4 q_scale; - if (q_scales_s != nullptr) { - // max_s = exp(local_max - new_max) / old_d_vf - VF4 max_s = hn::Mul(one_over_d, hn::Exp(df4, hn::Sub(local_max, new_max))); - - // Output the unquantize scale directly to array memory: - // Because we're capping out at 32767 / max_v_scale, the true scale goes up - // proportionately - hn::Store(hn::Mul(max_s, hn::Set(df4, max_v_scale / 32767.0f)), df4, - q_scales_s); - - // multiplier for x = 32767 * exp(new_max - local_max) / max_v_scale - auto max_s_gt_0 = hn::Gt(max_s, zero4); - float inv_max_v_scale = 1.0f / std::max(max_v_scale, 1e-10f); - VF4 mult = hn::Mul(hn::Set(df4, 32767.0f * inv_max_v_scale), - hn::Exp(df4, hn::Sub(new_max, local_max))); - q_scale = hn::IfThenElse(max_s_gt_0, mult, zero4); - } else { - q_scale = one_over_d; - } - HWY_ALIGN float tmp_one_over_d[4]; - hn::Store(q_scale, df4, tmp_one_over_d); - hn::BlendedStore(old_d_vf, changed_max, df4, old_d); - scale = hn::Mul(scale, one_over_d); - hn::BlendedStore(scale, changed_max, df4, scales); - // same as lambda - auto mul_or_zero = [&](VF& x_p0, VF& x_p1, int i) HWY_ATTR { - if (HWY_LIKELY(old_d[i] > 0.0f && scales[i] != 1.0f)) { - const VF one_over_d_i = hn::Set(df, tmp_one_over_d[i]); - x_p0 = hn::Mul(x_p0, one_over_d_i); - x_p1 = hn::Mul(x_p1, one_over_d_i); - } else { - x_p0 = zero; - x_p1 = zero; - } - }; - mul_or_zero(x_0_p0, x_0_p1, 0); - if constexpr (kNumQueries >= 2) { - mul_or_zero(x_1_p0, x_1_p1, 1); - } - if constexpr (kNumQueries >= 3) { - mul_or_zero(x_2_p0, x_2_p1, 2); - } - if constexpr (kNumQueries >= 4) { - mul_or_zero(x_3_p0, x_3_p1, 3); - } -} + HWY_DASSERT(kNumQueries <= 8); -template > -static HWY_INLINE void FlashAttentionTileStepAndApplySoftCap8( - DF df, float att_cap, float one_over_att_cap, VF& x_0_p0, VF& x_0_p1, - VF& x_1_p0, VF& x_1_p1, VF& x_2_p0, VF& x_2_p1, VF& x_3_p0, VF& x_3_p1, - VF& x_4_p0, VF& x_4_p1, VF& x_5_p0, VF& x_5_p1, VF& x_6_p0, VF& x_6_p1, - VF& x_7_p0, VF& x_7_p1, float* HWY_RESTRICT old_max, - float* HWY_RESTRICT old_d, float* HWY_RESTRICT scales, - float* HWY_RESTRICT q_scales_s = nullptr, float max_v_scale = 1.0f) { - using DF8 = hn::CappedTag; - const DF8 df8; - using VF8 = hn::Vec; - static_assert(kNumQueries >= 1 && kNumQueries <= 8); - VF8 new_max = hn::Set(df8, kMaskedLogitVal); - VF max_0, max_1, max_2, max_3, max_4, max_5, max_6, max_7 = hn::Zero(df); - max_0 = hn::Max(x_0_p0, x_0_p1); - if constexpr (kNumQueries >= 2) { - max_1 = hn::Max(x_1_p0, x_1_p1); - } - if constexpr (kNumQueries >= 3) { - max_2 = hn::Max(x_2_p0, x_2_p1); - } - if constexpr (kNumQueries >= 4) { - max_3 = hn::Max(x_3_p0, x_3_p1); - } - if constexpr (kNumQueries >= 5) { - max_4 = hn::Max(x_4_p0, x_4_p1); - } - if constexpr (kNumQueries >= 6) { - max_5 = hn::Max(x_5_p0, x_5_p1); - } - if constexpr (kNumQueries >= 7) { - max_6 = hn::Max(x_6_p0, x_6_p1); - } - if constexpr (kNumQueries >= 8) { - max_7 = hn::Max(x_7_p0, x_7_p1); - } + VI32 isum0_p0 = hn::Zero(di32), isum0_p1 = hn::Zero(di32); + VI32 isum1_p0 = hn::Zero(di32), isum1_p1 = hn::Zero(di32); + VI32 isum2_p0 = hn::Zero(di32), isum2_p1 = hn::Zero(di32); + VI32 isum3_p0 = hn::Zero(di32), isum3_p1 = hn::Zero(di32); + VI32 isum4_p0 = hn::Zero(di32), isum4_p1 = hn::Zero(di32); + VI32 isum5_p0 = hn::Zero(di32), isum5_p1 = hn::Zero(di32); + VI32 isum6_p0 = hn::Zero(di32), isum6_p1 = hn::Zero(di32); + VI32 isum7_p0 = hn::Zero(di32), isum7_p1 = hn::Zero(di32); - if constexpr (kNumQueries == 1) { - new_max = hn::InsertLane(new_max, 0, hn::ReduceMax(df, max_0)); - } else { - new_max = - Reduce8(df, max_0, max_1, max_2, max_3, max_4, max_5, max_6, max_7, - [](auto a, auto b) HWY_ATTR { return hn::Max(a, b); }); - } - if (att_cap > 0.0f) { - VF8 cap = hn::Set(df8, att_cap); - VF8 one_over_cap = hn::Set(df8, one_over_att_cap); - new_max = hn::Mul(cap, hn::FastTanh(df8, hn::Mul(new_max, one_over_cap))); - } - VF8 local_max = new_max; - VF8 old_max_vf = hn::Set(df8, kMaskedLogitVal); - old_max_vf = hn::LoadU(df8, old_max); - new_max = hn::Max(new_max, old_max_vf); - auto changed_max = hn::Gt(new_max, hn::Set(df8, kMaskedLogitVal)); - hn::StoreU(new_max, df8, old_max); - - auto apply_exp = [&](int i, VF& x_p0, VF& x_p1) HWY_ATTR { - const VF new_max_i = hn::Set(df, old_max[i]); - x_p0 = hn::FastExpMinusOrZero(df, hn::Sub(x_p0, new_max_i)); - x_p1 = hn::FastExpMinusOrZero(df, hn::Sub(x_p1, new_max_i)); - }; + const int32_t* q_base_i32 = HWY_RCAST_ALIGNED(const int32_t*, q_base); + const size_t q_stride = qkv_dim / 4; - if constexpr (kNumQueries >= 1) { - apply_exp(0, x_0_p0, x_0_p1); - } - if constexpr (kNumQueries >= 2) { - apply_exp(1, x_1_p0, x_1_p1); - } - if constexpr (kNumQueries >= 3) { - apply_exp(2, x_2_p0, x_2_p1); - } - if constexpr (kNumQueries >= 4) { - apply_exp(3, x_3_p0, x_3_p1); - } - if constexpr (kNumQueries >= 5) { - apply_exp(4, x_4_p0, x_4_p1); - } - if constexpr (kNumQueries >= 6) { - apply_exp(5, x_5_p0, x_5_p1); - } - if constexpr (kNumQueries >= 7) { - apply_exp(6, x_6_p0, x_6_p1); - } - if constexpr (kNumQueries >= 8) { - apply_exp(7, x_7_p0, x_7_p1); - } - VF8 old_d_vf = hn::Set(df8, 0.0f); - old_d_vf = hn::LoadU(df8, old_d); + for (size_t g = 0; g < qkv_dim / 4; g++) { + const int8_t* k_ptr = k_transposed_tile + g * kTileSize * 4; + auto k_vec_p0 = hn::LoadU(di8, k_ptr); + auto k_vec_p1 = hn::LoadU(di8, k_ptr + hn::Lanes(di8)); - VF8 x_sum = hn::Zero(df8); - if constexpr (kNumQueries == 1) { - x_sum = hn::Set(df8, hn::ReduceSum(df, x_0_p0) + hn::ReduceSum(df, x_0_p1)); - } else { - VF x_0_sum = hn::Add(x_0_p0, x_0_p1); - VF x_1_sum = hn::Add(x_1_p0, x_1_p1); - VF x_2_sum = hn::Add(x_2_p0, x_2_p1); - VF x_3_sum = hn::Add(x_3_p0, x_3_p1); - VF x_4_sum = hn::Add(x_4_p0, x_4_p1); - VF x_5_sum = hn::Add(x_5_p0, x_5_p1); - VF x_6_sum = hn::Add(x_6_p0, x_6_p1); - VF x_7_sum = hn::Add(x_7_p0, x_7_p1); - x_sum = Reduce8(df, x_0_sum, x_1_sum, x_2_sum, x_3_sum, x_4_sum, x_5_sum, - x_6_sum, x_7_sum, - [](auto a, auto b) HWY_ATTR { return hn::Add(a, b); }); - } - VF8 scale = hn::Mul( - old_d_vf, hn::FastExpMinusOrZero(df8, hn::Sub(old_max_vf, new_max))); - old_d_vf = hn::Add(scale, x_sum); - auto non_zero_mask = hn::Gt(old_d_vf, hn::Set(df8, 0.0f)); - const VF zero = hn::Zero(df); - const VF8 zero8 = hn::Zero(df8); - const VF8 one_over_d = - hn::MaskedDivOr(zero8, non_zero_mask, hn::Set(df8, 1.0f), old_d_vf); - VF8 q_scale; - if (q_scales_s != nullptr) { - VF8 max_s = hn::Mul(one_over_d, hn::Exp(df8, hn::Sub(local_max, new_max))); - hn::Store(hn::Mul(max_s, hn::Set(df8, max_v_scale / 32767.0f)), df8, - q_scales_s); - - auto max_s_gt_0 = hn::Gt(max_s, zero8); - float inv_max_v_scale = 1.0f / std::max(max_v_scale, 1e-10f); - VF8 mult = hn::Mul(hn::Set(df8, 32767.0f * inv_max_v_scale), - hn::Exp(df8, hn::Sub(new_max, local_max))); - q_scale = hn::IfThenElse(max_s_gt_0, mult, zero8); - } else { - q_scale = one_over_d; - } - HWY_ALIGN float tmp_one_over_d[8]; - hn::Store(q_scale, df8, tmp_one_over_d); - hn::BlendedStore(old_d_vf, changed_max, df8, old_d); - scale = hn::Mul(scale, one_over_d); - hn::BlendedStore(scale, changed_max, df8, scales); - auto mul_or_zero = [&](VF& x_p0, VF& x_p1, int i) HWY_ATTR { - if (HWY_LIKELY(old_d[i] > 0.0f && scales[i] != 1.0f)) { - const VF one_over_d_i = hn::Set(df, tmp_one_over_d[i]); - x_p0 = hn::Mul(x_p0, one_over_d_i); - x_p1 = hn::Mul(x_p1, one_over_d_i); - } else { - x_p0 = zero; - x_p1 = zero; + auto accumulate = [&](int32_t q_val, VI32& sum_p0, VI32& sum_p1) HWY_ATTR { + auto q_vec = hn::BitCast(du8, hn::Set(di32, q_val)); + + sum_p0 = hn::SumOfMulQuadAccumulate(di32, q_vec, k_vec_p0, sum_p0); + sum_p1 = hn::SumOfMulQuadAccumulate(di32, q_vec, k_vec_p1, sum_p1); + }; + + accumulate(q_base_i32[g], isum0_p0, isum0_p1); + if constexpr (kNumQueries >= 2) { + accumulate(q_base_i32[q_stride + g], isum1_p0, isum1_p1); } + if constexpr (kNumQueries >= 3) { + accumulate(q_base_i32[2 * q_stride + g], isum2_p0, isum2_p1); + } + if constexpr (kNumQueries >= 4) { + accumulate(q_base_i32[3 * q_stride + g], isum3_p0, isum3_p1); + } + if constexpr (kNumQueries >= 5) { + accumulate(q_base_i32[4 * q_stride + g], isum4_p0, isum4_p1); + } + if constexpr (kNumQueries >= 6) { + accumulate(q_base_i32[5 * q_stride + g], isum5_p0, isum5_p1); + } + if constexpr (kNumQueries >= 7) { + accumulate(q_base_i32[6 * q_stride + g], isum6_p0, isum6_p1); + } + if constexpr (kNumQueries >= 8) { + accumulate(q_base_i32[7 * q_stride + g], isum7_p0, isum7_p1); + } + } + + auto apply_correction = [&](const int32_t* k_sums_ptr, VI32& isum0, + VI32& isum1, VI32& isum2, VI32& isum3, + VI32& isum4, VI32& isum5, VI32& isum6, + VI32& isum7) HWY_ATTR { + auto k_sums_vec = hn::LoadU(di32, k_sums_ptr); + auto correction = hn::ShiftLeft<7>(k_sums_vec); + isum0 = hn::Sub(isum0, correction); + if constexpr (kNumQueries >= 2) isum1 = hn::Sub(isum1, correction); + if constexpr (kNumQueries >= 3) isum2 = hn::Sub(isum2, correction); + if constexpr (kNumQueries >= 4) isum3 = hn::Sub(isum3, correction); + if constexpr (kNumQueries >= 5) isum4 = hn::Sub(isum4, correction); + if constexpr (kNumQueries >= 6) isum5 = hn::Sub(isum5, correction); + if constexpr (kNumQueries >= 7) isum6 = hn::Sub(isum6, correction); + if constexpr (kNumQueries >= 8) isum7 = hn::Sub(isum7, correction); }; - mul_or_zero(x_0_p0, x_0_p1, 0); + + apply_correction(k_sums, isum0_p0, isum1_p0, isum2_p0, isum3_p0, isum4_p0, + isum5_p0, isum6_p0, isum7_p0); + apply_correction(k_sums + hn::Lanes(df), isum0_p1, isum1_p1, isum2_p1, + isum3_p1, isum4_p1, isum5_p1, isum6_p1, isum7_p1); + + sum0_p0 = hn::ConvertTo(df, isum0_p0); + sum0_p1 = hn::ConvertTo(df, isum0_p1); if constexpr (kNumQueries >= 2) { - mul_or_zero(x_1_p0, x_1_p1, 1); + sum1_p0 = hn::ConvertTo(df, isum1_p0); + sum1_p1 = hn::ConvertTo(df, isum1_p1); } if constexpr (kNumQueries >= 3) { - mul_or_zero(x_2_p0, x_2_p1, 2); + sum2_p0 = hn::ConvertTo(df, isum2_p0); + sum2_p1 = hn::ConvertTo(df, isum2_p1); } if constexpr (kNumQueries >= 4) { - mul_or_zero(x_3_p0, x_3_p1, 3); + sum3_p0 = hn::ConvertTo(df, isum3_p0); + sum3_p1 = hn::ConvertTo(df, isum3_p1); } if constexpr (kNumQueries >= 5) { - mul_or_zero(x_4_p0, x_4_p1, 4); + sum4_p0 = hn::ConvertTo(df, isum4_p0); + sum4_p1 = hn::ConvertTo(df, isum4_p1); } if constexpr (kNumQueries >= 6) { - mul_or_zero(x_5_p0, x_5_p1, 5); + sum5_p0 = hn::ConvertTo(df, isum5_p0); + sum5_p1 = hn::ConvertTo(df, isum5_p1); } if constexpr (kNumQueries >= 7) { - mul_or_zero(x_6_p0, x_6_p1, 6); + sum6_p0 = hn::ConvertTo(df, isum6_p0); + sum6_p1 = hn::ConvertTo(df, isum6_p1); } if constexpr (kNumQueries >= 8) { - mul_or_zero(x_7_p0, x_7_p1, 7); + sum7_p0 = hn::ConvertTo(df, isum7_p0); + sum7_p1 = hn::ConvertTo(df, isum7_p1); } } -template > -static HWY_INLINE void FlashAttentionTileStepAndApplySoftCap( - DF df, float att_cap, float one_over_att_cap, VF& x_0_p0, VF& x_0_p1, - VF& x_1_p0, VF& x_1_p1, VF& x_2_p0, VF& x_2_p1, VF& x_3_p0, VF& x_3_p1, - VF& x_4_p0, VF& x_4_p1, VF& x_5_p0, VF& x_5_p1, VF& x_6_p0, VF& x_6_p1, - VF& x_7_p0, VF& x_7_p1, float* HWY_RESTRICT old_max, - float* HWY_RESTRICT old_d, float* HWY_RESTRICT scales, size_t query_idx, - float* HWY_RESTRICT q_scales_s = nullptr, float max_v_scale = 1.0f) { - constexpr int kFirstHalfAmountOfQueries = std::min(kNumQueries, 4); - [[maybe_unused]] constexpr int kSecondHalfAmountOfQueries = - kNumQueries - kFirstHalfAmountOfQueries; - if constexpr (kNumQueries <= 4) { - FlashAttentionTileStepAndApplySoftCap4( - df, att_cap, one_over_att_cap, x_0_p0, x_0_p1, x_1_p0, x_1_p1, x_2_p0, - x_2_p1, x_3_p0, x_3_p1, old_max + query_idx, old_d + query_idx, scales, - q_scales_s, max_v_scale); - } else { -#if HWY_MAX_BYTES <= 16 - // Codepath if we have only 4 float lanes. - HWY_DASSERT(hn::Lanes(df) >= 4); - FlashAttentionTileStepAndApplySoftCap4<4>( - df, att_cap, one_over_att_cap, x_0_p0, x_0_p1, x_1_p0, x_1_p1, x_2_p0, - x_2_p1, x_3_p0, x_3_p1, old_max + query_idx, old_d + query_idx, scales, - q_scales_s, max_v_scale); - FlashAttentionTileStepAndApplySoftCap4( - df, att_cap, one_over_att_cap, x_4_p0, x_4_p1, x_5_p0, x_5_p1, x_6_p0, - x_6_p1, x_7_p0, x_7_p1, old_max + query_idx + 4, old_d + query_idx + 4, - scales + 4, q_scales_s == nullptr ? nullptr : q_scales_s + 4, - max_v_scale); -#else - FlashAttentionTileStepAndApplySoftCap8( - df, att_cap, one_over_att_cap, x_0_p0, x_0_p1, x_1_p0, x_1_p1, x_2_p0, - x_2_p1, x_3_p0, x_3_p1, x_4_p0, x_4_p1, x_5_p0, x_5_p1, x_6_p0, x_6_p1, - x_7_p0, x_7_p1, old_max + query_idx, old_d + query_idx, scales, - q_scales_s, max_v_scale); -#endif + +template > +HWY_INLINE HWY_MAYBE_UNUSED void MulByConstAndAddTileVNNI( + DF df, const float* HWY_RESTRICT scales_old, size_t actual_steps, + const uint8_t* HWY_RESTRICT step_cs_u8, + 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(); + constexpr size_t kMaxLanes = hn::MaxLanes(df); + HWY_LANES_CONSTEXPR size_t NF = hn::Lanes(df); + + using DI32 = hn::Repartition; + const DI32 di32; + using VI32 = hn::Vec; + using DI8 = hn::Repartition; + const DI8 di8; + using DU8 = hn::Repartition; + const DU8 du8; + + 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); + + for (size_t step_idx = 0; step_idx < actual_steps; ++step_idx) { + const uint8_t* cs_u8 = step_cs_u8 + step_idx * (N * 2 * kMaxLanes); + const int8_t* v_tile = step_v_tiles[step_idx]; + + for (size_t g = 0; g < (NF * 2) / 4; ++g) { + const int8_t* v_ptr0 = v_tile + g * qkv_dim * 4 + i * 4; + auto v_vec0 = hn::LoadU(di8, v_ptr0); + + const int8_t* v_ptr1 = v_tile + g * qkv_dim * 4 + (i + NF) * 4; + auto v_vec1 = hn::LoadU(di8, v_ptr1); + + auto mul_acc = [&](int j, VI32& a0, VI32& a1) HWY_ATTR { + const uint8_t* cs_u8_j = cs_u8 + j * 2 * kMaxLanes; + uint32_t s_32 = *HWY_RCAST_ALIGNED(const uint32_t*, cs_u8_j + 4 * g); + auto s_vec = + hn::BitCast(du8, hn::Set(di32, static_cast(s_32))); + + a0 = hn::SumOfMulQuadAccumulate(di32, s_vec, v_vec0, a0); + a1 = hn::SumOfMulQuadAccumulate(di32, s_vec, v_vec1, a1); + }; + + mul_acc(0, acc0_0, acc0_1); + if constexpr (N >= 2) mul_acc(1, acc1_0, acc1_1); + if constexpr (N >= 3) mul_acc(2, acc2_0, acc2_1); + if constexpr (N >= 4) mul_acc(3, acc3_0, acc3_1); + if constexpr (N >= 5) mul_acc(4, acc4_0, acc4_1); + if constexpr (N >= 6) mul_acc(5, acc5_0, acc5_1); + if constexpr (N >= 7) mul_acc(6, acc6_0, acc6_1); + if constexpr (N >= 8) mul_acc(7, acc7_0, acc7_1); + } + } + + const float* q_scales_s = step_q_scales_s[0]; + auto convert_and_add = [&](int j, VI32& a0, VI32& a1, VF& o0, VF& o1) + HWY_ATTR { + VF f0 = hn::ConvertTo(df, a0); + VF f1 = hn::ConvertTo(df, a1); + + VF scale_new = hn::Set(df, q_scales_s[j]); + o0 = hn::MulAdd(f0, scale_new, o0); + o1 = hn::MulAdd(f1, scale_new, o1); + }; + + 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); + + convert_and_add(0, acc0_0, acc0_1, out0_0, out0_1); + if constexpr (N >= 2) convert_and_add(1, acc1_0, acc1_1, out1_0, out1_1); + if constexpr (N >= 3) convert_and_add(2, acc2_0, acc2_1, out2_0, out2_1); + if constexpr (N >= 4) convert_and_add(3, acc3_0, acc3_1, out3_0, out3_1); + if constexpr (N >= 5) convert_and_add(4, acc4_0, acc4_1, out4_0, out4_1); + if constexpr (N >= 6) convert_and_add(5, acc5_0, acc5_1, out5_0, out5_1); + if constexpr (N >= 7) convert_and_add(6, acc6_0, acc6_1, out6_0, out6_1); + if constexpr (N >= 8) convert_and_add(7, acc7_0, acc7_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; } } @@ -1192,7 +1076,6 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidthBF16( #endif } - // Performs tiled flash attention for arbitrary number of queries // It depends on kv being tiled. // Runs 2 loops one over tiles, and inner one over queries(up to 4 at a time). @@ -1216,14 +1099,6 @@ 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. -#ifndef BENCHMARK_BLOCK_SIZE_BF16 -#define BENCHMARK_BLOCK_SIZE_BF16 256 -#endif - -#ifndef BENCHMARK_BLOCK_SIZE_INT8 -#define BENCHMARK_BLOCK_SIZE_INT8 64 -#endif - template static HWY_INLINE void ComputeQKMacroTile( @@ -1291,6 +1166,20 @@ static HWY_INLINE void ComputeQKMacroTile( 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 (IsInt8()) { + const int8_t* k_transposed_tile = + reinterpret_cast(tile_base); + const int32_t* k_sums = reinterpret_cast( + k_transposed_tile + qkv_dim * 2 * kTileSize + + kTileSize * 2 * sizeof(BF16)); + const int8_t* k_ptr = k_transposed_tile + pos_in_tile * 4; + const int32_t* k_sums_ptr = k_sums + pos_in_tile; + + QDotKTilexVNNI( + df, reinterpret_cast(q_group), k_ptr, qkv_dim, + k_sums_ptr, 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 + @@ -1362,6 +1251,9 @@ static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( const DI16 di16; using DBF = hn::Repartition; const DBF dbf; + using DU8 = hn::Repartition; + [[maybe_unused]] const DU8 du8; + const auto du8_half = hn::Half(); using D1 = hn::CappedTag; const D1 d1; @@ -1381,7 +1273,7 @@ static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( p0 = hn::Mul(p0, scales_p0); p1 = hn::Mul(p1, scales_p1); } - if constexpr (IsInt16()) { + if constexpr (IsInt16() || IsInt8()) { VF s = hn::Set(df, q_scale_val); p0 = hn::Mul(p0, s); p1 = hn::Mul(p1, s); @@ -1479,16 +1371,12 @@ static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( 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 block_max_val = 1e-10f; + if constexpr (IsInt16() || IsInt8()) { + 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); float step_max_v_scale = 1.0f; if constexpr (IsInt8()) { const PackedSpan scales_span = @@ -1499,18 +1387,37 @@ static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( } 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); + float eff_v_scale = std::max(step_max_v_scale, 1e-10f); + block_max_val = std::max(block_max_val, step_max_val * eff_v_scale); + } + } + + 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() || IsInt8()) { + constexpr float kMaxQuantVal = IsInt16() ? 32767.0f : 255.0f; + if (block_max_val > 1e-10f && q_scale > 0.0f) { + float scale_to_quant = kMaxQuantVal / block_max_val; + VF s_quant = hn::Set(df, scale_to_quant); + p0 = hn::Mul(p0, s_quant); + p1 = hn::Mul(p1, s_quant); + for (size_t s_idx = 0; s_idx < actual_steps; ++s_idx) { + step_q_scales_s_buf_ptr[s_idx * kNumQueries + q_offset] = + q_scale * block_max_val / kMaxQuantVal; + } } else { p0 = hn::Zero(df); p1 = hn::Zero(df); - step_q_scales_s_buf_ptr[step_idx * kNumQueries + q_offset] = 0.0f; + for (size_t s_idx = 0; s_idx < actual_steps; ++s_idx) { + step_q_scales_s_buf_ptr[s_idx * kNumQueries + q_offset] = 0.0f; + } } } @@ -1536,6 +1443,14 @@ static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( auto i0 = hn::OrderedDemote2To(di16, hn::NearestInt(p0), hn::NearestInt(p1)); hn::Store(i0, di16, dst); + } else if constexpr (IsInt8()) { + uint8_t* dst = ((uint8_t*)step_consts_buf) + + step_idx * (kNumQueries * 2 * kMaxLanes) + + q_offset * 2 * kMaxLanes; + auto i0 = + hn::OrderedDemote2To(di16, hn::NearestInt(p0), hn::NearestInt(p1)); + auto u0_half = hn::DemoteTo(du8_half, i0); + hn::Store(u0_half, du8_half, dst); } else { BF16* dst = ((BF16*)step_consts_buf) + step_idx * (kNumQueries * 2 * kMaxLanes) + @@ -1547,9 +1462,10 @@ static HWY_INLINE void UpdateOnlineSoftmaxAndPackSingleQuery( } template -using FlashStepBufT = - hwy::If(), float, - std::conditional_t(), int16_t, BF16>>; +using FlashStepBufT = hwy::If< + IsF32(), float, + std::conditional_t(), int16_t, + std::conditional_t(), uint8_t, BF16>>>; template @@ -1617,11 +1533,10 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle( step_microscaling_k_ptrs[step_idx] = microscaling_scales_k; step_microscaling_v_ptrs[step_idx] = microscaling_scales_v; } - if constexpr (IsInt16()) { + if constexpr (IsInt16() || IsInt8()) { step_q_scales_s_ptrs[step_idx] = step_q_scales_s_buf + step_idx * kNumQueries; } - } for (size_t q_offset = 0; q_offset < actual_q_count; ++q_offset) { @@ -1639,7 +1554,8 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle( 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, + step_consts_buf, + (IsInt16() || IsInt8()) ? 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); } @@ -1657,6 +1573,10 @@ static HWY_INLINE void ComputeSoftmaxAndSVBundle( MulByConstAndAddTileUpTo8_BF16_Int16( df, scales_old, actual_steps, step_consts_buf, step_v_tiles, group_out, step_q_scales_s_ptrs); + } else if constexpr (IsInt8()) { + MulByConstAndAddTileVNNI(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); @@ -1707,25 +1627,19 @@ struct TileFlashAttentionWorkspaceLayout { step_mk_bytes(hwy::RoundUpTo( (IsInt8() ? 256 : 1) * sizeof(const BF16*), HWY_ALIGNMENT)), step_qp_offset(step_mk_offset + step_mk_bytes), - step_qp_bytes(hwy::RoundUpTo( - (IsInt16() ? 256 : 1) * sizeof(const float*), HWY_ALIGNMENT)), + step_qp_bytes(hwy::RoundUpTo((IsInt16() || IsInt8()) + ? 256 * sizeof(const float*) + : sizeof(const float*), + HWY_ALIGNMENT)), step_qb_offset(step_qp_offset + step_qp_bytes), - step_qb_bytes(hwy::RoundUpTo( - (IsInt16() ? (256 * kNumQueriesPerLoop) : 1) * sizeof(float), - HWY_ALIGNMENT)), + step_qb_bytes( + hwy::RoundUpTo((IsInt16() || IsInt8()) + ? (256 * kNumQueriesPerLoop) * sizeof(float) + : sizeof(float), + HWY_ALIGNMENT)), 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, @@ -1745,9 +1659,11 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( 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 kBlockSizeBf16 = 256; + constexpr size_t kBlockSizeInt8 = + (HWY_ARCH_X86 && (HWY_TARGET >= HWY_AVX2)) ? 64 : 256; constexpr size_t kBlockSize = - IsInt8() ? BENCHMARK_BLOCK_SIZE_INT8 : BENCHMARK_BLOCK_SIZE_BF16; + IsInt8() ? kBlockSizeInt8 : kBlockSizeBf16; HWY_LANES_CONSTEXPR size_t step_size = 2 * kHTileSize; size_t smallest_start_pos = std::numeric_limits::max(); @@ -1901,7 +1817,7 @@ HWY_NOINLINE void TileFlashAttentionReturnExpSumsAndMaxLogits( MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, float* HWY_RESTRICT max_logits) { constexpr int kDefaultChunkSize = (HWY_REGISTERS >= 32) ? 8 : 4; - TileFlashAttentionReturnExpSumsAndMaxLogitsImpl( + 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); } @@ -1928,7 +1844,8 @@ void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsBF16( MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, float* HWY_RESTRICT max_logits) { CallUpcastedKVs(kvs, [&](const auto& kv_t) { - return TileFlashAttentionReturnExpSumsAndMaxLogits( + return TileFlashAttentionReturnExpSumsAndMaxLogits< + typename std::decay_t::T, BF16>( kv_t, q_count, q_base, {}, start_pos_per_query, last_pos_per_query, att_cap, att_out, exp_denominator_sums, max_logits); }); @@ -1952,6 +1869,24 @@ void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16( last_pos_per_query, att_cap, att_out, exp_denominator_sums, max_logits); } +void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( + hwy::Span kvs, size_t q_count, + const int8_t* HWY_RESTRICT q_base, const hwy::Span q_scales, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, float att_cap, + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT max_logits) { + for ([[maybe_unused]] auto&& mat : kvs) { + HWY_DASSERT(mat.GetType() == Type::kInt8); + } + auto matptrs = MakeMatPtrVec(kvs); + hwy::Span> matptrs_span(matptrs.data(), matptrs.size()); + + return TileFlashAttentionReturnExpSumsAndMaxLogits( + matptrs_span, q_count, q_base, q_scales, start_pos_per_query, + last_pos_per_query, att_cap, att_out, exp_denominator_sums, max_logits); +} + template MatT GetKVTypeHelper(const hwy::Span>&); diff --git a/gemma/flash_attention.h b/gemma/flash_attention.h index 602f15fc..e3e2ddbd 100644 --- a/gemma/flash_attention.h +++ b/gemma/flash_attention.h @@ -92,6 +92,14 @@ namespace gcpp { MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, \ float* HWY_RESTRICT max_logits); \ \ + void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( \ + hwy::Span kvs, size_t q_count, \ + const int8_t* HWY_RESTRICT q_base, 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); \ + \ /* NOLINTNEXTLINE(google-readability-namespace-comments) */ \ } // namespace NAMESPACE diff --git a/gemma/flash_attention_arm-inl.h b/gemma/flash_attention_arm-inl.h index 7e8ed08d..b6c6c834 100644 --- a/gemma/flash_attention_arm-inl.h +++ b/gemma/flash_attention_arm-inl.h @@ -32,19 +32,14 @@ #include "ops/ops-inl.h" #include "hwy/contrib/math/fast_math-inl.h" -#ifndef BENCHMARK_BLOCK_SIZE_BF16 -#define BENCHMARK_BLOCK_SIZE_BF16 128 -#endif - -#ifndef BENCHMARK_BLOCK_SIZE_INT8 -#define BENCHMARK_BLOCK_SIZE_INT8 512 -#endif - HWY_BEFORE_NAMESPACE(); namespace gcpp { namespace HWY_NAMESPACE { namespace hn = hwy::HWY_NAMESPACE; +constexpr size_t kArmBlockSizeBf16 = 128; +constexpr size_t kArmBlockSizeInt8 = 512; + struct TileAttentionGroupParams { size_t smallest_start_pos; size_t largest_last_pos; @@ -344,7 +339,7 @@ HWY_INLINE void QuantizeAndPackSoftmaxProbs( const size_t num_groups = actual_block_size / 8; const size_t num_qp = 4; - constexpr size_t kBlockSize = BENCHMARK_BLOCK_SIZE_INT8; + constexpr size_t kBlockSize = kArmBlockSizeInt8; // Process 2 queries at a time to reduce L1 working set and avoid large macros HWY_ALIGN float w_buf[2 * kBlockSize]; @@ -471,14 +466,14 @@ HWY_INLINE void TileFlashAttentionSVBlockInt8( size_t ch_g = ch_base / 8; // Initialize accumulators: up to 4 query pairs x 4 channel pairs - VI32 acc00 = hn::Zero(di32), acc01 = hn::Zero(di32), - acc02 = hn::Zero(di32), acc03 = hn::Zero(di32); - VI32 acc10 = hn::Zero(di32), acc11 = hn::Zero(di32), - acc12 = hn::Zero(di32), acc13 = hn::Zero(di32); - VI32 acc20 = hn::Zero(di32), acc21 = hn::Zero(di32), - acc22 = hn::Zero(di32), acc23 = hn::Zero(di32); - VI32 acc30 = hn::Zero(di32), acc31 = hn::Zero(di32), - acc32 = hn::Zero(di32), acc33 = hn::Zero(di32); + VI32 acc00 = hn::Zero(di32), acc01 = hn::Zero(di32), acc02 = hn::Zero(di32), + acc03 = hn::Zero(di32); + VI32 acc10 = hn::Zero(di32), acc11 = hn::Zero(di32), acc12 = hn::Zero(di32), + acc13 = hn::Zero(di32); + VI32 acc20 = hn::Zero(di32), acc21 = hn::Zero(di32), acc22 = hn::Zero(di32), + acc23 = hn::Zero(di32); + VI32 acc30 = hn::Zero(di32), acc31 = hn::Zero(di32), acc32 = hn::Zero(di32), + acc33 = hn::Zero(di32); // 3. Accumulate over steps in int32 for (size_t g = 0; g < num_groups; ++g) { @@ -497,14 +492,13 @@ HWY_INLINE void TileFlashAttentionSVBlockInt8( // bytes (4 query pairs) regardless of kNumQueries. const int8_t* q_w_ptr = q_weights_pre + g * 64; - auto load_A = [&](size_t idx) HWY_ATTR { - return hn::LoadU(di8, q_w_ptr + idx * 16); - }; + auto load_A = [&](size_t idx) + HWY_ATTR { return hn::LoadU(di8, q_w_ptr + idx * 16); }; - Accumulate4x4Grid(di32, load_A, B0, B1, B2, B3, acc00, - acc01, acc02, acc03, acc10, acc11, acc12, - acc13, acc20, acc21, acc22, acc23, acc30, - acc31, acc32, acc33); + Accumulate4x4Grid(di32, load_A, B0, B1, B2, B3, acc00, acc01, + acc02, acc03, acc10, acc11, acc12, acc13, + acc20, acc21, acc22, acc23, acc30, acc31, + acc32, acc33); } // 4. Dequantize and write back @@ -578,7 +572,7 @@ HWY_INLINE void TileFlashAttentionSVBlockBF16( using VF4 = hn::Vec; constexpr size_t kBf16Lanes = kRegBytes / sizeof(BF16); - constexpr size_t kBlockSize = BENCHMARK_BLOCK_SIZE_BF16; + constexpr size_t kBlockSize = kArmBlockSizeBf16; using KV_PTR_T = const KV_T*; // Dynamically scale the tile pointer buffer with the block size. @@ -986,7 +980,7 @@ HWY_ATTR void TileFlashAttentionReturnExpSumsAndMaxLogitsBF16_Impl( constexpr int kNumQueriesPerLoop = 8; constexpr size_t kTileSize = 32; constexpr size_t kBlockSize = - IsInt8() ? BENCHMARK_BLOCK_SIZE_INT8 : BENCHMARK_BLOCK_SIZE_BF16; + IsInt8() ? kArmBlockSizeInt8 : kArmBlockSizeBf16; const size_t kStepsPerTile = std::max(size_t(1), KVCache::kTileSize / kBf16Lanes); @@ -1063,7 +1057,7 @@ HWY_ATTR void TileFlashAttentionReturnExpSumsAndMaxLogitsBF16_Impl( size_t macro_tile_start_pos = position; auto inner_loop_qk = [&](size_t query_idx, - size_t step_idx) HWY_ATTR { + size_t step_idx) HWY_ATTR { size_t loop_idx = query_idx / kNumQueriesPerLoop; size_t step_pos = position + step_idx * kBf16Lanes; if (step_pos + kBf16Lanes <= min_start_pos_per_group[loop_idx] || diff --git a/gemma/flash_attention_test.cc b/gemma/flash_attention_test.cc index 36e2aeb6..ddbb6ab0 100644 --- a/gemma/flash_attention_test.cc +++ b/gemma/flash_attention_test.cc @@ -551,6 +551,7 @@ void RunTiledFlashAttentionTest(gcpp::KVEncoding kv_encoding, } hwy::Span kvs(&kv, 1); + if (attention_impl == AttentionImpl::kFlashTransposedQsBF16) { std::vector> bf16_queries(num_queries * qkv_dim); @@ -574,6 +575,18 @@ void RunTiledFlashAttentionTest(gcpp::KVEncoding kv_encoding, hwy::Span(start_pos_per_query), hwy::Span(last_pos_per_query), att_cap, att_out, exp_denominator_sums.data(), max_logits.data()); + } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + std::vector> int8_queries( + num_queries * qkv_dim); + AlignedFloatVector q_scales(num_queries); + CompressQueriesInt8Contiguous(q_all.data(), qkv_dim, num_queries, + int8_queries.data(), q_scales.data()); + + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( + kvs, num_queries, int8_queries.data(), q_scales, + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), att_cap, att_out, + exp_denominator_sums.data(), max_logits.data()); } else if (attention_impl == AttentionImpl::kInt8MatrixAccumulation) { size_t num_queries_rounded = hwy::RoundUpTo(num_queries, 2); hwy::AlignedVector int8_queries(num_queries_rounded * qkv_dim); @@ -667,9 +680,8 @@ void TestTiledFlashAttentionBF16() { } void TestTiledFlashAttentionInt8() { - RunTiledFlashAttentionTest(gcpp::KVEncoding::kInt8, - AttentionImpl::kFlash, 5e-3f, 2e-2f, - 1e-3f); + RunTiledFlashAttentionTest( + gcpp::KVEncoding::kInt8, AttentionImpl::kFlash, 5e-3f, 2e-2f, 1e-3f); } void TestTiledFlashAttentionInt8BF16() { @@ -684,6 +696,12 @@ void TestTiledFlashAttentionInt8Int16() { 5e-3f, 2e-2f, 1e-3f); } +void TestTiledFlashAttentionInt8VNNI() { + RunTiledFlashAttentionTest( + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashTransposedQsInt8, 5e-3f, 2e-2f, 1e-3f); +} + void TestTiledFlashAttentionBF16MatrixAccumulation() { const hn::ScalableTag dbf; if (hn::Lanes(dbf) > 32) { @@ -721,10 +739,12 @@ void TestTiledFlashAttentionInt8MatrixAccumulation() { } template -void RunTiledFlashAttentionDifferentialTest( - size_t kv_seq_len, float tol, float tol_exp, float tol_max, - gcpp::KVEncoding ref_encoding, gcpp::KVEncoding opt_encoding, - AttentionImpl opt_impl, const char* type_name) { +void RunTiledFlashAttentionDifferentialTest(size_t kv_seq_len, float tol, + float tol_exp, float tol_max, + gcpp::KVEncoding ref_encoding, + gcpp::KVEncoding opt_encoding, + AttentionImpl opt_impl, + const char* type_name) { const hn::ScalableTag dbf; if (hn::Lanes(dbf) > 32) { GTEST_SKIP() << "Skipping MatrixAccumulation test for target with register " @@ -777,11 +797,10 @@ void RunTiledFlashAttentionDifferentialTest( // Set up Optimized cache (Matrix Accumulation) size_t opt_tile_size_bytes = *gcpp::GetTileSizeBytes(opt_encoding, qkv_dim); size_t opt_tile_size_in_elements = opt_tile_size_bytes / sizeof(KV_T); - MatStorageT kv( - "kv", - Extents2D(padded_kv_seq_len / gcpp::KVCache::kTileSize, - opt_tile_size_in_elements), - ctx.allocator, MatPadding::kPacked); + MatStorageT kv("kv", + Extents2D(padded_kv_seq_len / gcpp::KVCache::kTileSize, + opt_tile_size_in_elements), + ctx.allocator, MatPadding::kPacked); PopulateTestKVCache(kv, opt_encoding, qkv_dim); AlignedFloatVector q_all = PopulateTestQueries(num_queries, qkv_dim); @@ -870,7 +889,8 @@ void RunTiledFlashAttentionDifferentialTest( if (failures < 5) { EXPECT_NEAR(exp_denominator_sums[i], exp_denominator_sums_ref[i], tol_exp) - << "i=" << i << " (Type: " << type_name << ", SeqLen: " << kv_seq_len << ")"; + << "i=" << i << " (Type: " << type_name + << ", SeqLen: " << kv_seq_len << ")"; } failures++; } @@ -879,7 +899,8 @@ void RunTiledFlashAttentionDifferentialTest( if (diff_max >= tol_max) { if (failures < 5) { EXPECT_NEAR(max_logits[i], max_logits_ref[i], tol_max) - << "i=" << i << " (Type: " << type_name << ", SeqLen: " << kv_seq_len << ")"; + << "i=" << i << " (Type: " << type_name + << ", SeqLen: " << kv_seq_len << ")"; } failures++; } @@ -896,7 +917,8 @@ void RunTiledFlashAttentionDifferentialTest( if (diff >= tol) { if (failures < 5) { EXPECT_NEAR(v_out, v_ref, tol) - << "i=" << i << " j=" << j << " (Type: " << type_name << ", SeqLen: " << kv_seq_len << ")"; + << "i=" << i << " j=" << j << " (Type: " << type_name + << ", SeqLen: " << kv_seq_len << ")"; } failures++; } @@ -913,16 +935,14 @@ void RunTiledFlashAttentionDifferentialTest( void TestTiledFlashAttentionBF16MatrixAccumulationLargeVerification() { RunTiledFlashAttentionDifferentialTest( - 2048, 1.0e-1f, 1.1e-1f, 1e-4f, - gcpp::KVEncoding::kBF16TwoTranspositions, + 2048, 1.0e-1f, 1.1e-1f, 1e-4f, gcpp::KVEncoding::kBF16TwoTranspositions, gcpp::KVEncoding::kBF16MatrixAccumulation, AttentionImpl::kFlashMatrixAccumulation, "BF16"); } void TestTiledFlashAttentionInt8MatrixAccumulationLargeVerification() { RunTiledFlashAttentionDifferentialTest( - 1024, 1.5e-1f, 5.0, 8.0e-2f, - gcpp::KVEncoding::kInt8TwoTranspositions, + 1024, 1.5e-1f, 5.0, 8.0e-2f, gcpp::KVEncoding::kInt8TwoTranspositions, gcpp::KVEncoding::kInt8MatrixAccumulation, AttentionImpl::kInt8MatrixAccumulation, "Int8"); } diff --git a/gemma/gemma.cc b/gemma/gemma.cc index a2f93a59..419ed91a 100644 --- a/gemma/gemma.cc +++ b/gemma/gemma.cc @@ -85,7 +85,10 @@ void Attention(LayerAttentionType type, const size_t num_tokens, const int kFlags = 0; if (activations.attention_impl == AttentionImpl::kFlashTransposedQs || activations.attention_impl == AttentionImpl::kFlashTransposedQsBF16 || - activations.attention_impl == AttentionImpl::kFlashTransposedQsInt16) { + activations.attention_impl == AttentionImpl::kFlashTransposedQsInt16 || + activations.attention_impl == AttentionImpl::kFlashTransposedQsInt8 || + activations.attention_impl == AttentionImpl::kInt8MatrixAccumulation || + activations.attention_impl == AttentionImpl::kFlashMatrixAccumulation) { TiledAttention(activations.attention_impl, num_tokens, layer_idx, layer, activations.attention, qbatch, env, kFlags); return; diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index d44996e6..3129e222 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -177,6 +177,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, if (runtime_config.attention_impl == AttentionImpl::kFlash || runtime_config.attention_impl == AttentionImpl::kFlashTransposedQs || runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsInt16 || + runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsInt8 || runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsBF16 || runtime_config.attention_impl == AttentionImpl::kFlashMatrixAccumulation || runtime_config.attention_impl == AttentionImpl::kInt8MatrixAccumulation @@ -211,6 +212,8 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, kv_cache_type = runtime_config.kv_cache_type.value_or(Type::kBF16); } else if (runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsInt16 || + runtime_config.attention_impl == + AttentionImpl::kFlashTransposedQsInt8 || runtime_config.attention_impl == AttentionImpl::kInt8MatrixAccumulation) { if (runtime_config.kv_cache_type.has_value() && @@ -231,6 +234,11 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, if (kv_cache_type == Type::kInt8) { // microscaling max_tile_length += 2 * sizeof(BF16) * kTileSize; + if (runtime_config.attention_impl == + AttentionImpl::kFlashTransposedQsInt8) { + // K sums + max_tile_length += sizeof(int32_t) * kTileSize; + } } auto num_tiles_per_head = [](size_t window_size, size_t prefill_tbatch_size, size_t max_seq_len) { @@ -314,6 +322,11 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, 2 * config.layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { layer_tile_length += 2 * sizeof(BF16) * kTileSize; + if (runtime_config.attention_impl == + AttentionImpl::kFlashTransposedQsInt8) { + // K sums + layer_tile_length += sizeof(int32_t) * kTileSize; + } } bool is_global = config.IsGlobalLayer(i); for (size_t kv = 0; kv < config.layer_configs[i].kv_heads; ++kv) { diff --git a/gemma/kv_transcoding.cc b/gemma/kv_transcoding.cc index 865c5998..bb1cd4e5 100644 --- a/gemma/kv_transcoding.cc +++ b/gemma/kv_transcoding.cc @@ -27,6 +27,10 @@ std::optional GetTileSizeBytes(gcpp::KVEncoding encoding, case gcpp::KVEncoding::kInt8MatrixAccumulation: return qkv_dim * kTileSize * 2 * sizeof(int8_t) + kTileSize * 2 * sizeof(gcpp::KV_microscale_t); + case gcpp::KVEncoding::kInt8VNNITwoTranspositions: + return qkv_dim * kTileSize * 2 * sizeof(int8_t) + + kTileSize * 2 * sizeof(gcpp::KV_microscale_t) + + kTileSize * sizeof(int32_t); case gcpp::KVEncoding::kBF16: case gcpp::KVEncoding::kBF16TwoTranspositions: case gcpp::KVEncoding::kBF16MatrixAccumulation: @@ -271,6 +275,92 @@ void EncodeTileInt8(bool transposed, size_t qkv_dim, const DecodedTile& decoded, }); } +void EncodeTileInt8VNNI(size_t qkv_dim, const DecodedTile& decoded, + hwy::Span out_encoded_tile_data) { + int8_t* k_data = HWY_RCAST_ALIGNED(int8_t*, out_encoded_tile_data.data()); + int8_t* v_data = k_data + qkv_dim * kTileSize; + gcpp::KV_microscale_t* scales = + HWY_RCAST_ALIGNED(gcpp::KV_microscale_t*, v_data + kTileSize * qkv_dim); + gcpp::KV_microscale_t* k_scales = scales; + gcpp::KV_microscale_t* v_scales = scales + kTileSize; + int32_t* k_sums = HWY_RCAST_ALIGNED(int32_t*, v_scales + kTileSize); + + AlignedFloatVector k_max_abs(kTileSize, 0.0f); + AlignedFloatVector v_max_abs(kTileSize, 0.0f); + + for (size_t token = 0; token < kTileSize; ++token) { + for (size_t dim = 0; dim < qkv_dim; ++dim) { + k_max_abs[token] = + std::max(k_max_abs[token], std::abs(decoded.k_elem(token, dim))); + v_max_abs[token] = + std::max(v_max_abs[token], std::abs(decoded.v_elem(token, dim))); + } + } + + AlignedFloatVector inv_scales_k(kTileSize); + AlignedFloatVector inv_scales_v(kTileSize); + for (size_t token = 0; token < kTileSize; ++token) { + float scale_k = k_max_abs[token] == 0.0f ? 1.0f : k_max_abs[token] / 127.0f; + k_scales[token] = hwy::ConvertScalarTo(scale_k); + inv_scales_k[token] = 1.0f / scale_k; + + float scale_v = v_max_abs[token] == 0.0f ? 1.0f : v_max_abs[token] / 127.0f; + v_scales[token] = hwy::ConvertScalarTo(scale_v); + inv_scales_v[token] = 1.0f / scale_v; + } + + for (size_t token = 0; token < kTileSize; ++token) { + k_sums[token] = 0; + } + + auto KOffset_VNNI = [&](size_t dim, size_t token) { + return (dim / 4) * kTileSize * 4 + token * 4 + (dim % 4); + }; + auto VOffset_VNNI = [&](size_t dim, size_t token) { + return (token / 4) * qkv_dim * 4 + dim * 4 + (token % 4); + }; + + EncodeTileWithFn( + qkv_dim, decoded, + [&](size_t dim, size_t token, float val) HWY_ATTR { + int8_t quantized = Quantize(val, inv_scales_k[token]); + k_data[KOffset_VNNI(dim, token)] = quantized; + k_sums[token] += quantized; + }, + [&](size_t dim, size_t token, float val) HWY_ATTR { + v_data[VOffset_VNNI(dim, token)] = Quantize(val, inv_scales_v[token]); + }); +} + +void DecodeTileInt8VNNI(size_t qkv_dim, hwy::Span encoded_tile_data, + DecodedTile* out) { + const int8_t* k_data = + HWY_RCAST_ALIGNED(const int8_t*, encoded_tile_data.data()); + const int8_t* v_data = k_data + qkv_dim * kTileSize; + const gcpp::KV_microscale_t* scales = HWY_RCAST_ALIGNED( + const gcpp::KV_microscale_t*, v_data + kTileSize * qkv_dim); + const gcpp::KV_microscale_t* k_scales = scales; + const gcpp::KV_microscale_t* v_scales = scales + kTileSize; + + auto KOffset_VNNI = [&](size_t dim, size_t token) { + return (dim / 4) * kTileSize * 4 + token * 4 + (dim % 4); + }; + auto VOffset_VNNI = [&](size_t dim, size_t token) { + return (token / 4) * qkv_dim * 4 + dim * 4 + (token % 4); + }; + + DecodeTileWithFn( + qkv_dim, out, + [&](size_t dim, size_t token) HWY_ATTR { + float scale = hwy::ConvertScalarTo(k_scales[token]); + return k_data[KOffset_VNNI(dim, token)] * scale; + }, + [&](size_t dim, size_t token) HWY_ATTR { + float scale = hwy::ConvertScalarTo(v_scales[token]); + return v_data[VOffset_VNNI(dim, token)] * scale; + }); +} + void DecodeTileF32(bool transposed, size_t qkv_dim, hwy::Span encoded_tile_data, DecodedTile* out) { const float* data = HWY_RCAST_ALIGNED(const float*, encoded_tile_data.data()); @@ -401,6 +491,7 @@ bool IsTransposed(KVEncoding encoding) { case KVEncoding::kF32TwoTranspositions: case KVEncoding::kBF16TwoTranspositions: case KVEncoding::kInt8TwoTranspositions: + case KVEncoding::kInt8VNNITwoTranspositions: return true; default: return false; @@ -443,6 +534,10 @@ bool DecodeTile(KVEncoding encoding, hwy::Span encoded_tile_data, DecodeTileInt8(transposed, qkv_dim, encoded_tile_data, out); return true; } + case gcpp::KVEncoding::kInt8VNNITwoTranspositions: { + DecodeTileInt8VNNI(qkv_dim, encoded_tile_data, out); + return true; + } case gcpp::KVEncoding::kInt8MatrixAccumulation: { DecodeTileInt8MatrixAccumulation(qkv_dim, encoded_tile_data, out); return true; @@ -481,6 +576,10 @@ bool EncodeTile(gcpp::KVEncoding encoding, const DecodedTile& decoded, EncodeTileInt8(transposed, qkv_dim, decoded, out_encoded_tile_data); return true; } + case gcpp::KVEncoding::kInt8VNNITwoTranspositions: { + EncodeTileInt8VNNI(qkv_dim, decoded, out_encoded_tile_data); + return true; + } case gcpp::KVEncoding::kInt8MatrixAccumulation: { EncodeTileInt8MatrixAccumulation(qkv_dim, decoded, out_encoded_tile_data); return true; diff --git a/gemma/kv_transcoding_test.cc b/gemma/kv_transcoding_test.cc index c5c93cd2..905fce35 100644 --- a/gemma/kv_transcoding_test.cc +++ b/gemma/kv_transcoding_test.cc @@ -102,6 +102,7 @@ INSTANTIATE_TEST_SUITE_P( EncodingTestCase{gcpp::KVEncoding::kBF16MatrixAccumulation, 0.05f}, EncodingTestCase{gcpp::KVEncoding::kInt8, 0.1f}, EncodingTestCase{gcpp::KVEncoding::kInt8TwoTranspositions, 0.1f}, + EncodingTestCase{gcpp::KVEncoding::kInt8VNNITwoTranspositions, 0.1f}, EncodingTestCase{gcpp::KVEncoding::kInt8MatrixAccumulation, 0.02f})); TEST(KVEncodingTest, ConvertTileFloat32ToBfloat16) { @@ -146,9 +147,13 @@ TEST(KVEncodingTest, PairwiseConversion) { constexpr size_t qkv_dim = 256; std::vector encodings = { - gcpp::KVEncoding::kF32, gcpp::KVEncoding::kF32TwoTranspositions, - gcpp::KVEncoding::kBF16, gcpp::KVEncoding::kBF16TwoTranspositions, - gcpp::KVEncoding::kInt8, gcpp::KVEncoding::kInt8TwoTranspositions}; + gcpp::KVEncoding::kF32, + gcpp::KVEncoding::kF32TwoTranspositions, + gcpp::KVEncoding::kBF16, + gcpp::KVEncoding::kBF16TwoTranspositions, + gcpp::KVEncoding::kInt8, + gcpp::KVEncoding::kInt8TwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions}; for (auto src : encodings) { for (auto dst : encodings) { diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index d790a52d..42cc2ab8 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -108,10 +108,8 @@ static HWY_INLINE void ComputeQKVTransposedTile( const size_t kv_heads = layer_config.kv_heads; // Resolve KV cache layer index and skip flag - const size_t kv_layer_idx = (layer_config.kv_share_layer_idx >= 0) - ? static_cast(layer_config.kv_share_layer_idx) - : layer_idx; - const bool skip_kv = (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV); + const bool skip_kv = + (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV); // The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim, // model_dim], which we reshaped to (heads + kv_heads * 2) * qkv_dim rows. @@ -140,7 +138,8 @@ static HWY_INLINE void ComputeQKVTransposedTile( bool is_transposed_qs = attention_impl == AttentionImpl::kFlashTransposedQsBF16 - || attention_impl == AttentionImpl::kFlashTransposedQsInt16; + || attention_impl == AttentionImpl::kFlashTransposedQsInt16 || + attention_impl == AttentionImpl::kFlashTransposedQsInt8; hn::ScalableTag df; static hwy::Divisor tile_size_divisor(KVCache::kTileSize); @@ -159,7 +158,7 @@ static HWY_INLINE void ComputeQKVTransposedTile( const bool is_global_layer = activations.config.IsGlobalLayer(layer_idx); std::vector kv_ptrs = qbatch.KV(query_idx).cache->GetPointers( - kv_layer_idx, kv_head, kv_heads, start_pos, is_global_layer); + layer_idx, kv_head, kv_heads, start_pos, is_global_layer); const size_t v_offset = qkv_dim * KVCache::kTileSize; const size_t tile_span_size = 2 * qkv_dim * KVCache::kTileSize; const size_t k_size = qkv_dim * KVCache::kTileSize; @@ -248,15 +247,34 @@ static HWY_INLINE void ComputeQKVTransposedTile( const hn::Vec v_inv_scale = hn::Set(df, inv_scale); const size_t lanes = hn::Lanes(df); + + const hn::Rebind di32; + auto sum_vec = hn::Zero(di32); + bool is_k = scale_idx < KVCache::kTileSize; + size_t i = 0; for (; i + lanes <= dim; i += lanes) { - hn::StoreU(hn::Mul(hn::LoadU(df, values + i), v_inv_scale), - df, values + i); + auto scaled = hn::Mul(hn::LoadU(df, values + i), v_inv_scale); + hn::StoreU(scaled, df, values + i); + if (is_k && + attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + sum_vec = hn::Add(sum_vec, hn::NearestInt(scaled)); + } } if (HWY_UNLIKELY(i < dim)) { - hn::StoreN( - hn::Mul(hn::LoadN(df, values + i, dim - i), v_inv_scale), - df, values + i, dim - i); + auto scaled = + hn::Mul(hn::LoadN(df, values + i, dim - i), v_inv_scale); + hn::StoreN(scaled, df, values + i, dim - i); + if (is_k && + attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + sum_vec = hn::Add(sum_vec, hn::NearestInt(scaled)); + } + } + if (is_k && + attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + int32_t* k_sums_ptr = reinterpret_cast( + scales_ptr + 2 * KVCache::kTileSize); + k_sums_ptr[scale_idx] = hn::ReduceSum(di32, sum_vec); } }; @@ -284,6 +302,36 @@ static HWY_INLINE void ComputeQKVTransposedTile( qkv_dim, in_tile_idx, dim); v_tile_vec[v_offset] = v_cache_values[dim]; } + } else if (attention_impl == + AttentionImpl::kFlashTransposedQsInt8) { + for (int dim = 0; dim < qkv_dim; ++dim) { + // K VNNI layout: [qkv_dim/4, kTileSize, 4] + size_t k_offset = (dim - dim % 4) * KVCache::kTileSize + + in_tile_idx * 4 + (dim % 4); + k_tile_vec[k_offset] = k_f32[dim]; + + // V VNNI layout: [kTileSize/4, qkv_dim, 4] + size_t v_offset_local = + (in_tile_idx - in_tile_idx % 4) * qkv_dim + dim * 4 + + (in_tile_idx % 4); + v_tile_vec[v_offset_local] = v_cache_values[dim]; + } + } else if (attention_impl == + AttentionImpl::kFlashTransposedQsBF16 && + std::is_same_v) { + for (int dim = 0; dim < qkv_dim; dim += 2) { + const int dim_mod_2 = dim % 2; + k_tile_vec[(dim - dim_mod_2) * KVCache::kTileSize + + in_tile_idx * 2] = k_f32[dim]; + k_tile_vec[(dim - dim_mod_2) * KVCache::kTileSize + + in_tile_idx * 2 + 1] = k_f32[dim + 1]; + } + for (int dim = 0; dim < qkv_dim; ++dim) { + size_t v_offset_local = + (in_tile_idx - in_tile_idx % 4) * qkv_dim + dim * 4 + + (in_tile_idx % 4); + v_tile_vec[v_offset_local] = v_cache_values[dim]; + } } else if (is_transposed_qs) { const int in_tile_idx_mod_2 = in_tile_idx % 2; for (int dim = 0; dim < qkv_dim; dim += 2) { @@ -417,6 +465,74 @@ void CompressQueriesInt16Contiguous(const float* HWY_RESTRICT input, scale); } +template +static HWY_INLINE void CompressSingleQueryInt8(DF df, const float* q_ptr, + int qkv_dim, int8_t* out_ptr, + float* scale_out) { + namespace hn = hwy::HWY_NAMESPACE; + const size_t lanes = hn::Lanes(df); + const hn::ScalableTag d_out_full; + const hn::ScalableTag d16; + + HWY_DASSERT(scale_out != nullptr); + float max_abs = AbsMaxOfSpan(hwy::Span(q_ptr, qkv_dim)); + float s = max_abs == 0.0f ? 1.0f : 127.0f / max_abs; + *scale_out = 1.0f / s; + const hn::Vec scale_vec = hn::Set(df, s); + + HWY_DASSERT(qkv_dim % (4 * lanes) == 0); + + for (size_t i = 0; i < qkv_dim; i += 4 * lanes) { + hn::Vec x0 = hn::LoadU(df, q_ptr + i); + hn::Vec x1 = hn::LoadU(df, q_ptr + i + lanes); + hn::Vec x2 = hn::LoadU(df, q_ptr + i + 2 * lanes); + hn::Vec x3 = hn::LoadU(df, q_ptr + i + 3 * lanes); + + x0 = hn::Mul(x0, scale_vec); + x1 = hn::Mul(x1, scale_vec); + x2 = hn::Mul(x2, scale_vec); + x3 = hn::Mul(x3, scale_vec); + + const hn::Vec demoted16_0 = + hn::OrderedDemote2To(d16, hn::NearestInt(x0), hn::NearestInt(x1)); + const hn::Vec demoted16_1 = + hn::OrderedDemote2To(d16, hn::NearestInt(x2), hn::NearestInt(x3)); + const hn::Vec demoted8 = + hn::OrderedDemote2To(d_out_full, demoted16_0, demoted16_1); + const hn::Vec biased8 = + hn::Add(demoted8, hn::Set(d_out_full, static_cast(-128))); + hn::StoreU(biased8, d_out_full, out_ptr + i); + } +} + +void CompressQueriesInt8(hwy::Span input, int qkv_dim, + int8_t* HWY_RESTRICT output, + float* HWY_RESTRICT scale) { + namespace hn = hwy::HWY_NAMESPACE; + using DF = hn::ScalableTag; + const DF df; + const size_t num_queries = input.size(); + + for (size_t q = 0; q < num_queries; ++q) { + CompressSingleQueryInt8(df, input[q], qkv_dim, output + q * qkv_dim, + scale + q); + } +} + +void CompressQueriesInt8Contiguous(const float* HWY_RESTRICT input, int qkv_dim, + size_t num_queries, + int8_t* HWY_RESTRICT output, + float* HWY_RESTRICT scale) { + namespace hn = hwy::HWY_NAMESPACE; + using DF = hn::ScalableTag; + const DF df; + + for (size_t q = 0; q < num_queries; ++q) { + CompressSingleQueryInt8(df, input + q * qkv_dim, qkv_dim, + output + q * qkv_dim, scale + q); + } +} + template static HWY_INLINE void MaybeResizeMatStorage(MatStorageT& mat_storage, int rows, int cols, @@ -685,6 +801,18 @@ void LocalAttentionForAllHeadsTokensAndBatch( activations.q_scales->size()) { activations.q_scales->resize(num_sub_tasks * max_queries_per_subtask); } + } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + if (activations.int8_queries != nullptr && + num_sub_tasks * max_queries_per_subtask * qkv_dim > + activations.int8_queries->size()) { + activations.int8_queries->resize(num_sub_tasks * max_queries_per_subtask * + qkv_dim); + } + if (activations.q_scales != nullptr && + num_sub_tasks * max_queries_per_subtask > + activations.q_scales->size()) { + activations.q_scales->resize(num_sub_tasks * max_queries_per_subtask); + } } else if (attention_impl == AttentionImpl::kInt8MatrixAccumulation) { if (activations.int8_queries != nullptr && num_sub_tasks * max_queries_per_subtask * qkv_dim > @@ -893,6 +1021,23 @@ void LocalAttentionForAllHeadsTokensAndBatch( hwy::Span(last_pos_per_query), activations.config.att_cap, att_out, exp_denominator_sums.data(), max_logits.data()); + } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + HWY_DASSERT(activations.int8_queries != nullptr); + HWY_DASSERT(activations.q_scales != nullptr); + int8_t* int8_queries_ptr = + activations.int8_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; + float* q_scales_ptr = + activations.q_scales->data() + task_idx * max_queries_per_subtask; + CompressQueriesInt8(queries_ptrs_span, qkv_dim, int8_queries_ptr, + q_scales_ptr); + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( + kv_ptrs, sub_num_queries, int8_queries_ptr, + hwy::Span(q_scales_ptr, sub_num_queries), + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), + activations.config.att_cap, att_out, exp_denominator_sums.data(), + max_logits.data()); } else { HWY_DASSERT(activations.float_queries != nullptr); float* contiguous_queries_ptr = diff --git a/gemma/tiled_attention.h b/gemma/tiled_attention.h index 5a08efba..38719f3d 100644 --- a/gemma/tiled_attention.h +++ b/gemma/tiled_attention.h @@ -43,6 +43,15 @@ namespace gcpp { int16_t* HWY_RESTRICT output, \ float* HWY_RESTRICT scale); \ \ + void CompressQueriesInt8(hwy::Span input, int qkv_dim, \ + int8_t* HWY_RESTRICT output, \ + float* HWY_RESTRICT scale); \ + \ + void CompressQueriesInt8Contiguous(const float* HWY_RESTRICT input, \ + int qkv_dim, size_t num_queries, \ + int8_t* HWY_RESTRICT output, \ + float* HWY_RESTRICT scale); \ + \ void CompressAndTransposeQueriesMatrixAccumulation(const float* raw_queries, \ BF16* packed_queries, \ size_t num_queries, \