@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d3ccbd6350
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, 18:30 UTC

Fix CUDA kernel index data type in vllm/csrc/quantization/fused_kernels/layernorm_utils.cuh +10 (#15159)

Signed-off-by: Lu Fang <lufang@fb.com> Co-authored-by: Richard Barnes <rbarnes@meta.com>

Lu Fang committed Mar 20, 2025 at 19:01 UTC d3ccbd6350bf8d601686e28fd059d1e6f53fcb4a
10 files changed +124 -124
csrc/quantization/fused_kernels/layernorm_utils.cuh
+6 -6
@@ -24,7 +24,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
24 // sum of squares
25 float ss = 0.0f;
26
27 - for (int32_t i = threadIdx.x; i < hidden_size; i += blockDim.x) {
27 + for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
28 float x = static_cast<float>(input[token_offset + i]);
29 if constexpr (has_residual) {
30 x += static_cast<float>(residual[token_offset + i]);
@@ -58,7 +58,7 @@ __device__ void compute_dynamic_per_token_scales(
58 constexpr scalar_out_t qmax{std::numeric_limits<scalar_out_t>::max()};
59
60 float block_absmax_val_maybe = 0.0f;
61 - for (int32_t i = threadIdx.x; i < hidden_size; i += blockDim.x) {
61 + for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
62 float x = static_cast<float>(input[token_offset + i]);
63 if constexpr (has_residual) {
64 x += static_cast<float>(residual[token_offset + i]);
@@ -103,7 +103,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
103 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
104 ;
105
106 - for (int32_t i = threadIdx.x; i < hidden_size; i += blockDim.x) {
106 + for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
107 float x = static_cast<float>(input[token_offset + i]);
108 if constexpr (has_residual) {
109 x += static_cast<float>(residual[token_offset + i]);
@@ -142,7 +142,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
142 int32_t const num_vec_elems = hidden_size >> 2;
143
144 #pragma unroll 4
145 - for (int32_t i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
145 + for (auto i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
146 vec4_t<scalar_t> in = vec_input[i];
147
148 vec4_t<float> x;
@@ -206,7 +206,7 @@ __device__ void compute_dynamic_per_token_scales(
206 float block_absmax_val_maybe = 0.0f;
207
208 #pragma unroll 4
209 - for (int32_t i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
209 + for (auto i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
210 vec4_t<scalar_t> in = vec_input[i];
211 vec4_t<scalar_t> const w = vec_weight[i];
212
@@ -286,7 +286,7 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
286 // TODO(luka/varun) extract into type-agnostic vectorized quant function to
287 // replace scaled_fp8_conversion_vec
288 #pragma unroll 4
289 - for (int32_t i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
289 + for (auto i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
290 vec4_t<scalar_t> const in = vec_input[i];
291 vec4_t<scalar_t> const w = vec_weight[i];
292
csrc/quantization/gguf/dequantize.cuh
+25 -25
@@ -101,10 +101,10 @@ static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __
101 template<typename dst_t>
102 static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
103
104 - const int i = blockIdx.x;
104 + const auto i = blockIdx.x;
105 const block_q2_K * x = (const block_q2_K *) vx;
106
107 - const int tid = threadIdx.x;
107 + const auto tid = threadIdx.x;
108 const int n = tid/32;
109 const int l = tid - 32*n;
110 const int is = 8*n + l/16;
@@ -123,10 +123,10 @@ static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t
123 template<typename dst_t>
124 static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
125
126 - const int i = blockIdx.x;
126 + const auto i = blockIdx.x;
127 const block_q3_K * x = (const block_q3_K *) vx;
128
129 - const int r = threadIdx.x/4;
129 + const auto r = threadIdx.x/4;
130 const int tid = r/2;
131 const int is0 = r%2;
132 const int l0 = 16*is0 + 4*(threadIdx.x%4);
@@ -164,10 +164,10 @@ template<typename dst_t>
164 static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
165 const block_q4_K * x = (const block_q4_K *) vx;
166
167 - const int i = blockIdx.x;
167 + const auto i = blockIdx.x;
168
169 // assume 32 threads
170 - const int tid = threadIdx.x;
170 + const auto tid = threadIdx.x;
171 const int il = tid/8;
172 const int ir = tid%8;
173 const int is = 2*il;
@@ -197,10 +197,10 @@ template<typename dst_t>
197 static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
198 const block_q5_K * x = (const block_q5_K *) vx;
199
200 - const int i = blockIdx.x;
200 + const auto i = blockIdx.x;
201
202 // assume 64 threads - this is very slightly better than the one below
203 - const int tid = threadIdx.x;
203 + const auto tid = threadIdx.x;
204 const int il = tid/16; // il is in 0...3
205 const int ir = tid%16; // ir is in 0...15
206 const int is = 2*il; // is is in 0...6
@@ -231,10 +231,10 @@ template<typename dst_t>
231 static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
232 const block_q6_K * x = (const block_q6_K *) vx;
233
234 - const int i = blockIdx.x;
234 + const auto i = blockIdx.x;
235
236 // assume 64 threads - this is very slightly better than the one below
237 - const int tid = threadIdx.x;
237 + const auto tid = threadIdx.x;
238 const int ip = tid/32; // ip is 0 or 1
239 const int il = tid - 32*ip; // 0...32
240 const int is = 8*ip + il/16;
@@ -256,10 +256,10 @@ static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t
256 template<typename dst_t>
257 static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
258
259 - const int i = blockIdx.x;
259 + const auto i = blockIdx.x;
260 const block_iq2_xxs * x = (const block_iq2_xxs *) vx;
261
262 - const int tid = threadIdx.x;
262 + const auto tid = threadIdx.x;
263 const int il = tid/8; // 0...3
264 const int ib = tid%8; // 0...7
265 dst_t * y = yy + i*QK_K + 32*ib + 8*il;
@@ -275,10 +275,10 @@ static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, ds
275 template<typename dst_t>
276 static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
277
278 - const int i = blockIdx.x;
278 + const auto i = blockIdx.x;
279 const block_iq2_xs * x = (const block_iq2_xs *) vx;
280
281 - const int tid = threadIdx.x;
281 + const auto tid = threadIdx.x;
282 const int il = tid/8; // 0...3
283 const int ib = tid%8; // 0...7
284 dst_t * y = yy + i*QK_K + 32*ib + 8*il;
@@ -293,10 +293,10 @@ static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst
293 template<typename dst_t>
294 static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
295
296 - const int i = blockIdx.x;
296 + const auto i = blockIdx.x;
297 const block_iq2_s * x = (const block_iq2_s *) vx;
298
299 - const int tid = threadIdx.x;
299 + const auto tid = threadIdx.x;
300 const int il = tid/8; // 0...3
301 const int ib = tid%8; // 0...7
302 dst_t * y = yy + i*QK_K + 32*ib + 8*il;
@@ -309,10 +309,10 @@ static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_
309 template<typename dst_t>
310 static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
311
312 - const int i = blockIdx.x;
312 + const auto i = blockIdx.x;
313 const block_iq3_xxs * x = (const block_iq3_xxs *) vx;
314
315 - const int tid = threadIdx.x;
315 + const auto tid = threadIdx.x;
316 const int il = tid/8; // 0...3
317 const int ib = tid%8; // 0...7
318 dst_t * y = yy + i*QK_K + 32*ib + 8*il;
@@ -332,10 +332,10 @@ static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, ds
332 template<typename dst_t>
333 static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
334
335 - const int i = blockIdx.x;
335 + const auto i = blockIdx.x;
336 const block_iq3_s * x = (const block_iq3_s *) vx;
337
338 - const int tid = threadIdx.x;
338 + const auto tid = threadIdx.x;
339 const int il = tid/8; // 0...3
340 const int ib = tid%8; // 0...7
341 dst_t * y = yy + i*QK_K + 32*ib + 8*il;
@@ -399,10 +399,10 @@ static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_
399 template<typename dst_t>
400 static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy) {
401
402 - const int i = blockIdx.x;
402 + const auto i = blockIdx.x;
403 const block_iq4_nl * x = (const block_iq4_nl *) vx + i*(QK_K/QK4_NL);
404
405 - const int tid = threadIdx.x;
405 + const auto tid = threadIdx.x;
406 const int il = tid/8; // 0...3
407 const int ib = tid%8; // 0...7
408 dst_t * y = yy + i*QK_K + 32*ib + 4*il;
@@ -417,10 +417,10 @@ static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst
417
418 template<typename dst_t>
419 static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
420 - const int i = blockIdx.x;
420 + const auto i = blockIdx.x;
421 const block_iq4_xs * x = (const block_iq4_xs *)vx;
422
423 - const int tid = threadIdx.x;
423 + const auto tid = threadIdx.x;
424 const int il = tid/8; // 0...3
425 const int ib = tid%8; // 0...7
426 dst_t * y = yy + i*QK_K + 32*ib + 4*il;
@@ -565,4 +565,4 @@ static to_fp16_cuda_t ggml_get_to_fp16_cuda(int64_t type) {
565 default:
566 return nullptr;
567 }
568 -}
\ No newline at end of file
568 +}
csrc/quantization/gguf/gguf_kernel.cu
+2 -2
@@ -19,11 +19,11 @@ template <typename scalar_t>
19 static __global__ void quantize_q8_1(const scalar_t* __restrict__ x,
20 void* __restrict__ vy, const int kx,
21 const int kx_padded) {
22 - const int ix = blockDim.x * blockIdx.x + threadIdx.x;
22 + const auto ix = blockDim.x * blockIdx.x + threadIdx.x;
23 if (ix >= kx_padded) {
24 return;
25 }
26 - const int iy = blockDim.y * blockIdx.y + threadIdx.y;
26 + const auto iy = blockDim.y * blockIdx.y + threadIdx.y;
27 const int i_padded = iy * kx_padded + ix;
28
29 block_q8_1* y = (block_q8_1*)vy;
csrc/quantization/gguf/mmq.cuh
+6 -6
@@ -14,10 +14,10 @@ static __device__ __forceinline__ void mul_mat_q(
14
15 const int & ncols_dst = ncols_y;
16
17 - const int row_dst_0 = blockIdx.x*mmq_y;
17 + const auto row_dst_0 = blockIdx.x*mmq_y;
18 const int & row_x_0 = row_dst_0;
19
20 - const int col_dst_0 = blockIdx.y*mmq_x;
20 + const auto col_dst_0 = blockIdx.y*mmq_x;
21 const int & col_y_0 = col_dst_0;
22
23 int * tile_x_ql = nullptr;
@@ -39,7 +39,7 @@ static __device__ __forceinline__ void mul_mat_q(
39
40 #pragma unroll
41 for (int ir = 0; ir < qr && ib0 + ir * blocks_per_warp/qr < blocks_per_row_x; ++ir) {
42 - const int kqs = ir*WARP_SIZE_GGUF + threadIdx.x;
42 + const auto kqs = ir*WARP_SIZE_GGUF + threadIdx.x;
43 const int kbxd = kqs / QI8_1;
44
45 #pragma unroll
@@ -53,7 +53,7 @@ static __device__ __forceinline__ void mul_mat_q(
53 #pragma unroll
54 for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) {
55 const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF/QI8_1)) % mmq_x;
56 - const int kby = threadIdx.x % (WARP_SIZE_GGUF/QI8_1);
56 + const auto kby = threadIdx.x % (WARP_SIZE_GGUF/QI8_1);
57 const int col_y_eff = min(col_y_0 + ids, ncols_y-1);
58
59 // if the sum is not needed it's faster to transform the scale to f32 ahead of time
@@ -87,14 +87,14 @@ static __device__ __forceinline__ void mul_mat_q(
87
88 #pragma unroll
89 for (int j = 0; j < mmq_x; j += nwarps) {
90 - const int col_dst = col_dst_0 + j + threadIdx.y;
90 + const auto col_dst = col_dst_0 + j + threadIdx.y;
91 if (col_dst >= ncols_dst) {
92 return;
93 }
94
95 #pragma unroll
96 for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
97 - const int row_dst = row_dst_0 + threadIdx.x + i;
97 + const auto row_dst = row_dst_0 + threadIdx.x + i;
98 if (row_dst >= nrows_dst) {
99 continue;
100 }
csrc/quantization/gguf/mmvq.cuh
+2 -2
@@ -1,7 +1,7 @@
1 // copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu
2 template <typename scalar_t, int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>
3 static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, const int ncols, const int nrows) {
4 - const int row = blockIdx.x*blockDim.y + threadIdx.y;
4 + const auto row = blockIdx.x*blockDim.y + threadIdx.y;
5
6 if (row >= nrows) {
7 return;
@@ -16,7 +16,7 @@ static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void *
16 const block_q_t * x = (const block_q_t *) vx;
17 const block_q8_1 * y = (const block_q8_1 *) vy;
18
19 - for (int i = threadIdx.x / (qi/vdr); i < blocks_per_row; i += blocks_per_warp) {
19 + for (auto i = threadIdx.x / (qi/vdr); i < blocks_per_row; i += blocks_per_warp) {
20 const int ibx = row*blocks_per_row + i; // x block index
21
22 const int iby = i * (qk/QK8_1); // y block index that aligns with ibx
csrc/quantization/gguf/moe.cuh
+5 -5
@@ -19,10 +19,10 @@ static __device__ __forceinline__ void moe_q(
19
20 const int ncols_dst = ncols_y * top_k;
21
22 - const int row_dst_0 = blockIdx.x * mmq_y;
22 + const auto row_dst_0 = blockIdx.x * mmq_y;
23 const int& row_x_0 = row_dst_0;
24
25 - const int col_dst_0 = blockIdx.y * mmq_x;
25 + const auto col_dst_0 = blockIdx.y * mmq_x;
26
27 int token_offs[mmq_x / nwarps];
28 for (int i = 0; i < mmq_x; i += nwarps) {
@@ -56,7 +56,7 @@ static __device__ __forceinline__ void moe_q(
56 const int n_per_r = ((qk * blocks_per_warp) / qr);
57 #pragma unroll
58 for (int ir = 0; ir < qr && ib0 * qk + ir * n_per_r < ncols_x; ++ir) {
59 - const int kqs = ir * WARP_SIZE_GGUF + threadIdx.x;
59 + const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x;
60 const int kbxd = kqs / QI8_1;
61
62 #pragma unroll
@@ -73,7 +73,7 @@ static __device__ __forceinline__ void moe_q(
73 }
74
75 if (threadIdx.x < n_per_r / QK8_1) {
76 - const int kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1);
76 + const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1);
77 const int col_y_eff = token_offs[threadIdx.y] / top_k;
78 const int block_x =
79 ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby;
@@ -119,7 +119,7 @@ static __device__ __forceinline__ void moe_q(
119
120 #pragma unroll
121 for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
122 - const int row_dst = row_dst_0 + threadIdx.x + i;
122 + const auto row_dst = row_dst_0 + threadIdx.x + i;
123 if (row_dst >= nrows_dst) {
124 continue;
125 }
csrc/quantization/gptq/q_gemm.cu
+53 -53
@@ -199,12 +199,12 @@ __global__ void gemm_half_q_half_gptq_4bit_kernel(
199 MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
200 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
201
202 - int t = threadIdx.x;
202 + auto t = threadIdx.x;
203
204 // Block
205 - int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
206 - int offset_m = blockIdx.y * m_count;
207 - int offset_k = blockIdx.z * BLOCK_KN_SIZE;
205 + auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
206 + auto offset_m = blockIdx.y * m_count;
207 + auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
208
209 [[maybe_unused]] int end_n = min(offset_n + BLOCK_KN_SIZE * 4, size_n);
210 [[maybe_unused]] int end_m = min(offset_m + m_count, size_m);
@@ -337,12 +337,12 @@ __global__ void gemm_half_q_half_gptq_2bit_kernel(
337 MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
338 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
339
340 - int t = threadIdx.x;
340 + auto t = threadIdx.x;
341
342 // Block
343 - int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
344 - int offset_m = blockIdx.y * m_count;
345 - int offset_k = blockIdx.z * BLOCK_KN_SIZE;
343 + auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
344 + auto offset_m = blockIdx.y * m_count;
345 + auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
346
347 [[maybe_unused]] int end_n = min(offset_n + BLOCK_KN_SIZE * 4, size_n);
348 [[maybe_unused]] int end_m = min(offset_m + m_count, size_m);
@@ -458,12 +458,12 @@ __global__ void gemm_half_q_half_gptq_3bit_kernel(
458 MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
459 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
460
461 - int t = threadIdx.x;
461 + auto t = threadIdx.x;
462
463 // Block
464 - int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
465 - int offset_m = blockIdx.y * m_count;
466 - int offset_k = blockIdx.z * BLOCK_KN_SIZE;
464 + auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
465 + auto offset_m = blockIdx.y * m_count;
466 + auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
467
468 [[maybe_unused]] int end_n = min(offset_n + BLOCK_KN_SIZE * 4, size_n);
469 [[maybe_unused]] int end_m = min(offset_m + m_count, size_m);
@@ -586,12 +586,12 @@ __global__ void gemm_half_q_half_gptq_8bit_kernel(
586 MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
587 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
588
589 - int t = threadIdx.x;
589 + auto t = threadIdx.x;
590
591 // Block
592 - int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
593 - int offset_m = blockIdx.y * m_count;
594 - int offset_k = blockIdx.z * BLOCK_KN_SIZE;
592 + auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
593 + auto offset_m = blockIdx.y * m_count;
594 + auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
595
596 [[maybe_unused]] int end_n = min(offset_n + BLOCK_KN_SIZE * 4, size_n);
597 [[maybe_unused]] int end_m = min(offset_m + m_count, size_m);
@@ -765,14 +765,14 @@ __global__ void reconstruct_exllama_8bit_kernel(
765 MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
766 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
767
768 - int offset_k = BLOCK_KN_SIZE * blockIdx.y;
769 - int offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
768 + auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
769 + auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
770
771 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
772
773 // Preload remapping table
774 __shared__ int perm[BLOCK_KN_SIZE];
775 - int t = threadIdx.x;
775 + auto t = threadIdx.x;
776
777 if (b_q_perm) {
778 if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
@@ -862,14 +862,14 @@ __global__ void reconstruct_exllama_4bit_kernel(
862 MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
863 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
864
865 - int offset_k = BLOCK_KN_SIZE * blockIdx.y;
866 - int offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
865 + auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
866 + auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
867
868 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
869
870 // Preload remapping table
871 __shared__ int perm[BLOCK_KN_SIZE];
872 - int t = threadIdx.x;
872 + auto t = threadIdx.x;
873
874 if (b_q_perm) {
875 if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
@@ -967,14 +967,14 @@ __global__ void reconstruct_exllama_3bit_kernel(
967 MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
968 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
969
970 - int offset_k = BLOCK_KN_SIZE * blockIdx.y;
971 - int offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
970 + auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
971 + auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
972
973 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
974
975 // Preload remapping table
976 __shared__ int perm[BLOCK_KN_SIZE];
977 - int t = threadIdx.x;
977 + auto t = threadIdx.x;
978
979 if (b_q_perm) {
980 if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
@@ -1065,14 +1065,14 @@ __global__ void reconstruct_exllama_2bit_kernel(
1065 MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
1066 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
1067
1068 - int offset_k = BLOCK_KN_SIZE * blockIdx.y;
1069 - int offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
1068 + auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
1069 + auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
1070
1071 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
1072
1073 // Preload remapping table
1074 __shared__ int perm[BLOCK_KN_SIZE];
1075 - int t = threadIdx.x;
1075 + auto t = threadIdx.x;
1076
1077 if (b_q_perm) {
1078 if (offset_k + t < size_k) perm[t] = b_q_perm[offset_k + t];
@@ -1181,11 +1181,11 @@ __global__ void gemm_half_q_half_alt_4bit_kernel(
1181 int zero_width = width / 8;
1182 int vec_height = height * 4;
1183 const int blockwidth2 = BLOCK_KN_SIZE / 2;
1184 - int b = blockIdx.y * BLOCK_M_SIZE_MAX;
1184 + auto b = blockIdx.y * BLOCK_M_SIZE_MAX;
1185 int b_end = min(BLOCK_M_SIZE_MAX, batch - b);
1186 - int h = BLOCK_KN_SIZE * blockIdx.z / 8;
1186 + auto h = BLOCK_KN_SIZE * blockIdx.z / 8;
1187 int h_end = min(BLOCK_KN_SIZE / 8, height - h) * 4;
1188 - int w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1188 + auto w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1189
1190 __shared__ half2 blockvec[BLOCK_M_SIZE_MAX][blockwidth2];
1191 if (threadIdx.x < h_end) {
@@ -1197,8 +1197,8 @@ __global__ void gemm_half_q_half_alt_4bit_kernel(
1197 }
1198
1199 __shared__ half2 deq2[256][8];
1200 - int val = threadIdx.x / 8;
1201 - int off = threadIdx.x % 8;
1200 + auto val = threadIdx.x / 8;
1201 + auto off = threadIdx.x % 8;
1202 for (; val < 256; val += BLOCK_KN_SIZE / 8) {
1203 deq2[val][off] =
1204 __halves2half2(__int2half_rn(val & 0xF), __int2half_rn(val >> 4));
@@ -1280,11 +1280,11 @@ __global__ void gemm_half_q_half_alt_8bit_kernel(
1280 int zero_width = width / 4;
1281 int vec_height = height * 2;
1282 const int blockwidth2 = BLOCK_KN_SIZE / 2;
1283 - int b = blockIdx.y * BLOCK_M_SIZE_MAX;
1283 + auto b = blockIdx.y * BLOCK_M_SIZE_MAX;
1284 int b_end = min(BLOCK_M_SIZE_MAX, batch - b);
1285 - int h = BLOCK_KN_SIZE * blockIdx.z / 4;
1285 + auto h = BLOCK_KN_SIZE * blockIdx.z / 4;
1286 int h_end = min(BLOCK_KN_SIZE / 4, height - h) * 2;
1287 - int w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1287 + auto w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1288
1289 __shared__ half2 blockvec[BLOCK_M_SIZE_MAX][blockwidth2];
1290 if (threadIdx.x < h_end) {
@@ -1393,8 +1393,8 @@ __global__ void reconstruct_gptq_kernel(const uint32_t* __restrict__ w,
1393 half* __restrict__ out) {
1394 // Start of block
1395
1396 - int column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1397 - int row = blockIdx.y * 32 / bit;
1396 + auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1397 + auto row = blockIdx.y * 32 / bit;
1398 if (column >= width) return;
1399
1400 // Views
@@ -1425,8 +1425,8 @@ __global__ void reconstruct_gptq_3bit_kernel(
1425 const int height, const int width, const int group,
1426 half* __restrict__ out) {
1427 // Start of block
1428 - int column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1429 - int row = blockIdx.y * 32;
1428 + auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1429 + auto row = blockIdx.y * 32;
1430 if (column >= width) return;
1431
1432 // Views
@@ -1542,7 +1542,7 @@ void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a,
1542
1543 __global__ void shuffle_4bit_kernel(uint32_t* __restrict__ b_q_weight,
1544 const int size_k, const int size_n) {
1545 - int n = blockIdx.x * THREADS_X + threadIdx.x;
1545 + auto n = blockIdx.x * THREADS_X + threadIdx.x;
1546 if (n >= size_n) return;
1547 int k = 0;
1548 uint32_t* b_ptr = b_q_weight + n;
@@ -1555,7 +1555,7 @@ __global__ void shuffle_4bit_kernel(uint32_t* __restrict__ b_q_weight,
1555
1556 __global__ void shuffle_8bit_kernel(uint32_t* __restrict__ b_q_weight,
1557 const int size_k, const int size_n) {
1558 - int n = blockIdx.x * THREADS_X + threadIdx.x;
1558 + auto n = blockIdx.x * THREADS_X + threadIdx.x;
1559 if (n >= size_n) return;
1560 int k = 0;
1561 uint32_t* b_ptr = b_q_weight + n;
@@ -1568,7 +1568,7 @@ __global__ void shuffle_8bit_kernel(uint32_t* __restrict__ b_q_weight,
1568
1569 __global__ void shuffle_2bit_kernel(uint32_t* __restrict__ b_q_weight,
1570 const int size_k, const int size_n) {
1571 - int n = blockIdx.x * THREADS_X + threadIdx.x;
1571 + auto n = blockIdx.x * THREADS_X + threadIdx.x;
1572 if (n >= size_n) return;
1573 int k = 0;
1574 uint32_t* b_ptr = b_q_weight + n;
@@ -1581,7 +1581,7 @@ __global__ void shuffle_2bit_kernel(uint32_t* __restrict__ b_q_weight,
1581
1582 __global__ void shuffle_3bit_kernel(uint32_t* __restrict__ b_q_weight,
1583 const int size_k, const int size_n) {
1584 - int n = blockIdx.x * THREADS_X + threadIdx.x;
1584 + auto n = blockIdx.x * THREADS_X + threadIdx.x;
1585 if (n >= size_n) return;
1586 int k = 0;
1587 uint32_t* b_ptr = b_q_weight + n;
@@ -1599,9 +1599,9 @@ __global__ void make_sequential_4bit_kernel(const uint32_t* __restrict__ w,
1599 const uint64_t* w2 = (uint64_t*)w;
1600 uint64_t* w_new2 = (uint64_t*)w_new;
1601 int w2_stride = w_width >> 1;
1602 - int w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1602 + auto w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1603 if (w2_column >= w2_stride) return;
1604 - int w_new2_row = blockIdx.y;
1604 + auto w_new2_row = blockIdx.y;
1605 int q_perm_idx = w_new2_row << 3;
1606 uint64_t dst = 0;
1607
@@ -1630,9 +1630,9 @@ __global__ void make_sequential_2bit_kernel(const uint32_t* __restrict__ w,
1630 const uint64_t* w2 = (uint64_t*)w;
1631 uint64_t* w_new2 = (uint64_t*)w_new;
1632 int w2_stride = w_width >> 1;
1633 - int w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1633 + auto w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1634 if (w2_column >= w2_stride) return;
1635 - int w_new2_row = blockIdx.y;
1635 + auto w_new2_row = blockIdx.y;
1636 int q_perm_idx = w_new2_row << 4;
1637 uint64_t dst = 0;
1638
@@ -1658,10 +1658,10 @@ __global__ void make_sequential_3bit_kernel(const uint32_t* __restrict__ w,
1658 uint32_t* __restrict__ w_new,
1659 const int* __restrict__ q_perm,
1660 const int w_width) {
1661 - int w_column = THREADS_X * blockIdx.x + threadIdx.x;
1661 + auto w_column = THREADS_X * blockIdx.x + threadIdx.x;
1662 if (w_column >= w_width) return;
1663 - int w_new_row = blockIdx.y * 3;
1664 - int q_perm_idx = blockIdx.y << 5;
1663 + auto w_new_row = blockIdx.y * 3;
1664 + auto q_perm_idx = blockIdx.y << 5;
1665 uint32_t dst[3] = {0, 0, 0};
1666
1667 #pragma unroll
@@ -1744,9 +1744,9 @@ __global__ void make_sequential_8bit_kernel(const uint32_t* __restrict__ w,
1744 const uint64_t* w2 = (uint64_t*)w;
1745 uint64_t* w_new2 = (uint64_t*)w_new;
1746 int w2_stride = w_width >> 1;
1747 - int w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1747 + auto w2_column = THREADS_X * blockIdx.x + threadIdx.x;
1748 if (w2_column >= w2_stride) return;
1749 - int w_new2_row = blockIdx.y;
1749 + auto w_new2_row = blockIdx.y;
1750 int q_perm_idx = w_new2_row << 2;
1751 uint64_t dst = 0;
1752
csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu
+19 -19
@@ -55,11 +55,11 @@ struct GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
55 this_block_B_base_ptr = params.B_ptr + blockIdx.y * Ntile * params.K +
56 blockIdx.z * params.SplitK * 4;
57
58 - const int lane_id = threadIdx.x % WARP_SIZE;
58 + const auto lane_id = threadIdx.x % WARP_SIZE;
59
60 // For matrix A, a block load/store Mtile(row) x 32(col) elements in
61 // multiple iters, 8x4 warp load/store 8(row) x 32(col) elements per iter
62 - const int Aldg_row_base_idx = threadIdx.x / 4;
62 + const auto Aldg_row_base_idx = threadIdx.x / 4;
63 Aldg_col_idx = (threadIdx.x % 4) * LDG_ELEMENT_CNT_A;
64 const int Aldg_base_offset = Aldg_row_base_idx * params.K + Aldg_col_idx;
65
@@ -67,7 +67,7 @@ struct GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
67 // elements of N32K16 packing in multiple iters, 4x8 warp load/store 4(row)
68 // * 128(col) per iter
69 Bldg_col_idx = (threadIdx.x % 8) * LDG_ELEMENT_CNT_B;
70 - const int Bldg_row_base_idx = threadIdx.x / 8;
70 + const auto Bldg_row_base_idx = threadIdx.x / 8;
71 const int Bldg_base_offset =
72 Bldg_row_base_idx * params.K * 4 + Bldg_col_idx;
73
@@ -89,7 +89,7 @@ struct GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
89 B_ldg_guard = 0;
90 #pragma unroll
91 for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
92 - int m_idx = blockIdx.x * Mtile + Aldg_row_base_idx + i * M_SIZE_ONE_LOAD;
92 + auto m_idx = blockIdx.x * Mtile + Aldg_row_base_idx + i * M_SIZE_ONE_LOAD;
93 if (m_idx < params.M) {
94 A_ldg_guard |= (1u << i);
95 }
@@ -98,8 +98,8 @@ struct GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
98 const int N_padded = (params.N + 31) / 32 * 32;
99 #pragma unroll
100 for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
101 - int n_idx = blockIdx.y * Ntile + (Bldg_row_base_idx / 8) * 32 +
102 - i * N_SIZE_ONE_LOAD;
101 + auto n_idx = blockIdx.y * Ntile + (Bldg_row_base_idx / 8) * 32 +
102 + i * N_SIZE_ONE_LOAD;
103 if (n_idx < N_padded) {
104 B_ldg_guard |= (1u << i);
105 }
@@ -355,7 +355,7 @@ struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
355 __device__ void fused_splitk_reduce() {
356 // need splitk-reduce if enable splitk
357 if (gridDim.z > 1) {
358 - int blk_red_idx = blockIdx.x * gridDim.y + blockIdx.y;
358 + auto blk_red_idx = blockIdx.x * gridDim.y + blockIdx.y;
359 // Wait for all previous blocks in the splitk direction to accumulate the
360 // results into C_tmp
361 if (threadIdx.x == 0) {
@@ -371,7 +371,7 @@ struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
371 }
372 __syncthreads();
373
374 - int C_tmp_base_offset = blk_red_idx * Mtile * Ntile + threadIdx.x * 4;
374 + auto C_tmp_base_offset = blk_red_idx * Mtile * Ntile + threadIdx.x * 4;
375 if (blockIdx.z != 0) {
376 // expecting that temporary register here reuses the previous A&B frag
377 // register
@@ -456,7 +456,7 @@ struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
456
457 FType* C_base_ptr = this_block_C_base_ptr + store_c_base_offset;
458 // C_tile lds and stg
459 - int m_base_idx = store_c_row_base_idx + blockIdx.x * Mtile;
459 + auto m_base_idx = store_c_row_base_idx + blockIdx.x * Mtile;
460 bool n_guard = (store_c_col_idx + blockIdx.y * Ntile) < params.N;
461 if (WARP_NTILE == 32) {
462 int lds_c_base_offset = warp_id * Mtile * WARP_NTILE +
@@ -580,9 +580,9 @@ __global__ void __launch_bounds__(BLOCK)
580 int sts_stage_idx = 0;
581 int lds_stage_idx = 0;
582
583 - int tb_k_slice = blockIdx.z * params.SplitK + params.SplitK <= params.K
584 - ? params.SplitK
585 - : params.K - blockIdx.z * params.SplitK;
583 + auto tb_k_slice = blockIdx.z * params.SplitK + params.SplitK <= params.K
584 + ? params.SplitK
585 + : params.K - blockIdx.z * params.SplitK;
586 int k_tiles = (tb_k_slice + 31) / 32;
587 int first_k_tile = tb_k_slice - (k_tiles - 1) * 32;
588
@@ -777,13 +777,13 @@ __global__ void restore_N32_K16_dequantize_rhs_w8a16_perc_kernel(
777 const QT* qdata, const FT* scales, const FT* zeros, FT* fdata,
778 const int N_32align, const int N, const int K) {
779 __shared__ FT smem[64 * 32];
780 - int warp_id = threadIdx.x / 32;
781 - int lane_id = threadIdx.x % 32;
782 - const int src_row_idx = blockIdx.x * 8 + lane_id / 4;
780 + auto warp_id = threadIdx.x / 32;
781 + auto lane_id = threadIdx.x % 32;
782 + const auto src_row_idx = blockIdx.x * 8 + lane_id / 4;
783 const int src_col_idx =
784 blockIdx.y * 64 * 4 + warp_id * 16 * 4 + (lane_id % 4) * 16;
785 const int src_offset = src_row_idx * K * 4 + src_col_idx;
786 - int params_nidx = blockIdx.x * 32 + (lane_id / 4) * 4;
786 + auto params_nidx = blockIdx.x * 32 + (lane_id / 4) * 4;
787
788 QT qval_reg[16];
789 const QT* pdata = qdata + src_offset;
@@ -829,8 +829,8 @@ __global__ void restore_N32_K16_dequantize_rhs_w8a16_perc_kernel(
829 *reinterpret_cast<uint4*>(smem + lds_base_offset + i * 32 * 32);
830 }
831
832 - const int dst_row_base_kidx = blockIdx.y * 64 + threadIdx.x / 4;
833 - const int dst_col_nidx = blockIdx.x * 32 + (threadIdx.x % 4) * 8;
832 + const auto dst_row_base_kidx = blockIdx.y * 64 + threadIdx.x / 4;
833 + const auto dst_col_nidx = blockIdx.x * 32 + (threadIdx.x % 4) * 8;
834 #pragma unroll
835 for (int i = 0; i < 2; ++i) {
836 int dst_row_kidx = dst_row_base_kidx + i * 32;
@@ -1008,4 +1008,4 @@ torch::Tensor allspark_w8a16_gemm(
1008
1009 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
1010 m.impl("allspark_w8a16_gemm", &allspark_w8a16_gemm);
1011 -}
\ No newline at end of file
1011 +}
csrc/quantization/gptq_allspark/allspark_repack.cu
+4 -4
@@ -13,8 +13,8 @@ __global__ void __launch_bounds__(128)
13 const uint8_t* B, const FType* B_scale, const FType* B_zero,
14 uint8_t* B_result, FType* B_scale_result, FType* B_zero_result,
15 const int K, const int N, const int N_32align) {
16 - const int lane_id = threadIdx.x % 32;
17 - const int warp_id = threadIdx.x / 32;
16 + const auto lane_id = threadIdx.x % 32;
17 + const auto warp_id = threadIdx.x / 32;
18
19 if (blockIdx.x != gridDim.x - 1) {
20 // Load B
@@ -50,7 +50,7 @@ __global__ void __launch_bounds__(128)
50 }
51
52 // Store B
53 - const int dst_row_base_idx = blockIdx.y * (128 / 4) + (lane_id / 8) * 8;
53 + const auto dst_row_base_idx = blockIdx.y * (128 / 4) + (lane_id / 8) * 8;
54 const int dst_col_idx =
55 blockIdx.x * (64 * 4) + warp_id * 64 + (lane_id % 8) * 8;
56 for (int i = 0; i < 8; ++i) {
@@ -65,7 +65,7 @@ __global__ void __launch_bounds__(128)
65 } else {
66 // Load B_scale and B_zero
67 FType b_scale_reg, b_zero_reg;
68 - int src_offset = blockIdx.y * 128 + threadIdx.x;
68 + auto src_offset = blockIdx.y * 128 + threadIdx.x;
69 ldg16_cg_0(b_scale_reg, B_scale + src_offset, src_offset < N);
70 if (B_zero != nullptr)
71 ldg16_cg_0(b_zero_reg, B_zero + src_offset, src_offset < N);
csrc/quantization/gptq_allspark/allspark_utils.cuh
+2 -2
@@ -62,7 +62,7 @@ template <typename FType, int BLOCK, int N_MATRIX>
62 __global__ void f16_gemm_splitk_reduce_kernel(const FType* C_split, FType* C,
63 uint32_t n, uint32_t n_matrix,
64 uint32_t matrix_size) {
65 - int idx = blockIdx.x * BLOCK + threadIdx.x;
65 + auto idx = blockIdx.x * BLOCK + threadIdx.x;
66
67 if (idx >= matrix_size) {
68 return;
@@ -407,4 +407,4 @@ static __device__ half2 inline num2num2(const half x) {
407 return __half2half2(x);
408 }
409
410 -} // namespace allspark
\ No newline at end of file
410 +} // namespace allspark