KBaba7/llama.cpp

1#include "quantize.cuh"2#include <cstdint>3 4static __global__ void quantize_q8_1(const float * __restrict__ x, void * __restrict__ vy, const int64_t kx, const int64_t kx0_padded) {5    const int64_t ix0 = (int64_t)blockDim.x*blockIdx.x + threadIdx.x;6 7    if (ix0 >= kx0_padded) {8        return;9    }10 11    const int64_t ix1 = blockIdx.y;12 13    const int64_t i_padded = ix1*kx0_padded + ix0;14 15    block_q8_1 * y = (block_q8_1 *) vy;16 17    const int64_t ib = i_padded / QK8_1; // block index18    const int64_t iqs = i_padded % QK8_1; // quant index19 20    const float xi = ix0 < kx ? x[ix1*kx + ix0] : 0.0f;21    float amax = fabsf(xi);22    float sum = xi;23 24    amax = warp_reduce_max(amax);25    sum = warp_reduce_sum(sum);26 27    const float d = amax / 127;28    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);29 30    y[ib].qs[iqs] = q;31 32    if (iqs > 0) {33        return;34    }35 36    reinterpret_cast<half&>(y[ib].ds.x) = d;37    reinterpret_cast<half&>(y[ib].ds.y) = sum;38}39 40template <mmq_q8_1_ds_layout ds_layout>41static __global__ void quantize_mmq_q8_1(42    const float * __restrict__ x, void * __restrict__ vy, const int64_t kx0, const int64_t kx1, const int64_t kx0_padded) {43 44    constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;45    constexpr int vals_per_sum   = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;46 47    const int64_t ix0 = ((int64_t)blockDim.x*blockIdx.x + threadIdx.x)*4;48 49    if (ix0 >= kx0_padded) {50        return;51    }52 53    const float4 * x4 = (const float4 *) x;54 55    const int64_t ix1 = kx1*blockIdx.z + blockIdx.y;56 57    block_q8_1_mmq * y = (block_q8_1_mmq *) vy;58 59    const int64_t ib0 = blockIdx.z*((int64_t)gridDim.y*gridDim.x*blockDim.x/QK8_1); // first block of channel60    const int64_t ib  = ib0 + (ix0 / (4*QK8_1))*kx1 + blockIdx.y;                   // block index in channel61    const int64_t iqs = ix0 % (4*QK8_1);                                            // quant index in block62 63    // Load 4 floats per thread and calculate max. abs. value between them:64    const float4 xi = ix0 < kx0 ? x4[(ix1*kx0 + ix0)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);65    float amax = fabsf(xi.x);66    amax = fmaxf(amax, fabsf(xi.y));67    amax = fmaxf(amax, fabsf(xi.z));68    amax = fmaxf(amax, fabsf(xi.w));69 70    // Exchange max. abs. value between vals_per_scale/4 threads.71#pragma unroll72    for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) {73        amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));74    }75 76    float sum;77    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {78        sum = xi.x + xi.y + xi.z + xi.w;79 80        // Exchange calculate sum across vals_per_sum/4 threads.81#pragma unroll82        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {83            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);84        }85    }86 87    const float d_inv = 127.0f / amax;88    char4 q;89    q.x = roundf(xi.x*d_inv);90    q.y = roundf(xi.y*d_inv);91    q.z = roundf(xi.z*d_inv);92    q.w = roundf(xi.w*d_inv);93 94    // Write back 4 int8 values as a single 32 bit value for better memroy bandwidth:95    char4 * yqs4 = (char4 *) y[ib].qs;96    yqs4[iqs/4] = q;97 98    if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {99        if (iqs % 16 != 0 || iqs >= 96) {100            return;101        }102 103        y[ib].d2s6[2 + iqs/16] = sum;104 105        if (iqs % 64 != 0) {106            return;107        }108 109        const float d = 1.0f / d_inv;110 111        y[ib].d2s6[iqs/64] = d;112 113        return;114    }115 116    if (iqs % 32 != 0) {117        return;118    }119 120    const float d = 1.0f / d_inv;121 122    if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {123        y[ib].ds4[iqs/32] = make_half2(d, sum);124    } else {125        y[ib].d4[iqs/32]  = d;126    }127}128 129void quantize_row_q8_1_cuda(130    const float * x, void * vy, const int64_t kx0, const int64_t kx1, const int64_t channels,131    const int64_t kx0_padded, const ggml_type type_x, cudaStream_t stream) {132 133    GGML_ASSERT(kx0_padded % QK8_1 == 0);134 135    const int64_t block_num_x = (kx0_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;136    const dim3 num_blocks(block_num_x, kx1*channels, 1);137    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1);138    quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(x, vy, kx0, kx0_padded);139 140    GGML_UNUSED(type_x);141}142 143void quantize_mmq_q8_1_cuda(144    const float * x, void * vy, const int64_t kx0, const int64_t kx1, const int64_t channels,145    const int64_t kx0_padded, const ggml_type type_x, cudaStream_t stream) {146 147    GGML_ASSERT(kx0_padded % (4*QK8_1) == 0);148 149    const int64_t block_num_x = (kx0_padded + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);150    const dim3 num_blocks(block_num_x, kx1, channels);151    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);152    switch (mmq_get_q8_1_ds_layout(type_x)) {153        case MMQ_Q8_1_DS_LAYOUT_D4:154            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>155                <<<num_blocks, block_size, 0, stream>>>(x, vy, kx0, kx1, kx0_padded);156            break;157        case MMQ_Q8_1_DS_LAYOUT_DS4:158            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>159                <<<num_blocks, block_size, 0, stream>>>(x, vy, kx0, kx1, kx0_padded);160            break;161        case MMQ_Q8_1_DS_LAYOUT_D2S6:162            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6>163                <<<num_blocks, block_size, 0, stream>>>(x, vy, kx0, kx1, kx0_padded);164            break;165        default:166            GGML_ABORT("fatal error");167            break;168    }169}170