diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 260ed7cd417..907f8895682 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -417,8 +417,10 @@ class AttentionScheduler { has_decode_request = has_decode_request || (q_token_num == 1); decode_only_batch = decode_only_batch && (q_token_num == 1); } - int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv; - const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter; + const int32_t original_q_head_per_kv = + input.num_heads_q / input.num_heads_kv; + int32_t q_head_per_kv = original_q_head_per_kv; + const bool supports_gqa = original_q_head_per_kv <= max_num_q_per_iter; const bool use_gqa_fast_path = supports_gqa && decode_only_batch; const bool use_gqa_scratchpad = supports_gqa && has_decode_request; if (!use_gqa_scratchpad) { @@ -671,22 +673,62 @@ class AttentionScheduler { metadata_ptr->effective_thread_num = effective_thread_num; { - // when q_tile_size = max_num_q_per_iter, requires max - // attention_scratchpad_size AttentionScratchPad sc(0, *metadata_ptr, 0x0); - int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q( - cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size, - input.logits_buffer_elem_size, input.output_buffer_elem_size, - max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true); - sc.update(input.head_dim, input.q_buffer_elem_size, - input.logits_buffer_elem_size, input.output_buffer_elem_size, - max_num_q_per_iter, max_num_q_per_iter, n); + int64_t max_attention_scratchpad_size = 0; + + for (const AttentionWorkItemGroup& item : workitems) { + const bool curr_use_gqa = + use_gqa_fast_path || (supports_gqa && item.q_token_num == 1); + const int32_t curr_q_heads_per_kv = + curr_use_gqa ? original_q_head_per_kv : 1; + const int32_t curr_default_q_tile_token_num = + default_tile_size / curr_q_heads_per_kv; + + for (int32_t q_token_offset = 0; q_token_offset < item.q_token_num; + q_token_offset += curr_default_q_tile_token_num) { + const int32_t actual_q_token_num = std::min( + curr_default_q_tile_token_num, item.q_token_num - q_token_offset); + const int32_t q_head_tile_size = + actual_q_token_num * curr_q_heads_per_kv; + const int32_t rounded_q_head_tile_size = + ((q_head_tile_size + max_num_q_per_iter - 1) / + max_num_q_per_iter) * + max_num_q_per_iter; + + const int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q( + cache_size, input.head_dim, input.elem_size, + input.q_buffer_elem_size, input.logits_buffer_elem_size, + input.output_buffer_elem_size, max_num_q_per_iter, + kv_len_alignment, rounded_q_head_tile_size, + rounded_q_head_tile_size <= max_num_q_per_iter); + + sc.update(input.head_dim, input.q_buffer_elem_size, + input.logits_buffer_elem_size, + input.output_buffer_elem_size, max_num_q_per_iter, + rounded_q_head_tile_size, n); + + max_attention_scratchpad_size = std::max( + max_attention_scratchpad_size, sc.get_thread_scratchpad_size()); + } + } + metadata_ptr->attention_scratchpad_size_per_thread = - ((sc.get_thread_scratchpad_size() + 63) / 64) * 64; + ((max_attention_scratchpad_size + 63) / 64) * 64; + + int32_t max_reduction_q_head_tile_size = 0; + for (const ReductionWorkItemGroup& item : reduce_workitems) { + const bool curr_use_gqa = + use_gqa_fast_path || (supports_gqa && item.q_token_id_num == 1); + const int32_t curr_q_heads_per_kv = + curr_use_gqa ? original_q_head_per_kv : 1; + + max_reduction_q_head_tile_size = + std::max(max_reduction_q_head_tile_size, + item.q_token_id_num * curr_q_heads_per_kv); + } sc.update(0, metadata_ptr->reduction_split_num, input.head_dim, - q_head_per_kv * split_kv_q_token_num_threshold, - input.output_buffer_elem_size); + max_reduction_q_head_tile_size, input.output_buffer_elem_size); metadata_ptr->reduction_scratchpad_size_per_kv_head = ((sc.get_reduction_scratchpad_size() + 63) / 64) * 64; } diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 056107c364d..6b2d202d3f1 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -141,7 +141,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata] self.vllm_config = vllm_config parallel_config = vllm_config.parallel_config - self.num_kv_heads = vllm_config.model_config.get_num_kv_heads(parallel_config) + self.num_kv_heads = kv_cache_spec.num_kv_heads self.num_heads = vllm_config.model_config.get_num_attention_heads( parallel_config )