diff --git a/gemma/gemma_args.h b/gemma/gemma_args.h index beabea00..0a191c21 100644 --- a/gemma/gemma_args.h +++ b/gemma/gemma_args.h @@ -211,6 +211,7 @@ struct InferenceArgs : public ArgsBase { Path prompt_file; std::string eot_line; std::string attention_impl; + std::string kv_cache_type; template void ForEach(const Visitor& visitor) { @@ -266,6 +267,10 @@ struct InferenceArgs : public ArgsBase { 2); visitor(attention_impl, "attention_impl", std::string("flash"), "Attention implementation to use. See configs.cc for options.", 2); + visitor(kv_cache_type, "kv_cache_type", std::string(""), + "KV cache data type (f32, bf16, int8). If empty, deduced from " + "attention_impl.", + 2); } void CopyTo(RuntimeConfig& runtime_config) const { @@ -288,6 +293,15 @@ struct InferenceArgs : public ArgsBase { runtime_config.temperature = temperature; runtime_config.top_k = top_k; runtime_config.attention_impl = GetAttentionImpl(attention_impl); + if (!kv_cache_type.empty()) { + if (kv_cache_type == "int8" || kv_cache_type == "i8") { + runtime_config.kv_cache_type = Type::kInt8; + } else if (kv_cache_type == "bf16") { + runtime_config.kv_cache_type = Type::kBF16; + } else if (kv_cache_type == "f32" || kv_cache_type == "float") { + runtime_config.kv_cache_type = Type::kF32; + } + } } }; diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index 59fab29e..8284cac5 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -605,8 +605,11 @@ void LocalAttentionForAllHeadsTokensAndBatch( const size_t layer_idx, const LayerWeightsPtrs& layer, AttentionActivationsPtrs& activations, QBatch& qbatch, ThreadingContext& ctx) { + constexpr int kQueriesPerSubtask = 32; const size_t heads_per_kv_head = layer.layer_config.heads / layer.layer_config.kv_heads; + const hwy::Divisor div_heads_per_kv_head(heads_per_kv_head); + const hwy::Divisor div_kv_heads(layer.layer_config.kv_heads); int core_count = ctx.pools.MaxWorkers(); int task_multiplier = 1; @@ -630,9 +633,13 @@ void LocalAttentionForAllHeadsTokensAndBatch( // amount of work is small. task_multiplier = 1; } - [[maybe_unused]] int num_tasks = qbatch.Size() * layer.layer_config.kv_heads; + int num_queries = num_query_tokens * heads_per_kv_head; + int num_query_tasks = hwy::DivCeil(num_queries, kQueriesPerSubtask); + [[maybe_unused]] int num_tasks = + qbatch.Size() * layer.layer_config.kv_heads * num_query_tasks; [[maybe_unused]] int num_sub_tasks = - qbatch.Size() * layer.layer_config.kv_heads * task_multiplier; + qbatch.Size() * layer.layer_config.kv_heads * num_query_tasks * + task_multiplier; HWY_DASSERT_M(activations.q.Rows() == num_query_tokens * qbatch.Size(), "qbatch size mismatch"); int qkv_dim = layer.layer_config.qkv_dim; @@ -643,30 +650,46 @@ void LocalAttentionForAllHeadsTokensAndBatch( activations.sub_task_exp_denominator_sums->resize(num_sub_tasks); activations.sub_task_max_logits->resize(num_sub_tasks); } - int num_queries = num_query_tokens * heads_per_kv_head; - if (activations.bf16_queries != nullptr && - num_sub_tasks * num_queries * qkv_dim > - activations.bf16_queries->size()) { - activations.bf16_queries->resize(num_sub_tasks * num_queries * qkv_dim); - } - if (activations.int16_queries != nullptr && - num_sub_tasks * num_queries * qkv_dim > - activations.int16_queries->size()) { - activations.int16_queries->resize(num_sub_tasks * num_queries * qkv_dim); - } - if (activations.int8_queries != nullptr && - num_sub_tasks * num_queries * qkv_dim > - activations.int8_queries->size()) { - activations.int8_queries->resize(num_sub_tasks * num_queries * qkv_dim); - } - if (activations.float_queries != nullptr && - num_sub_tasks * num_queries * qkv_dim > - activations.float_queries->size()) { - activations.float_queries->resize(num_sub_tasks * num_queries * qkv_dim); - } - if (activations.q_scales != nullptr && - num_sub_tasks * num_queries > activations.q_scales->size()) { - activations.q_scales->resize(num_sub_tasks * num_queries); + int max_queries_per_subtask = std::min(num_queries, kQueriesPerSubtask); + if (attention_impl == AttentionImpl::kFlashTransposedQsBF16 || + attention_impl == AttentionImpl::kFlashMatrixAccumulation) { + if (activations.bf16_queries != nullptr && + num_sub_tasks * max_queries_per_subtask * qkv_dim > + activations.bf16_queries->size()) { + activations.bf16_queries->resize(num_sub_tasks * max_queries_per_subtask * + qkv_dim); + } + } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt16) { + if (activations.int16_queries != nullptr && + num_sub_tasks * max_queries_per_subtask * qkv_dim > + activations.int16_queries->size()) { + activations.int16_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 > + 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 (activations.float_queries != nullptr && + num_sub_tasks * max_queries_per_subtask * qkv_dim > + activations.float_queries->size()) { + activations.float_queries->resize(num_sub_tasks * max_queries_per_subtask * + qkv_dim); + } } std::vector skip_sub_task(num_sub_tasks, 0); @@ -679,9 +702,12 @@ void LocalAttentionForAllHeadsTokensAndBatch( [&](size_t task_idx, size_t worker) HWY_ATTR { size_t main_task_idx = task_idx / task_multiplier; size_t sub_task_idx = task_idx % task_multiplier; + size_t query_task_idx = main_task_idx % num_query_tasks; + size_t qbatch_and_kv_head_idx = main_task_idx / num_query_tasks; size_t current_qbatch_idx = - main_task_idx / layer.layer_config.kv_heads; - size_t kv_head_idx = main_task_idx % layer.layer_config.kv_heads; + div_kv_heads.Divide(qbatch_and_kv_head_idx); + size_t kv_head_idx = + div_kv_heads.Remainder(qbatch_and_kv_head_idx); // First and last context token we will attend to. size_t global_start_context_pos = StartPos( qbatch.Pos(current_qbatch_idx), activations.config, layer_idx); @@ -714,17 +740,18 @@ void LocalAttentionForAllHeadsTokensAndBatch( std::min(last_context_pos, start_context_pos + context_tokens_per_sub_task - 1); // pre-initialize memory [to avoid racy resizes laters]. - int num_queries = num_query_tokens * heads_per_kv_head; + int query_start_idx = static_cast(query_task_idx) * kQueriesPerSubtask; + int query_end_idx = std::min(num_queries, query_start_idx + kQueriesPerSubtask); + int sub_num_queries = query_end_idx - query_start_idx; std::vector queries_ptrs; - queries_ptrs.reserve(num_queries); - for (int token_idx = 0; token_idx < num_query_tokens; ++token_idx) { - for (int q_head_idx = 0; q_head_idx < heads_per_kv_head; - ++q_head_idx) { - queries_ptrs.push_back( - activations.q.Row(token_idx * qbatch.Size() + - current_qbatch_idx) + - (kv_head_idx * heads_per_kv_head + q_head_idx) * qkv_dim); - } + queries_ptrs.reserve(sub_num_queries); + for (int q_idx = query_start_idx; q_idx < query_end_idx; ++q_idx) { + int token_idx = div_heads_per_kv_head.Divide(q_idx); + int q_head_idx = div_heads_per_kv_head.Remainder(q_idx); + queries_ptrs.push_back( + activations.q.Row(token_idx * qbatch.Size() + + current_qbatch_idx) + + (kv_head_idx * heads_per_kv_head + q_head_idx) * qkv_dim); } hwy::Span queries_ptrs_span(queries_ptrs.data(), queries_ptrs.size()); @@ -735,14 +762,14 @@ void LocalAttentionForAllHeadsTokensAndBatch( activations.sub_task_exp_denominator_sums->at(task_idx); AlignedFloatVector& max_logits = activations.sub_task_max_logits->at(task_idx); - MaybeResizeMatStorage(att_out, num_queries, qkv_dim, "att_out", + MaybeResizeMatStorage(att_out, sub_num_queries, qkv_dim, "att_out", ctx.allocator); - for (int i = 0; i < num_queries; ++i) { + for (int i = 0; i < sub_num_queries; ++i) { hwy::ZeroBytes(att_out.Row(i), att_out.Cols() * sizeof(decltype(att_out.Row(i)[0]))); } - int num_queries_rounded_to_8 = hwy::RoundUpTo(num_queries, 8); + int num_queries_rounded_to_8 = hwy::RoundUpTo(sub_num_queries, 8); exp_denominator_sums.resize(num_queries_rounded_to_8); max_logits.resize(num_queries_rounded_to_8); for (int i = 0; i < num_queries_rounded_to_8; ++i) { @@ -761,14 +788,15 @@ void LocalAttentionForAllHeadsTokensAndBatch( std::vector> start_pos_per_query; std::vector> last_pos_per_query; - start_pos_per_query.reserve(num_queries); - last_pos_per_query.reserve(num_queries); + start_pos_per_query.reserve(sub_num_queries); + last_pos_per_query.reserve(sub_num_queries); // Position of the first token in the first tile whose pointer was // returned above. Allows for handling of token positions relative to // the KV tiles returned above. size_t rounded_down_global_start_pos = hwy::RoundDownTo(global_start_context_pos, KVCache::kTileSize); - for (int token_idx = 0; token_idx < num_query_tokens; ++token_idx) { + for (int q_idx = query_start_idx; q_idx < query_end_idx; ++q_idx) { + int token_idx = q_idx / heads_per_kv_head; int64_t global_query_pos = qbatch.Pos(current_qbatch_idx) + token_idx; // Intersect context to attend to for this specific query token @@ -787,20 +815,18 @@ void LocalAttentionForAllHeadsTokensAndBatch( // Turn token position into KV-tile relative token positions. query_last_context_pos -= rounded_down_global_start_pos; query_start_context_pos -= rounded_down_global_start_pos; - for (int q_head_idx = 0; q_head_idx < heads_per_kv_head; - ++q_head_idx) { - start_pos_per_query.push_back(query_start_context_pos); - last_pos_per_query.push_back(query_last_context_pos); - } + start_pos_per_query.push_back(query_start_context_pos); + last_pos_per_query.push_back(query_last_context_pos); } if (attention_impl == AttentionImpl::kFlashTransposedQsBF16) { HWY_DASSERT(activations.bf16_queries != nullptr); - BF16* bf16_queries_ptr = activations.bf16_queries->data() + - task_idx * num_queries * qkv_dim; + BF16* bf16_queries_ptr = + activations.bf16_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; CompressQueriesBF16(queries_ptrs_span, qkv_dim, bf16_queries_ptr); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsBF16( - kv_ptrs, num_queries, bf16_queries_ptr, + kv_ptrs, sub_num_queries, bf16_queries_ptr, hwy::Span(start_pos_per_query), hwy::Span(last_pos_per_query), activations.config.att_cap, att_out, exp_denominator_sums.data(), @@ -809,27 +835,29 @@ void LocalAttentionForAllHeadsTokensAndBatch( } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt16) { HWY_DASSERT(activations.int16_queries != nullptr); HWY_DASSERT(activations.q_scales != nullptr); - int16_t* int16_queries_ptr = activations.int16_queries->data() + - task_idx * num_queries * qkv_dim; - float* q_scales_ptr = - activations.q_scales->data() + task_idx * num_queries; + int16_t* int16_queries_ptr = + activations.int16_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; + float* q_scales_ptr = activations.q_scales->data() + + task_idx * max_queries_per_subtask; CompressQueriesInt16(queries_ptrs_span, qkv_dim, int16_queries_ptr, q_scales_ptr); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16( - kv_ptrs, num_queries, int16_queries_ptr, - hwy::Span(q_scales_ptr, num_queries), + kv_ptrs, sub_num_queries, int16_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 if (attention_impl == AttentionImpl::kFlashMatrixAccumulation) { HWY_DASSERT(activations.bf16_queries != nullptr); - BF16* bf16_queries_ptr = activations.bf16_queries->data() + - task_idx * num_queries * qkv_dim; + BF16* bf16_queries_ptr = + activations.bf16_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; CompressAndTransposeQueriesMatrixAccumulationNonContiguous( queries_ptrs_span, bf16_queries_ptr, qkv_dim); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsMatrixAccumulation( - kv_ptrs, num_queries, bf16_queries_ptr, + kv_ptrs, sub_num_queries, bf16_queries_ptr, hwy::Span(start_pos_per_query), hwy::Span(last_pos_per_query), activations.config.att_cap, att_out, exp_denominator_sums.data(), @@ -837,32 +865,34 @@ void LocalAttentionForAllHeadsTokensAndBatch( } else if (attention_impl == AttentionImpl::kInt8MatrixAccumulation) { HWY_DASSERT(activations.int8_queries != nullptr); HWY_DASSERT(activations.q_scales != nullptr); - int8_t* int8_queries_ptr = activations.int8_queries->data() + - task_idx * num_queries * qkv_dim; - float* q_scales_ptr = - activations.q_scales->data() + task_idx * num_queries; + 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; CompressAndQuantizeQueriesMatrixAccumulationInt8NonContiguous( queries_ptrs_span, int8_queries_ptr, q_scales_ptr, qkv_dim); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsMatrixAccumulationInt8( - kv_ptrs, num_queries, int8_queries_ptr, - hwy::Span(q_scales_ptr, num_queries), + 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 = activations.float_queries->data() + - task_idx * num_queries * qkv_dim; - for (int i = 0; i < num_queries; ++i) { + float* contiguous_queries_ptr = + activations.float_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; + for (int i = 0; i < sub_num_queries; ++i) { hwy::CopyBytes(queries_ptrs_span[i], contiguous_queries_ptr + i * qkv_dim, qkv_dim * sizeof(float)); } DispatchTileFlashAttentionReturnExpSumsAndMaxLogits( - kv_ptrs, num_queries, contiguous_queries_ptr, + kv_ptrs, sub_num_queries, contiguous_queries_ptr, hwy::Span(start_pos_per_query), hwy::Span(last_pos_per_query), activations.config.att_cap, att_out, exp_denominator_sums.data(), @@ -876,56 +906,59 @@ void LocalAttentionForAllHeadsTokensAndBatch( Parallelism::kFlat, num_tasks, ctx, /*cluster_idx=*/0, Callers::kFlashAttention, [&](size_t main_task_idx, size_t worker) HWY_ATTR { - size_t current_qbatch_idx = main_task_idx / layer.layer_config.kv_heads; - size_t kv_head_idx = main_task_idx % layer.layer_config.kv_heads; - for (int token_idx = 0; token_idx < num_query_tokens; ++token_idx) { - for (int head_in_group_idx = 0; head_in_group_idx < heads_per_kv_head; - ++head_in_group_idx) { - const size_t batch_index = - current_qbatch_idx * num_query_tokens + token_idx; - const size_t q_head_idx = - kv_head_idx * heads_per_kv_head + head_in_group_idx; - const size_t att_out_row_idx = - token_idx * heads_per_kv_head + head_in_group_idx; - const size_t activations_att_out_start_idx = q_head_idx * qkv_dim; - auto& att_out_0 = activations.sub_task_att_out->at( - main_task_idx * task_multiplier + 0); - auto& exp_denominator_sums_0 = - activations.sub_task_exp_denominator_sums->at( - main_task_idx * task_multiplier + 0); - auto& max_logits_0 = activations.sub_task_max_logits->at( - main_task_idx * task_multiplier + 0); - - hwy::CopyBytes(att_out_0.Row(att_out_row_idx), - activations.att_out.Row(batch_index) + - activations_att_out_start_idx, - qkv_dim * sizeof(float)); - activations.softmax_d.Row(batch_index)[q_head_idx] = - exp_denominator_sums_0[token_idx * heads_per_kv_head + - head_in_group_idx]; - activations.softmax_max.Row(batch_index)[q_head_idx] = - max_logits_0[token_idx * heads_per_kv_head + head_in_group_idx]; - for (int sub_task_idx = 1; sub_task_idx < task_multiplier; - ++sub_task_idx) { - int task_idx = main_task_idx * task_multiplier + sub_task_idx; - if (skip_sub_task[task_idx] == 1) { - continue; - } - auto& att_out = activations.sub_task_att_out->at(task_idx); - auto& exp_denominator_sums = - activations.sub_task_exp_denominator_sums->at(task_idx); - auto& max_logits = activations.sub_task_max_logits->at(task_idx); - MergeOnlineSoftmax( - att_out.Row(att_out_row_idx), - max_logits[token_idx * heads_per_kv_head + head_in_group_idx], - exp_denominator_sums[token_idx * heads_per_kv_head + - head_in_group_idx], - qkv_dim, - activations.att_out.Row(batch_index) + - activations_att_out_start_idx, - activations.softmax_max.Row(batch_index)[q_head_idx], - activations.softmax_d.Row(batch_index)[q_head_idx]); + size_t query_task_idx = main_task_idx % num_query_tasks; + size_t qbatch_and_kv_head_idx = main_task_idx / num_query_tasks; + size_t current_qbatch_idx = + div_kv_heads.Divide(qbatch_and_kv_head_idx); + size_t kv_head_idx = + div_kv_heads.Remainder(qbatch_and_kv_head_idx); + + int query_start_idx = static_cast(query_task_idx) * kQueriesPerSubtask; + int query_end_idx = std::min(num_queries, query_start_idx + kQueriesPerSubtask); + + for (int q_idx = query_start_idx; q_idx < query_end_idx; ++q_idx) { + int sub_q_idx = q_idx - query_start_idx; + int token_idx = div_heads_per_kv_head.Divide(q_idx); + int head_in_group_idx = div_heads_per_kv_head.Remainder(q_idx); + + const size_t batch_index = + current_qbatch_idx * num_query_tokens + token_idx; + const size_t q_head_idx = + kv_head_idx * heads_per_kv_head + head_in_group_idx; + const size_t activations_att_out_start_idx = q_head_idx * qkv_dim; + auto& att_out_0 = activations.sub_task_att_out->at( + main_task_idx * task_multiplier + 0); + auto& exp_denominator_sums_0 = + activations.sub_task_exp_denominator_sums->at( + main_task_idx * task_multiplier + 0); + auto& max_logits_0 = activations.sub_task_max_logits->at( + main_task_idx * task_multiplier + 0); + + hwy::CopyBytes(att_out_0.Row(sub_q_idx), + activations.att_out.Row(batch_index) + + activations_att_out_start_idx, + qkv_dim * sizeof(float)); + activations.softmax_d.Row(batch_index)[q_head_idx] = + exp_denominator_sums_0[sub_q_idx]; + activations.softmax_max.Row(batch_index)[q_head_idx] = + max_logits_0[sub_q_idx]; + for (int sub_task_idx = 1; sub_task_idx < task_multiplier; + ++sub_task_idx) { + int task_idx = main_task_idx * task_multiplier + sub_task_idx; + if (skip_sub_task[task_idx] == 1) { + continue; } + auto& att_out = activations.sub_task_att_out->at(task_idx); + auto& exp_denominator_sums = + activations.sub_task_exp_denominator_sums->at(task_idx); + auto& max_logits = activations.sub_task_max_logits->at(task_idx); + MergeOnlineSoftmax( + att_out.Row(sub_q_idx), max_logits[sub_q_idx], + exp_denominator_sums[sub_q_idx], qkv_dim, + activations.att_out.Row(batch_index) + + activations_att_out_start_idx, + activations.softmax_max.Row(batch_index)[q_head_idx], + activations.softmax_d.Row(batch_index)[q_head_idx]); } } });