KBaba7/llama.cpp

1#include "getrows.cuh"2#include "dequantize.cuh"3 4template<int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>5static __global__ void k_get_rows(6        const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,7        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/8        /*const int64_t ne10, const int64_t ne11,*/ const int64_t ne12, /*const int64_t ne13,*/9        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,10        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,11        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {12 13    const int i00 = (blockIdx.x*blockDim.x + threadIdx.x)*2;14    const int i10 =  blockDim.y*blockIdx.y + threadIdx.y;15    const int i11 = (blockIdx.z*blockDim.z + threadIdx.z)/ne12;16    const int i12 = (blockIdx.z*blockDim.z + threadIdx.z)%ne12;17 18    if (i00 >= ne00) {19        return;20    }21 22    const int i01 = src1[i10*s10 + i11*s11 + i12*s12];23 24    dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;25    const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;26 27    const int ib   =  i00/qk;      // block index28    const int iqs  = (i00%qk)/qr;  // quant index29    const int iybs = i00 - i00%qk; // dst block start index30    const int y_offset = qr == 1 ? 1 : qk/2;31 32    // dequantize33    dfloat2 v;34    dequantize_kernel(src0_row, ib, iqs, v);35 36    dst_row[iybs + iqs + 0]        = v.x;37    dst_row[iybs + iqs + y_offset] = v.y;38}39 40template<typename src0_t, typename dst_t>41static __global__ void k_get_rows_float(42        const src0_t * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,43        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/44        /*const int64_t ne10, const int64_t ne11,*/ const int64_t ne12, /*const int64_t ne13,*/45        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,46        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,47        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {48 49    const int i00 =  blockIdx.x*blockDim.x + threadIdx.x;50    const int i10 =  blockDim.y*blockIdx.y + threadIdx.y;51    const int i11 = (blockIdx.z*blockDim.z + threadIdx.z)/ne12;52    const int i12 = (blockIdx.z*blockDim.z + threadIdx.z)%ne12;53 54    if (i00 >= ne00) {55        return;56    }57 58    const int i01 = src1[i10*s10 + i11*s11 + i12*s12];59 60    dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;61    const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);62 63    dst_row[i00] = src0_row[i00];64}65 66template<typename grad_t, typename dst_t>67static __global__ void k_get_rows_back_float(68        const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst, const int64_t ncols, const int64_t nrows_grad) {69    const int col = blockIdx.x*blockDim.x + threadIdx.x;70 71    if (col >= ncols) {72        return;73    }74 75    const int dst_row = blockIdx.y*blockDim.y + threadIdx.y;76 77    float sum = 0.0f;78 79    for (int64_t i = 0; i < nrows_grad; ++i) {80        if (rows[i] != dst_row) {81            continue;82        }83        sum += grad[i*ncols + col];84    }85 86    dst[dst_row*ncols + col] = sum;87}88 89template<int qk, int qr, dequantize_kernel_t dq>90static void get_rows_cuda(91        const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,92        const void * src0_dd, const int32_t * src1_dd, float * dst_dd, cudaStream_t stream) {93 94    GGML_TENSOR_BINARY_OP_LOCALS95 96    const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);97    const int block_num_x = (ne00 + 2*CUDA_GET_ROWS_BLOCK_SIZE - 1) / (2*CUDA_GET_ROWS_BLOCK_SIZE);98    const dim3 block_nums(block_num_x, ne10, ne11*ne12);99 100    // strides in elements101    //const size_t s0 = nb0 / ggml_element_size(dst);102    const size_t s1 = nb1 / ggml_element_size(dst);103    const size_t s2 = nb2 / ggml_element_size(dst);104    const size_t s3 = nb3 / ggml_element_size(dst);105 106    const size_t s10 = nb10 / ggml_element_size(src1);107    const size_t s11 = nb11 / ggml_element_size(src1);108    const size_t s12 = nb12 / ggml_element_size(src1);109    //const size_t s13 = nb13 / ggml_element_size(src1);110 111    GGML_ASSERT(ne00 % 2 == 0);112 113    k_get_rows<qk, qr, dq><<<block_nums, block_dims, 0, stream>>>(114        src0_dd, src1_dd, dst_dd,115        ne00, /*ne01, ne02, ne03,*/116        /*ne10, ne11,*/ ne12, /*ne13,*/117        /* s0,*/ s1, s2, s3,118        /* nb00,*/ nb01, nb02, nb03,119        s10, s11, s12/*, s13*/);120 121    GGML_UNUSED(dst);122}123 124template<typename src0_t>125static void get_rows_cuda_float(126        const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,127        const src0_t * src0_dd, const int32_t * src1_dd, float * dst_dd, cudaStream_t stream) {128 129    GGML_TENSOR_BINARY_OP_LOCALS130 131    GGML_ASSERT(ne13 == 1);132 133    const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);134    const int block_num_x = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;135    const dim3 block_nums(block_num_x, ne10, ne11*ne12);136 137    // strides in elements138    //const size_t s0 = nb0 / ggml_element_size(dst);139    const size_t s1 = nb1 / ggml_element_size(dst);140    const size_t s2 = nb2 / ggml_element_size(dst);141    const size_t s3 = nb3 / ggml_element_size(dst);142 143    const size_t s10 = nb10 / ggml_element_size(src1);144    const size_t s11 = nb11 / ggml_element_size(src1);145    const size_t s12 = nb12 / ggml_element_size(src1);146    //const size_t s13 = nb13 / ggml_element_size(src1);147 148    k_get_rows_float<<<block_nums, block_dims, 0, stream>>>(149        src0_dd, src1_dd, dst_dd,150        ne00, /*ne01, ne02, ne03,*/151        /*ne10, ne11,*/ ne12, /*ne13,*/152        /* s0,*/ s1, s2, s3,153        /* nb00,*/ nb01, nb02, nb03,154        s10, s11, s12/*, s13*/);155 156    GGML_UNUSED(dst);157}158 159void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {160    const ggml_tensor * src0 = dst->src[0];161    const ggml_tensor * src1 = dst->src[1];162 163    const void    * src0_d = (const void    *) src0->data;164    const int32_t * src1_d = (const int32_t *) src1->data;165    float         * dst_d  = (float         *) dst->data;166 167    cudaStream_t stream = ctx.stream();168 169    GGML_ASSERT(src1->type == GGML_TYPE_I32);170    GGML_ASSERT(dst->type  == GGML_TYPE_F32);171 172    GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type));173    GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type));174    GGML_ASSERT(dst->nb[0]  == ggml_type_size(dst->type));175 176    switch (src0->type) {177        case GGML_TYPE_F16:178            get_rows_cuda_float(src0, src1, dst, (const half *) src0_d, src1_d, dst_d, stream);179            break;180        case GGML_TYPE_F32:181            get_rows_cuda_float(src0, src1, dst, (const float *) src0_d, src1_d, dst_d, stream);182            break;183        case GGML_TYPE_Q4_0:184            get_rows_cuda<QK4_0, QR4_0, dequantize_q4_0>(src0, src1, dst, src0_d, src1_d, dst_d, stream);185            break;186        case GGML_TYPE_Q4_1:187            get_rows_cuda<QK4_1, QR4_1, dequantize_q4_1>(src0, src1, dst, src0_d, src1_d, dst_d, stream);188            break;189        case GGML_TYPE_Q5_0:190            get_rows_cuda<QK5_0, QR5_0, dequantize_q5_0>(src0, src1, dst, src0_d, src1_d, dst_d, stream);191            break;192        case GGML_TYPE_Q5_1:193            get_rows_cuda<QK5_1, QR5_1, dequantize_q5_1>(src0, src1, dst, src0_d, src1_d, dst_d, stream);194            break;195        case GGML_TYPE_Q8_0:196            get_rows_cuda<QK8_0, QR8_0, dequantize_q8_0>(src0, src1, dst, src0_d, src1_d, dst_d, stream);197            break;198        default:199            // TODO: k-quants200            GGML_ABORT("%s: unsupported type: %s\n", __func__, ggml_type_name(src0->type));201            break;202    }203}204 205void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {206    const ggml_tensor * src0 = dst->src[0]; // gradients of forward pass output207    const ggml_tensor * src1 = dst->src[1]; // src1 in forward pass208 209    GGML_TENSOR_BINARY_OP_LOCALS210 211    const float   * src0_d = (const float   *) src0->data;212    const int32_t * src1_d = (const int32_t *) src1->data;213    float         * dst_d  = (float         *) dst->data;214 215    cudaStream_t stream = ctx.stream();216 217    GGML_ASSERT(src0->type == GGML_TYPE_F32);218    GGML_ASSERT(src1->type == GGML_TYPE_I32);219    GGML_ASSERT(dst->type  == GGML_TYPE_F32);220 221    GGML_ASSERT(ggml_is_contiguous(src0));222    GGML_ASSERT(ggml_is_contiguous(src1));223    GGML_ASSERT(ggml_is_contiguous(dst));224 225    GGML_ASSERT(ne02*ne03 == 1);226    GGML_ASSERT(ne12*ne13 == 1);227    GGML_ASSERT(ne2*ne3 == 1);228 229    const dim3 block_dims(CUDA_GET_ROWS_BACK_BLOCK_SIZE, 1, 1);230    const int block_num_x = (ne00 + CUDA_GET_ROWS_BACK_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BACK_BLOCK_SIZE;231    const dim3 block_nums(block_num_x, ne1, 1);232 233    k_get_rows_back_float<<<block_nums, block_dims, 0, stream>>>(src0_d, src1_d, dst_d, ne00, ne10);234}235