char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
a094f3814
char/aliceai-t5-moe
1#include "add-id.cuh" 2 3static __global__ void add_id_kernel( 4 const float * src0, const float * src1, const int32_t * src2, float * dst, 5 int64_t ne0, int64_t ne1, 6 size_t nb01, size_t nb02, 7 size_t nb11, 8 size_t nb21 9 ) { 10 11 const int64_t i1 = blockIdx.x; 12 const int64_t i2 = blockIdx.y; 13 14 const int i11 = *(const int32_t *) ((const char *) src2 + i1*sizeof(int32_t) + i2*nb21); 15 16 const size_t nb1 = ne0 * sizeof(float); 17 const size_t nb2 = ne1 * nb1; 18 19 float * dst_row = (float *)((char *)dst + i1*nb1 + i2*nb2); 20 const float * src0_row = (const float *)((const char *)src0 + i1*nb01 + i2*nb02); 21 const float * src1_row = (const float *)((const char *)src1 + i11*nb11); 22 23 for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) { 24 dst_row[i0] = src0_row[i0] + src1_row[i0]; 25 } 26} 27 28void ggml_cuda_op_add_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 29 const ggml_tensor * src0 = dst->src[0]; 30 const ggml_tensor * src1 = dst->src[1]; 31 const ggml_tensor * src2 = dst->src[2]; 32 33 GGML_TENSOR_TERNARY_OP_LOCALS 34 35 GGML_ASSERT(dst->type == GGML_TYPE_F32); 36 GGML_ASSERT(src0->type == GGML_TYPE_F32); 37 GGML_ASSERT(src1->type == GGML_TYPE_F32); 38 GGML_ASSERT(src2->type == GGML_TYPE_I32); 39 40 GGML_ASSERT(nb00 == sizeof(float)); 41 GGML_ASSERT(nb10 == sizeof(float)); 42 GGML_ASSERT(nb20 == sizeof(int32_t)); 43 44 const float * src0_d = (const float *)src0->data; 45 const float * src1_d = (const float *)src1->data; 46 const int32_t * src2_d = (const int32_t *)src2->data; 47 float * dst_d = (float *)dst->data; 48 49 int threads = std::min((int)ne00, 768); // cols 50 dim3 blocks(ne01, ne02); // n_experts_used, n_tokens 51 add_id_kernel<<<blocks, threads, 0, ctx.stream()>>>( 52 src0_d, src1_d, src2_d, dst_d, 53 ne0, ne1, 54 nb01, nb02, 55 nb11, 56 nb21 57 ); 58}