@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e73ff24e31

Fork 0

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, 16:00 UTC

[ROCM][KERNEL] Paged attention for V1 (#15720)

Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Signed-off-by: root <root@banff-cyxtera-s65-4.amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: root <root@banff-cyxtera-s65-4.amd.com>

Aleksandr Malyshev committed Apr 2, 2025 at 19:48 UTC e73ff24e31d29b095e24e160d3b0651d2436d977

11 files changed +220 -110

csrc/rocm/attention.cu

+73 -32

@@ -272,6 +272,7 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(

272 const float scale,

273 const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]

274 const int* __restrict__ context_lens, // [num_seqs]

275 + const int* __restrict__ query_start_loc_ptr, // [num_seqs]

276 const int max_num_blocks_per_seq,

277 const float* __restrict__ alibi_slopes, // [num_heads]

278 const int q_stride,

@@ -291,6 +292,13 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(

292 const int rowid = laneid / 16;

293

294 const auto seq_idx = blockIdx.x;

295 + // NOTE queries with sequence len > 1 are prefills and taken care by another

296 + // kernel.

297 + if (query_start_loc_ptr != nullptr &&

298 + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) {

299 + return;

300 + }

301 +

302 const auto partition_idx = blockIdx.y;

303

304 constexpr int T_PAR_SIZE = 256; // token partition size set to 256

@@ -377,9 +385,10 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(

385 // fetch Q in shared across warps and then write to registers

386 const int local_qhead_idx = 4 * warpid + rowid;

387 const int global_qhead_idx = wg_start_head_idx + local_qhead_idx;

380 - const int64_t seq_idx64 = static_cast<int64_t>(seq_idx);

388 + const int64_t query_start_off = static_cast<int64_t>(

389 + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);

390 const scalar_t* q_ptr =

382 - q + seq_idx64 * q_stride + global_qhead_idx * HEAD_SIZE;

391 + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE;

392

393 const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B;

394 if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) {

@@ -777,6 +786,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(

786 const float scale,

787 const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]

788 const int* __restrict__ context_lens, // [num_seqs]

789 + const int* __restrict__ query_start_loc_ptr, // [num_seqs]

790 const int max_num_blocks_per_seq,

791 const float* __restrict__ alibi_slopes, // [num_heads]

792 const int q_stride,

@@ -794,6 +804,12 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(

804 const int lane4id = laneid % 4;

805

806 const auto seq_idx = blockIdx.x;

807 + // NOTE queries with sequence len > 1 are prefills and taken care by another

808 + // kernel.

809 + if (query_start_loc_ptr != nullptr &&

810 + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) {

811 + return;

812 + }

813 const auto partition_idx = blockIdx.y;

814 const auto partition_size = blockDim.x;

815 const auto max_num_partitions = gridDim.y;

@@ -882,9 +898,11 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(

898 }

899

900 // fetch q elements

885 - // every 4 lanes fetch 8 elems, so warp fetches 8*16 = 128 elems

901 + // every 4 lanes fetch 8 elems, so warp fetches 8*16 = 128 elemsc

902 + const int64_t query_start_off = static_cast<int64_t>(

903 + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);

904 const scalar_t* q_ptr =

887 - q + seq_idx * q_stride + wg_start_head_idx * HEAD_SIZE;

905 + q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE;

906 const _B16x8* q_ptrh8 = reinterpret_cast<const _B16x8*>(q_ptr);

907 const int qhead_elemh8 = laneid / 4;

908

@@ -1267,10 +1285,19 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(

1285 const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,

1286 // max_num_partitions, head_size]

1287 const int* __restrict__ context_lens, // [num_seqs]

1288 + const int* __restrict__ query_start_loc_ptr, // [num_seqs]

1289 const int max_num_partitions) {

1290 const auto num_heads = gridDim.x;

1291 const auto head_idx = blockIdx.x;

1292 const auto seq_idx = blockIdx.y;

1293 +

1294 + // NOTE queries with sequence len > 1 are prefills and taken care by another

1295 + // kernel.

1296 + if (query_start_loc_ptr != nullptr &&

1297 + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) {

1298 + return;

1299 + }

1300 +

1301 const int context_len = context_lens[seq_idx];

1302 const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);

1303 [[maybe_unused]] constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;

@@ -1439,7 +1466,9 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(

1466 __fdividef(1.0f, shared_global_exp_sum + 1e-6f);

1467 acc *= inv_global_exp_sum;

1468

1442 - OUTT* out_ptr = out + static_cast<int64_t>(seq_idx) * num_heads * HEAD_SIZE +

1469 + const int64_t query_start_off = static_cast<int64_t>(

1470 + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx);

1471 + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE +

1472 static_cast<int64_t>(head_idx) * HEAD_SIZE;

1473 if constexpr (std::is_same<OUTT, bit8_t>::value) {

1474 out_ptr[threadIdx.x] =

@@ -1466,6 +1495,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(

1495 const float scale,

1496 const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]

1497 const int* __restrict__ context_lens, // [num_seqs]

1498 + const int* __restrict__ query_start_loc_ptr, // [num_seqs]

1499 const int max_num_blocks_per_seq,

1500 const float* __restrict__ alibi_slopes, // [num_heads]

1501 const int q_stride,

@@ -1492,6 +1522,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(

1522 const float scale,

1523 const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]

1524 const int* __restrict__ context_lens, // [num_seqs]

1525 + const int* __restrict__ query_start_loc_ptr, // [num_seqs]

1526 const int max_num_blocks_per_seq,

1527 const float* __restrict__ alibi_slopes, // [num_heads]

1528 const int q_stride,

@@ -1515,6 +1546,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(

1546 const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]

1547 const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]

1548 const int* __restrict__ context_lens, // [num_seqs]

1549 + const int* __restrict__ query_start_loc_ptr, // [num_seqs]

1550 const int max_num_partitions) {

1551 UNREACHABLE_CODE

1552 }

@@ -1522,34 +1554,34 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(

1554

1555 #endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support

1556

1525 -#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \

1526 - paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \

1527 - HEAD_SIZE, NTHR, ALIBI_ENABLED, \

1528 - GQA_RATIO> \

1529 - <<<grid, block, 0, stream>>>( \

1530 - query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \

1531 - block_tables_ptr, context_lens_ptr, max_num_blocks_per_seq, \

1532 - alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \

1533 - exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, max_ctx_blocks, \

1534 - k_scale_ptr, v_scale_ptr);

1535 -

1536 -#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \

1537 - paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \

1538 - HEAD_SIZE, NTHR, ALIBI_ENABLED, \

1539 - GQA_RATIO> \

1540 - <<<grid, block, 0, stream>>>( \

1541 - query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \

1542 - block_tables_ptr, context_lens_ptr, max_num_blocks_per_seq, \

1543 - alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \

1544 - exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, max_ctx_blocks, \

1545 - k_scale_ptr, v_scale_ptr);

1557 +#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \

1558 + paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \

1559 + HEAD_SIZE, NTHR, ALIBI_ENABLED, \

1560 + GQA_RATIO> \

1561 + <<<grid, block, 0, stream>>>( \

1562 + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \

1563 + block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \

1564 + max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \

1565 + kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \

1566 + max_ctx_blocks, k_scale_ptr, v_scale_ptr);

1567 +

1568 +#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \

1569 + paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \

1570 + HEAD_SIZE, NTHR, ALIBI_ENABLED, \

1571 + GQA_RATIO> \

1572 + <<<grid, block, 0, stream>>>( \

1573 + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \

1574 + block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \

1575 + max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \

1576 + kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \

1577 + max_ctx_blocks, k_scale_ptr, v_scale_ptr);

1578

1579 #define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \

1580 paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \

1581 PARTITION_SIZE, NPAR_LOOPS> \

1582 <<<reduce_grid, reduce_block, 0, stream>>>( \

1583 out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, \

1552 - context_lens_ptr, max_num_partitions);

1584 + context_lens_ptr, query_start_loc_ptr, max_num_partitions);

1585

1586 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,

1587 int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,

@@ -1559,9 +1591,10 @@ void paged_attention_custom_launcher(

1591 torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,

1592 torch::Tensor& value_cache, const int num_kv_heads, float scale,

1593 torch::Tensor& block_tables, torch::Tensor& context_lens,

1562 - int max_context_len, const std::optional<torch::Tensor>& alibi_slopes,

1563 - torch::Tensor& k_scale, torch::Tensor& v_scale) {

1564 - int num_seqs = query.size(0);

1594 + const std::optional<torch::Tensor>& query_start_loc, int max_context_len,

1595 + const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,

1596 + torch::Tensor& v_scale) {

1597 + int num_seqs = block_tables.size(0);

1598 int num_heads = query.size(1);

1599 int head_size = query.size(2);

1600 int max_num_blocks_per_seq = block_tables.size(1);

@@ -1569,6 +1602,13 @@ void paged_attention_custom_launcher(

1602 int kv_block_stride = key_cache.stride(0);

1603 int kv_head_stride = key_cache.stride(1);

1604

1605 + // NOTE: query start location is optional for V0 decode should not be used.

1606 + // If batch contains mix of prefills and decode, prefills should be skipped.

1607 + const int* query_start_loc_ptr =

1608 + query_start_loc

1609 + ? reinterpret_cast<const int*>(query_start_loc.value().data_ptr())

1610 + : nullptr;

1611 +

1612 // NOTE: alibi_slopes is optional.

1613 const float* alibi_slopes_ptr =

1614 alibi_slopes

@@ -1700,8 +1740,8 @@ void paged_attention_custom_launcher(

1740 paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \

1741 PSIZE, ALIBI_ENABLED>( \

1742 out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \

1703 - num_kv_heads, scale, block_tables, context_lens, max_context_len, \

1704 - alibi_slopes, k_scale, v_scale);

1743 + num_kv_heads, scale, block_tables, context_lens, query_start_loc, \

1744 + max_context_len, alibi_slopes, k_scale, v_scale);

1745

1746 #define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \

1747 PSIZE) \

@@ -1750,6 +1790,7 @@ void paged_attention(

1790 double scale,

1791 torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]

1792 torch::Tensor& context_lens, // [num_seqs]

1793 + const std::optional<torch::Tensor>& query_start_loc, // [num_seqs]

1794 int64_t block_size, int64_t max_context_len,

1795 const std::optional<torch::Tensor>& alibi_slopes,

1796 const std::string& kv_cache_dtype, torch::Tensor& k_scale,

csrc/rocm/ops.h

+3 -2

@@ -7,8 +7,9 @@ void paged_attention(torch::Tensor& out, torch::Tensor& exp_sums,

7 torch::Tensor& query, torch::Tensor& key_cache,

8 torch::Tensor& value_cache, int64_t num_kv_heads,

9 double scale, torch::Tensor& block_tables,

10 - torch::Tensor& context_lens, int64_t block_size,

11 - int64_t max_context_len,

10 + torch::Tensor& context_lens,

11 + const std::optional<torch::Tensor>& query_start_loc,

12 + int64_t block_size, int64_t max_context_len,

13 const std::optional<torch::Tensor>& alibi_slopes,

14 const std::string& kv_cache_dtype, torch::Tensor& k_scale,

15 torch::Tensor& v_scale);

csrc/rocm/torch_bindings.cpp

+3 -1

@@ -23,7 +23,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {

23 " Tensor query, Tensor key_cache,"

24 " Tensor value_cache, int num_kv_heads,"

25 " float scale, Tensor block_tables,"

26 - " Tensor context_lens, int block_size,"

26 + " Tensor context_lens,"

27 + " Tensor? query_start_loc,"

28 + " int block_size,"

29 " int max_context_len,"

30 " Tensor? alibi_slopes,"

31 " str kv_cache_dtype,"

tests/kernels/test_prefix_prefill.py

@@ -164,6 +164,7 @@ def test_contexted_kv_attention(

164 block_table,

165 b_start_loc,

166 b_seq_len,

167 + MAX_CTX_LEN,

168 max_input_len,

169 k_scale,

170 v_scale,

@@ -180,6 +181,7 @@ def test_contexted_kv_attention(

181 block_table,

182 b_start_loc,

183 b_seq_len,

184 + MAX_CTX_LEN,

185 max_input_len,

186 k_scale,

187 v_scale,

@@ -397,6 +399,7 @@ def test_contexted_kv_attention_alibi(

399 block_table,

400 b_start_loc,

401 b_seq_len,

402 + MAX_CTX_LEN,

403 max_input_len,

404 k_scale,

405 v_scale,

@@ -413,6 +416,7 @@ def test_contexted_kv_attention_alibi(

416 block_table,

417 b_start_loc,

418 b_seq_len,

419 + MAX_CTX_LEN,

420 max_input_len,

421 k_scale,

422 v_scale,

vllm/_custom_ops.py

+4 -2

@@ -110,6 +110,7 @@ def paged_attention_rocm(

110 scale: float,

111 block_tables: torch.Tensor,

112 seq_lens: torch.Tensor,

113 + query_start_loc: Optional[torch.Tensor],

114 block_size: int,

115 max_seq_len: int,

116 alibi_slopes: Optional[torch.Tensor],

@@ -120,8 +121,9 @@ def paged_attention_rocm(

121 torch.ops._rocm_C.paged_attention(out, exp_sum, max_logits, tmp_out, query,

122 key_cache, value_cache, num_kv_heads,

123 scale, block_tables, seq_lens,

123 - block_size, max_seq_len, alibi_slopes,

124 - kv_cache_dtype, k_scale, v_scale)

124 + query_start_loc, block_size, max_seq_len,

125 + alibi_slopes, kv_cache_dtype, k_scale,

126 + v_scale)

127

128

129 def mla_decode_kvcache_cpu(

vllm/attention/backends/rocm_flash_attn.py

+6 -18

@@ -17,16 +17,13 @@ from vllm.attention.ops.paged_attn import (PagedAttention,

17 PagedAttentionMetadata)

18 from vllm.logger import init_logger

19 from vllm.platforms import current_platform

20 +from vllm.platforms.rocm import use_rocm_custom_paged_attention

22 if TYPE_CHECKING:

23 from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata

25 logger = init_logger(__name__)

25 -

26 _PARTITION_SIZE_ROCM = 256

27 -_GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName

28 -_ON_NAVI = "gfx1" in _GPU_ARCH

29 -_ON_MI250_MI300 = any(arch in _GPU_ARCH for arch in ["gfx90a", "gfx942"])

29 class ROCmFlashAttentionBackend(AttentionBackend):

@@ -790,9 +787,9 @@ class ROCmFlashAttentionImpl(AttentionImpl):

787 num_seqs, num_heads, head_size = decode_query.shape

788 block_size = value_cache.shape[3]

789 gqa_ratio = num_heads // self.num_kv_heads

793 - use_custom = _use_rocm_custom_paged_attention(

790 + use_custom = use_rocm_custom_paged_attention(

791 decode_query.dtype, head_size, block_size, gqa_ratio,

795 - decode_meta.max_decode_seq_len)

792 + decode_meta.max_decode_seq_len, self.sliding_window)

793 if use_custom:

794 max_seq_len = (decode_meta.max_decode_seq_len if self.attn_type

795 != AttentionType.ENCODER_DECODER else

@@ -817,6 +814,8 @@ class ROCmFlashAttentionImpl(AttentionImpl):

814 out = output[num_prefill_tokens:]

815 else:

816 out = output

817 +

818 + query_start_loc = None

819 ops.paged_attention_rocm(

820 out,

821 exp_sums,

@@ -833,6 +832,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):

832 decode_meta.seq_lens_tensor

833 if self.attn_type != AttentionType.ENCODER_DECODER else

834 decode_meta.encoder_seq_lens_tensor,

835 + query_start_loc,

836 block_size,

837 max_seq_len,

838 self.alibi_slopes,

@@ -898,15 +898,3 @@ def _sdpa_attention(

898 start = end

899

900 return output

901 -

902 -

903 -def _use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,

904 - block_size: int, gqa_ratio: int,

905 - max_seq_len: int) -> bool:

906 - # rocm custom page attention not support on navi (gfx1*)

907 - return (_ON_MI250_MI300 and not _ON_NAVI

908 - and (qtype == torch.half or qtype == torch.bfloat16)

909 - and (head_size == 64 or head_size == 128)

910 - and (block_size == 16 or block_size == 32)

911 - and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768

912 - and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)

vllm/attention/ops/chunked_prefill_paged_decode.py

+103 -54

@@ -10,6 +10,9 @@ import torch

10 import triton

11 import triton.language as tl

13 +from vllm import _custom_ops as ops

14 +from vllm.platforms.rocm import use_rocm_custom_paged_attention

15 +

16 from .prefix_prefill import context_attention_fwd

@@ -33,26 +36,26 @@ def kernel_paged_attention_2d(

36 num_query_heads: tl.constexpr, # int

37 num_queries_per_kv: tl.constexpr, # int

38 num_queries_per_kv_padded: tl.constexpr, # int

36 - block_table_stride: tl.constexpr, # int

37 - query_stride_0: tl.constexpr, # int

38 - query_stride_1: tl.constexpr, # int, should be equal to head_size

39 - output_stride_0: tl.constexpr, # int

40 - output_stride_1: tl.constexpr, # int, should be equal to head_size

39 + block_table_stride: tl.int64, # int

40 + query_stride_0: tl.int64, # int

41 + query_stride_1: tl.int64, # int, should be equal to head_size

42 + output_stride_0: tl.int64, # int

43 + output_stride_1: tl.int64, # int, should be equal to head_size

44 BLOCK_SIZE: tl.constexpr, # int

45 HEAD_SIZE: tl.constexpr, # int

46 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2

47 USE_ALIBI_SLOPES: tl.constexpr, # bool

48 SLIDING_WINDOW: tl.constexpr, # int

49 x: tl.constexpr, # int

47 - stride_k_cache_0: tl.constexpr, # int

48 - stride_k_cache_1: tl.constexpr, # int

49 - stride_k_cache_2: tl.constexpr, # int

50 - stride_k_cache_3: tl.constexpr, # int

51 - stride_k_cache_4: tl.constexpr, # int

52 - stride_v_cache_0: tl.constexpr, # int

53 - stride_v_cache_1: tl.constexpr, # int

54 - stride_v_cache_2: tl.constexpr, # int

55 - stride_v_cache_3: tl.constexpr, # int

50 + stride_k_cache_0: tl.int64, # int

51 + stride_k_cache_1: tl.int64, # int

52 + stride_k_cache_2: tl.int64, # int

53 + stride_k_cache_3: tl.int64, # int

54 + stride_k_cache_4: tl.int64, # int

55 + stride_v_cache_0: tl.int64, # int

56 + stride_v_cache_1: tl.int64, # int

57 + stride_v_cache_2: tl.int64, # int

58 + stride_v_cache_3: tl.int64, # int

59 filter_by_query_len: tl.constexpr, # bool

60 query_start_len_ptr, # [num_seqs+1]

61 ):

@@ -212,6 +215,7 @@ def chunked_prefill_paged_decode(

215 block_table,

216 query_start_loc,

217 seq_lens,

218 + max_seq_len,

219 max_query_len,

220 k_scale,

221 v_scale,

@@ -240,6 +244,7 @@ def chunked_prefill_paged_decode(

244 b_loc=block_table,

245 b_start_loc=query_start_loc,

246 b_seq_len=seq_lens,

247 + max_seq_len=max_seq_len,

248 max_input_len=max_query_len,

249 k_scale=k_scale,

250 v_scale=v_scale,

@@ -275,43 +280,87 @@ def chunked_prefill_paged_decode(

280 num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv),

281 16)

282

278 - kernel_paged_attention_2d[(

279 - num_seqs,

280 - num_kv_heads,

281 - )](

282 - output_ptr=output,

283 - query_ptr=query,

284 - key_cache_ptr=key_cache,

285 - value_cache_ptr=value_cache,

286 - block_tables_ptr=block_table,

287 - seq_lens_ptr=seq_lens,

288 - alibi_slopes_ptr=alibi_slopes,

289 - scale=sm_scale,

290 - k_scale=k_scale,

291 - v_scale=v_scale,

292 - num_query_heads=num_query_heads,

293 - num_queries_per_kv=num_queries_per_kv,

294 - num_queries_per_kv_padded=num_queries_per_kv_padded,

295 - block_table_stride=block_table.stride(0),

296 - query_stride_0=query.stride(0),

297 - query_stride_1=query.stride(1),

298 - output_stride_0=output.stride(0),

299 - output_stride_1=output.stride(1),

300 - BLOCK_SIZE=block_size,

301 - HEAD_SIZE=head_size,

302 - HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

303 - USE_ALIBI_SLOPES=use_alibi_slopes,

304 - SLIDING_WINDOW=sliding_window,

305 - x=key_cache.shape[4],

306 - stride_k_cache_0=key_cache.stride(0),

307 - stride_k_cache_1=key_cache.stride(1),

308 - stride_k_cache_2=key_cache.stride(2),

309 - stride_k_cache_3=key_cache.stride(3),

310 - stride_k_cache_4=key_cache.stride(4),

311 - stride_v_cache_0=value_cache.stride(0),

312 - stride_v_cache_1=value_cache.stride(1),

313 - stride_v_cache_2=value_cache.stride(2),

314 - stride_v_cache_3=value_cache.stride(3),

315 - filter_by_query_len=True,

316 - query_start_len_ptr=query_start_loc,

317 - )

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)

287 + if use_custom:

288 + _PARTITION_SIZE_ROCM = 256

289 + max_num_partitions = ((max_seq_len + _PARTITION_SIZE_ROCM - 1) //

290 + _PARTITION_SIZE_ROCM)

291 + assert _PARTITION_SIZE_ROCM % block_size == 0

292 + total_num_seq = query.shape[0]

293 + tmp_output = torch.empty(

294 + size=(total_num_seq, num_query_heads, max_num_partitions,

295 + head_size),

296 + dtype=output.dtype,

297 + device=output.device,

298 + )

299 + exp_sums = torch.empty(

300 + size=(total_num_seq, num_query_heads, max_num_partitions),

301 + dtype=torch.float32,

302 + device=output.device,

303 + )

304 + max_logits = torch.empty_like(exp_sums)

305 +

306 + ops.paged_attention_rocm(

307 + output,

308 + exp_sums,

309 + max_logits,

310 + tmp_output,

311 + query,

312 + key_cache,

313 + value_cache,

314 + num_kv_heads,

315 + scale=sm_scale,

316 + block_tables=block_table,

317 + seq_lens=seq_lens,

318 + query_start_loc=query_start_loc,

319 + block_size=block_size,

320 + max_seq_len=max_seq_len,

321 + alibi_slopes=alibi_slopes,

322 + kv_cache_dtype=kv_cache_dtype,

323 + k_scale=k_scale,

324 + v_scale=v_scale,

325 + )

326 + else:

327 + kernel_paged_attention_2d[(

328 + num_seqs,

329 + num_kv_heads,

330 + )](

331 + output_ptr=output,

332 + query_ptr=query,

333 + key_cache_ptr=key_cache,

334 + value_cache_ptr=value_cache,

335 + block_tables_ptr=block_table,

336 + seq_lens_ptr=seq_lens,

337 + alibi_slopes_ptr=alibi_slopes,

338 + scale=sm_scale,

339 + k_scale=k_scale,

340 + v_scale=v_scale,

341 + num_query_heads=num_query_heads,

342 + num_queries_per_kv=num_queries_per_kv,

343 + num_queries_per_kv_padded=num_queries_per_kv_padded,

344 + block_table_stride=block_table.stride(0),

345 + query_stride_0=query.stride(0),

346 + query_stride_1=query.stride(1),

347 + output_stride_0=output.stride(0),

348 + output_stride_1=output.stride(1),

349 + BLOCK_SIZE=block_size,

350 + HEAD_SIZE=head_size,

351 + HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

352 + USE_ALIBI_SLOPES=use_alibi_slopes,

353 + SLIDING_WINDOW=sliding_window,

354 + x=key_cache.shape[4],

355 + stride_k_cache_0=key_cache.stride(0),

356 + stride_k_cache_1=key_cache.stride(1),

357 + stride_k_cache_2=key_cache.stride(2),

358 + stride_k_cache_3=key_cache.stride(3),

359 + stride_k_cache_4=key_cache.stride(4),

360 + stride_v_cache_0=value_cache.stride(0),

361 + stride_v_cache_1=value_cache.stride(1),

362 + stride_v_cache_2=value_cache.stride(2),

363 + stride_v_cache_3=value_cache.stride(3),

364 + filter_by_query_len=True,

365 + query_start_len_ptr=query_start_loc,

366 + )

vllm/attention/ops/paged_attn.py

@@ -209,6 +209,7 @@ class PagedAttention:

209 v_scale: torch.Tensor,

210 ) -> torch.Tensor:

211 output = torch.empty_like(query)

212 + max_seq_len = None

213 context_attention_fwd(

214 query,

215 key,

@@ -221,6 +222,7 @@ class PagedAttention:

222 # query_start_loc is (batch_size + 1,)

223 query_start_loc,

224 seq_lens_tensor,

225 + max_seq_len,

226 max_query_len,

227 k_scale,

228 v_scale,

vllm/attention/ops/prefix_prefill.py

@@ -725,6 +725,7 @@ if triton.__version__ >= "2.1.0":

725 b_loc,

726 b_start_loc,

727 b_seq_len,

728 + max_seq_len,

729 max_input_len,

730 k_scale: torch.Tensor,

731 v_scale: torch.Tensor,

vllm/platforms/rocm.py

+20 -1

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import os

4 -from functools import lru_cache, wraps

4 +from functools import cache, lru_cache, wraps

5 from typing import TYPE_CHECKING, Dict, List, Optional

7 import torch

@@ -98,6 +98,25 @@ def device_id_to_physical_device_id(device_id: int) -> int:

98 return device_id

100

101 +@cache

102 +def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,

103 + block_size: int, gqa_ratio: int,

104 + max_seq_len: int,

105 + sliding_window: int) -> bool:

106 +

107 + GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName

108 + ON_NAVI = "gfx1" in GPU_ARCH

109 + ON_MI250_MI300 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])

110 +

111 + # rocm custom page attention not support on navi (gfx1*)

112 + return (ON_MI250_MI300 and not ON_NAVI and (sliding_window == 0)

113 + and (qtype == torch.half or qtype == torch.bfloat16)

114 + and (head_size == 64 or head_size == 128)

115 + and (block_size == 16 or block_size == 32)

116 + and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768

117 + and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)

118 +

119 +

120 class RocmPlatform(Platform):

121 _enum = PlatformEnum.ROCM

122 device_name: str = "rocm"

vllm/v1/attention/backends/triton_attn.py

@@ -168,6 +168,7 @@ class TritonAttentionImpl(AttentionImpl):

168 block_table=attn_metadata.block_table,

169 query_start_loc=attn_metadata.query_start_loc,

170 seq_lens=attn_metadata.seq_lens,

171 + max_seq_len=attn_metadata.max_seq_len,

172 max_query_len=attn_metadata.max_query_len,

173 k_scale=layer._k_scale,

174 v_scale=layer._v_scale,