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 14, 05:30 UTC
[ROCm][Misc] Rename the context_len to seq_len in ROCm custom paged attention kernel (#22097)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
Aug 9, 2025 at 01:15 UTC
b7c0942b65380ab8c53ecf2657121e1c21150672
3 files changed
+91
-96
csrc/rocm/attention.cu
+87
-92
@@ -270,7 +270,7 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
270
const int num_kv_heads,
271
const float scale,
272
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
273
- const int* __restrict__ context_lens, // [num_seqs]
273
+ const int* __restrict__ seq_lens, // [num_seqs]
274
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
275
const int max_num_blocks_per_seq,
276
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -304,12 +304,12 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
304
305
const auto max_num_partitions = gridDim.y;
306
307
- const int context_len = context_lens[seq_idx];
307
+ const int seq_len = seq_lens[seq_idx];
308
309
const int partition_start_token_idx =
310
partition_idx * T_PAR_SIZE; // partition_size;
311
// exit if partition is out of context for seq
312
- if (partition_start_token_idx >= context_len) {
312
+ if (partition_start_token_idx >= seq_len) {
313
return;
314
}
315
@@ -361,8 +361,8 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
361
// output layout from QKmfma : QH16xT4x4 16 qheads across 16 lanes, 16 tokens
362
// across 4 rows x 4 tokens per lane
363
364
- const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
365
- const int last_ctx_block = num_context_blocks - 1;
364
+ const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
365
+ const int last_seq_block = num_seq_blocks - 1;
366
367
const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
368
@@ -373,9 +373,9 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
373
const int klocal_token_idx =
374
TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
375
const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
376
- const int kblock_idx = (kglobal_token_idx < context_len)
376
+ const int kblock_idx = (kglobal_token_idx < seq_len)
377
? kglobal_token_idx / BLOCK_SIZE
378
- : last_ctx_block;
378
+ : last_seq_block;
379
kphysical_block_number[token_depth] = block_table_seq[kblock_idx];
380
}
381
@@ -476,9 +476,9 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
476
// tokens
477
const int vglobal_token_idx =
478
partition_start_token_idx + vlocal_token_idx;
479
- const int vblock_idx = (vglobal_token_idx < context_len)
479
+ const int vblock_idx = (vglobal_token_idx < seq_len)
480
? vglobal_token_idx / BLOCK_SIZE
481
- : last_ctx_block;
481
+ : last_seq_block;
482
vphysical_block_number[vtoken_depth][vblock_depth] =
483
block_table_seq[vblock_idx];
484
}
@@ -554,7 +554,7 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
554
if constexpr (ALIBI_ENABLED) {
555
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
556
const int local_token_idx = qkout_token_idx + token_depth * 16;
557
- const int alibi_offset = local_token_idx - context_len + 1;
557
+ const int alibi_offset = local_token_idx - seq_len + 1;
558
for (int i = 0; i < 4; i++) {
559
d_out[token_depth][i] += alibi_slope * (alibi_offset + i);
560
}
@@ -568,9 +568,8 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
568
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
569
const int local_token_idx = qkout_token_idx + token_depth * 16;
570
for (int i = 0; i < 4; i++) {
571
- const float tmp = (local_token_idx + i < context_len)
572
- ? d_out[token_depth][i]
573
- : -FLT_MAX;
571
+ const float tmp =
572
+ (local_token_idx + i < seq_len) ? d_out[token_depth][i] : -FLT_MAX;
573
qk_max = fmaxf(qk_max, tmp);
574
}
575
}
@@ -582,7 +581,7 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
581
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
582
const int local_token_idx = qkout_token_idx + token_depth * 16;
583
for (int i = 0; i < 4; i++) {
585
- const float tmp = (local_token_idx + i < context_len)
584
+ const float tmp = (local_token_idx + i < seq_len)
585
? __expf(d_out[token_depth][i] - qk_max)
586
: 0.0f;
587
d_out[token_depth][i] = tmp;
@@ -780,7 +779,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
779
const int num_kv_heads,
780
const float scale,
781
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
783
- const int* __restrict__ context_lens, // [num_seqs]
782
+ const int* __restrict__ seq_lens, // [num_seqs]
783
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
784
const int max_num_blocks_per_seq,
785
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -809,10 +808,10 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
808
const auto partition_size = blockDim.x;
809
const auto max_num_partitions = gridDim.y;
810
812
- const int context_len = context_lens[seq_idx];
811
+ const int seq_len = seq_lens[seq_idx];
812
const int partition_start_token_idx = partition_idx * partition_size;
813
// exit if partition is out of context for seq
815
- if (partition_start_token_idx >= context_len) {
814
+ if (partition_start_token_idx >= seq_len) {
815
return;
816
}
817
// every 4 lanes fetch 4 different qheads
@@ -855,7 +854,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
854
const int warp_start_token_idx =
855
partition_start_token_idx + warpid * WARP_SIZE;
856
858
- if (warp_start_token_idx >= context_len) { // warp out of context
857
+ if (warp_start_token_idx >= seq_len) { // warp out of context
858
#pragma unroll
859
for (int h = 0; h < GQA_RATIO4; h++) {
860
shared_qk_max[warpid][h] = -FLT_MAX;
@@ -863,8 +862,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
862
}
863
} else { // warp within context
864
866
- const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
867
- const int last_ctx_block = num_context_blocks - 1;
865
+ const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
866
+ const int last_seq_block = num_seq_blocks - 1;
867
868
const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
869
// token id within partition
@@ -873,9 +872,9 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
872
const int global_token_idx = partition_start_token_idx + local_token_idx;
873
874
// fetch block number for k
876
- const int block_idx = (global_token_idx < context_len)
875
+ const int block_idx = (global_token_idx < seq_len)
876
? global_token_idx / BLOCK_SIZE
878
- : last_ctx_block;
877
+ : last_seq_block;
878
879
// fetch k physical block number
880
// int32 physical_block_number leads to overflow when multiplied with
@@ -888,7 +887,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
887
for (int b = 0; b < VBLOCKS; b++) {
888
const int vblock_idx = warp_start_block_idx + b;
889
const int vblock_idx_ctx =
891
- (vblock_idx <= last_ctx_block) ? vblock_idx : last_ctx_block;
890
+ (vblock_idx <= last_seq_block) ? vblock_idx : last_seq_block;
891
vphysical_blocks[b] = block_table[vblock_idx_ctx];
892
}
893
@@ -1057,7 +1056,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1056
const int lane4_token_idx = 4 * (global_token_idx >> 2);
1057
1058
if constexpr (ALIBI_ENABLED) {
1060
- const int alibi_offset = lane4_token_idx - context_len + 1;
1059
+ const int alibi_offset = lane4_token_idx - seq_len + 1;
1060
for (int h = 0; h < QHLOOP; h++) {
1061
for (int i = 0; i < 4; i++) {
1062
d_out[h][i] += alibi_slope[h] * (alibi_offset + i);
@@ -1070,7 +1069,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1069
for (int h = 0; h < QHLOOP; h++) {
1070
qk_max[h] = -FLT_MAX;
1071
for (int i = 0; i < 4; i++) {
1073
- qk_max[h] = (lane4_token_idx + i < context_len)
1072
+ qk_max[h] = (lane4_token_idx + i < seq_len)
1073
? fmaxf(qk_max[h], d_out[h][i])
1074
: qk_max[h];
1075
}
@@ -1101,7 +1100,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1100
for (int h = 0; h < QHLOOP; h++) {
1101
exp_sum[h] = 0.0f;
1102
for (int i = 0; i < 4; i++) {
1104
- d_out[h][i] = (lane4_token_idx + i < context_len)
1103
+ d_out[h][i] = (lane4_token_idx + i < seq_len)
1104
? __expf(d_out[h][i] - qk_max[h])
1105
: 0.0f;
1106
exp_sum[h] += d_out[h][i];
@@ -1181,7 +1180,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1180
}
1181
}
1182
1184
- if (warp_start_token_idx >= context_len) { // warp out of context
1183
+ if (warp_start_token_idx >= seq_len) { // warp out of context
1184
for (int qh = 0; qh < QHLOOP; qh++) {
1185
for (int vh = 0; vh < VHELOOP; vh++) {
1186
vout_shared[qh][vh][laneid][warpid] = {0};
@@ -1279,7 +1278,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1278
// max_num_partitions]
1279
const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
1280
// max_num_partitions, head_size]
1282
- const int* __restrict__ context_lens, // [num_seqs]
1281
+ const int* __restrict__ seq_lens, // [num_seqs]
1282
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
1283
const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
1284
const auto num_heads = gridDim.x;
@@ -1293,8 +1292,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1292
return;
1293
}
1294
1296
- const int context_len = context_lens[seq_idx];
1297
- const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
1295
+ const int seq_len = seq_lens[seq_idx];
1296
+ const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
1297
const auto warpid = threadIdx.x / WARP_SIZE;
1298
1299
__shared__ float shared_global_exp_sum;
@@ -1581,7 +1580,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1580
// head_size, block_size]
1581
const int num_kv_heads, const float scale,
1582
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
1584
- const int* __restrict__ context_lens, // [num_seqs]
1583
+ const int* __restrict__ seq_lens, // [num_seqs]
1584
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
1585
const int max_num_blocks_per_seq,
1586
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -1615,11 +1614,11 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1614
1615
const int max_num_partitions = gridDim.y;
1616
1618
- const int context_len = context_lens[seq_idx]; // length of a seq
1617
+ const int seq_len = seq_lens[seq_idx]; // length of a seq
1618
1619
const int partition_start_token_idx = partition_idx * T_PAR_SIZE;
1620
// exit if partition is out of context for seq
1622
- if (partition_start_token_idx >= context_len) {
1621
+ if (partition_start_token_idx >= seq_len) {
1622
return;
1623
}
1624
@@ -1715,8 +1714,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1714
}
1715
}
1716
1718
- const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
1719
- const int last_ctx_block = num_context_blocks - 1;
1717
+ const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
1718
+ const int last_seq_block = num_seq_blocks - 1;
1719
1720
const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
1721
@@ -1727,9 +1726,9 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1726
const int klocal_token_idx =
1727
TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
1728
const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
1730
- const int kblock_idx = (kglobal_token_idx < context_len)
1729
+ const int kblock_idx = (kglobal_token_idx < seq_len)
1730
? kglobal_token_idx / BLOCK_SIZE
1732
- : last_ctx_block;
1731
+ : last_seq_block;
1732
kphysical_block_number[token_depth] = block_table_seq[kblock_idx];
1733
}
1734
@@ -1781,9 +1780,9 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1780
vblock_depth * BLOCK_SIZE;
1781
const int vglobal_token_idx =
1782
partition_start_token_idx + vlocal_token_idx;
1784
- const int vblock_idx = (vglobal_token_idx < context_len)
1783
+ const int vblock_idx = (vglobal_token_idx < seq_len)
1784
? vglobal_token_idx / BLOCK_SIZE
1786
- : last_ctx_block;
1785
+ : last_seq_block;
1786
vphysical_block_number[vtoken_depth][vblock_depth] =
1787
block_table_seq[vblock_idx];
1788
}
@@ -1836,9 +1835,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1835
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1836
const int local_token_idx = qkout_token_idx + token_depth * 16;
1837
for (int i = 0; i < 8; i++) {
1839
- const float tmp = (local_token_idx + 2 * i < context_len)
1840
- ? dout[token_depth][i]
1841
- : -FLT_MAX;
1838
+ const float tmp =
1839
+ (local_token_idx + 2 * i < seq_len) ? dout[token_depth][i] : -FLT_MAX;
1840
qk_max = fmaxf(qk_max, tmp);
1841
}
1842
}
@@ -1848,7 +1846,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1846
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
1847
const int local_token_idx = qkout_token_idx + token_depth * 16;
1848
for (int i = 0; i < 8; i++) {
1851
- const float tmp = (local_token_idx + 2 * i < context_len)
1849
+ const float tmp = (local_token_idx + 2 * i < seq_len)
1850
? __expf(dout[token_depth][i] - qk_max)
1851
: 0.0f;
1852
dout[token_depth][i] = tmp;
@@ -2019,7 +2017,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
2017
// head_size, block_size]
2018
const int num_kv_heads, const float scale,
2019
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2022
- const int* __restrict__ context_lens, // [num_seqs]
2020
+ const int* __restrict__ seq_lens, // [num_seqs]
2021
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2022
const int max_num_blocks_per_seq,
2023
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -2046,7 +2044,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
2044
// max_num_partitions]
2045
const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
2046
// max_num_partitions, head_size]
2049
- const int* __restrict__ context_lens, // [num_seqs]
2047
+ const int* __restrict__ seq_lens, // [num_seqs]
2048
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2049
const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
2050
const auto num_heads = gridDim.x;
@@ -2060,8 +2058,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
2058
return;
2059
}
2060
2063
- const int context_len = context_lens[seq_idx];
2064
- const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
2061
+ const int seq_len = seq_lens[seq_idx];
2062
+ const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
2063
const int warpid = threadIdx.x / WARP_SIZE;
2064
2065
__shared__ float shared_global_exp_sum;
@@ -2349,7 +2347,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2347
// head_size, block_size]
2348
const int num_kv_heads, const float scale,
2349
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2352
- const int* __restrict__ context_lens, // [num_seqs]
2350
+ const int* __restrict__ seq_lens, // [num_seqs]
2351
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2352
const int max_num_blocks_per_seq,
2353
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -2382,11 +2380,11 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2380
2381
const int max_num_partitions = gridDim.y;
2382
2385
- const int context_len = context_lens[seq_idx]; // length of a seq
2383
+ const int seq_len = seq_lens[seq_idx]; // length of a seq
2384
2385
const int partition_start_token_idx = partition_idx * T_PAR_SIZE;
2386
// exit if partition is out of context for seq
2389
- if (partition_start_token_idx >= context_len) {
2387
+ if (partition_start_token_idx >= seq_len) {
2388
return;
2389
}
2390
@@ -2482,8 +2480,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2480
}
2481
}
2482
2485
- const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
2486
- const int last_ctx_block = num_context_blocks - 1;
2483
+ const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
2484
+ const int last_seq_block = num_seq_blocks - 1;
2485
2486
const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
2487
@@ -2494,9 +2492,9 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2492
const int klocal_token_idx =
2493
TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
2494
const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
2497
- const int kblock_idx = (kglobal_token_idx < context_len)
2495
+ const int kblock_idx = (kglobal_token_idx < seq_len)
2496
? kglobal_token_idx / BLOCK_SIZE
2499
- : last_ctx_block;
2497
+ : last_seq_block;
2498
kphysical_block_number[token_depth] = block_table_seq[kblock_idx];
2499
}
2500
@@ -2548,9 +2546,9 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2546
rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE;
2547
const int vglobal_token_idx =
2548
partition_start_token_idx + vlocal_token_idx;
2551
- const int vblock_idx = (vglobal_token_idx < context_len)
2549
+ const int vblock_idx = (vglobal_token_idx < seq_len)
2550
? vglobal_token_idx / BLOCK_SIZE
2553
- : last_ctx_block;
2551
+ : last_seq_block;
2552
vphysical_block_number[vtoken_depth][vblock_depth] =
2553
block_table_seq[vblock_idx];
2554
}
@@ -2604,7 +2602,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2602
const int local_token_idx = qkout_token_idx + token_depth * 16;
2603
for (int i = 0; i < 8; i++) {
2604
const float tmp =
2607
- (local_token_idx + i < context_len) ? dout[token_depth][i] : -FLT_MAX;
2605
+ (local_token_idx + i < seq_len) ? dout[token_depth][i] : -FLT_MAX;
2606
qk_max = fmaxf(qk_max, tmp);
2607
}
2608
}
@@ -2614,7 +2612,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2612
for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
2613
const int local_token_idx = qkout_token_idx + token_depth * 16;
2614
for (int i = 0; i < 8; i++) {
2617
- const float tmp = (local_token_idx + i < context_len)
2615
+ const float tmp = (local_token_idx + i < seq_len)
2616
? __expf(dout[token_depth][i] - qk_max)
2617
: 0.0f;
2618
dout[token_depth][i] = tmp;
@@ -2751,7 +2749,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
2749
// head_size, block_size]
2750
const int num_kv_heads, const float scale,
2751
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2754
- const int* __restrict__ context_lens, // [num_seqs]
2752
+ const int* __restrict__ seq_lens, // [num_seqs]
2753
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2754
const int max_num_blocks_per_seq,
2755
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -2778,7 +2776,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
2776
// max_num_partitions]
2777
const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
2778
// max_num_partitions, head_size]
2781
- const int* __restrict__ context_lens, // [num_seqs]
2779
+ const int* __restrict__ seq_lens, // [num_seqs]
2780
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2781
const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
2782
const auto num_heads = gridDim.x;
@@ -2792,8 +2790,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
2790
return;
2791
}
2792
2795
- const int context_len = context_lens[seq_idx];
2796
- const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
2793
+ const int seq_len = seq_lens[seq_idx];
2794
+ const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
2795
const int warpid = threadIdx.x / WARP_SIZE;
2796
2797
__shared__ float shared_global_exp_sum;
@@ -2980,7 +2978,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
2978
const int num_kv_heads,
2979
const float scale,
2980
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
2983
- const int* __restrict__ context_lens, // [num_seqs]
2981
+ const int* __restrict__ seq_lens, // [num_seqs]
2982
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
2983
const int max_num_blocks_per_seq,
2984
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -3007,7 +3005,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
3005
const int num_kv_heads,
3006
const float scale,
3007
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
3010
- const int* __restrict__ context_lens, // [num_seqs]
3008
+ const int* __restrict__ seq_lens, // [num_seqs]
3009
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
3010
const int max_num_blocks_per_seq,
3011
const float* __restrict__ alibi_slopes, // [num_heads]
@@ -3031,7 +3029,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3029
const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
3030
const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
3031
const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
3034
- const int* __restrict__ context_lens, // [num_seqs]
3032
+ const int* __restrict__ seq_lens, // [num_seqs]
3033
const int* __restrict__ query_start_loc_ptr, // [num_seqs]
3034
const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) {
3035
UNREACHABLE_CODE
@@ -3046,7 +3044,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3044
GQA_RATIO> \
3045
<<<grid, block, 0, stream>>>( \
3046
query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
3049
- block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \
3047
+ block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \
3048
max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
3049
kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
3050
max_ctx_blocks, k_scale_ptr, v_scale_ptr);
@@ -3057,18 +3055,17 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3055
GQA_RATIO> \
3056
<<<grid, block, 0, stream>>>( \
3057
query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
3060
- block_tables_ptr, context_lens_ptr, query_start_loc_ptr, \
3058
+ block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \
3059
max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \
3060
kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \
3061
max_ctx_blocks, k_scale_ptr, v_scale_ptr);
3062
3065
-#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \
3066
- paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \
3067
- PARTITION_SIZE, NPAR_LOOPS> \
3068
- <<<reduce_grid, reduce_block, 0, stream>>>( \
3069
- out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, \
3070
- context_lens_ptr, query_start_loc_ptr, max_num_partitions, \
3071
- fp8_out_scale_ptr);
3063
+#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \
3064
+ paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \
3065
+ PARTITION_SIZE, NPAR_LOOPS> \
3066
+ <<<reduce_grid, reduce_block, 0, stream>>>( \
3067
+ out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \
3068
+ query_start_loc_ptr, max_num_partitions, fp8_out_scale_ptr);
3069
3070
template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
3071
int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
@@ -3077,8 +3074,8 @@ void paged_attention_custom_launcher(
3074
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
3075
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
3076
torch::Tensor& value_cache, const int num_kv_heads, float scale,
3080
- torch::Tensor& block_tables, torch::Tensor& context_lens,
3081
- const std::optional<torch::Tensor>& query_start_loc, int max_context_len,
3077
+ torch::Tensor& block_tables, torch::Tensor& seq_lens,
3078
+ const std::optional<torch::Tensor>& query_start_loc, int max_seq_len,
3079
const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
3080
torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale) {
3081
int num_seqs = block_tables.size(0);
@@ -3109,7 +3106,7 @@ void paged_attention_custom_launcher(
3106
KVT* key_cache_ptr = reinterpret_cast<KVT*>(key_cache.data_ptr());
3107
KVT* value_cache_ptr = reinterpret_cast<KVT*>(value_cache.data_ptr());
3108
int* block_tables_ptr = block_tables.data_ptr<int>();
3112
- int* context_lens_ptr = context_lens.data_ptr<int>();
3109
+ int* seq_lens_ptr = seq_lens.data_ptr<int>();
3110
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
3111
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
3112
// NOTE: fp8_out_scale is optional.
@@ -3119,13 +3116,12 @@ void paged_attention_custom_launcher(
3116
: nullptr;
3117
OUTT* out_ptr = reinterpret_cast<OUTT*>(out.data_ptr());
3118
3122
- const int max_ctx_blocks = DIVIDE_ROUND_UP(max_context_len, BLOCK_SIZE);
3119
+ const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE);
3120
3121
// partition size is fixed at 256 since both mfma4 and mfma16 kernels support
3122
// it mfma4 kernel also supports partition size 512
3123
constexpr int PARTITION_SIZE = 256;
3127
- const int max_num_partitions =
3128
- DIVIDE_ROUND_UP(max_context_len, PARTITION_SIZE);
3124
+ const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
3125
const int gqa_ratio = num_heads / num_kv_heads;
3126
assert(num_heads % num_kv_heads == 0);
3127
assert(head_size == HEAD_SIZE);
@@ -3234,8 +3230,8 @@ void paged_attention_custom_launcher_navi(
3230
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
3231
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
3232
torch::Tensor& value_cache, const int num_kv_heads, float scale,
3237
- torch::Tensor& block_tables, torch::Tensor& context_lens,
3238
- const std::optional<torch::Tensor>& query_start_loc, int max_context_len,
3233
+ torch::Tensor& block_tables, torch::Tensor& seq_lens,
3234
+ const std::optional<torch::Tensor>& query_start_loc, int max_seq_len,
3235
const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
3236
torch::Tensor& v_scale) {
3237
int num_seqs = block_tables.size(0);
@@ -3263,7 +3259,7 @@ void paged_attention_custom_launcher_navi(
3259
KVT* key_cache_ptr = reinterpret_cast<KVT*>(key_cache.data_ptr());
3260
KVT* value_cache_ptr = reinterpret_cast<KVT*>(value_cache.data_ptr());
3261
int* block_tables_ptr = block_tables.data_ptr<int>();
3266
- int* context_lens_ptr = context_lens.data_ptr<int>();
3262
+ int* seq_lens_ptr = seq_lens.data_ptr<int>();
3263
3264
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
3265
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
@@ -3271,11 +3267,10 @@ void paged_attention_custom_launcher_navi(
3267
const auto fp8_out_scale_ptr = nullptr;
3268
OUTT* out_ptr = reinterpret_cast<OUTT*>(out.data_ptr());
3269
3274
- const int max_ctx_blocks = DIVIDE_ROUND_UP(max_context_len, BLOCK_SIZE);
3270
+ const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE);
3271
3272
constexpr int PARTITION_SIZE = 256;
3277
- const int max_num_partitions =
3278
- DIVIDE_ROUND_UP(max_context_len, PARTITION_SIZE);
3273
+ const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
3274
const int gqa_ratio = num_heads / num_kv_heads;
3275
assert(num_heads % num_kv_heads == 0);
3276
assert(head_size == HEAD_SIZE);
@@ -3407,14 +3402,14 @@ void paged_attention_custom_launcher_navi(
3402
paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3403
OUTT, PSIZE, ALIBI_ENABLED>( \
3404
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
3410
- num_kv_heads, scale, block_tables, context_lens, query_start_loc, \
3411
- max_context_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
3405
+ num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
3406
+ max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
3407
} else { \
3408
paged_attention_custom_launcher_navi< \
3409
T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, ALIBI_ENABLED>( \
3410
out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
3416
- num_kv_heads, scale, block_tables, context_lens, query_start_loc, \
3417
- max_context_len, alibi_slopes, k_scale, v_scale); \
3411
+ num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
3412
+ max_seq_len, alibi_slopes, k_scale, v_scale); \
3413
}
3414
3415
#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
@@ -3502,9 +3497,9 @@ void paged_attention(
3497
int64_t num_kv_heads,
3498
double scale,
3499
torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
3505
- torch::Tensor& context_lens, // [num_seqs]
3500
+ torch::Tensor& seq_lens, // [num_seqs]
3501
const std::optional<torch::Tensor>& query_start_loc, // [num_seqs]
3507
- int64_t block_size, int64_t max_context_len,
3502
+ int64_t block_size, int64_t max_seq_len,
3503
const std::optional<torch::Tensor>& alibi_slopes,
3504
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
3505
torch::Tensor& v_scale,
csrc/rocm/ops.h
+2
-2
@@ -15,8 +15,8 @@ void paged_attention(
15
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
16
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
17
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
18
- torch::Tensor& block_tables, torch::Tensor& context_lens,
18
+ torch::Tensor& block_tables, torch::Tensor& seq_lens,
19
const std::optional<torch::Tensor>& query_start_loc, int64_t block_size,
20
- int64_t max_context_len, const std::optional<torch::Tensor>& alibi_slopes,
20
+ int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
21
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
22
torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale);
csrc/rocm/torch_bindings.cpp
+2
-2
@@ -41,10 +41,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
41
" Tensor query, Tensor key_cache,"
42
" Tensor value_cache, int num_kv_heads,"
43
" float scale, Tensor block_tables,"
44
- " Tensor context_lens,"
44
+ " Tensor seq_lens,"
45
" Tensor? query_start_loc,"
46
" int block_size,"
47
- " int max_context_len,"
47
+ " int max_seq_len,"
48
" Tensor? alibi_slopes,"
49
" str kv_cache_dtype,"
50
" Tensor k_scale, Tensor v_scale,"