small refactor

Author: JohannesGaessler

ggml-cuda.cu

@@ -1415,24 +1415,25 @@ static __device__ void convert_f16(const void * vx, const int ib, const int iqs,
 }
 
 static __global__ void quantize_q8_1(
-    const float * __restrict__ x, void * __restrict__ vy, const int kx, const int kx_padded, const int nchannels) {
+    const float * __restrict__ x, void * __restrict__ vy, const int kx, const int kx_padded, const int ky) {
 
     const int ix = blockDim.x*blockIdx.x + threadIdx.x;
 
     if (ix >= kx_padded) {
         return;
     }
 
-    const int iy = blockDim.y*blockIdx.y + threadIdx.y;
+    const int iy      = blockDim.y*blockIdx.y + threadIdx.y;
+    const int channel = blockDim.z*blockIdx.z + threadIdx.z;
 
-    const int i_padded = iy*kx_padded + ix;
+    const int i_padded = channel*ky*kx_padded + iy*kx_padded + ix;
 
     block_q8_1 * y = (block_q8_1 *) vy;
 
     const int ib = i_padded / QK8_1; // block index
     const int iqs = i_padded % QK8_1; // quant index
 
-    const float xi = ix < kx ? x[iy*kx + ix] : 0.0f;
+    const float xi = ix < kx ? x[channel*ky*kx + iy*kx + ix] : 0.0f;
     float amax = fabsf(xi);
     float sum = xi;
 
@@ -4298,9 +4299,9 @@ static void quantize_row_q8_1_cuda(
     const float * x, void * vy, const int kx, const int ky, const int kx_padded, const int nchannels, cudaStream_t stream) {
 
     const int block_num_x = (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
-    const dim3 num_blocks(block_num_x, ky*nchannels, 1);
+    const dim3 num_blocks(block_num_x, ky, nchannels);
     const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1);
-    quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(x, vy, kx, kx_padded, nchannels);
+    quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(x, vy, kx, kx_padded, ky);
 }
 
 static void dequantize_row_q4_0_cuda(const void * vx, float * y, const int k, cudaStream_t stream) {