A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 20:45 UTC
[CPU] Enable non-divisible GQA for decode workitems in mixed batches (#43032)
Signed-off-by: zhejiangxiaomai <zhenhui.zhao@intel.com>
zhao, zhenhui committed
May 26, 2026 at 14:15 UTC
771e1e48b153df0e5e4334061deb1cb2ebde8fab
1 file changed
+81
-48
csrc/cpu/cpu_attn_impl.hpp
+81
-48
@@ -408,9 +408,19 @@ class AttentionScheduler {
408
const int64_t cache_size = cpu_utils::get_available_l2_size();
409
const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
410
const int32_t kv_len_alignment = input.kv_block_alignment;
411
+ bool has_decode_request = false;
412
+ bool decode_only_batch = true;
413
+ for (int32_t req_id = 0; req_id < input.num_reqs; ++req_id) {
414
+ const int32_t q_token_num =
415
+ input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
416
+ has_decode_request = has_decode_request || (q_token_num == 1);
417
+ decode_only_batch = decode_only_batch && (q_token_num == 1);
418
+ }
419
int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
412
- const bool use_gqa = (max_num_q_per_iter % q_head_per_kv == 0);
413
- if (!use_gqa) {
420
+ const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
421
+ const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
422
+ const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
423
+ if (!use_gqa_scratchpad) {
424
q_head_per_kv = 1; // fallback to MHA
425
}
426
const int32_t min_split_kv_len =
@@ -680,7 +690,7 @@ class AttentionScheduler {
690
metadata_ptr->attention_scratchpad_size_per_thread *
691
metadata_ptr->thread_num +
692
metadata_ptr->reduction_scratchpad_size_per_kv_head *
683
- (use_gqa ? input.num_heads_kv : input.num_heads_q);
693
+ (use_gqa_fast_path ? input.num_heads_kv : input.num_heads_q);
694
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(
695
scratchpad_size);
696
@@ -1409,13 +1419,24 @@ class AttentionMainLoop {
1419
const int32_t q_head_num = input->num_heads;
1420
const int32_t kv_head_num = input->num_kv_heads;
1421
const int32_t q_heads_per_kv = q_head_num / kv_head_num;
1412
- const bool use_gqa =
1413
- (max_q_head_num_per_iter % q_heads_per_kv == 0) ? true : false;
1414
- const int32_t actual_kv_head_num = use_gqa ? kv_head_num : q_head_num;
1415
- const int32_t actual_q_heads_per_kv = use_gqa ? q_heads_per_kv : 1;
1422
+ AttentionWorkItemGroup* const workitem_groups =
1423
+ metadata.workitem_groups_ptr;
1424
+ const int32_t* cu_workitem_num_per_thread =
1425
+ metadata.cu_workitem_num_per_thread;
1426
+ ReductionWorkItemGroup* const reduction_items =
1427
+ metadata.reduction_items_ptr;
1428
+ const bool supports_gqa = q_heads_per_kv <= max_q_head_num_per_iter;
1429
+ bool decode_only_batch = true;
1430
+ for (int32_t i = 0; i < metadata.workitem_group_num; ++i) {
1431
+ decode_only_batch =
1432
+ decode_only_batch && (workitem_groups[i].q_token_num == 1);
1433
+ }
1434
+ const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
1435
+ const int32_t actual_kv_head_num =
1436
+ use_gqa_fast_path ? kv_head_num : q_head_num;
1437
+ const int32_t actual_q_heads_per_kv =
1438
+ use_gqa_fast_path ? q_heads_per_kv : 1;
1439
TORCH_CHECK_LE(actual_q_heads_per_kv, max_q_head_num_per_iter);
1417
- const int32_t max_q_token_num_per_iter =
1418
- max_q_head_num_per_iter / actual_q_heads_per_kv;
1440
const int64_t q_token_num_stride = input->query_num_tokens_stride;
1441
const int64_t q_head_num_stride = input->query_num_heads_stride;
1442
const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
@@ -1461,15 +1482,6 @@ class AttentionMainLoop {
1482
sizeof(q_buffer_t), sizeof(logits_buffer_t),
1483
sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
1484
max_q_head_num_per_iter);
1464
- const int32_t default_q_tile_token_num =
1465
- default_tile_size / actual_q_heads_per_kv;
1466
-
1467
- AttentionWorkItemGroup* const workitem_groups =
1468
- metadata.workitem_groups_ptr;
1469
- const int32_t* cu_workitem_num_per_thread =
1470
- metadata.cu_workitem_num_per_thread;
1471
- ReductionWorkItemGroup* const reduction_items =
1472
- metadata.reduction_items_ptr;
1485
1486
const int32_t effective_thread_num = metadata.effective_thread_num;
1487
const int32_t reduction_item_num = metadata.reduction_item_num;
@@ -1513,8 +1525,6 @@ class AttentionMainLoop {
1525
cu_workitem_num_per_thread[thread_offset + 1] -
1526
cu_workitem_num_per_thread[thread_offset];
1527
1516
- const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
1517
-
1528
for (int32_t workitem_group_idx = 0;
1529
workitem_group_idx < curr_workitem_groups_num;
1530
++workitem_group_idx) {
@@ -1529,6 +1539,21 @@ class AttentionMainLoop {
1539
const int32_t q_token_id_start =
1540
current_workitem_group->q_token_id_start;
1541
const int32_t q_token_num = current_workitem_group->q_token_num;
1542
+ const bool curr_use_gqa =
1543
+ use_gqa_fast_path || (supports_gqa && q_token_num == 1);
1544
+ if (!use_gqa_fast_path && curr_use_gqa &&
1545
+ kv_head_idx % q_heads_per_kv != 0) {
1546
+ continue;
1547
+ }
1548
+ const int32_t curr_q_heads_per_kv =
1549
+ curr_use_gqa ? q_heads_per_kv : 1;
1550
+ const int32_t curr_max_q_token_num_per_iter =
1551
+ max_q_head_num_per_iter / curr_q_heads_per_kv;
1552
+ const int32_t curr_default_q_tile_token_num =
1553
+ default_tile_size / curr_q_heads_per_kv;
1554
+ const int32_t q_head_start_idx =
1555
+ use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv)
1556
+ : kv_head_idx;
1557
1558
// taskgroup general information
1559
const int32_t q_end = input->query_start_loc[current_group_idx + 1];
@@ -1542,7 +1567,7 @@ class AttentionMainLoop {
1567
current_workitem_group->local_split_id == 0);
1568
1569
for (int32_t q_token_offset = 0; q_token_offset < q_token_num;
1545
- q_token_offset += default_q_tile_token_num) {
1570
+ q_token_offset += curr_default_q_tile_token_num) {
1571
bool first_iter_flag[AttentionScheduler::MaxQTileIterNum];
1572
for (int32_t i = 0; i < AttentionScheduler::MaxQTileIterNum;
1573
++i) {
@@ -1552,9 +1577,9 @@ class AttentionMainLoop {
1577
const int32_t q_token_start_idx =
1578
q_start + q_token_offset + q_token_id_start;
1579
const int32_t actual_q_token_num = std::min(
1555
- default_q_tile_token_num, q_token_num - q_token_offset);
1580
+ curr_default_q_tile_token_num, q_token_num - q_token_offset);
1581
const int32_t q_head_tile_size =
1557
- actual_q_token_num * actual_q_heads_per_kv;
1582
+ actual_q_token_num * curr_q_heads_per_kv;
1583
const int32_t rounded_q_head_tile_size =
1584
((q_head_tile_size + max_q_head_num_per_iter - 1) /
1585
max_q_head_num_per_iter) *
@@ -1591,10 +1616,9 @@ class AttentionMainLoop {
1616
AttentionScheduler::align_kv_tile_pos(
1617
kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
1618
1594
- int32_t curr_kv_head_idx =
1595
- use_gqa ? kv_head_idx
1596
- : (kv_head_idx /
1597
- q_heads_per_kv); // for GQA disabled case
1619
+ const int32_t curr_kv_head_idx =
1620
+ use_gqa_fast_path ? kv_head_idx
1621
+ : (kv_head_idx / q_heads_per_kv);
1622
1623
// std::printf("thread_id: %d, req_id: %d, q_token_start: %d,
1624
// q_token_end: %d, q_head_start: %d, q_head_end: %d, kv_head_idx:
@@ -1629,12 +1653,12 @@ class AttentionMainLoop {
1653
(s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
1654
1655
// copy the Q tile to q_buffer, the logical layout of q_buffer is
1632
- // [actual_q_token_num, actual_q_heads_per_kv, head_dim]
1656
+ // [actual_q_token_num, curr_q_heads_per_kv, head_dim]
1657
{
1658
attn_impl.copy_q_heads_tile(
1659
q_tile_ptr, q_buffer, actual_q_token_num,
1636
- actual_q_heads_per_kv, q_token_num_stride,
1637
- q_head_num_stride, scale);
1660
+ curr_q_heads_per_kv, q_token_num_stride, q_head_num_stride,
1661
+ scale);
1662
}
1663
1664
if (use_sink) {
@@ -1648,29 +1672,29 @@ class AttentionMainLoop {
1672
float* __restrict__ curr_max_buffer = max_buffer;
1673
for (int32_t token_idx = 0; token_idx < actual_q_token_num;
1674
++token_idx) {
1651
- for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
1675
+ for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
1676
++head_idx) {
1677
curr_sum_buffer[head_idx] = 1.0f;
1678
curr_max_buffer[head_idx] = s_aux_fp32[head_idx];
1679
}
1680
1657
- curr_sum_buffer += actual_q_heads_per_kv;
1658
- curr_max_buffer += actual_q_heads_per_kv;
1681
+ curr_sum_buffer += curr_q_heads_per_kv;
1682
+ curr_max_buffer += curr_q_heads_per_kv;
1683
}
1684
} else {
1685
float* __restrict__ curr_sum_buffer = sum_buffer;
1686
float* __restrict__ curr_max_buffer = max_buffer;
1687
for (int32_t token_idx = 0; token_idx < actual_q_token_num;
1688
++token_idx) {
1665
- for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
1689
+ for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
1690
++head_idx) {
1691
curr_sum_buffer[head_idx] = 0.0f;
1692
curr_max_buffer[head_idx] =
1693
std::numeric_limits<float>::lowest();
1694
}
1695
1672
- curr_sum_buffer += actual_q_heads_per_kv;
1673
- curr_max_buffer += actual_q_heads_per_kv;
1696
+ curr_sum_buffer += curr_q_heads_per_kv;
1697
+ curr_max_buffer += curr_q_heads_per_kv;
1698
}
1699
}
1700
@@ -1683,16 +1707,17 @@ class AttentionMainLoop {
1707
kv_tile_pos_left + kv_tile_size, rounded_kv_tile_end_pos);
1708
for (int32_t q_head_tile_token_offset = 0;
1709
q_head_tile_token_offset < actual_q_token_num;
1686
- q_head_tile_token_offset += max_q_token_num_per_iter) {
1710
+ q_head_tile_token_offset +=
1711
+ curr_max_q_token_num_per_iter) {
1712
const int32_t q_tile_pos_left =
1713
q_tile_start_pos + q_head_tile_token_offset;
1714
const int32_t q_tile_token_num =
1690
- std::min(max_q_token_num_per_iter,
1715
+ std::min(curr_max_q_token_num_per_iter,
1716
actual_q_token_num - q_head_tile_token_offset);
1717
const int32_t q_tile_head_offset =
1693
- q_head_tile_token_offset * actual_q_heads_per_kv;
1718
+ q_head_tile_token_offset * curr_q_heads_per_kv;
1719
const int32_t q_tile_head_num =
1695
- q_tile_token_num * actual_q_heads_per_kv;
1720
+ q_tile_token_num * curr_q_heads_per_kv;
1721
const int32_t q_tile_pos_right =
1722
q_tile_pos_left + q_tile_token_num;
1723
const auto [actual_kv_tile_pos_left,
@@ -1702,7 +1727,7 @@ class AttentionMainLoop {
1727
q_tile_pos_right, sliding_window_left,
1728
sliding_window_right);
1729
const int32_t q_iter_idx =
1705
- q_head_tile_token_offset / max_q_token_num_per_iter;
1730
+ q_head_tile_token_offset / curr_max_q_token_num_per_iter;
1731
1732
if (actual_kv_tile_pos_right <= actual_kv_tile_pos_left) {
1733
continue;
@@ -1768,7 +1793,7 @@ class AttentionMainLoop {
1793
aligned_actual_kv_tile_pos_left,
1794
aligned_actual_kv_tile_pos_right, actual_kv_token_num,
1795
kv_cache_block_num_stride, q_tile_head_num,
1771
- q_tile_token_num, q_tile_pos_left, actual_q_heads_per_kv,
1796
+ q_tile_token_num, q_tile_pos_left, curr_q_heads_per_kv,
1797
block_size, sliding_window_left, sliding_window_right,
1798
scale, softcap_scale, curr_alibi_slopes,
1799
first_iter_flag[q_iter_idx], use_sink, debug_info);
@@ -1782,11 +1807,11 @@ class AttentionMainLoop {
1807
final_output(partial_q_buffer,
1808
reinterpret_cast<query_t*>(input->output) +
1809
output_buffer_offset,
1785
- sum_buffer, actual_q_heads_per_kv,
1810
+ sum_buffer, curr_q_heads_per_kv,
1811
actual_q_token_num, q_head_num, output_v_scale);
1812
} else {
1813
const int32_t stride =
1789
- actual_q_heads_per_kv * split_kv_q_token_num_threshold;
1814
+ curr_q_heads_per_kv * split_kv_q_token_num_threshold;
1815
buffer_manager.update(kv_head_idx, total_reduction_split_num,
1816
head_dim, stride, sizeof(float));
1817
volatile bool* split_flag_buffer =
@@ -1822,18 +1847,26 @@ class AttentionMainLoop {
1847
const int32_t curr_split_id = curr_workitem_groups->split_start_id;
1848
const int32_t curr_split_num = curr_workitem_groups->split_num;
1849
const int32_t current_group_idx = curr_workitem_groups->req_id;
1850
+ const bool curr_use_gqa =
1851
+ use_gqa_fast_path || (supports_gqa && curr_output_token_num == 1);
1852
+ if (!use_gqa_fast_path && curr_use_gqa &&
1853
+ kv_head_idx % q_heads_per_kv != 0) {
1854
+ continue;
1855
+ }
1856
+ const int32_t curr_q_heads_per_kv = curr_use_gqa ? q_heads_per_kv : 1;
1857
const int32_t curr_output_head_num =
1826
- curr_output_token_num * actual_q_heads_per_kv;
1858
+ curr_output_token_num * curr_q_heads_per_kv;
1859
1860
const int32_t q_start = input->query_start_loc[current_group_idx];
1861
const int32_t q_token_start_idx = q_start + curr_output_token_idx;
1830
- const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
1862
+ const int32_t q_head_start_idx =
1863
+ use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv) : kv_head_idx;
1864
size_t output_buffer_offset =
1865
q_token_start_idx * q_head_num * head_dim +
1866
q_head_start_idx * head_dim;
1867
1868
const int32_t stride =
1836
- actual_q_heads_per_kv * split_kv_q_token_num_threshold;
1869
+ curr_q_heads_per_kv * split_kv_q_token_num_threshold;
1870
buffer_manager.update(kv_head_idx, total_reduction_split_num,
1871
head_dim, stride, sizeof(float));
1872
volatile bool* split_flag_buffer =
@@ -1852,7 +1885,7 @@ class AttentionMainLoop {
1885
final_output(
1886
split_output_buffer,
1887
reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
1855
- split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
1888
+ split_sum_buffer, curr_q_heads_per_kv, curr_output_token_num,
1889
q_head_num, output_v_scale);
1890
}
1891
}