[Bugfix] Fix CPU split-KV scratchpad sizing (#45844)
Signed-off-by: Gauri Sahnan <gauri.sahnan@arm.com>
gausah01 committed
Jul 4, 2026 at 07:47 UTC
26eb87204d58d9d8a25c378fbaec9da23db3fb85
2 files changed
+57
-15
csrc/cpu/cpu_attn_impl.hpp
+56
-14
@@ -417,8 +417,10 @@ class AttentionScheduler {
417
has_decode_request = has_decode_request || (q_token_num == 1);
418
decode_only_batch = decode_only_batch && (q_token_num == 1);
419
}
420
- int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
421
- const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
420
+ const int32_t original_q_head_per_kv =
421
+ input.num_heads_q / input.num_heads_kv;
422
+ int32_t q_head_per_kv = original_q_head_per_kv;
423
+ const bool supports_gqa = original_q_head_per_kv <= max_num_q_per_iter;
424
const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
425
const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
426
if (!use_gqa_scratchpad) {
@@ -671,22 +673,62 @@ class AttentionScheduler {
673
metadata_ptr->effective_thread_num = effective_thread_num;
674
675
{
674
- // when q_tile_size = max_num_q_per_iter, requires max
675
- // attention_scratchpad_size
676
AttentionScratchPad sc(0, *metadata_ptr, 0x0);
677
- int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
678
- cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size,
679
- input.logits_buffer_elem_size, input.output_buffer_elem_size,
680
- max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true);
681
- sc.update(input.head_dim, input.q_buffer_elem_size,
682
- input.logits_buffer_elem_size, input.output_buffer_elem_size,
683
- max_num_q_per_iter, max_num_q_per_iter, n);
677
+ int64_t max_attention_scratchpad_size = 0;
678
+
679
+ for (const AttentionWorkItemGroup& item : workitems) {
680
+ const bool curr_use_gqa =
681
+ use_gqa_fast_path || (supports_gqa && item.q_token_num == 1);
682
+ const int32_t curr_q_heads_per_kv =
683
+ curr_use_gqa ? original_q_head_per_kv : 1;
684
+ const int32_t curr_default_q_tile_token_num =
685
+ default_tile_size / curr_q_heads_per_kv;
686
+
687
+ for (int32_t q_token_offset = 0; q_token_offset < item.q_token_num;
688
+ q_token_offset += curr_default_q_tile_token_num) {
689
+ const int32_t actual_q_token_num = std::min(
690
+ curr_default_q_tile_token_num, item.q_token_num - q_token_offset);
691
+ const int32_t q_head_tile_size =
692
+ actual_q_token_num * curr_q_heads_per_kv;
693
+ const int32_t rounded_q_head_tile_size =
694
+ ((q_head_tile_size + max_num_q_per_iter - 1) /
695
+ max_num_q_per_iter) *
696
+ max_num_q_per_iter;
697
+
698
+ const int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
699
+ cache_size, input.head_dim, input.elem_size,
700
+ input.q_buffer_elem_size, input.logits_buffer_elem_size,
701
+ input.output_buffer_elem_size, max_num_q_per_iter,
702
+ kv_len_alignment, rounded_q_head_tile_size,
703
+ rounded_q_head_tile_size <= max_num_q_per_iter);
704
+
705
+ sc.update(input.head_dim, input.q_buffer_elem_size,
706
+ input.logits_buffer_elem_size,
707
+ input.output_buffer_elem_size, max_num_q_per_iter,
708
+ rounded_q_head_tile_size, n);
709
+
710
+ max_attention_scratchpad_size = std::max(
711
+ max_attention_scratchpad_size, sc.get_thread_scratchpad_size());
712
+ }
713
+ }
714
+
715
metadata_ptr->attention_scratchpad_size_per_thread =
685
- ((sc.get_thread_scratchpad_size() + 63) / 64) * 64;
716
+ ((max_attention_scratchpad_size + 63) / 64) * 64;
717
+
718
+ int32_t max_reduction_q_head_tile_size = 0;
719
+ for (const ReductionWorkItemGroup& item : reduce_workitems) {
720
+ const bool curr_use_gqa =
721
+ use_gqa_fast_path || (supports_gqa && item.q_token_id_num == 1);
722
+ const int32_t curr_q_heads_per_kv =
723
+ curr_use_gqa ? original_q_head_per_kv : 1;
724
+
725
+ max_reduction_q_head_tile_size =
726
+ std::max(max_reduction_q_head_tile_size,
727
+ item.q_token_id_num * curr_q_heads_per_kv);
728
+ }
729
730
sc.update(0, metadata_ptr->reduction_split_num, input.head_dim,
688
- q_head_per_kv * split_kv_q_token_num_threshold,
689
- input.output_buffer_elem_size);
731
+ max_reduction_q_head_tile_size, input.output_buffer_elem_size);
732
metadata_ptr->reduction_scratchpad_size_per_kv_head =
733
((sc.get_reduction_scratchpad_size() + 63) / 64) * 64;
734
}
vllm/v1/attention/backends/cpu_attn.py
+1
-1
@@ -141,7 +141,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
141
self.vllm_config = vllm_config
142
143
parallel_config = vllm_config.parallel_config
144
- self.num_kv_heads = vllm_config.model_config.get_num_kv_heads(parallel_config)
144
+ self.num_kv_heads = kv_cache_spec.num_kv_heads
145
self.num_heads = vllm_config.model_config.get_num_attention_heads(
146
parallel_config
147
)