#define WARP_SIZE 32

typedef struct {

half d;

int8_t qs[QK8_0];

} block_q8_0;

const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;

const dim3 block_nums(1, block_num_y, 1);

const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);

dequantize_mul_mat_vec<QK8_0, QR8_0, dequantize_q8_0>

<<<block_nums, block_dims, 0, stream>>>(vx, y, dst, ncols, nrows);

}

static __device__ __forceinline__ void dequantize_q8_0(const void *vx, const int ib, const int iqs, dfloat2 & v) {

// 均匀对称量化

// dequantize is int8 * scale

const block_q8_0* x = (const block_q8_0*) vx;

const dfloat2 d = x[ib].d; // scale

v.x = x[ib].qs[iqs+0];

v.y = x[ib].qs[iqs+1];

#ifdef GGML_CUDA_FP16

// FP16

v = __hmul2(v, {d, d});

#else

// FP32

v.x *= d;

v.y *= d;

#endif

}

template<int qk, int qr, dequantize_kernel_t dequantize_kernel>

static __global__ void dequantize_mul_mat_vec(const void * __restrict__ vx, const dfloat *__restrict__ y,

float* __restrict__ dst, const int ncols, const int nrows)

{

// qk = quantized weights per x block

// qr = number of quantized weights per data value in x block

const int row = blockIdx.y * blockDim.y + threadIdx.x;

if (row >= nrows) return;

const int tid = threadIdx.x;

const int iter_stride = 2 * GGML_CUDA_DMMV_X; // 2*32

const int vals_per_iter = iter_stride / WARP_SIZE;

const int y_offset = qr == 1 ? 1 : qk/2;

#ifdef GGML_CUDA_FP16

half2 tmp = {0.0f, 0.0f};

#else

float tmp = 0.0f;

#endif

// 32 threads process 4096-set data

for (int i =0; i < ncols; i += iter_stride) {

const int col = i + vals_per_iter * tid;

const int ib = (row * ncols + col) / qk;

const int iqs = (col % qk) / qr;

const int iybs = col - col % qk;

for (int j = 0; j < vals_per_iter; j+= 2) {

// 2 vals per j iter

// dequantize

// for qr = 2 the iqs needs to increase by 1 per j iter because 2 weights per data val

dfloat2 v;

dequantize_kernel(vx, ib, iqs + j / qr, v);

#ifdef GGML_CUDA_FP16

tmp += __hmul2(v, {

y[iybs + iqs + j/qr + 0],

y[iybs + iqs + j/qr + y_offset]

});

#else

tmp += v.x * y[iybs + iqs + j / qr + 0];

tmp += v.y * y[iybs + iqs + j / qr + y_offset];

#endif

}

}

#pragma unroll

for (int mask = 16; mask > 0; mask >>= 1) {

tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32);

}

if (tid == 0) {

#ifdef GGML_CUDA_FP16

dst[row] = tmp.x + tmp.y;

#else

dst[row] = tmp;

#endif

}

#define VDR_Q8_0_Q8_1_MMVQ 2

static void mul_mat_vec_q8_0_q8_1_cuda(const void * vx, const void * vy, float * dst,

const int ncols, const int nrows, cudaStream_t stream)

{

GGML_ASSERT(ncols % QK8_0 == 0);

const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;

const dim3 block_nums(1, block_num_y, 1);

const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);

// QK8_0 = 32, QI8_0 = 8

mul_mat_vec_q<QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1>

<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);

}

template<int vdr>

static __device__ __forceinline__ float vec_dot_q8_0_q8_1_impl(

const int *v, const int *u, const float & d8_0, const float & d8_1)

{

#if __CUDA_ARCH__ >= MIN_CC_DP4A

int sumi = 0;

#pragma unroll

for (int i = 0; i < vdr; ++i) {

sumi = __dp4a(v[i], u[i], sumi);

}

return d8_0 * d8_1 * sumi;

#else

assert(false);

return 0.0f;

#endif

}

static __device__ __forceinline__ float vec_dot_q8_0_q8_1(

const void* __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs)

{

const block_q8_0* bq8_0 = (const block_q8_0 *) vbq;

int v[VDR_Q8_0_Q8_1_MMVQ];

int u[VDR_Q8_0_Q8_1_MMVQ];

#pragma unroll

for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) {

v[i] = get_int_from_int8(bq8_0->qs, iqs + i);

u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs+i);

}

return vec_dot_q8_0_q8_1_impl<VDR_Q8_0_Q8_1_MMVQ>(v, u, bq8_0->d, bq8_1->ds.x);

}

template <int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>

static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy,

float* __restrict__ dst, const int ncols, const int nrows)

{

const int row = blockIdx.y * blockDim.y + threadIdx.y;

if (row >= nrows) return;

const int blocks_per_row = ncols / qk;

const int blocks_per_warp = vdr * WARP_SIZE / qi;

float tmp = 0.0f;

const block_q_t * x = (const block_q_t * ) vx;

const block_q8_1 * y = (const block_q8_1 *) vy;

for (int i = 0; i < blocks_per_row; i += blocks_per_warp) {

const int ibx = row * blocks_per_row + i + threadIdx.x / (qi / vdr);

const int iby = (i + threadIdx.x / (qi / vdr)) * (qk / QK8_1);

const int iqs = vdr * (threadIdx.x % (qi / vdr));

tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs);

}

#pragma unroll

for (int mask = 16; mask > 0; mask >>= 1) {

tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32);

}

if (threadIdx.x == 0)

dst[row] = tmp;

}

#define QK8_1 32

typedef struct {

half2 ds;

int8_t qs[QK8_0];

} block_q8_1;

static __global__ void quantize_q8_1(const float* __restrict__ x, void* __restrict__ vy, const int kx,

const int kx_padded)

{

const int ix = blockDim.x * blockIdx.x + threadIdx.x; // 0-4096

if (ix >= kx_padded) return;

const int iy = blockDimx.y * blockIdx.y + threadIdx.y; // 0

const int i_padded = iy * ky_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;

float amax = fabsf(xi);

float sum = xi;

#pragma unroll

for (int mask = 16; mask > 0; mask >>= 1) {

amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, mask, 32));

sum += __shfl_xor_sync(0xffffffff, sum, mask, 32);

}

// q = round(clip(r_i / scale, Q_{min}, Q_{max}))

// scale = fmax - fmin / qmax - qmin

const float d = amax / 127;

const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);

y[ib].qs[iqs] = q;

if (iqs > 0) return;

y[ib].ds.x = d;

y[ib].ds.y = sum;

}