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 15, 09:45 UTC
[ROCm][Kernel][V1] Enable AMD Radeon GPU Custom Paged Attention on v1 (#17004)
Signed-off-by: Hosang Yoon <hosang.yoon@amd.com>
Hosang committed
May 21, 2025 at 11:35 UTC
dd5fa7e04f7544dca276701816453e8cc31fb7de
6 files changed
+1930
-189
benchmarks/kernels/benchmark_paged_attention.py
+5
-1
@@ -84,7 +84,10 @@ def main(
84
if version == "v2":
85
if current_platform.is_rocm():
86
global PARTITION_SIZE
87
- PARTITION_SIZE = 1024 if not args.custom_paged_attn else PARTITION_SIZE_ROCM
87
+ if not args.custom_paged_attn and not current_platform.is_navi():
88
+ PARTITION_SIZE = 1024
89
+ else:
90
+ PARTITION_SIZE = PARTITION_SIZE_ROCM
91
num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
92
tmp_output = torch.empty(
93
size=(num_seqs, num_query_heads, num_partitions, head_size),
@@ -159,6 +162,7 @@ def main(
162
scale,
163
block_tables,
164
seq_lens,
165
+ None,
166
block_size,
167
max_seq_len,
168
alibi_slopes,
csrc/rocm/attention.cu
+1880
-171
@@ -30,6 +30,14 @@
30
#define __HIP__GFX9__
31
#endif
32
33
+#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__))
34
+ #define __HIP__GFX11__
35
+#endif
36
+
37
+#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
38
+ #define __HIP__GFX12__
39
+#endif
40
+
41
#if defined(NDEBUG)
42
#undef NDEBUG
43
#include <assert.h>
@@ -43,7 +51,7 @@
51
#define MIN(a, b) ((a) < (b) ? (a) : (b))
52
#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
53
46
-#if defined(__HIP__GFX9__) // TODO: Add NAVI support
54
+#if defined(__HIP__GFX9__)
55
56
#define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32
57
#define GCN_MFMA_INSTR __builtin_amdgcn_mfma_f32_4x4x4f16
@@ -1482,191 +1490,1690 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1490
}
1491
}
1492
1485
-#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1493
+#elif defined(__HIP__GFX11__)
1494
1487
-// clang-format off
1488
-template <typename scalar_t, typename cache_t,
1489
- vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
1490
- int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
1491
- int GQA_RATIO>
1492
-__global__
1493
-__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
1494
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
1495
- const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
1496
- const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
1497
- const int num_kv_heads,
1498
- const float scale,
1499
- const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
1500
- const int* __restrict__ context_lens, // [num_seqs]
1501
- const int* __restrict__ query_start_loc_ptr, // [num_seqs]
1502
- const int max_num_blocks_per_seq,
1503
- const float* __restrict__ alibi_slopes, // [num_heads]
1504
- const int q_stride,
1505
- const int kv_block_stride,
1506
- const int kv_head_stride,
1507
- float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1508
- float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
1509
- scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
1510
- OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
1511
- int max_ctx_blocks, const float* k_scale, const float* v_scale) {
1512
- UNREACHABLE_CODE
1495
+using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
1496
+
1497
+using bit16_t = uint16_t;
1498
+using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t;
1499
+typedef bit16x4 _B16x4;
1500
+
1501
+using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t;
1502
+union b16x8_u {
1503
+ bit16x8 u16x8;
1504
+ _B16x4 xy[2];
1505
+};
1506
+typedef b16x8_u _B16x8;
1507
+
1508
+using bit16x16 =
1509
+ __attribute__((__vector_size__(16 * sizeof(uint16_t)))) uint16_t;
1510
+union b16x16_u {
1511
+ bit16x16 u16x16;
1512
+ _B16x8 xy[2];
1513
+};
1514
+typedef b16x16_u _B16x16;
1515
+
1516
+using _B8x8 = uint2;
1517
+using bit8_t = uint8_t;
1518
+
1519
+typedef struct _B8x16 {
1520
+ _B8x8 xy[2];
1521
+} _B8x16;
1522
+
1523
+template <typename T, int absz, int cbid, int blgp>
1524
+__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x16& inpA,
1525
+ const bit16x16& inpB,
1526
+ const floatx8& inpC) {
1527
+ if constexpr (std::is_same<T, _Float16>::value) {
1528
+ return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(inpA, inpB, inpC);
1529
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
1530
+ return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(inpA, inpB, inpC);
1531
+ } else {
1532
+ static_assert(false, "unsupported 16b dtype");
1533
+ }
1534
+}
1535
+
1536
+template <typename T>
1537
+__device__ __forceinline__ float to_float(const T& inp) {
1538
+ if constexpr (std::is_same<T, _Float16>::value) {
1539
+ return (float)inp;
1540
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
1541
+ return __bfloat162float(inp);
1542
+ } else {
1543
+ static_assert(false, "unsupported 16b dtype");
1544
+ }
1545
+}
1546
+
1547
+template <typename T>
1548
+__device__ __forceinline__ T from_float(const float& inp) {
1549
+ if constexpr (std::is_same<T, _Float16>::value) {
1550
+ return (_Float16)inp;
1551
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
1552
+ return __float2bfloat16(inp);
1553
+ } else {
1554
+ static_assert(false, "unsupported 16b dtype");
1555
+ }
1556
+}
1557
+
1558
+template <typename T>
1559
+__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
1560
+ if constexpr (std::is_same<T, _Float16>::value) {
1561
+ union h2cvt {
1562
+ __half2 h2[4];
1563
+ _B16x8 b16x8;
1564
+ } u;
1565
+ u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1]));
1566
+ u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3]));
1567
+ u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5]));
1568
+ u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7]));
1569
+ return u.b16x8;
1570
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
1571
+ union b2cvt {
1572
+ __hip_bfloat162 b2[4];
1573
+ _B16x8 b16x8;
1574
+ } u;
1575
+
1576
+ u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1]));
1577
+ u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3]));
1578
+ u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5]));
1579
+ u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7]));
1580
+
1581
+ return u.b16x8;
1582
+ } else {
1583
+ static_assert(false, "unsupported 16b dtype");
1584
+ }
1585
}
1586
1587
+// clang-format off
1588
template <typename scalar_t, typename cache_t,
1589
vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
1517
- int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
1518
- int GQA_RATIO>
1590
+ int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
1591
__global__
1520
-__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1521
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
1522
- const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
1523
- const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
1524
- const int num_kv_heads,
1525
- const float scale,
1526
- const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
1527
- const int* __restrict__ context_lens, // [num_seqs]
1528
- const int* __restrict__ query_start_loc_ptr, // [num_seqs]
1592
+__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1593
+ const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
1594
+ const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
1595
+ // head_size/x, block_size, x]
1596
+ const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
1597
+ // head_size, block_size]
1598
+ const int num_kv_heads, const float scale,
1599
+ const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
1600
+ const int* __restrict__ context_lens, // [num_seqs]
1601
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
1602
const int max_num_blocks_per_seq,
1603
const float* __restrict__ alibi_slopes, // [num_heads]
1531
- const int q_stride,
1532
- const int kv_block_stride,
1533
- const int kv_head_stride,
1534
- float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1535
- float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
1536
- scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
1537
- OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
1604
+ const int q_stride, const int kv_block_stride, const int kv_head_stride,
1605
+ float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1606
+ float* __restrict__ max_logits, // [num_seqs, num_heads,
1607
+ // max_num_partitions]
1608
+ scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
1609
+ // head_size]
1610
+ OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
1611
int max_ctx_blocks, const float* k_scale, const float* v_scale) {
1539
- UNREACHABLE_CODE
1540
-}
1612
+ // clang-format on
1613
+ constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11
1614
+ const int warpid = threadIdx.x / WARP_SIZE;
1615
+ const int laneid = threadIdx.x % WARP_SIZE;
1616
+ const int lane2id = laneid % 2;
1617
+ const int lane4id = laneid % 4;
1618
+ const int lane16id = laneid % 16;
1619
+ const int rowid = laneid / 16;
1620
1542
-// Grid: (num_heads, num_seqs).
1543
-template <typename scalar_t, typename OUTT, int HEAD_SIZE, int NUM_THREADS,
1544
- int PARTITION_SIZE, int NPAR_LOOPS>
1545
-__global__
1546
-__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1547
- OUTT* __restrict__ out, // [num_seqs, num_heads, head_size]
1548
- const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1549
- const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
1550
- const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
1551
- const int* __restrict__ context_lens, // [num_seqs]
1552
- const int* __restrict__ query_start_loc_ptr, // [num_seqs]
1553
- const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
1554
- UNREACHABLE_CODE
1555
-}
1556
-// clang-format on
1621
+ const int seq_idx = blockIdx.x;
1622
+ // NOTE queries with sequence len > 1 are prefills and taken care by another
1623
+ // kernel.
1624
+ if (query_start_loc_ptr != nullptr &&
1625
+ (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) {
1626
+ return;
1627
+ }
1628
1558
-#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1629
+ const int partition_idx = blockIdx.y;
1630
1560
-#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \
1561
- paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
1562
- HEAD_SIZE, NTHR, ALIBI_ENABLED, \
1563
- GQA_RATIO> \
1564
- <<<grid, block, 0, stream>>>( \
1565
- query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
1566
- block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \
1567
- max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
1568
- kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
1569
- max_ctx_blocks, k_scale_ptr, v_scale_ptr);
1631
+ constexpr int T_PAR_SIZE = 256; // token partition size set to 256
1632
1571
-#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \
1572
- paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
1573
- HEAD_SIZE, NTHR, ALIBI_ENABLED, \
1574
- GQA_RATIO> \
1575
- <<<grid, block, 0, stream>>>( \
1576
- query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
1577
- block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \
1578
- max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
1579
- kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
1580
- max_ctx_blocks, k_scale_ptr, v_scale_ptr);
1633
+ const int max_num_partitions = gridDim.y;
1634
1582
-#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \
1583
- paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \
1584
- PARTITION_SIZE, NPAR_LOOPS> \
1585
- <<<reduce_grid, reduce_block, 0, stream>>>( \
1586
- out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, \
1587
- context_lens_ptr, query_start_loc_ptr, max_num_partitions, \
1588
- fp8_out_scale_ptr);
1635
+ const int context_len = context_lens[seq_idx]; // length of a seq
1636
1590
-template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
1591
- int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
1592
- bool ALIBI_ENABLED>
1593
-void paged_attention_custom_launcher(
1594
- torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
1595
- torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
1596
- torch::Tensor& value_cache, const int num_kv_heads, float scale,
1597
- torch::Tensor& block_tables, torch::Tensor& context_lens,
1598
- const std::optional<torch::Tensor>& query_start_loc, int max_context_len,
1599
- const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
1600
- torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale) {
1601
- int num_seqs = block_tables.size(0);
1602
- int num_heads = query.size(1);
1603
- int head_size = query.size(2);
1604
- int max_num_blocks_per_seq = block_tables.size(1);
1605
- int q_stride = query.stride(0);
1606
- int kv_block_stride = key_cache.stride(0);
1607
- int kv_head_stride = key_cache.stride(1);
1637
+ const int partition_start_token_idx = partition_idx * T_PAR_SIZE;
1638
+ // exit if partition is out of context for seq
1639
+ if (partition_start_token_idx >= context_len) {
1640
+ return;
1641
+ }
1642
1609
- // NOTE: query start location is optional for V0 decode should not be used.
1610
- // If batch contains mix of prefills and decode, prefills should be skipped.
1611
- const int* query_start_loc_ptr =
1612
- query_start_loc
1613
- ? reinterpret_cast<const int*>(query_start_loc.value().data_ptr())
1614
- : nullptr;
1643
+ constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2);
1644
1616
- // NOTE: alibi_slopes is optional.
1617
- const float* alibi_slopes_ptr =
1618
- alibi_slopes
1619
- ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
1620
- : nullptr;
1645
+ __shared__ float shared_qk_max[NWARPS][16 + 1];
1646
+ __shared__ float shared_exp_sum[NWARPS][16 + 1];
1647
+ // shared_logits is used for multiple purposes
1648
+ __shared__ _B16x16 shared_logits[NWARPS][2][16][2];
1649
1622
- float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
1623
- float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
1624
- T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
1625
- T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
1626
- KVT* key_cache_ptr = reinterpret_cast<KVT*>(key_cache.data_ptr());
1627
- KVT* value_cache_ptr = reinterpret_cast<KVT*>(value_cache.data_ptr());
1628
- int* block_tables_ptr = block_tables.data_ptr<int>();
1629
- int* context_lens_ptr = context_lens.data_ptr<int>();
1630
- const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
1631
- const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
1632
- // NOTE: fp8_out_scale is optional.
1633
- const auto fp8_out_scale_ptr =
1634
- fp8_out_scale
1635
- ? static_cast<const float*>(fp8_out_scale.value().data_ptr())
1636
- : nullptr;
1637
- OUTT* out_ptr = reinterpret_cast<OUTT*>(out.data_ptr());
1650
+ // for QK wmma16x16, layout is QHead/Tokenx16 across every 16 lanes,
1651
+ // 32 Bytes HeadElements in each lane, 2x16B HeadElements across a row of warp
1652
+ constexpr int ROWS_PER_WARP =
1653
+ WARP_SIZE / 16 / 2; // rows refers to 16 lanes; refer dpp terminology
1654
+ constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD =
1655
+ 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types
1656
+ constexpr int QKHE_PER_FETCH =
1657
+ CONTIGUOUS_KV_ELEMS_16B_LOAD *
1658
+ ROWS_PER_WARP; // each fetch across a warp fetches these many elements
1659
+ constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across
1660
+ // warp
1661
1639
- const int max_ctx_blocks = DIVIDE_ROUND_UP(max_context_len, BLOCK_SIZE);
1662
+ _B16x16 Qlocal[QKHELOOP / 2]; // note that 16 contiguous elements of Q should
1663
+ // be fetched per lane for 16 bit cache types
1664
1641
- // partition size is fixed at 256 since both mfma4 and mfma16 kernels support
1642
- // it mfma4 kernel also supports partition size 512
1643
- constexpr int PARTITION_SIZE = 256;
1644
- const int max_num_partitions =
1645
- DIVIDE_ROUND_UP(max_context_len, PARTITION_SIZE);
1646
- const int gqa_ratio = num_heads / num_kv_heads;
1647
- assert(num_heads % num_kv_heads == 0);
1648
- assert(head_size == HEAD_SIZE);
1665
+ constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t);
1666
1650
- constexpr int NTHR = 256;
1651
- dim3 grid(num_seqs, max_num_partitions, num_kv_heads);
1652
- dim3 block(NTHR);
1653
- const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
1654
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1667
+ constexpr int TOKENS_PER_WARP =
1668
+ T_PAR_SIZE /
1669
+ NWARPS; // sub partition of tokens per warp for qk calculation
1670
+ constexpr int TLOOP =
1671
+ TOKENS_PER_WARP /
1672
+ 16; // each wmma16x16x16 instruction processes 16 tokens
1673
1656
- // mfma4 kernel is faster than mfma16 for gqa_ratio <= 4
1657
- switch (gqa_ratio) {
1658
- case 1:
1659
- LAUNCH_CUSTOM_ATTENTION_MFMA4(1);
1660
- break;
1661
- case 2:
1662
- LAUNCH_CUSTOM_ATTENTION_MFMA4(2);
1663
- break;
1664
- case 3:
1665
- LAUNCH_CUSTOM_ATTENTION_MFMA4(3);
1666
- break;
1667
- case 4:
1668
- LAUNCH_CUSTOM_ATTENTION_MFMA4(4);
1669
- break;
1674
+ _B16x16 Klocal[TLOOP]
1675
+ [QKHELOOP / 2]; // can be interpreted as B8x16 for 8 bit types
1676
+
1677
+ const int wg_start_head_idx = blockIdx.z * GQA_RATIO;
1678
+ const int wg_start_kv_head_idx = blockIdx.z;
1679
+ const int total_num_heads = gridDim.z * GQA_RATIO;
1680
+
1681
+ // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps
1682
+ // each wmma takes QH16xT16x16HE across warp
1683
+ // repeat wmma across QKHELOOP dimension
1684
+ // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens
1685
+ // across 2 rows x 8 tokens per lane
1686
+
1687
+ const int64_t query_start_off = static_cast<int64_t>(
1688
+ query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);
1689
+
1690
+ if (GQA_RATIO == 1) {
1691
+ const int local_qhead_idx = lane16id % GQA_RATIO;
1692
+ const int global_qhead_idx = wg_start_head_idx + local_qhead_idx;
1693
+ const scalar_t* q_ptr =
1694
+ q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE;
1695
+ if (lane16id < GQA_RATIO) {
1696
+ #pragma unroll
1697
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) {
1698
+ const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH * 2;
1699
+ const _B16x16* q_fetch_ptr_32B =
1700
+ reinterpret_cast<const _B16x16*>(q_fetch_ptr);
1701
+ Qlocal[qkhe_depth] = *q_fetch_ptr_32B;
1702
+ }
1703
+ }
1704
+ } else {
1705
+ // fetch Q in shared across warps and then write to registers
1706
+ const int local_qhead_idx = 2 * warpid + rowid;
1707
+ const int global_qhead_idx = wg_start_head_idx + local_qhead_idx;
1708
+ const scalar_t* q_ptr =
1709
+ q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE;
1710
+
1711
+ const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B;
1712
+ if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) {
1713
+ const scalar_t* q_fetch_ptr = q_ptr + qhead_element;
1714
+ const _B16x8* q_fetch_ptr_16B =
1715
+ reinterpret_cast<const _B16x8*>(q_fetch_ptr);
1716
+ _B16x8 tmp = *q_fetch_ptr_16B;
1717
+
1718
+ const int offset1 =
1719
+ lane16id /
1720
+ 2; // 16 contiguous chunks of head elems are spread across 8x2lanes
1721
+ shared_logits[offset1][lane2id][local_qhead_idx][0].xy[0] = tmp;
1722
+ }
1723
+
1724
+ __syncthreads();
1725
+
1726
+ #pragma unroll
1727
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) {
1728
+ Qlocal[qkhe_depth].xy[0] =
1729
+ shared_logits[qkhe_depth][0][lane16id % GQA_RATIO][0].xy[0];
1730
+ Qlocal[qkhe_depth].xy[1] =
1731
+ shared_logits[qkhe_depth][1][lane16id % GQA_RATIO][0].xy[0];
1732
+ }
1733
+ }
1734
+
1735
+ const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
1736
+ const int last_ctx_block = num_context_blocks - 1;
1737
+
1738
+ const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
1739
+
1740
+ int kphysical_block_number[TLOOP];
1741
+
1742
+ // fetch k physical block numbers
1743
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1744
+ const int klocal_token_idx =
1745
+ TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
1746
+ const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
1747
+ const int kblock_idx = (kglobal_token_idx < context_len)
1748
+ ? kglobal_token_idx / BLOCK_SIZE
1749
+ : last_ctx_block;
1750
+ kphysical_block_number[token_depth] = block_table_seq[kblock_idx];
1751
+ }
1752
+
1753
+ constexpr int KX = 16 / sizeof(cache_t);
1754
+ const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride;
1755
+
1756
+ const int row_head_elem = 0;
1757
+
1758
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1759
+ const int64_t kblock_number =
1760
+ static_cast<int64_t>(kphysical_block_number[token_depth]);
1761
+ const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride;
1762
+ const int klocal_token_idx =
1763
+ TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
1764
+ const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
1765
+ const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE;
1766
+ const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX;
1767
+
1768
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
1769
+ const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH;
1770
+ const int offset1 = head_elem / KX;
1771
+ const int offset2 = head_elem % KX;
1772
+ const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2;
1773
+ const _B16x8* k_fetch_ptr_16B =
1774
+ reinterpret_cast<const _B16x8*>(k_fetch_ptr);
1775
+ Klocal[token_depth][qkhe_depth / 2].xy[qkhe_depth % 2] = *k_fetch_ptr_16B;
1776
+ }
1777
+ }
1778
+
1779
+ constexpr int VTOKENS_PER_LANE =
1780
+ TOKENS_PER_WARP / ROWS_PER_WARP; // 32/1 = 32 vtokens per lane
1781
+ constexpr int VBLOCKS_PER_LANE = 2; // assumes block size >=16
1782
+ constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps
1783
+ constexpr int VTLANELOOP = DIVIDE_ROUND_UP(
1784
+ VTOKENS_PER_LANE,
1785
+ CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes
1786
+ // minimum block size is 16
1787
+ constexpr int VHELOOP = DIVIDE_ROUND_UP(
1788
+ (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each
1789
+ // wmma instr works on 16 head elements
1790
+
1791
+ int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE];
1792
+
1793
+ // fetch v physical block numbers
1794
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
1795
+ for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE;
1796
+ vblock_depth++) {
1797
+ const int vlocal_token_idx =
1798
+ vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
1799
+ vblock_depth * BLOCK_SIZE;
1800
+ const int vglobal_token_idx =
1801
+ partition_start_token_idx + vlocal_token_idx;
1802
+ const int vblock_idx = (vglobal_token_idx < context_len)
1803
+ ? vglobal_token_idx / BLOCK_SIZE
1804
+ : last_ctx_block;
1805
+ vphysical_block_number[vtoken_depth][vblock_depth] =
1806
+ block_table_seq[vblock_idx];
1807
+ }
1808
+ }
1809
+
1810
+ _B16x16 Vlocal[VTLOOP][VHELOOP]
1811
+ [VTLANELOOP / 2]; // this can be interpreted as B8x16 too
1812
+
1813
+ const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride;
1814
+ // v fetches are 16head elems across lanes x (16x2) tokens per lane
1815
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
1816
+ const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
1817
+ const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
1818
+
1819
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
1820
+ for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
1821
+ const int64_t vblock_number = static_cast<int64_t>(
1822
+ vphysical_block_number[vtoken_depth]
1823
+ [vfetch_depth / VBLOCKS_PER_LANE]);
1824
+ const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
1825
+
1826
+ const cache_t* v_fetch_ptr =
1827
+ v_ptr3 +
1828
+ (vfetch_depth % VBLOCKS_PER_LANE) * CONTIGUOUS_KV_ELEMS_16B_LOAD;
1829
+ const _B16x8* v_fetch_ptr_16B =
1830
+ reinterpret_cast<const _B16x8*>(v_fetch_ptr);
1831
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth / 2].xy[vfetch_depth % 2] =
1832
+ *v_fetch_ptr_16B;
1833
+ }
1834
+ }
1835
+ }
1836
+
1837
+ floatx8 dout[TLOOP];
1838
+ // qk wmma
1839
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1840
+ dout[token_depth] = {0};
1841
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) {
1842
+ dout[token_depth] = gcn_wmma16x16x16_instr<scalar_t, 0, 0, 0>(
1843
+ Klocal[token_depth][qkhe_depth].u16x16, Qlocal[qkhe_depth].u16x16,
1844
+ dout[token_depth]);
1845
+ }
1846
+ dout[token_depth] *= scale;
1847
+ }
1848
+
1849
+ // calculate qk_max and exp_sum per warp and write to shared memory
1850
+ float qk_max = -FLT_MAX;
1851
+ float exp_sum = 0.0f;
1852
+ const int qkout_token_idx =
1853
+ partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid;
1854
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1855
+ const int local_token_idx = qkout_token_idx + token_depth * 16;
1856
+ for (int i = 0; i < 8; i++) {
1857
+ const float tmp = (local_token_idx + 2 * i < context_len)
1858
+ ? dout[token_depth][i]
1859
+ : -FLT_MAX;
1860
+ qk_max = fmaxf(qk_max, tmp);
1861
+ }
1862
+ }
1863
+
1864
+ qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16));
1865
+
1866
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1867
+ const int local_token_idx = qkout_token_idx + token_depth * 16;
1868
+ for (int i = 0; i < 8; i++) {
1869
+ const float tmp = (local_token_idx + 2 * i < context_len)
1870
+ ? __expf(dout[token_depth][i] - qk_max)
1871
+ : 0.0f;
1872
+ dout[token_depth][i] = tmp;
1873
+ exp_sum += tmp;
1874
+ }
1875
+ }
1876
+
1877
+ exp_sum += __shfl_xor(exp_sum, 16);
1878
+
1879
+ __syncthreads();
1880
+
1881
+ if (laneid < 16) {
1882
+ shared_qk_max[warpid][lane16id] = qk_max;
1883
+ shared_exp_sum[warpid][lane16id] = exp_sum;
1884
+ }
1885
+
1886
+ __syncthreads();
1887
+
1888
+ // calculate partition qk_max and exp_sum
1889
+ float partition_qk_max = -FLT_MAX;
1890
+ float warp_qk_max_exp[NWARPS];
1891
+ float partition_exp_sum = 0.0f;
1892
+
1893
+ #pragma unroll
1894
+ for (int w = 0; w < NWARPS; w++) {
1895
+ warp_qk_max_exp[w] = shared_qk_max[w][lane16id];
1896
+ partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]);
1897
+ }
1898
+
1899
+ for (int w = 0; w < NWARPS; w++) {
1900
+ warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max);
1901
+ partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w];
1902
+ }
1903
+
1904
+ const float inv_sum_scale =
1905
+ __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid];
1906
+
1907
+ __syncthreads();
1908
+
1909
+ // write logits to shared mem
1910
+ #pragma unroll
1911
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1912
+ dout[token_depth] *= inv_sum_scale;
1913
+ shared_logits[warpid][token_depth][lane16id][0].xy[rowid] =
1914
+ from_floatx8<scalar_t>(dout[token_depth]);
1915
+ }
1916
+ __syncthreads();
1917
+
1918
+ _B16x8 swp_buf[TLOOP][2];
1919
+ #pragma unroll
1920
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1921
+ swp_buf[token_depth][0] =
1922
+ shared_logits[warpid][token_depth][lane16id][0].xy[0];
1923
+ swp_buf[token_depth][1] =
1924
+ shared_logits[warpid][token_depth][lane16id][0].xy[1];
1925
+ }
1926
+
1927
+ #pragma unroll
1928
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1929
+ #pragma unroll
1930
+ for (int i = 0; i < 8; i++) {
1931
+ shared_logits[warpid][token_depth][lane16id][0].xy[rowid].u16x8[i] =
1932
+ swp_buf[token_depth][i % 2].u16x8[4 * rowid + (i / 2)];
1933
+ }
1934
+ }
1935
+
1936
+ // write out partition max_logits and exp_sum
1937
+ if (threadIdx.x < GQA_RATIO) {
1938
+ const int qhead_idx = lane16id;
1939
+ const int offset = seq_idx * total_num_heads * max_num_partitions +
1940
+ (wg_start_head_idx + qhead_idx) * max_num_partitions +
1941
+ partition_idx;
1942
+ max_logits[offset] = partition_qk_max;
1943
+ exp_sums[offset] = partition_exp_sum;
1944
+ }
1945
+
1946
+ __syncthreads();
1947
+
1948
+ _B16x8 outelems[VHELOOP];
1949
+ // Softmax V wmma
1950
+ // v layout: 16he across lanes x (16x2) tokens per lane
1951
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
1952
+ floatx8 tmp_out = {0};
1953
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
1954
+ for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP / 2;
1955
+ vfetch_depth++) {
1956
+ const int offset = vfetch_depth;
1957
+ // if output format is 16 qheads across 16 lanes, 16 head elems spread
1958
+ // across rows
1959
+ tmp_out = gcn_wmma16x16x16_instr<scalar_t, 0, 0, 0>(
1960
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x16,
1961
+ shared_logits[vtoken_depth][offset][lane16id][0].u16x16, tmp_out);
1962
+ }
1963
+ }
1964
+ outelems[vhe_depth] = from_floatx8<scalar_t>(tmp_out);
1965
+ }
1966
+
1967
+ __syncthreads();
1968
+
1969
+ #pragma unroll
1970
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
1971
+ shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid] =
1972
+ outelems[vhe_depth]; // lane16 id head dimension; rowid head element
1973
+ // dimension
1974
+ }
1975
+
1976
+ __syncthreads();
1977
+
1978
+ #pragma unroll
1979
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
1980
+ swp_buf[vhe_depth][0] = shared_logits[warpid][vhe_depth][lane16id][0].xy[0];
1981
+ swp_buf[vhe_depth][1] = shared_logits[warpid][vhe_depth][lane16id][0].xy[1];
1982
+ }
1983
+
1984
+ #pragma unroll
1985
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
1986
+ #pragma unroll
1987
+ for (int i = 0; i < 8; i++) {
1988
+ shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid].u16x8[i] =
1989
+ swp_buf[vhe_depth][i % 2].u16x8[4 * rowid + (i / 2)];
1990
+ }
1991
+ }
1992
+
1993
+ __syncthreads();
1994
+
1995
+ // write to tmp_out with coalesced writes after reading from shared mem
1996
+ if (warpid == 0) {
1997
+ _B16x8 vout[GQA_RATIO2];
1998
+ // each lane writes out 16Bytes of tmp_out along head elem dimension
1999
+ const int head_elem_idx = lane16id * 8;
2000
+ if (head_elem_idx < HEAD_SIZE) {
2001
+ for (int h = 0; h < GQA_RATIO2; h++) {
2002
+ const int local_head_idx = 2 * h + rowid;
2003
+ const int offset1 = (head_elem_idx / 16) % NWARPS;
2004
+ const int offset2 = head_elem_idx / 16 / NWARPS;
2005
+ const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row
2006
+ vout[h] =
2007
+ shared_logits[offset1][offset2][local_head_idx][0].xy[offset3];
2008
+ }
2009
+
2010
+ const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions;
2011
+ scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult +
2012
+ partition_idx * HEAD_SIZE;
2013
+ for (int h = 0; h < GQA_RATIO2; h++) {
2014
+ const int local_head_idx = 2 * h + rowid;
2015
+ if (local_head_idx < GQA_RATIO) {
2016
+ const int out_head_idx = wg_start_head_idx + local_head_idx;
2017
+ scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult;
2018
+ scalar_t* out_ptr3 = out_ptr2 + head_elem_idx;
2019
+ _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3);
2020
+ *out_ptr_B16x8 = vout[h];
2021
+ }
2022
+ }
2023
+ }
2024
+ }
2025
+}
2026
+
2027
+template <typename scalar_t, typename cache_t,
2028
+ vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
2029
+ int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
2030
+ int GQA_RATIO>
2031
+__global__
2032
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
2033
+ const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
2034
+ const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
2035
+ // head_size/x, block_size, x]
2036
+ const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
2037
+ // head_size, block_size]
2038
+ const int num_kv_heads, const float scale,
2039
+ const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2040
+ const int* __restrict__ context_lens, // [num_seqs]
2041
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2042
+ const int max_num_blocks_per_seq,
2043
+ const float* __restrict__ alibi_slopes, // [num_heads]
2044
+ const int q_stride, const int kv_block_stride, const int kv_head_stride,
2045
+ float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
2046
+ float* __restrict__ max_logits, // [num_seqs, num_heads,
2047
+ // max_num_partitions]
2048
+ scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
2049
+ // head_size]
2050
+ OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
2051
+ int max_ctx_blocks, const float* k_scale, const float* v_scale) {
2052
+ UNREACHABLE_CODE
2053
+}
2054
+
2055
+// Grid: (num_heads, num_seqs).
2056
+template <typename scalar_t, typename OUTT, int HEAD_SIZE, int NUM_THREADS,
2057
+ int PARTITION_SIZE, int NPAR_LOOPS>
2058
+__global__
2059
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
2060
+ OUTT* __restrict__ out, // [num_seqs, num_heads, head_size]
2061
+ const float* __restrict__ exp_sums, // [num_seqs, num_heads,
2062
+ // max_num_partitions]
2063
+ const float* __restrict__ max_logits, // [num_seqs, num_heads,
2064
+ // max_num_partitions]
2065
+ const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
2066
+ // max_num_partitions, head_size]
2067
+ const int* __restrict__ context_lens, // [num_seqs]
2068
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2069
+ const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
2070
+ const auto num_heads = gridDim.x;
2071
+ const auto head_idx = blockIdx.x;
2072
+ const auto seq_idx = blockIdx.y;
2073
+
2074
+ // NOTE queries with sequence len > 1 are prefills and taken care by another
2075
+ // kernel.
2076
+ if (query_start_loc_ptr != nullptr &&
2077
+ (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) {
2078
+ return;
2079
+ }
2080
+
2081
+ const int context_len = context_lens[seq_idx];
2082
+ const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
2083
+ [[maybe_unused]] constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
2084
+ const int warpid = threadIdx.x / WARP_SIZE;
2085
+ [[maybe_unused]] const int laneid = threadIdx.x % WARP_SIZE;
2086
+
2087
+ __shared__ float shared_global_exp_sum;
2088
+ // max num partitions supported is warp_size * NPAR_LOOPS
2089
+ __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE];
2090
+
2091
+ if (warpid == 0) {
2092
+ const float* max_logits_ptr = max_logits +
2093
+ seq_idx * num_heads * max_num_partitions +
2094
+ head_idx * max_num_partitions;
2095
+
2096
+ // valid partition is the last valid partition in case threadid > num
2097
+ // partitions
2098
+ int valid_partition[NPAR_LOOPS];
2099
+ float reg_max_logit[NPAR_LOOPS];
2100
+ const int last_valid_partition = num_partitions - 1;
2101
+
2102
+ #pragma unroll
2103
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2104
+ const int partition_no = i * WARP_SIZE + threadIdx.x;
2105
+ valid_partition[i] =
2106
+ (partition_no < num_partitions) ? partition_no : last_valid_partition;
2107
+ }
2108
+ #pragma unroll
2109
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2110
+ reg_max_logit[i] = max_logits_ptr[valid_partition[i]];
2111
+ }
2112
+ float max_logit = reg_max_logit[0];
2113
+ #pragma unroll
2114
+ for (int i = 1; i < NPAR_LOOPS; i++) {
2115
+ max_logit = fmaxf(max_logit, reg_max_logit[i]);
2116
+ }
2117
+
2118
+ #pragma unroll
2119
+ for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
2120
+ max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask));
2121
+ }
2122
+
2123
+ const float* exp_sums_ptr = exp_sums +
2124
+ seq_idx * num_heads * max_num_partitions +
2125
+ head_idx * max_num_partitions;
2126
+
2127
+ float rescaled_exp_sum[NPAR_LOOPS];
2128
+ #pragma unroll
2129
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2130
+ rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]];
2131
+ }
2132
+ #pragma unroll
2133
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2134
+ const int partition_no = i * WARP_SIZE + threadIdx.x;
2135
+ rescaled_exp_sum[i] *= (partition_no < num_partitions)
2136
+ ? expf(reg_max_logit[i] - max_logit)
2137
+ : 0.0f;
2138
+ }
2139
+ float global_exp_sum = rescaled_exp_sum[0];
2140
+ #pragma unroll
2141
+ for (int i = 1; i < NPAR_LOOPS; i++) {
2142
+ global_exp_sum += rescaled_exp_sum[i];
2143
+ }
2144
+ #pragma unroll
2145
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2146
+ const int partition_no = i * WARP_SIZE + threadIdx.x;
2147
+ shared_exp_sums[partition_no] = rescaled_exp_sum[i];
2148
+ }
2149
+
2150
+ #pragma unroll
2151
+ for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
2152
+ global_exp_sum += __shfl_xor(global_exp_sum, mask);
2153
+ }
2154
+ if (threadIdx.x == 0) {
2155
+ shared_global_exp_sum = global_exp_sum;
2156
+ }
2157
+ } // warpid == 0
2158
+ const scalar_t* tmp_out_ptr =
2159
+ tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
2160
+ head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x;
2161
+ constexpr int MAX_NPAR = 32;
2162
+ scalar_t tmps[MAX_NPAR];
2163
+ const float dzero = 0.0f;
2164
+ #pragma unroll
2165
+ for (int j = 0; j < MAX_NPAR; j++) {
2166
+ tmps[j] = from_float<scalar_t>(dzero);
2167
+ }
2168
+ const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE;
2169
+ const int num_partition_offset = (num_partitions)*HEAD_SIZE;
2170
+ int idx = 0;
2171
+
2172
+ constexpr int JCHUNK = 16;
2173
+
2174
+ #pragma unroll
2175
+ for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) {
2176
+ // lastj is last valid partition
2177
+ const int lastj_offset =
2178
+ (j < num_partition_offset) ? j : last_partition_offset;
2179
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2180
+ idx++;
2181
+ }
2182
+ __syncthreads();
2183
+
2184
+ if (num_partitions > JCHUNK) {
2185
+ #pragma unroll
2186
+ for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE;
2187
+ j += HEAD_SIZE) {
2188
+ const int lastj_offset =
2189
+ (j < num_partition_offset) ? j : last_partition_offset;
2190
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2191
+ idx++;
2192
+ }
2193
+
2194
+ if (num_partitions > 2 * JCHUNK) {
2195
+ #pragma unroll
2196
+ for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE;
2197
+ j += HEAD_SIZE) {
2198
+ const int lastj_offset =
2199
+ (j < num_partition_offset) ? j : last_partition_offset;
2200
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2201
+ idx++;
2202
+ }
2203
+ }
2204
+ } // num_partitions > JCHUNK
2205
+
2206
+ // Aggregate tmp_out to out.
2207
+ float acc = 0.0f;
2208
+ #pragma unroll
2209
+ for (int j = 0; j < JCHUNK; j++) {
2210
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j];
2211
+ }
2212
+ if (num_partitions > JCHUNK) {
2213
+ #pragma unroll
2214
+ for (int j = JCHUNK; j < 2 * JCHUNK; j++) {
2215
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j];
2216
+ }
2217
+ if (num_partitions > 2 * JCHUNK) {
2218
+ #pragma unroll
2219
+ for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) {
2220
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j];
2221
+ }
2222
+ }
2223
+ }
2224
+
2225
+ for (int p = 1; p < NPAR_LOOPS; p++) {
2226
+ if (num_partitions > p * MAX_NPAR) {
2227
+ idx = 0;
2228
+ #pragma unroll
2229
+ for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE;
2230
+ j += HEAD_SIZE) {
2231
+ // lastj is last valid partition
2232
+ const int lastj_offset =
2233
+ (j < num_partition_offset) ? j : last_partition_offset;
2234
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2235
+ idx++;
2236
+ }
2237
+
2238
+ #pragma unroll
2239
+ for (int j = 0; j < MAX_NPAR; j++) {
2240
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR];
2241
+ }
2242
+ }
2243
+ }
2244
+
2245
+ const float inv_global_exp_sum =
2246
+ __fdividef(1.0f, shared_global_exp_sum + 1e-6f);
2247
+ acc *= inv_global_exp_sum;
2248
+
2249
+ const int64_t query_start_off = static_cast<int64_t>(
2250
+ query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);
2251
+ OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE +
2252
+ static_cast<int64_t>(head_idx) * HEAD_SIZE;
2253
+ out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
2254
+}
2255
+
2256
+#elif defined(__HIP__GFX12__)
2257
+
2258
+using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
2259
+
2260
+using bit16_t = uint16_t;
2261
+using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t;
2262
+typedef bit16x4 _B16x4;
2263
+
2264
+using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t;
2265
+union b16x8_u {
2266
+ bit16x8 u16x8;
2267
+ _B16x4 xy[2];
2268
+};
2269
+typedef b16x8_u _B16x8;
2270
+
2271
+using _B8x8 = uint2;
2272
+using bit8_t = uint8_t;
2273
+
2274
+typedef struct _B8x16 {
2275
+ _B8x8 xy[2];
2276
+} _B8x16;
2277
+
2278
+template <typename T, int absz, int cbid, int blgp>
2279
+__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA,
2280
+ const bit16x8& inpB,
2281
+ const floatx8& inpC) {
2282
+ if constexpr (std::is_same<T, _Float16>::value) {
2283
+ return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC);
2284
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
2285
+ return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(inpA, inpB, inpC);
2286
+ } else {
2287
+ static_assert(false, "unsupported 16b dtype");
2288
+ }
2289
+}
2290
+
2291
+template <typename T>
2292
+__device__ __forceinline__ float to_float(const T& inp) {
2293
+ if constexpr (std::is_same<T, _Float16>::value) {
2294
+ return (float)inp;
2295
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
2296
+ return __bfloat162float(inp);
2297
+ } else {
2298
+ static_assert(false, "unsupported 16b dtype");
2299
+ }
2300
+}
2301
+
2302
+template <typename T>
2303
+__device__ __forceinline__ float to_float_b16(const bit16_t& inp) {
2304
+ union tmpcvt {
2305
+ bit16_t u;
2306
+ _Float16 f;
2307
+ __hip_bfloat16 b;
2308
+ } t16;
2309
+ t16.u = inp;
2310
+ if constexpr (std::is_same<T, _Float16>::value) {
2311
+ return (float)t16.f;
2312
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
2313
+ return __bfloat162float(t16.b);
2314
+ } else {
2315
+ static_assert(false, "unsupported 16b dtype");
2316
+ }
2317
+}
2318
+
2319
+template <typename T>
2320
+__device__ __forceinline__ T from_float(const float& inp) {
2321
+ if constexpr (std::is_same<T, _Float16>::value) {
2322
+ return (_Float16)inp;
2323
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
2324
+ return __float2bfloat16(inp);
2325
+ } else {
2326
+ static_assert(false, "unsupported 16b dtype");
2327
+ }
2328
+}
2329
+
2330
+template <typename T>
2331
+__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
2332
+ if constexpr (std::is_same<T, _Float16>::value) {
2333
+ union h2cvt {
2334
+ __half2 h2[4];
2335
+ _B16x8 b16x8;
2336
+ } u;
2337
+ u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1]));
2338
+ u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3]));
2339
+ u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5]));
2340
+ u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7]));
2341
+ return u.b16x8;
2342
+ } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
2343
+ union b2cvt {
2344
+ __hip_bfloat162 b2[4];
2345
+ _B16x8 b16x8;
2346
+ } u;
2347
+
2348
+ u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1]));
2349
+ u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3]));
2350
+ u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5]));
2351
+ u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7]));
2352
+
2353
+ return u.b16x8;
2354
+ } else {
2355
+ static_assert(false, "unsupported 16b dtype");
2356
+ }
2357
+}
2358
+
2359
+// clang-format off
2360
+template <typename scalar_t, typename cache_t,
2361
+ vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
2362
+ int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
2363
+__global__
2364
+__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2365
+ const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
2366
+ const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
2367
+ // head_size/x, block_size, x]
2368
+ const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
2369
+ // head_size, block_size]
2370
+ const int num_kv_heads, const float scale,
2371
+ const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2372
+ const int* __restrict__ context_lens, // [num_seqs]
2373
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2374
+ const int max_num_blocks_per_seq,
2375
+ const float* __restrict__ alibi_slopes, // [num_heads]
2376
+ const int q_stride, const int kv_block_stride, const int kv_head_stride,
2377
+ float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
2378
+ float* __restrict__ max_logits, // [num_seqs, num_heads,
2379
+ // max_num_partitions]
2380
+ scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
2381
+ // head_size]
2382
+ OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
2383
+ int max_ctx_blocks, const float* k_scale, const float* v_scale) {
2384
+ // clang-format on
2385
+ constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11
2386
+ const int warpid = threadIdx.x / WARP_SIZE;
2387
+ const int laneid = threadIdx.x % WARP_SIZE;
2388
+ const int lane2id = laneid % 2;
2389
+ const int lane4id = laneid % 4;
2390
+ const int lane16id = laneid % 16;
2391
+ const int rowid = laneid / 16;
2392
+
2393
+ const int seq_idx = blockIdx.x;
2394
+ // NOTE queries with sequence len > 1 are prefills and taken care by another
2395
+ // kernel.
2396
+ if (query_start_loc_ptr != nullptr &&
2397
+ (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) {
2398
+ return;
2399
+ }
2400
+ const int partition_idx = blockIdx.y;
2401
+
2402
+ constexpr int T_PAR_SIZE = 256; // token partition size set to 256
2403
+
2404
+ const int max_num_partitions = gridDim.y;
2405
+
2406
+ const int context_len = context_lens[seq_idx]; // length of a seq
2407
+
2408
+ const int partition_start_token_idx = partition_idx * T_PAR_SIZE;
2409
+ // exit if partition is out of context for seq
2410
+ if (partition_start_token_idx >= context_len) {
2411
+ return;
2412
+ }
2413
+
2414
+ constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2);
2415
+
2416
+ __shared__ float shared_qk_max[NWARPS][16 + 1];
2417
+ __shared__ float shared_exp_sum[NWARPS][16 + 1];
2418
+ // shared_logits is used for multiple purposes
2419
+ __shared__ _B16x8 shared_logits[NWARPS][2][16][2];
2420
+
2421
+ // for QK wmma16x16_gfx12, layout is QHead/Tokenx16 across every 16 lanes,
2422
+ // 16 Bytes HeadElements in each lane, 2x16B HeadElements across 2 rows of
2423
+ // warp
2424
+ constexpr int ROWS_PER_WARP =
2425
+ WARP_SIZE / 16; // rows refers to 16 lanes; refer dpp terminology
2426
+ constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD =
2427
+ 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types
2428
+ constexpr int QKHE_PER_FETCH =
2429
+ CONTIGUOUS_KV_ELEMS_16B_LOAD *
2430
+ ROWS_PER_WARP; // each fetch across a warp fetches these many elements
2431
+ constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across
2432
+ // warp
2433
+
2434
+ _B16x8 Qlocal[QKHELOOP]; // note that 16 contiguous elements of Q should
2435
+ // be fetched per lane for 16 bit cache types
2436
+
2437
+ constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t);
2438
+
2439
+ constexpr int TOKENS_PER_WARP =
2440
+ T_PAR_SIZE /
2441
+ NWARPS; // sub partition of tokens per warp for qk calculation
2442
+ constexpr int TLOOP =
2443
+ TOKENS_PER_WARP /
2444
+ 16; // each wmma16x16x16 instruction processes 16 tokens
2445
+
2446
+ _B16x8 Klocal[TLOOP]
2447
+ [QKHELOOP]; // can be interpreted as B8x16 for 8 bit types
2448
+
2449
+ const int wg_start_head_idx = blockIdx.z * GQA_RATIO;
2450
+ const int wg_start_kv_head_idx = blockIdx.z;
2451
+ const int total_num_heads = gridDim.z * GQA_RATIO;
2452
+
2453
+ // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps
2454
+ // each wmma takes QH16xT16x16HE across warp
2455
+ // repeat wmma across QKHELOOP dimension
2456
+ // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens
2457
+ // across 2 rows x 8 tokens per lane
2458
+
2459
+ const int64_t query_start_off = static_cast<int64_t>(
2460
+ query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);
2461
+
2462
+ if (GQA_RATIO == 1) {
2463
+ const int local_qhead_idx = lane16id % GQA_RATIO;
2464
+ const int global_qhead_idx = wg_start_head_idx + local_qhead_idx;
2465
+ const scalar_t* q_ptr = q + query_start_off * q_stride +
2466
+ global_qhead_idx * HEAD_SIZE +
2467
+ rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD;
2468
+ if (lane16id < GQA_RATIO) {
2469
+ #pragma unroll
2470
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
2471
+ const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH;
2472
+ const _B16x8* q_fetch_ptr_16B =
2473
+ reinterpret_cast<const _B16x8*>(q_fetch_ptr);
2474
+ Qlocal[qkhe_depth] = *q_fetch_ptr_16B;
2475
+ }
2476
+ }
2477
+ } else {
2478
+ // fetch Q in shared across warps and then write to registers
2479
+ const int local_qhead_idx = 2 * warpid + rowid;
2480
+ const int global_qhead_idx = wg_start_head_idx + local_qhead_idx;
2481
+ const scalar_t* q_ptr =
2482
+ q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE;
2483
+
2484
+ const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B;
2485
+ if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) {
2486
+ const scalar_t* q_fetch_ptr = q_ptr + qhead_element;
2487
+ const _B16x8* q_fetch_ptr_16B =
2488
+ reinterpret_cast<const _B16x8*>(q_fetch_ptr);
2489
+ _B16x8 tmp = *q_fetch_ptr_16B;
2490
+
2491
+ const int offset1 =
2492
+ lane16id /
2493
+ 2; // 16 contiguous chunks of head elems are spread across 8x2lanes
2494
+ shared_logits[offset1][lane2id][local_qhead_idx][0] = tmp;
2495
+ }
2496
+
2497
+ __syncthreads();
2498
+
2499
+ #pragma unroll
2500
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
2501
+ Qlocal[qkhe_depth] =
2502
+ shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO][0];
2503
+ }
2504
+ }
2505
+
2506
+ const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
2507
+ const int last_ctx_block = num_context_blocks - 1;
2508
+
2509
+ const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
2510
+
2511
+ int kphysical_block_number[TLOOP];
2512
+
2513
+ // fetch k physical block numbers
2514
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2515
+ const int klocal_token_idx =
2516
+ TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
2517
+ const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
2518
+ const int kblock_idx = (kglobal_token_idx < context_len)
2519
+ ? kglobal_token_idx / BLOCK_SIZE
2520
+ : last_ctx_block;
2521
+ kphysical_block_number[token_depth] = block_table_seq[kblock_idx];
2522
+ }
2523
+
2524
+ constexpr int KX = 16 / sizeof(cache_t);
2525
+ const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride;
2526
+
2527
+ const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD;
2528
+
2529
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2530
+ const int64_t kblock_number =
2531
+ static_cast<int64_t>(kphysical_block_number[token_depth]);
2532
+ const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride;
2533
+ const int klocal_token_idx =
2534
+ TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
2535
+ const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
2536
+ const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE;
2537
+ const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX;
2538
+
2539
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
2540
+ const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH;
2541
+ const int offset1 = head_elem / KX;
2542
+ const int offset2 = head_elem % KX;
2543
+ const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2;
2544
+ const _B16x8* k_fetch_ptr_16B =
2545
+ reinterpret_cast<const _B16x8*>(k_fetch_ptr);
2546
+ Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B;
2547
+ }
2548
+ }
2549
+
2550
+ constexpr int VTOKENS_PER_LANE =
2551
+ TOKENS_PER_WARP / ROWS_PER_WARP; // 32/2 = 16 vtokens per lane
2552
+ constexpr int VBLOCKS_PER_LANE = 1; // assumes block size >=16
2553
+ constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps
2554
+ constexpr int VTLANELOOP = DIVIDE_ROUND_UP(
2555
+ VTOKENS_PER_LANE,
2556
+ CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes
2557
+ // minimum block size is 16
2558
+ constexpr int VHELOOP = DIVIDE_ROUND_UP(
2559
+ (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each
2560
+ // wmma instr works on 16 head elements
2561
+
2562
+ int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE];
2563
+
2564
+ // fetch v physical block numbers
2565
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
2566
+ for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE;
2567
+ vblock_depth++) {
2568
+ const int vlocal_token_idx =
2569
+ vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
2570
+ rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE;
2571
+ const int vglobal_token_idx =
2572
+ partition_start_token_idx + vlocal_token_idx;
2573
+ const int vblock_idx = (vglobal_token_idx < context_len)
2574
+ ? vglobal_token_idx / BLOCK_SIZE
2575
+ : last_ctx_block;
2576
+ vphysical_block_number[vtoken_depth][vblock_depth] =
2577
+ block_table_seq[vblock_idx];
2578
+ }
2579
+ }
2580
+
2581
+ _B16x8 Vlocal[VTLOOP][VHELOOP]
2582
+ [VTLANELOOP]; // this can be interpreted as B8x16 too
2583
+
2584
+ const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride +
2585
+ ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE);
2586
+
2587
+ // v fetches are 16head elems across lanes x 16 tokens per lane
2588
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
2589
+ const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
2590
+ const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
2591
+
2592
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
2593
+ for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
2594
+ const int vblock_depth = 0;
2595
+ const int64_t vblock_number = static_cast<int64_t>(
2596
+ vphysical_block_number[vtoken_depth][vblock_depth]);
2597
+ const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
2598
+
2599
+ const cache_t* v_fetch_ptr =
2600
+ v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
2601
+ const _B16x8* v_fetch_ptr_16B =
2602
+ reinterpret_cast<const _B16x8*>(v_fetch_ptr);
2603
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B;
2604
+ }
2605
+ }
2606
+ }
2607
+
2608
+ floatx8 dout[TLOOP];
2609
+ // qk wmma
2610
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2611
+ dout[token_depth] = {0};
2612
+ for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
2613
+ dout[token_depth] = gcn_wmma16x16x16_instr<scalar_t, 0, 0, 0>(
2614
+ Klocal[token_depth][qkhe_depth].u16x8, Qlocal[qkhe_depth].u16x8,
2615
+ dout[token_depth]);
2616
+ }
2617
+ dout[token_depth] *= scale;
2618
+ }
2619
+
2620
+ // calculate qk_max and exp_sum per warp and write to shared memory
2621
+ float qk_max = -FLT_MAX;
2622
+ float exp_sum = 0.0f;
2623
+ const int qkout_token_idx =
2624
+ partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 8;
2625
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2626
+ const int local_token_idx = qkout_token_idx + token_depth * 16;
2627
+ for (int i = 0; i < 8; i++) {
2628
+ const float tmp =
2629
+ (local_token_idx + i < context_len) ? dout[token_depth][i] : -FLT_MAX;
2630
+ qk_max = fmaxf(qk_max, tmp);
2631
+ }
2632
+ }
2633
+
2634
+ qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16));
2635
+
2636
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2637
+ const int local_token_idx = qkout_token_idx + token_depth * 16;
2638
+ for (int i = 0; i < 8; i++) {
2639
+ const float tmp = (local_token_idx + i < context_len)
2640
+ ? __expf(dout[token_depth][i] - qk_max)
2641
+ : 0.0f;
2642
+ dout[token_depth][i] = tmp;
2643
+ exp_sum += tmp;
2644
+ }
2645
+ }
2646
+
2647
+ exp_sum += __shfl_xor(exp_sum, 16);
2648
+
2649
+ __syncthreads();
2650
+
2651
+ if (laneid < 16) {
2652
+ shared_qk_max[warpid][lane16id] = qk_max;
2653
+ shared_exp_sum[warpid][lane16id] = exp_sum;
2654
+ }
2655
+
2656
+ __syncthreads();
2657
+
2658
+ // calculate partition qk_max and exp_sum
2659
+ float partition_qk_max = -FLT_MAX;
2660
+ float warp_qk_max_exp[NWARPS];
2661
+ float partition_exp_sum = 0.0f;
2662
+
2663
+ #pragma unroll
2664
+ for (int w = 0; w < NWARPS; w++) {
2665
+ warp_qk_max_exp[w] = shared_qk_max[w][lane16id];
2666
+ partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]);
2667
+ }
2668
+
2669
+ for (int w = 0; w < NWARPS; w++) {
2670
+ warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max);
2671
+ partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w];
2672
+ }
2673
+
2674
+ const float inv_sum_scale =
2675
+ __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid];
2676
+
2677
+ __syncthreads();
2678
+
2679
+ // write logits to shared mem
2680
+ #pragma unroll
2681
+ for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2682
+ dout[token_depth] *= inv_sum_scale;
2683
+ shared_logits[warpid][token_depth][lane16id][rowid] =
2684
+ from_floatx8<scalar_t>(dout[token_depth]);
2685
+ }
2686
+
2687
+ // write out partition max_logits and exp_sum
2688
+ if (threadIdx.x < GQA_RATIO) {
2689
+ const int qhead_idx = lane16id;
2690
+ const int offset = seq_idx * total_num_heads * max_num_partitions +
2691
+ (wg_start_head_idx + qhead_idx) * max_num_partitions +
2692
+ partition_idx;
2693
+ max_logits[offset] = partition_qk_max;
2694
+ exp_sums[offset] = partition_exp_sum;
2695
+ }
2696
+
2697
+ __syncthreads();
2698
+
2699
+ _B16x8 outelems[VHELOOP];
2700
+ // Softmax V wmma
2701
+ // v layout: 16he across lanes x 16 tokens per lane
2702
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
2703
+ floatx8 tmp_out = {0};
2704
+
2705
+ for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
2706
+ for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
2707
+ const int offset = rowid * VTLANELOOP + vfetch_depth;
2708
+ const int offset1 = offset % ROWS_PER_WARP;
2709
+ const int offset2 = offset / ROWS_PER_WARP;
2710
+ // if output format is 16 qheads across 16 lanes, 16 head elems spread
2711
+ // across rows
2712
+ tmp_out = gcn_wmma16x16x16_instr<scalar_t, 0, 0, 0>(
2713
+ Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x8,
2714
+ shared_logits[vtoken_depth][offset2][lane16id][offset1].u16x8,
2715
+ tmp_out);
2716
+ }
2717
+ }
2718
+ outelems[vhe_depth] = from_floatx8<scalar_t>(tmp_out);
2719
+ }
2720
+
2721
+ __syncthreads();
2722
+
2723
+ #pragma unroll
2724
+ for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
2725
+ shared_logits[warpid][vhe_depth][lane16id][rowid] =
2726
+ outelems[vhe_depth]; // lane16 id head dimension; rowid head element
2727
+ // dimension
2728
+ }
2729
+
2730
+ __syncthreads();
2731
+
2732
+ // write to tmp_out with coalesced writes after reading from shared mem
2733
+ if (warpid == 0) {
2734
+ _B16x8 vout[GQA_RATIO2];
2735
+ // each lane writes out 16Bytes of tmp_out along head elem dimension
2736
+ const int head_elem_idx = lane16id * 8;
2737
+ if (head_elem_idx < HEAD_SIZE) {
2738
+ for (int h = 0; h < GQA_RATIO2; h++) {
2739
+ const int local_head_idx = 2 * h + rowid;
2740
+ const int offset1 = (head_elem_idx / 16) % NWARPS;
2741
+ const int offset2 = head_elem_idx / 16 / NWARPS;
2742
+ const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row
2743
+ vout[h] = shared_logits[offset1][offset2][local_head_idx][offset3];
2744
+ }
2745
+
2746
+ const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions;
2747
+ scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult +
2748
+ partition_idx * HEAD_SIZE;
2749
+ for (int h = 0; h < GQA_RATIO2; h++) {
2750
+ const int local_head_idx = 2 * h + rowid;
2751
+ if (local_head_idx < GQA_RATIO) {
2752
+ const int out_head_idx = wg_start_head_idx + local_head_idx;
2753
+ scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult;
2754
+ scalar_t* out_ptr3 = out_ptr2 + head_elem_idx;
2755
+ _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3);
2756
+ *out_ptr_B16x8 = vout[h];
2757
+ }
2758
+ }
2759
+ }
2760
+ }
2761
+}
2762
+
2763
+template <typename scalar_t, typename cache_t,
2764
+ vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
2765
+ int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
2766
+ int GQA_RATIO>
2767
+__global__
2768
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
2769
+ const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
2770
+ const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
2771
+ // head_size/x, block_size, x]
2772
+ const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
2773
+ // head_size, block_size]
2774
+ const int num_kv_heads, const float scale,
2775
+ const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2776
+ const int* __restrict__ context_lens, // [num_seqs]
2777
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2778
+ const int max_num_blocks_per_seq,
2779
+ const float* __restrict__ alibi_slopes, // [num_heads]
2780
+ const int q_stride, const int kv_block_stride, const int kv_head_stride,
2781
+ float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
2782
+ float* __restrict__ max_logits, // [num_seqs, num_heads,
2783
+ // max_num_partitions]
2784
+ scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
2785
+ // head_size]
2786
+ OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
2787
+ int max_ctx_blocks, const float* k_scale, const float* v_scale) {
2788
+ UNREACHABLE_CODE
2789
+}
2790
+
2791
+// Grid: (num_heads, num_seqs).
2792
+template <typename scalar_t, typename OUTT, int HEAD_SIZE, int NUM_THREADS,
2793
+ int PARTITION_SIZE, int NPAR_LOOPS>
2794
+__global__
2795
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
2796
+ OUTT* __restrict__ out, // [num_seqs, num_heads, head_size]
2797
+ const float* __restrict__ exp_sums, // [num_seqs, num_heads,
2798
+ // max_num_partitions]
2799
+ const float* __restrict__ max_logits, // [num_seqs, num_heads,
2800
+ // max_num_partitions]
2801
+ const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
2802
+ // max_num_partitions, head_size]
2803
+ const int* __restrict__ context_lens, // [num_seqs]
2804
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2805
+ const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
2806
+ const auto num_heads = gridDim.x;
2807
+ const auto head_idx = blockIdx.x;
2808
+ const auto seq_idx = blockIdx.y;
2809
+
2810
+ // NOTE queries with sequence len > 1 are prefills and taken care by another
2811
+ // kernel.
2812
+ if (query_start_loc_ptr != nullptr &&
2813
+ (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) {
2814
+ return;
2815
+ }
2816
+
2817
+ const int context_len = context_lens[seq_idx];
2818
+ const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
2819
+ [[maybe_unused]] constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
2820
+ const int warpid = threadIdx.x / WARP_SIZE;
2821
+ [[maybe_unused]] const int laneid = threadIdx.x % WARP_SIZE;
2822
+
2823
+ __shared__ float shared_global_exp_sum;
2824
+ // max num partitions supported is warp_size * NPAR_LOOPS
2825
+ __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE];
2826
+
2827
+ if (warpid == 0) {
2828
+ const float* max_logits_ptr = max_logits +
2829
+ seq_idx * num_heads * max_num_partitions +
2830
+ head_idx * max_num_partitions;
2831
+
2832
+ // valid partition is the last valid partition in case threadid > num
2833
+ // partitions
2834
+ int valid_partition[NPAR_LOOPS];
2835
+ float reg_max_logit[NPAR_LOOPS];
2836
+ const int last_valid_partition = num_partitions - 1;
2837
+
2838
+ #pragma unroll
2839
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2840
+ const int partition_no = i * WARP_SIZE + threadIdx.x;
2841
+ valid_partition[i] =
2842
+ (partition_no < num_partitions) ? partition_no : last_valid_partition;
2843
+ }
2844
+ #pragma unroll
2845
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2846
+ reg_max_logit[i] = max_logits_ptr[valid_partition[i]];
2847
+ }
2848
+ float max_logit = reg_max_logit[0];
2849
+ #pragma unroll
2850
+ for (int i = 1; i < NPAR_LOOPS; i++) {
2851
+ max_logit = fmaxf(max_logit, reg_max_logit[i]);
2852
+ }
2853
+
2854
+ #pragma unroll
2855
+ for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
2856
+ max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask));
2857
+ }
2858
+
2859
+ const float* exp_sums_ptr = exp_sums +
2860
+ seq_idx * num_heads * max_num_partitions +
2861
+ head_idx * max_num_partitions;
2862
+
2863
+ float rescaled_exp_sum[NPAR_LOOPS];
2864
+ #pragma unroll
2865
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2866
+ rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]];
2867
+ }
2868
+ #pragma unroll
2869
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2870
+ const int partition_no = i * WARP_SIZE + threadIdx.x;
2871
+ rescaled_exp_sum[i] *= (partition_no < num_partitions)
2872
+ ? expf(reg_max_logit[i] - max_logit)
2873
+ : 0.0f;
2874
+ }
2875
+ float global_exp_sum = rescaled_exp_sum[0];
2876
+ #pragma unroll
2877
+ for (int i = 1; i < NPAR_LOOPS; i++) {
2878
+ global_exp_sum += rescaled_exp_sum[i];
2879
+ }
2880
+ #pragma unroll
2881
+ for (int i = 0; i < NPAR_LOOPS; i++) {
2882
+ const int partition_no = i * WARP_SIZE + threadIdx.x;
2883
+ shared_exp_sums[partition_no] = rescaled_exp_sum[i];
2884
+ }
2885
+
2886
+ #pragma unroll
2887
+ for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
2888
+ global_exp_sum += __shfl_xor(global_exp_sum, mask);
2889
+ }
2890
+ if (threadIdx.x == 0) {
2891
+ shared_global_exp_sum = global_exp_sum;
2892
+ }
2893
+ } // warpid == 0
2894
+ const scalar_t* tmp_out_ptr =
2895
+ tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
2896
+ head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x;
2897
+ constexpr int MAX_NPAR = 32;
2898
+ scalar_t tmps[MAX_NPAR];
2899
+ const float dzero = 0.0f;
2900
+ #pragma unroll
2901
+ for (int j = 0; j < MAX_NPAR; j++) {
2902
+ tmps[j] = from_float<scalar_t>(dzero);
2903
+ }
2904
+ const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE;
2905
+ const int num_partition_offset = (num_partitions)*HEAD_SIZE;
2906
+ int idx = 0;
2907
+
2908
+ constexpr int JCHUNK = 16;
2909
+
2910
+ #pragma unroll
2911
+ for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) {
2912
+ // lastj is last valid partition
2913
+ const int lastj_offset =
2914
+ (j < num_partition_offset) ? j : last_partition_offset;
2915
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2916
+ idx++;
2917
+ }
2918
+ __syncthreads();
2919
+
2920
+ if (num_partitions > JCHUNK) {
2921
+ #pragma unroll
2922
+ for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE;
2923
+ j += HEAD_SIZE) {
2924
+ const int lastj_offset =
2925
+ (j < num_partition_offset) ? j : last_partition_offset;
2926
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2927
+ idx++;
2928
+ }
2929
+
2930
+ if (num_partitions > 2 * JCHUNK) {
2931
+ #pragma unroll
2932
+ for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE;
2933
+ j += HEAD_SIZE) {
2934
+ const int lastj_offset =
2935
+ (j < num_partition_offset) ? j : last_partition_offset;
2936
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2937
+ idx++;
2938
+ }
2939
+ }
2940
+ } // num_partitions > JCHUNK
2941
+
2942
+ // Aggregate tmp_out to out.
2943
+ float acc = 0.0f;
2944
+ #pragma unroll
2945
+ for (int j = 0; j < JCHUNK; j++) {
2946
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j];
2947
+ }
2948
+ if (num_partitions > JCHUNK) {
2949
+ #pragma unroll
2950
+ for (int j = JCHUNK; j < 2 * JCHUNK; j++) {
2951
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j];
2952
+ }
2953
+ if (num_partitions > 2 * JCHUNK) {
2954
+ #pragma unroll
2955
+ for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) {
2956
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j];
2957
+ }
2958
+ }
2959
+ }
2960
+
2961
+ for (int p = 1; p < NPAR_LOOPS; p++) {
2962
+ if (num_partitions > p * MAX_NPAR) {
2963
+ idx = 0;
2964
+ #pragma unroll
2965
+ for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE;
2966
+ j += HEAD_SIZE) {
2967
+ // lastj is last valid partition
2968
+ const int lastj_offset =
2969
+ (j < num_partition_offset) ? j : last_partition_offset;
2970
+ tmps[idx] = tmp_out_ptr[lastj_offset];
2971
+ idx++;
2972
+ }
2973
+
2974
+ #pragma unroll
2975
+ for (int j = 0; j < MAX_NPAR; j++) {
2976
+ acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR];
2977
+ }
2978
+ }
2979
+ }
2980
+
2981
+ const float inv_global_exp_sum =
2982
+ __fdividef(1.0f, shared_global_exp_sum + 1e-6f);
2983
+ acc *= inv_global_exp_sum;
2984
+
2985
+ const int64_t query_start_off = static_cast<int64_t>(
2986
+ query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);
2987
+ OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE +
2988
+ static_cast<int64_t>(head_idx) * HEAD_SIZE;
2989
+ out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
2990
+}
2991
+
2992
+#else
2993
+
2994
+// clang-format off
2995
+template <typename scalar_t, typename cache_t,
2996
+ vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
2997
+ int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
2998
+ int GQA_RATIO>
2999
+__global__
3000
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
3001
+ const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
3002
+ const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
3003
+ const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
3004
+ const int num_kv_heads,
3005
+ const float scale,
3006
+ const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
3007
+ const int* __restrict__ context_lens, // [num_seqs]
3008
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
3009
+ const int max_num_blocks_per_seq,
3010
+ const float* __restrict__ alibi_slopes, // [num_heads]
3011
+ const int q_stride,
3012
+ const int kv_block_stride,
3013
+ const int kv_head_stride,
3014
+ float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
3015
+ float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
3016
+ scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
3017
+ OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
3018
+ int max_ctx_blocks, const float* k_scale, const float* v_scale) {
3019
+ UNREACHABLE_CODE
3020
+}
3021
+
3022
+template <typename scalar_t, typename cache_t,
3023
+ vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
3024
+ int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
3025
+ int GQA_RATIO>
3026
+__global__
3027
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
3028
+ const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
3029
+ const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
3030
+ const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
3031
+ const int num_kv_heads,
3032
+ const float scale,
3033
+ const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
3034
+ const int* __restrict__ context_lens, // [num_seqs]
3035
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
3036
+ const int max_num_blocks_per_seq,
3037
+ const float* __restrict__ alibi_slopes, // [num_heads]
3038
+ const int q_stride,
3039
+ const int kv_block_stride,
3040
+ const int kv_head_stride,
3041
+ float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
3042
+ float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
3043
+ scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
3044
+ OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
3045
+ int max_ctx_blocks, const float* k_scale, const float* v_scale) {
3046
+ UNREACHABLE_CODE
3047
+}
3048
+
3049
+// Grid: (num_heads, num_seqs).
3050
+template <typename scalar_t, typename OUTT, int HEAD_SIZE, int NUM_THREADS,
3051
+ int PARTITION_SIZE, int NPAR_LOOPS>
3052
+__global__
3053
+__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3054
+ OUTT* __restrict__ out, // [num_seqs, num_heads, head_size]
3055
+ const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
3056
+ const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
3057
+ const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
3058
+ const int* __restrict__ context_lens, // [num_seqs]
3059
+ const int* __restrict__ query_start_loc_ptr, // [num_seqs]
3060
+ const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
3061
+ UNREACHABLE_CODE
3062
+}
3063
+// clang-format on
3064
+
3065
+#endif
3066
+
3067
+#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \
3068
+ paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
3069
+ HEAD_SIZE, NTHR, ALIBI_ENABLED, \
3070
+ GQA_RATIO> \
3071
+ <<<grid, block, 0, stream>>>( \
3072
+ query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
3073
+ block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \
3074
+ max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
3075
+ kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
3076
+ max_ctx_blocks, k_scale_ptr, v_scale_ptr);
3077
+
3078
+#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \
3079
+ paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
3080
+ HEAD_SIZE, NTHR, ALIBI_ENABLED, \
3081
+ GQA_RATIO> \
3082
+ <<<grid, block, 0, stream>>>( \
3083
+ query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
3084
+ block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \
3085
+ max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
3086
+ kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
3087
+ max_ctx_blocks, k_scale_ptr, v_scale_ptr);
3088
+
3089
+#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \
3090
+ paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \
3091
+ PARTITION_SIZE, NPAR_LOOPS> \
3092
+ <<<reduce_grid, reduce_block, 0, stream>>>( \
3093
+ out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, \
3094
+ context_lens_ptr, query_start_loc_ptr, max_num_partitions, \
3095
+ fp8_out_scale_ptr);
3096
+
3097
+template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
3098
+ int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
3099
+ bool ALIBI_ENABLED>
3100
+void paged_attention_custom_launcher(
3101
+ torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
3102
+ torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
3103
+ torch::Tensor& value_cache, const int num_kv_heads, float scale,
3104
+ torch::Tensor& block_tables, torch::Tensor& context_lens,
3105
+ const std::optional<torch::Tensor>& query_start_loc, int max_context_len,
3106
+ const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
3107
+ torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale) {
3108
+ int num_seqs = block_tables.size(0);
3109
+ int num_heads = query.size(1);
3110
+ int head_size = query.size(2);
3111
+ int max_num_blocks_per_seq = block_tables.size(1);
3112
+ int q_stride = query.stride(0);
3113
+ int kv_block_stride = key_cache.stride(0);
3114
+ int kv_head_stride = key_cache.stride(1);
3115
+
3116
+ // NOTE: query start location is optional for V0 decode should not be used.
3117
+ // If batch contains mix of prefills and decode, prefills should be skipped.
3118
+ const int* query_start_loc_ptr =
3119
+ query_start_loc
3120
+ ? reinterpret_cast<const int*>(query_start_loc.value().data_ptr())
3121
+ : nullptr;
3122
+
3123
+ // NOTE: alibi_slopes is optional.
3124
+ const float* alibi_slopes_ptr =
3125
+ alibi_slopes
3126
+ ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
3127
+ : nullptr;
3128
+
3129
+ float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
3130
+ float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
3131
+ T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
3132
+ T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
3133
+ KVT* key_cache_ptr = reinterpret_cast<KVT*>(key_cache.data_ptr());
3134
+ KVT* value_cache_ptr = reinterpret_cast<KVT*>(value_cache.data_ptr());
3135
+ int* block_tables_ptr = block_tables.data_ptr<int>();
3136
+ int* context_lens_ptr = context_lens.data_ptr<int>();
3137
+ const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
3138
+ const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
3139
+ // NOTE: fp8_out_scale is optional.
3140
+ const auto fp8_out_scale_ptr =
3141
+ fp8_out_scale
3142
+ ? static_cast<const float*>(fp8_out_scale.value().data_ptr())
3143
+ : nullptr;
3144
+ OUTT* out_ptr = reinterpret_cast<OUTT*>(out.data_ptr());
3145
+
3146
+ const int max_ctx_blocks = DIVIDE_ROUND_UP(max_context_len, BLOCK_SIZE);
3147
+
3148
+ // partition size is fixed at 256 since both mfma4 and mfma16 kernels support
3149
+ // it mfma4 kernel also supports partition size 512
3150
+ constexpr int PARTITION_SIZE = 256;
3151
+ const int max_num_partitions =
3152
+ DIVIDE_ROUND_UP(max_context_len, PARTITION_SIZE);
3153
+ const int gqa_ratio = num_heads / num_kv_heads;
3154
+ assert(num_heads % num_kv_heads == 0);
3155
+ assert(head_size == HEAD_SIZE);
3156
+
3157
+ constexpr int NTHR = 256;
3158
+ dim3 grid(num_seqs, max_num_partitions, num_kv_heads);
3159
+ dim3 block(NTHR);
3160
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
3161
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
3162
+
3163
+ // mfma4 kernel is faster than mfma16 for gqa_ratio <= 4
3164
+ switch (gqa_ratio) {
3165
+ case 1:
3166
+ LAUNCH_CUSTOM_ATTENTION_MFMA4(1);
3167
+ break;
3168
+ case 2:
3169
+ LAUNCH_CUSTOM_ATTENTION_MFMA4(2);
3170
+ break;
3171
+ case 3:
3172
+ LAUNCH_CUSTOM_ATTENTION_MFMA4(3);
3173
+ break;
3174
+ case 4:
3175
+ LAUNCH_CUSTOM_ATTENTION_MFMA4(4);
3176
+ break;
3177
case 5:
3178
LAUNCH_CUSTOM_ATTENTION_MFMA16(5);
3179
break;
@@ -1744,13 +3251,195 @@ void paged_attention_custom_launcher(
3251
}
3252
}
3253
1747
-#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
1748
- PSIZE, ALIBI_ENABLED) \
1749
- paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
1750
- PSIZE, ALIBI_ENABLED>( \
1751
- out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
1752
- num_kv_heads, scale, block_tables, context_lens, query_start_loc, \
1753
- max_context_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);
3254
+template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
3255
+ int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
3256
+ bool ALIBI_ENABLED>
3257
+void paged_attention_custom_launcher_navi(
3258
+ torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
3259
+ torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
3260
+ torch::Tensor& value_cache, const int num_kv_heads, float scale,
3261
+ torch::Tensor& block_tables, torch::Tensor& context_lens,
3262
+ const std::optional<torch::Tensor>& query_start_loc, int max_context_len,
3263
+ const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
3264
+ torch::Tensor& v_scale) {
3265
+ int num_seqs = block_tables.size(0);
3266
+ int num_heads = query.size(1);
3267
+ int head_size = query.size(2);
3268
+ int max_num_blocks_per_seq = block_tables.size(1);
3269
+ int q_stride = query.stride(0);
3270
+ int kv_block_stride = key_cache.stride(0);
3271
+ int kv_head_stride = key_cache.stride(1);
3272
+
3273
+ // NOTE: query start location is optional for V0 decode should not be used.
3274
+ // If batch contains mix of prefills and decode, prefills should be skipped.
3275
+ const int* query_start_loc_ptr =
3276
+ query_start_loc
3277
+ ? reinterpret_cast<const int*>(query_start_loc.value().data_ptr())
3278
+ : nullptr;
3279
+
3280
+ // NOTE: Navi does not support alibi_slopes.
3281
+ const float* alibi_slopes_ptr = nullptr;
3282
+
3283
+ float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
3284
+ float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
3285
+ T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
3286
+ T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
3287
+ KVT* key_cache_ptr = reinterpret_cast<KVT*>(key_cache.data_ptr());
3288
+ KVT* value_cache_ptr = reinterpret_cast<KVT*>(value_cache.data_ptr());
3289
+ int* block_tables_ptr = block_tables.data_ptr<int>();
3290
+ int* context_lens_ptr = context_lens.data_ptr<int>();
3291
+
3292
+ const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
3293
+ const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
3294
+ // NOTE: Navi does not support fp8.
3295
+ const auto fp8_out_scale_ptr = nullptr;
3296
+ OUTT* out_ptr = reinterpret_cast<OUTT*>(out.data_ptr());
3297
+
3298
+ const int max_ctx_blocks = DIVIDE_ROUND_UP(max_context_len, BLOCK_SIZE);
3299
+
3300
+ constexpr int PARTITION_SIZE = 256;
3301
+ const int max_num_partitions =
3302
+ DIVIDE_ROUND_UP(max_context_len, PARTITION_SIZE);
3303
+ const int gqa_ratio = num_heads / num_kv_heads;
3304
+ assert(num_heads % num_kv_heads == 0);
3305
+ assert(head_size == HEAD_SIZE);
3306
+
3307
+ constexpr int NTHR = 256;
3308
+ dim3 grid(num_seqs, max_num_partitions, num_kv_heads);
3309
+ dim3 block(NTHR);
3310
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
3311
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
3312
+
3313
+ switch (gqa_ratio) {
3314
+ case 1:
3315
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(1);
3316
+ break;
3317
+ case 2:
3318
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(2);
3319
+ break;
3320
+ case 3:
3321
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(3);
3322
+ break;
3323
+ case 4:
3324
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(4);
3325
+ break;
3326
+ case 5:
3327
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(5);
3328
+ break;
3329
+ case 6:
3330
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(6);
3331
+ break;
3332
+ case 7:
3333
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(7);
3334
+ break;
3335
+ case 8:
3336
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(8);
3337
+ break;
3338
+ case 9:
3339
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(9);
3340
+ break;
3341
+ case 10:
3342
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(10);
3343
+ break;
3344
+ case 11:
3345
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(11);
3346
+ break;
3347
+ case 12:
3348
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(12);
3349
+ break;
3350
+ case 13:
3351
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(13);
3352
+ break;
3353
+ case 14:
3354
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(14);
3355
+ break;
3356
+ case 15:
3357
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(15);
3358
+ break;
3359
+ case 16:
3360
+ LAUNCH_CUSTOM_ATTENTION_MFMA16(16);
3361
+ break;
3362
+ default:
3363
+ TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio);
3364
+ break;
3365
+ }
3366
+
3367
+ dim3 reduce_grid(num_heads, num_seqs);
3368
+ dim3 reduce_block(head_size);
3369
+ const int warp_size = 32;
3370
+ const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, warp_size);
3371
+ // reduction kernel supports upto 16 NPAR_loops * 32 (warp_size) * 256
3372
+ // (partition size) = 128K context length
3373
+ switch (npar_loops) {
3374
+ case 1:
3375
+ LAUNCH_CUSTOM_REDUCTION(1);
3376
+ break;
3377
+ case 2:
3378
+ LAUNCH_CUSTOM_REDUCTION(2);
3379
+ break;
3380
+ case 3:
3381
+ LAUNCH_CUSTOM_REDUCTION(3);
3382
+ break;
3383
+ case 4:
3384
+ LAUNCH_CUSTOM_REDUCTION(4);
3385
+ break;
3386
+ case 5:
3387
+ LAUNCH_CUSTOM_REDUCTION(5);
3388
+ break;
3389
+ case 6:
3390
+ LAUNCH_CUSTOM_REDUCTION(6);
3391
+ break;
3392
+ case 7:
3393
+ LAUNCH_CUSTOM_REDUCTION(7);
3394
+ break;
3395
+ case 8:
3396
+ LAUNCH_CUSTOM_REDUCTION(8);
3397
+ break;
3398
+ case 9:
3399
+ LAUNCH_CUSTOM_REDUCTION(9);
3400
+ break;
3401
+ case 10:
3402
+ LAUNCH_CUSTOM_REDUCTION(10);
3403
+ break;
3404
+ case 11:
3405
+ LAUNCH_CUSTOM_REDUCTION(11);
3406
+ break;
3407
+ case 12:
3408
+ LAUNCH_CUSTOM_REDUCTION(12);
3409
+ break;
3410
+ case 13:
3411
+ LAUNCH_CUSTOM_REDUCTION(13);
3412
+ break;
3413
+ case 14:
3414
+ LAUNCH_CUSTOM_REDUCTION(14);
3415
+ break;
3416
+ case 15:
3417
+ LAUNCH_CUSTOM_REDUCTION(15);
3418
+ break;
3419
+ case 16:
3420
+ LAUNCH_CUSTOM_REDUCTION(16);
3421
+ break;
3422
+ default:
3423
+ TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops);
3424
+ break;
3425
+ }
3426
+}
3427
+
3428
+#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
3429
+ PSIZE, ALIBI_ENABLED) \
3430
+ if (!is_navi) { \
3431
+ paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3432
+ OUTT, PSIZE, ALIBI_ENABLED>( \
3433
+ out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
3434
+ num_kv_heads, scale, block_tables, context_lens, query_start_loc, \
3435
+ max_context_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
3436
+ } else { \
3437
+ paged_attention_custom_launcher_navi< \
3438
+ T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, ALIBI_ENABLED>( \
3439
+ out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
3440
+ num_kv_heads, scale, block_tables, context_lens, query_start_loc, \
3441
+ max_context_len, alibi_slopes, k_scale, v_scale); \
3442
+ }
3443
3444
#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3445
OUTT, PSIZE) \
@@ -1807,6 +3496,24 @@ void paged_attention_custom_launcher(
3496
break; \
3497
}
3498
3499
+bool is_navi_gpu() {
3500
+ static bool is_cached = false;
3501
+ static bool result;
3502
+
3503
+ if (!is_cached) {
3504
+ int device_id;
3505
+ hipDeviceProp_t deviceProp;
3506
+ hipGetDevice(&device_id);
3507
+ hipGetDeviceProperties(&deviceProp, device_id);
3508
+
3509
+ std::string arch = deviceProp.gcnArchName;
3510
+ result = arch.find("gfx11") == 0 || arch.find("gfx12") == 0;
3511
+ is_cached = true;
3512
+ }
3513
+
3514
+ return result;
3515
+}
3516
+
3517
// clang-format off
3518
void paged_attention(
3519
torch::Tensor& out, // [num_seqs, num_heads, head_size]
@@ -1827,6 +3534,8 @@ void paged_attention(
3534
torch::Tensor& v_scale,
3535
const std::optional<torch::Tensor>& fp8_out_scale) {
3536
// clang-format on
3537
+ bool is_navi = is_navi_gpu();
3538
+
3539
const int head_size = query.size(2);
3540
if (kv_cache_dtype == "auto") {
3541
if (query.dtype() == at::ScalarType::Half) {
tests/kernels/attention/test_attention.py
+7
-1
@@ -148,6 +148,11 @@ def test_paged_attention(
148
or (version == "rocm" and head_size not in (64, 128))):
149
pytest.skip()
150
151
+ if (version == "rocm" and current_platform.is_navi()
152
+ and (kv_cache_dtype == "fp8" or head_size != 128
153
+ or block_size != 16 or use_alibi)):
154
+ pytest.skip()
155
+
156
global PARTITION_SIZE
157
158
current_platform.seed_everything(seed)
@@ -275,6 +280,7 @@ def test_paged_attention(
280
scale,
281
block_tables,
282
seq_lens,
283
+ None,
284
block_size,
285
max_seq_len,
286
alibi_slopes,
@@ -286,7 +292,7 @@ def test_paged_attention(
292
opcheck(torch.ops._rocm_C.paged_attention,
293
(output, exp_sums, max_logits, tmp_output, query,
294
key_cache, value_cache, num_kv_heads, scale, block_tables,
289
- seq_lens, block_size, max_seq_len, alibi_slopes,
295
+ seq_lens, None, block_size, max_seq_len, alibi_slopes,
296
kv_cache_dtype, k_scale, v_scale),
297
cond=(head_size == HEAD_SIZES[0]
298
and block_size == BLOCK_SIZES[0]))
vllm/attention/backends/rocm_flash_attn.py
+2
-1
@@ -861,7 +861,8 @@ class ROCmFlashAttentionImpl(AttentionImpl):
861
gqa_ratio = num_heads // self.num_kv_heads
862
use_custom = use_rocm_custom_paged_attention(
863
decode_query.dtype, head_size, block_size, gqa_ratio,
864
- decode_meta.max_decode_seq_len, self.sliding_window)
864
+ decode_meta.max_decode_seq_len, self.sliding_window,
865
+ self.kv_cache_dtype, self.alibi_slopes)
866
if use_custom:
867
max_seq_len = (decode_meta.max_decode_seq_len if self.attn_type
868
!= AttentionType.ENCODER_DECODER else
vllm/attention/ops/chunked_prefill_paged_decode.py
+2
-1
@@ -283,7 +283,8 @@ def chunked_prefill_paged_decode(
283
use_custom = use_rocm_custom_paged_attention(query.dtype, head_size,
284
block_size,
285
num_queries_per_kv,
286
- max_seq_len, sliding_window)
286
+ max_seq_len, sliding_window,
287
+ kv_cache_dtype, alibi_slopes)
288
if use_custom:
289
_PARTITION_SIZE_ROCM = 256
290
max_num_partitions = ((max_seq_len + _PARTITION_SIZE_ROCM - 1) //
vllm/platforms/rocm.py
+34
-14
@@ -102,26 +102,42 @@ def on_mi250_mi300() -> bool:
102
103
104
@cache
105
-def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
106
- block_size: int, gqa_ratio: int,
107
- max_seq_len: int,
108
- sliding_window: int) -> bool:
105
+def use_rocm_custom_paged_attention(
106
+ qtype: torch.dtype,
107
+ head_size: int,
108
+ block_size: int,
109
+ gqa_ratio: int,
110
+ max_seq_len: int,
111
+ sliding_window: int,
112
+ kv_cache_dtype: str,
113
+ alibi_slopes: Optional[torch.Tensor] = None) -> bool:
114
115
GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
116
ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
117
+ ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
118
113
- # rocm custom page attention not support on gfx1*
119
# custom paged attn always supported on V0. On V1, requires sliding window
120
# disabled due to observed numerical discrepancy.
116
- return (ON_GFX9 and (not envs.VLLM_USE_V1 or sliding_window == 0
117
- or sliding_window == (-1, -1))
118
- and (qtype == torch.half or qtype == torch.bfloat16)
119
- and (head_size == 64 or head_size == 128)
120
- and (block_size == 16 or block_size == 32)
121
- and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768
122
- and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
123
- and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
124
- and envs.VLLM_ROCM_USE_AITER))
121
+ if ON_GFX9:
122
+ return ((not envs.VLLM_USE_V1 or sliding_window == 0
123
+ or sliding_window == (-1, -1))
124
+ and (qtype == torch.half or qtype == torch.bfloat16)
125
+ and (head_size == 64 or head_size == 128)
126
+ and (block_size == 16 or block_size == 32)
127
+ and (gqa_ratio >= 1 and gqa_ratio <= 16)
128
+ and max_seq_len <= 32768 and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
129
+ and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
130
+ and envs.VLLM_ROCM_USE_AITER))
131
+
132
+ else:
133
+ return (ON_GFX11_GFX12 and (not envs.VLLM_USE_V1 or sliding_window == 0
134
+ or sliding_window == (-1, -1))
135
+ and (qtype == torch.half or qtype == torch.bfloat16)
136
+ and head_size == 128 and block_size == 16
137
+ and (gqa_ratio >= 3 and gqa_ratio <= 16)
138
+ and max_seq_len <= 32768 and alibi_slopes is None
139
+ and kv_cache_dtype == "auto"
140
+ and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
141
142
143
class RocmPlatform(Platform):
@@ -362,3 +378,7 @@ class RocmPlatform(Platform):
378
def get_cu_count(cls, device_id: int = 0) -> int:
379
return torch.cuda.get_device_properties(
380
device_id).multi_processor_count
381
+
382
+ @classmethod
383
+ def is_navi(cls) -> bool:
384
+ return 'gfx1' in torch.cuda.get_device_properties(0).gcnArchName