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, 17:15 UTC
Fix CUDA kernel index data type in vllm/csrc/quantization/gptq_marlin/awq_marlin_repack.cu +10 (#15160)
Signed-off-by: Lu Fang <lufang@fb.com> Co-authored-by: Richard Barnes <rbarnes@meta.com>
Lu Fang committed
Mar 25, 2025 at 00:36 UTC
051da7efe39591f256dea30c286ed7c920c0f7d2
7 files changed
+73
-73
csrc/quantization/gptq_marlin/awq_marlin_repack.cu
+6
-6
@@ -14,7 +14,7 @@ __global__ void awq_marlin_repack_kernel(
14
int n_tiles = size_n / tile_n_size;
15
int block_k_tiles = div_ceil(k_tiles, gridDim.x);
16
17
- int start_k_tile = blockIdx.x * block_k_tiles;
17
+ auto start_k_tile = blockIdx.x * block_k_tiles;
18
if (start_k_tile >= k_tiles) {
19
return;
20
}
@@ -51,8 +51,8 @@ __global__ void awq_marlin_repack_kernel(
51
int4* sh_ptr = sh + stage_size * pipe;
52
53
if (threadIdx.x < stage_size) {
54
- int k_id = threadIdx.x / stage_n_threads;
55
- int n_id = threadIdx.x % stage_n_threads;
54
+ auto k_id = threadIdx.x / stage_n_threads;
55
+ auto n_id = threadIdx.x % stage_n_threads;
56
57
int first_k = k_tile_id * tile_k_size;
58
@@ -70,8 +70,8 @@ __global__ void awq_marlin_repack_kernel(
70
return;
71
}
72
73
- int warp_id = threadIdx.x / 32;
74
- int th_id = threadIdx.x % 32;
73
+ auto warp_id = threadIdx.x / 32;
74
+ auto th_id = threadIdx.x % 32;
75
76
if (warp_id >= 4) {
77
return;
@@ -265,4 +265,4 @@ TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
265
266
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, Meta, m) {
267
m.impl("awq_marlin_repack", &awq_marlin_repack_meta);
268
-}
\ No newline at end of file
268
+}
csrc/quantization/gptq_marlin/gptq_marlin.cu
+14
-14
@@ -460,7 +460,7 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
460
int const* __restrict__ perm_int_ptr,
461
int4* __restrict__ out_int4_ptr, int size_m,
462
int size_k, int lda, int block_rows) {
463
- int start_row = block_rows * blockIdx.x;
463
+ auto start_row = block_rows * blockIdx.x;
464
int finish_row = start_row + block_rows;
465
if (finish_row > size_m) {
466
finish_row = size_m;
@@ -484,7 +484,7 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
484
int base_k = 0;
485
486
for (int i = 0; i < iters; i++) {
487
- int cur_k = base_k + threadIdx.x;
487
+ auto cur_k = base_k + threadIdx.x;
488
int src_pos = perm_int_ptr[cur_k];
489
490
out_half[cur_k] = a_row_half[src_pos];
@@ -494,7 +494,7 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
494
495
if (rest) {
496
if (threadIdx.x < rest) {
497
- int cur_k = base_k + threadIdx.x;
497
+ auto cur_k = base_k + threadIdx.x;
498
int src_pos = perm_int_ptr[cur_k];
499
500
out_half[cur_k] = a_row_half[src_pos];
@@ -723,8 +723,8 @@ __global__ void Marlin(
723
(threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
724
b_gl_rd += b_sh_stride * slice_col;
725
b_gl_rd += b_gl_rd_delta_o * slice_row;
726
- int b_sh_wr = threadIdx.x * b_thread_vecs;
727
- int b_sh_rd = threadIdx.x * b_thread_vecs;
726
+ auto b_sh_wr = threadIdx.x * b_thread_vecs;
727
+ auto b_sh_rd = threadIdx.x * b_thread_vecs;
728
729
// For act_order
730
constexpr int k_iter_size = tb_k / b_sh_wr_iters;
@@ -743,7 +743,7 @@ __global__ void Marlin(
743
s_sh_stride * slice_col + threadIdx.x;
744
}
745
}
746
- int s_sh_wr = threadIdx.x;
746
+ auto s_sh_wr = threadIdx.x;
747
bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
748
749
// Zero-points
@@ -756,7 +756,7 @@ __global__ void Marlin(
756
zp_sh_stride * slice_col + threadIdx.x;
757
}
758
}
759
- int zp_sh_wr = threadIdx.x;
759
+ auto zp_sh_wr = threadIdx.x;
760
bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
761
762
// We use a different scale layout for grouped and column-wise quantization as
@@ -1047,7 +1047,7 @@ __global__ void Marlin(
1047
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1048
reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1049
} else {
1050
- int warp_id = threadIdx.x / 32;
1050
+ auto warp_id = threadIdx.x / 32;
1051
int n_warps = thread_n_blocks / 4;
1052
1053
int warp_row = warp_id / n_warps;
@@ -1085,7 +1085,7 @@ __global__ void Marlin(
1085
1086
// Determine "position" inside the thread-block (based on warp and
1087
// thread-id)
1088
- int warp_id = threadIdx.x / 32;
1088
+ auto warp_id = threadIdx.x / 32;
1089
int n_warps =
1090
thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
1091
@@ -1094,7 +1094,7 @@ __global__ void Marlin(
1094
1095
cur_k += warp_row * 16;
1096
1097
- int th_id = threadIdx.x % 32;
1097
+ auto th_id = threadIdx.x % 32;
1098
cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
1099
1100
int s_col_shift =
@@ -1159,7 +1159,7 @@ __global__ void Marlin(
1159
(reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1160
}
1161
} else {
1162
- int warp_id = threadIdx.x / 32;
1162
+ auto warp_id = threadIdx.x / 32;
1163
int n_warps = thread_n_blocks / 4;
1164
1165
int warp_row = warp_id / n_warps;
@@ -1197,7 +1197,7 @@ __global__ void Marlin(
1197
(pipe / (group_blocks / thread_k_blocks)));
1198
reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd];
1199
} else {
1200
- int warp_id = threadIdx.x / 32;
1200
+ auto warp_id = threadIdx.x / 32;
1201
int n_warps = thread_n_blocks / 4;
1202
1203
int warp_row = warp_id / n_warps;
@@ -1323,7 +1323,7 @@ __global__ void Marlin(
1323
auto thread_block_reduce = [&]() {
1324
constexpr int red_off = threads / b_sh_stride_threads / 2;
1325
if (red_off >= 1) {
1326
- int red_idx = threadIdx.x / b_sh_stride_threads;
1326
+ auto red_idx = threadIdx.x / b_sh_stride_threads;
1327
constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1328
constexpr int red_sh_delta = b_sh_stride_threads;
1329
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
@@ -1390,7 +1390,7 @@ __global__ void Marlin(
1390
4 * (threadIdx.x / 32) + threadIdx.x % 4;
1391
c_gl_wr += (2 * thread_n_blocks) * slice_col;
1392
constexpr int c_sh_wr_delta = active_threads;
1393
- int c_sh_wr = threadIdx.x;
1393
+ auto c_sh_wr = threadIdx.x;
1394
1395
int row = (threadIdx.x % 32) / 4;
1396
csrc/quantization/gptq_marlin/gptq_marlin_repack.cu
+8
-8
@@ -15,7 +15,7 @@ __global__ void gptq_marlin_repack_kernel(
15
int n_tiles = size_n / tile_n_size;
16
int block_k_tiles = div_ceil(k_tiles, gridDim.x);
17
18
- int start_k_tile = blockIdx.x * block_k_tiles;
18
+ auto start_k_tile = blockIdx.x * block_k_tiles;
19
if (start_k_tile >= k_tiles) {
20
return;
21
}
@@ -71,8 +71,8 @@ __global__ void gptq_marlin_repack_kernel(
71
72
if constexpr (has_perm) {
73
if (threadIdx.x < stage_size) {
74
- int k_id = threadIdx.x / stage_n_threads;
75
- int n_id = threadIdx.x % stage_n_threads;
74
+ auto k_id = threadIdx.x / stage_n_threads;
75
+ auto n_id = threadIdx.x % stage_n_threads;
76
77
uint32_t const* sh_perm_int_ptr =
78
reinterpret_cast<uint32_t const*>(sh_perm_ptr);
@@ -88,8 +88,8 @@ __global__ void gptq_marlin_repack_kernel(
88
89
} else {
90
if (threadIdx.x < stage_size) {
91
- int k_id = threadIdx.x / stage_n_threads;
92
- int n_id = threadIdx.x % stage_n_threads;
91
+ auto k_id = threadIdx.x / stage_n_threads;
92
+ auto n_id = threadIdx.x % stage_n_threads;
93
94
int first_k = k_tile_id * tile_k_size;
95
int first_k_packed = first_k / pack_factor;
@@ -109,8 +109,8 @@ __global__ void gptq_marlin_repack_kernel(
109
return;
110
}
111
112
- int warp_id = threadIdx.x / 32;
113
- int th_id = threadIdx.x % 32;
112
+ auto warp_id = threadIdx.x / 32;
113
+ auto th_id = threadIdx.x % 32;
114
115
if (warp_id >= 4) {
116
return;
@@ -339,4 +339,4 @@ TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
339
340
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, Meta, m) {
341
m.impl("gptq_marlin_repack", &gptq_marlin_repack_meta);
342
-}
\ No newline at end of file
342
+}
csrc/quantization/marlin/dense/marlin_cuda_kernel.cu
+5
-5
@@ -277,12 +277,12 @@ __global__ void Marlin(
277
b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
278
b_gl_rd += b_sh_stride * slice_col;
279
b_gl_rd += b_gl_rd_delta_o * slice_row;
280
- int b_sh_wr = threadIdx.x;
281
- int b_sh_rd = threadIdx.x;
280
+ auto b_sh_wr = threadIdx.x;
281
+ auto b_sh_rd = threadIdx.x;
282
283
int s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
284
s_sh_stride * slice_col + threadIdx.x;
285
- int s_sh_wr = threadIdx.x;
285
+ auto s_sh_wr = threadIdx.x;
286
int s_sh_rd;
287
// We use a different scale layout for grouped and column-wise quantization as
288
// we scale a `half2` tile in column-major layout in the former and in
@@ -455,7 +455,7 @@ __global__ void Marlin(
455
auto thread_block_reduce = [&]() {
456
constexpr int red_off = threads / b_sh_stride / 2;
457
if (red_off >= 1) {
458
- int red_idx = threadIdx.x / b_sh_stride;
458
+ auto red_idx = threadIdx.x / b_sh_stride;
459
constexpr int red_sh_stride = b_sh_stride * 4 * 2;
460
constexpr int red_sh_delta = b_sh_stride;
461
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride) +
@@ -522,7 +522,7 @@ __global__ void Marlin(
522
4 * (threadIdx.x / 32) + threadIdx.x % 4;
523
c_gl_wr += (2 * thread_n_blocks) * slice_col;
524
constexpr int c_sh_wr_delta = active_threads;
525
- int c_sh_wr = threadIdx.x;
525
+ auto c_sh_wr = threadIdx.x;
526
527
int row = (threadIdx.x % 32) / 4;
528
csrc/quantization/marlin/qqq/marlin_qqq_gemm_kernel.cu
+7
-7
@@ -353,10 +353,10 @@ __global__ void Marlin(
353
b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
354
b_gl_rd += b_sh_stride * slice_col;
355
b_gl_rd += b_gl_rd_delta_o * slice_row;
356
- int b_sh_wr = threadIdx.x;
357
- int b_sh_rd = threadIdx.x;
356
+ auto b_sh_wr = threadIdx.x;
357
+ auto b_sh_rd = threadIdx.x;
358
359
- int s_tok_gl_rd = threadIdx.x;
359
+ auto s_tok_gl_rd = threadIdx.x;
360
// NOTE(HandH1998): activation scale s_tok need shuffle to [0, 8, 1, 9, 2, 10,
361
// 3, 11, 4, 12, 5, 13, 6, 14, 7, 15] for example, 0, 8 row scales serve for
362
// thread 0, 1, 2, 3. For more details, refer to mma operand A layout as
@@ -368,8 +368,8 @@ __global__ void Marlin(
368
int s_tok_sh_rd = (threadIdx.x % 32) / 4;
369
bool s_tok_sh_wr_pred = threadIdx.x < prob_m;
370
371
- int s_ch_gl_rd = s_ch_sh_stride * slice_col + threadIdx.x;
372
- int s_ch_sh_wr = threadIdx.x;
371
+ auto s_ch_gl_rd = s_ch_sh_stride * slice_col + threadIdx.x;
372
+ auto s_ch_sh_wr = threadIdx.x;
373
int s_ch_sh_rd = 16 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
374
2 * ((threadIdx.x % 32) % 4);
375
bool s_ch_sh_wr_pred = threadIdx.x < s_ch_sh_stride;
@@ -558,7 +558,7 @@ __global__ void Marlin(
558
auto thread_block_reduce = [&]() {
559
constexpr int red_off = threads / b_sh_stride / 2;
560
if (red_off >= 1) {
561
- int red_idx = threadIdx.x / b_sh_stride;
561
+ auto red_idx = threadIdx.x / b_sh_stride;
562
constexpr int red_sh_stride = b_sh_stride * 4 * 2;
563
constexpr int red_sh_delta = b_sh_stride;
564
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride) +
@@ -628,7 +628,7 @@ __global__ void Marlin(
628
8 * (threadIdx.x / 32) + (threadIdx.x % 4) * 2;
629
c_gl_wr += (4 * thread_n_blocks) * slice_col;
630
constexpr int c_sh_wr_delta = active_threads * 2;
631
- int c_sh_wr = 2 * threadIdx.x;
631
+ auto c_sh_wr = 2 * threadIdx.x;
632
633
int row = (threadIdx.x % 32) / 4;
634
csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu
+7
-7
@@ -273,15 +273,15 @@ __global__ void Marlin_24(
273
(threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
274
b_gl_rd += b_sh_stride * slice_col;
275
b_gl_rd += b_gl_rd_delta_o * slice_row;
276
- int b_sh_wr = threadIdx.x * b_thread_vecs;
277
- int b_sh_rd = threadIdx.x * b_thread_vecs;
276
+ auto b_sh_wr = threadIdx.x * b_thread_vecs;
277
+ auto b_sh_rd = threadIdx.x * b_thread_vecs;
278
279
int m_gl_rd = m_gl_stride * (threadIdx.x / (m_sh_stride)) +
280
(threadIdx.x % (m_sh_stride));
281
m_gl_rd += (m_sh_stride)*slice_col;
282
m_gl_rd += m_gl_rd_delta_o * slice_row;
283
- int m_sh_wr = threadIdx.x;
284
- int m_sh_rd = threadIdx.x % 16 + (threadIdx.x / 32) * 16;
283
+ auto m_sh_wr = threadIdx.x;
284
+ auto m_sh_rd = threadIdx.x % 16 + (threadIdx.x / 32) * 16;
285
286
int s_gl_rd;
287
if constexpr (group_blocks == -1) {
@@ -291,7 +291,7 @@ __global__ void Marlin_24(
291
s_sh_stride * slice_col + threadIdx.x;
292
}
293
294
- int s_sh_wr = threadIdx.x;
294
+ auto s_sh_wr = threadIdx.x;
295
int s_sh_rd;
296
// We use a different scale layout for grouped and column-wise quantization as
297
// we scale a `half2` tile in column-major layout in the former and in
@@ -516,7 +516,7 @@ __global__ void Marlin_24(
516
auto thread_block_reduce = [&]() {
517
constexpr int red_off = threads / b_sh_stride_threads / 2;
518
if (red_off >= 1) {
519
- int red_idx = threadIdx.x / b_sh_stride_threads;
519
+ auto red_idx = threadIdx.x / b_sh_stride_threads;
520
constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
521
constexpr int red_sh_delta = b_sh_stride_threads;
522
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
@@ -583,7 +583,7 @@ __global__ void Marlin_24(
583
8 * (threadIdx.x / 32) + (threadIdx.x % 32) / 4;
584
c_gl_wr += (2 * thread_n_blocks) * slice_col;
585
constexpr int c_sh_wr_delta = active_threads;
586
- int c_sh_wr = threadIdx.x;
586
+ auto c_sh_wr = threadIdx.x;
587
588
int col = 2 * ((threadIdx.x % 32) % 4);
589
csrc/rocm/attention.cu
+26
-26
@@ -284,18 +284,18 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
284
int max_ctx_blocks, const float* k_scale, const float* v_scale) {
285
// clang-format on
286
constexpr int NWARPS = NUM_THREADS / WARP_SIZE;
287
- const int warpid = threadIdx.x / WARP_SIZE;
288
- const int laneid = threadIdx.x % WARP_SIZE;
287
+ const auto warpid = threadIdx.x / WARP_SIZE;
288
+ const auto laneid = threadIdx.x % WARP_SIZE;
289
const int lane4id = laneid % 4;
290
const int lane16id = laneid % 16;
291
const int rowid = laneid / 16;
292
293
- const int seq_idx = blockIdx.x;
294
- const int partition_idx = blockIdx.y;
293
+ const auto seq_idx = blockIdx.x;
294
+ const auto partition_idx = blockIdx.y;
295
296
constexpr int T_PAR_SIZE = 256; // token partition size set to 256
297
298
- const int max_num_partitions = gridDim.y;
298
+ const auto max_num_partitions = gridDim.y;
299
300
const int context_len = context_lens[seq_idx];
301
@@ -346,9 +346,9 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
346
// can be interpreted as B8x16 for 8 bit types
347
_B16x8 Klocal[TLOOP][QKHELOOP];
348
349
- const int wg_start_head_idx = blockIdx.z * GQA_RATIO;
350
- const int wg_start_kv_head_idx = blockIdx.z;
351
- const int total_num_heads = gridDim.z * GQA_RATIO;
349
+ const auto wg_start_head_idx = blockIdx.z * GQA_RATIO;
350
+ const auto wg_start_kv_head_idx = blockIdx.z;
351
+ const auto total_num_heads = gridDim.z * GQA_RATIO;
352
353
// for QK mfma, tokens in multiples of TOKENS_PER_WARP are spread across warps
354
// each mfma takes QH16xT16x16HE across warp
@@ -789,14 +789,14 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
789
int max_ctx_blocks, const float* k_scale, const float* v_scale) {
790
// clang-format on
791
constexpr int NWARPS = NUM_THREADS / WARP_SIZE;
792
- const int warpid = threadIdx.x / WARP_SIZE;
793
- const int laneid = threadIdx.x % WARP_SIZE;
792
+ const auto warpid = threadIdx.x / WARP_SIZE;
793
+ const auto laneid = threadIdx.x % WARP_SIZE;
794
const int lane4id = laneid % 4;
795
796
- const int seq_idx = blockIdx.x;
797
- const int partition_idx = blockIdx.y;
798
- const int partition_size = blockDim.x;
799
- const int max_num_partitions = gridDim.y;
796
+ const auto seq_idx = blockIdx.x;
797
+ const auto partition_idx = blockIdx.y;
798
+ const auto partition_size = blockDim.x;
799
+ const auto max_num_partitions = gridDim.y;
800
801
const int context_len = context_lens[seq_idx];
802
const int partition_start_token_idx = partition_idx * partition_size;
@@ -838,8 +838,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
838
qk_max[h] = -FLT_MAX;
839
}
840
841
- const int wg_start_head_idx = blockIdx.z * GQA_RATIO;
842
- const int wg_start_kv_head_idx = blockIdx.z;
841
+ const auto wg_start_head_idx = blockIdx.z * GQA_RATIO;
842
+ const auto wg_start_kv_head_idx = blockIdx.z;
843
844
const int warp_start_token_idx =
845
partition_start_token_idx + warpid * WARP_SIZE;
@@ -857,7 +857,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
857
858
const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
859
// token id within partition
860
- const int local_token_idx = threadIdx.x;
860
+ const auto local_token_idx = threadIdx.x;
861
// token id within sequence
862
const int global_token_idx = partition_start_token_idx + local_token_idx;
863
@@ -1126,7 +1126,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1126
1127
__syncthreads();
1128
1129
- const int num_heads = gridDim.z * GQA_RATIO;
1129
+ const auto num_heads = gridDim.z * GQA_RATIO;
1130
float* max_logits_ptr =
1131
max_logits + seq_idx * num_heads * max_num_partitions + partition_idx;
1132
float* exp_sums_ptr =
@@ -1268,14 +1268,14 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1268
// max_num_partitions, head_size]
1269
const int* __restrict__ context_lens, // [num_seqs]
1270
const int max_num_partitions) {
1271
- const int num_heads = gridDim.x;
1272
- const int head_idx = blockIdx.x;
1273
- const int seq_idx = blockIdx.y;
1271
+ const auto num_heads = gridDim.x;
1272
+ const auto head_idx = blockIdx.x;
1273
+ const auto seq_idx = blockIdx.y;
1274
const int context_len = context_lens[seq_idx];
1275
const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
1276
[[maybe_unused]] constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
1277
- const int warpid = threadIdx.x / WARP_SIZE;
1278
- [[maybe_unused]] const int laneid = threadIdx.x % WARP_SIZE;
1277
+ const auto warpid = threadIdx.x / WARP_SIZE;
1278
+ [[maybe_unused]] const auto laneid = threadIdx.x % WARP_SIZE;
1279
1280
__shared__ float shared_global_exp_sum;
1281
// max num partitions supported is warp_size * NPAR_LOOPS
@@ -1294,7 +1294,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1294
1295
#pragma unroll
1296
for (int i = 0; i < NPAR_LOOPS; i++) {
1297
- const int partition_no = i * WARP_SIZE + threadIdx.x;
1297
+ const auto partition_no = i * WARP_SIZE + threadIdx.x;
1298
valid_partition[i] =
1299
(partition_no < num_partitions) ? partition_no : last_valid_partition;
1300
}
@@ -1324,7 +1324,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1324
}
1325
#pragma unroll
1326
for (int i = 0; i < NPAR_LOOPS; i++) {
1327
- const int partition_no = i * WARP_SIZE + threadIdx.x;
1327
+ const auto partition_no = i * WARP_SIZE + threadIdx.x;
1328
rescaled_exp_sum[i] *= (partition_no < num_partitions)
1329
? expf(reg_max_logit[i] - max_logit)
1330
: 0.0f;
@@ -1336,7 +1336,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1336
}
1337
#pragma unroll
1338
for (int i = 0; i < NPAR_LOOPS; i++) {
1339
- const int partition_no = i * WARP_SIZE + threadIdx.x;
1339
+ const auto partition_no = i * WARP_SIZE + threadIdx.x;
1340
shared_exp_sums[partition_no] = rescaled_exp_sum[i];
1341
}
1342