From 3da71dcda2bbd115d390256c08160427a6d777bf Mon Sep 17 00:00:00 2001 From: Iwan Kawrakow Date: Fri, 24 Oct 2025 13:07:54 +0300 Subject: [PATCH] Fused ffn_up*unary_op(ffn_gate) for MMVQ (with bias) --- ggml/src/ggml-cuda.cu | 140 +++++++++++---------------------- ggml/src/ggml-cuda/mmvq-args.h | 3 + ggml/src/ggml-cuda/mmvq.cu | 73 +++++++++++------ ggml/src/ggml-cuda/mmvq.cuh | 1 + 4 files changed, 101 insertions(+), 116 deletions(-) diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index f507012a..41d4ba3a 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -2618,8 +2618,6 @@ static bool ggml_cuda_moe_up_gate_unary(ggml_backend_cuda_context & ctx, ggml_te // Fast TG path const int64_t n_ids = ids->ne[0]; auto stream = ctx.stream(device_id, 0); - ggml_cuda_pool_alloc dst_up_contiguous(ctx.pool(), sizeof(float)*dst->ne[0]*n_ids); - ggml_cuda_pool_alloc dst_gate_contiguous(ctx.pool(), sizeof(float)*dst->ne[0]*n_ids); auto local_dst = *dst; local_dst.ne[2] = n_ids; @@ -2651,100 +2649,54 @@ static bool ggml_cuda_moe_up_gate_unary(ggml_backend_cuda_context & ctx, ggml_te ggml_backend_buffer_is_cuda(next->buffer) && ((ggml_backend_cuda_buffer_context *)next->buffer->context)->device == device_id; - if (!dst->src[4] && !dst->src[5]) { - local_dst.data = dst_gate_contiguous.get(); - auto the_destination = fuse_next ? &local_dst : dst; - auto unary_op = (ggml_unary_op)dst->op_params[0]; - ggml_cuda_op_fused_mul_mat_vec_q_id(ctx, src0_1, &local_src1, ids, the_destination, - (const char *)src0_1->data, (const char *)src0_2->data, (const float *)src1->data, src1_quantized.get(), - (float *)dst_gate_contiguous.get(), - 0, src0_1->ne[1], 1, src1_padded_col_size, unary_op, stream); - CUDA_CHECK(cudaGetLastError()); - - if (!fuse_next) return false; - } - else { - - local_dst.data = dst_up_contiguous.get(); - ggml_cuda_op_mul_mat_vec_q_id(ctx, src0_1, &local_src1, ids, &local_dst, - (const char *)src0_1->data, (const float *)src1->data, src1_quantized.get(), (float *)dst_up_contiguous.get(), - 0, src0_1->ne[1], 1, src1_padded_col_size, stream); - CUDA_CHECK(cudaGetLastError()); - - if (dst->src[4]) { - ggml_cuda_add_id((const float *)local_dst.data, (const float *)dst->src[4]->data, - (const int32_t *)ids->data, (float *)local_dst.data, - local_dst.ne[0], local_dst.ne[2], local_dst.ne[1], local_dst.ne[0], local_dst.ne[2], - local_dst.nb[1], local_dst.nb[2], dst->src[4]->nb[1], ids->nb[2], stream); - } - - local_dst.data = dst_gate_contiguous.get(); - ggml_cuda_op_mul_mat_vec_q_id(ctx, src0_2, &local_src1, ids, &local_dst, - (const char *)src0_2->data, (const float *)src1->data, src1_quantized.get(), (float *)dst_gate_contiguous.get(), - 0, src0_2->ne[1], 1, src1_padded_col_size, stream); - CUDA_CHECK(cudaGetLastError()); - - if (dst->src[5]) { - ggml_cuda_add_id((const float *)local_dst.data, (const float *)dst->src[5]->data, - (const int32_t *)ids->data, (float *)local_dst.data, - local_dst.ne[0], local_dst.ne[2], local_dst.ne[1], local_dst.ne[0], local_dst.ne[2], - local_dst.nb[1], local_dst.nb[2], dst->src[5]->nb[1], ids->nb[2], stream); - } - } - + auto the_destination = dst; + ggml_cuda_pool_alloc dst_gate_contiguous(ctx.pool()); if (fuse_next) { - - auto unary_op = (ggml_unary_op)dst->op_params[0]; - if (unary_op == GGML_UNARY_OP_SWIGLU_OAI) { - ggml_swiglu_oai_cuda_f32((const float *)dst_gate_contiguous.get(), (const float *)dst_up_contiguous.get(), - (float *)dst_gate_contiguous.get(), dst->ne[0]*n_ids, dst->ne[0], dst->ne[0], dst->ne[0], 1.702f, 7.0f, stream); - CUDA_CHECK(cudaGetLastError()); - } - - const int64_t dst_padded_col_size = GGML_PAD(dst->ne[0], MATRIX_ROW_PADDING); - GGML_ASSERT(dst->ne[0] % QK8_1 == 0); - auto dst_row_size = dst_padded_col_size*sizeof(block_q8_1)/QK8_1; - auto dst_ddq_size = n_ids*dst_row_size; - ggml_cuda_pool_alloc dst_quantized(ctx.pool(), dst_ddq_size); - quantize_row_q8_1_cuda((const float *)dst_gate_contiguous.get(), (void *)dst_quantized.get(), dst->ne[0], n_ids, 1, - dst_padded_col_size, next->src[0]->type, stream); - CUDA_CHECK(cudaGetLastError()); - - local_dst.ne[2] = 1; - - auto local_next = *next; - local_next.ne[2] = local_next.ne[1]; - local_next.ne[1] = local_next.ne[3] = 1; - local_next.nb[2] = local_next.nb[1]; - - local_src1 = *next->src[1]; - local_src1.ne[1] = local_src1.ne[2] = local_src1.ne[3] = 1; - local_src1.nb[1] = local_src1.nb[2] = local_src1.nb[3] = dst_row_size; - - auto local_src0 = *next->src[0]; - local_src0.ne[2] = local_src0.ne[3] = 1; - - CUDA_CHECK(cudaMemsetAsync(next->data, 0, ggml_nbytes(next), stream)); - - ggml_cuda_op_mul_mat_vec_q_id(ctx, &local_src0, &local_src1, ids, &local_next, - (const char *)next->src[0]->data, nullptr, dst_quantized.get(), (float *)next->data, - 0, next->src[0]->ne[1], 1, dst_padded_col_size, stream); - CUDA_CHECK(cudaGetLastError()); - - return true; - } else { - CUDA_CHECK(cudaMemsetAsync(dst->data, 0, ggml_nbytes(dst), stream)); - auto unary_op = (ggml_unary_op)dst->op_params[0]; - if (unary_op == GGML_UNARY_OP_SWIGLU_OAI) { - ggml_swiglu_oai_cuda_f32((const float *)dst_gate_contiguous.get(), (const float *)dst_up_contiguous.get(), - (float *)dst->data, dst->ne[0]*n_ids, dst->ne[0], dst->ne[0], dst->ne[0], 1.702f, 7.0f, stream); - } else { - ggml_fused_mul_unary(ctx, unary_op, ggml_nelements(dst), - (const float *)dst_gate_contiguous.get(), (const float *)dst_up_contiguous.get(), (float *)dst->data); - } - CUDA_CHECK(cudaGetLastError()); - return false; + dst_gate_contiguous.alloc(sizeof(float)*dst->ne[0]*n_ids); + local_dst.data = dst_gate_contiguous.get(); + the_destination = &local_dst; } + auto unary_op = (ggml_unary_op)dst->op_params[0]; + ggml_cuda_op_fused_mul_mat_vec_q_id(ctx, src0_1, &local_src1, ids, the_destination, + dst->src[4], dst->src[5], + (const char *)src0_1->data, (const char *)src0_2->data, (const float *)src1->data, src1_quantized.get(), + (float *)dst_gate_contiguous.get(), + 0, src0_1->ne[1], 1, src1_padded_col_size, unary_op, stream); + CUDA_CHECK(cudaGetLastError()); + + if (!fuse_next) return false; + + const int64_t dst_padded_col_size = GGML_PAD(dst->ne[0], MATRIX_ROW_PADDING); + GGML_ASSERT(dst->ne[0] % QK8_1 == 0); + auto dst_row_size = dst_padded_col_size*sizeof(block_q8_1)/QK8_1; + auto dst_ddq_size = n_ids*dst_row_size; + ggml_cuda_pool_alloc dst_quantized(ctx.pool(), dst_ddq_size); + quantize_row_q8_1_cuda((const float *)dst_gate_contiguous.get(), (void *)dst_quantized.get(), dst->ne[0], n_ids, 1, + dst_padded_col_size, next->src[0]->type, stream); + CUDA_CHECK(cudaGetLastError()); + + local_dst.ne[2] = 1; + + auto local_next = *next; + local_next.ne[2] = local_next.ne[1]; + local_next.ne[1] = local_next.ne[3] = 1; + local_next.nb[2] = local_next.nb[1]; + + local_src1 = *next->src[1]; + local_src1.ne[1] = local_src1.ne[2] = local_src1.ne[3] = 1; + local_src1.nb[1] = local_src1.nb[2] = local_src1.nb[3] = dst_row_size; + + auto local_src0 = *next->src[0]; + local_src0.ne[2] = local_src0.ne[3] = 1; + + CUDA_CHECK(cudaMemsetAsync(next->data, 0, ggml_nbytes(next), stream)); + + ggml_cuda_op_mul_mat_vec_q_id(ctx, &local_src0, &local_src1, ids, &local_next, + (const char *)next->src[0]->data, nullptr, dst_quantized.get(), (float *)next->data, + 0, next->src[0]->ne[1], 1, dst_padded_col_size, stream); + CUDA_CHECK(cudaGetLastError()); + + return true; } } diff --git a/ggml/src/ggml-cuda/mmvq-args.h b/ggml/src/ggml-cuda/mmvq-args.h index 3cf9189b..dac4227d 100644 --- a/ggml/src/ggml-cuda/mmvq-args.h +++ b/ggml/src/ggml-cuda/mmvq-args.h @@ -5,6 +5,8 @@ struct mmvq_args { const void * vx_u; const void * vx_g; + const void * bias_u; + const void * bias_g; const void * vy; float * dst; const char * ids_data; @@ -18,6 +20,7 @@ struct mmvq_args { const uint64_t nb12; const uint64_t nb2; const uint64_t ids_nb0; + const uint64_t bias_nb1; ggml_unary_op unary_op; }; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index d60edcbd..e6f5c4d2 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -145,7 +145,9 @@ static __device__ void mul_mat_vec_q( template static __device__ void fused_mul_mat_vec_q( - const void * __restrict__ vup, const void * __restrict__ vgate, const void * __restrict__ vy, float * __restrict__ dst, + const void * __restrict__ vup, const void * __restrict__ vgate, + const float * __restrict__ bias_u, const float * __restrict__ bias_g, + const void * __restrict__ vy, float * __restrict__ dst, const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst, ggml_unary_op unary_op) { constexpr int qk = ggml_cuda_type_traits::qk; @@ -230,11 +232,20 @@ static __device__ void fused_mul_mat_vec_q( case GGML_UNARY_OP_SILU: r = u*g/(1 + expf(-g)); break; case GGML_UNARY_OP_RELU: r = fmaxf(g, 0.0f) * u; break; // we assume that the supported ops have been checked by the caller - default: { + case GGML_UNARY_OP_GELU: { constexpr float GELU_COEF_A = 0.044715f; constexpr float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; r = 0.5f*g*u*(1.0f + tanhf(SQRT_2_OVER_PI*g*(1.0f + GELU_COEF_A*g*g))); } break; + default: { + constexpr float alpha = 1.702f; + constexpr float limit = 7.0f; + g += bias_g[j*nrows_dst + row0 + threadIdx.x]; + u += bias_u[j*nrows_dst + row0 + threadIdx.x]; + g = fminf(g, limit); + u = fmaxf(fminf(u, limit), -limit); + r = g / (1.0f + expf(-g * alpha)) * (1.0f + u); + } break; } dst[j*nrows_dst + row0 + threadIdx.x] = r; } @@ -270,6 +281,7 @@ __launch_bounds__(nwarps*WARP_SIZE, 1) static __global__ void fused_mul_mat_vec_q( const void * __restrict__ vup, const void * __restrict__ vgate, const void * __restrict__ vy, float * __restrict__ dst, const char * __restrict__ ids_data, + const void * __restrict__ bias_u, const void * __restrict__ bias_g, const uint64_t bias_nb1, const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst, const uint64_t nb02, const uint64_t nb12, const uint64_t nb2, const int64_t ids_nb0, ggml_unary_op unary_op) { @@ -281,8 +293,10 @@ static __global__ void fused_mul_mat_vec_q( } const char * cx_u = (const char *)vup + i02*nb02; const char * cx_g = (const char *)vgate + i02*nb02; + const float * cx_u_b = bias_u ? (const float *)((const char *)bias_u + i02*bias_nb1) : nullptr; + const float * cx_g_b = bias_g ? (const float *)((const char *)bias_g + i02*bias_nb1) : nullptr; const char * cy = (const char *)vy + i2*nb12; - fused_mul_mat_vec_q(cx_u, cx_g, cy, (float *)cdst, ncols_x, nrows_x, nrows_y, nrows_dst, unary_op); + fused_mul_mat_vec_q(cx_u, cx_g, cx_u_b, cx_g_b, cy, (float *)cdst, ncols_x, nrows_x, nrows_y, nrows_dst, unary_op); } template @@ -304,42 +318,42 @@ static void mul_mat_vec_q_cuda_T(const mmvq_args & args, cudaStream_t stream) { switch (args.ncols_y) { case 1: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 2: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 3: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 4: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 5: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 6: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 7: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; case 8: fused_mul_mat_vec_q<<>>(args.vx_u, args.vx_g, args.vy, - args.dst, args.ids_data, + args.dst, args.ids_data, args.bias_u, args.bias_g, args.bias_nb1, args.ncols_x, args.nrows_x, args.nrows_y, args.nrows_dst, args.nb02, args.nb12, args.nb2, args.ids_nb0, args.unary_op); break; default: @@ -492,8 +506,9 @@ static void mul_mat_vec_iq3_s_q8_1_cuda(const mmvq_args & args, cudaStream_t str static void ggml_cuda_op_mul_mat_vec_q_impl(ggml_backend_cuda_context & ctx, ggml_type type, const int64_t ne00, const int64_t ne0, const int64_t ne2, - const int64_t nb02, const int64_t nb12, const int64_t nb2, const int64_t ids_nb0, + const int64_t nb02, const int64_t nb12, const int64_t nb2, const int64_t ids_nb0, const int64_t bias_nb1, const char * src0_dd_u, const char * src0_dd_g, const char * src1_ddq_i, float * dst_dd_i, const char * ids_data, + const void * bias_u, const void * bias_g, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, ggml_unary_op unary_op, cudaStream_t stream) { @@ -524,6 +539,8 @@ static void ggml_cuda_op_mul_mat_vec_q_impl(ggml_backend_cuda_context & ctx, ggm //}; mmvq_args args{/* vx_u */ src0_dd_u, /* vx_g */ src0_dd_g, + /* bias_u */ bias_u, + /* bias_g */ bias_g, /* vy */ src1_ddq_i, /* dst */ dst_dd_i, /* ids_data */ ids_data, @@ -537,6 +554,7 @@ static void ggml_cuda_op_mul_mat_vec_q_impl(ggml_backend_cuda_context & ctx, ggm /* nb12 */ uint64_t(nb12), /* nb2 */ uint64_t(nb2), /* ids_nb0 */ uint64_t(ids_nb0), + /* bias_nb1 */ uint64_t(bias_nb1), /* unary_op */ unary_op }; @@ -656,8 +674,8 @@ void ggml_cuda_op_mul_mat_vec_q_3D( ggml_cuda_op_mul_mat_vec_q_impl(ctx, src0->type, ne00, ne0, dst->ne[2], - src0->nb[2], src1_row_size, dst->nb[2], 0, - src0_dd_i, nullptr, src1_ddq_i, dst_dd_i, nullptr, + src0->nb[2], src1_row_size, dst->nb[2], 0, 0, + src0_dd_i, nullptr, src1_ddq_i, dst_dd_i, nullptr, nullptr, nullptr, row_low, row_high, src1_ncols, src1_padded_row_size, GGML_UNARY_OP_COUNT, stream); @@ -677,8 +695,8 @@ void ggml_cuda_op_mul_mat_vec_q( const int64_t ne0 = dst->ne[0]; ggml_cuda_op_mul_mat_vec_q_impl(ctx, src0->type, - ne00, ne0, 1, 0, 0, 0, 0, - src0_dd_i, nullptr, src1_ddq_i, dst_dd_i, nullptr, + ne00, ne0, 1, 0, 0, 0, 0, 0, + src0_dd_i, nullptr, src1_ddq_i, dst_dd_i, nullptr, nullptr, nullptr, row_low, row_high, src1_ncols, src1_padded_row_size, GGML_UNARY_OP_COUNT, stream); @@ -703,8 +721,8 @@ void ggml_cuda_op_mul_mat_vec_q_id( ggml_cuda_op_mul_mat_vec_q_impl(ctx, src0->type, ne00, ne0, dst->ne[2], - src0->nb[2], src1->nb[2], dst->nb[2], ids->nb[0], - src0_dd_i, nullptr, src1_ddq_i, dst_dd_i, (const char *)ids->data, + src0->nb[2], src1->nb[2], dst->nb[2], ids->nb[0], 0, + src0_dd_i, nullptr, src1_ddq_i, dst_dd_i, (const char *)ids->data, nullptr, nullptr, row_low, row_high, src1_ncols, src1_padded_row_size, GGML_UNARY_OP_COUNT, stream); @@ -713,13 +731,23 @@ void ggml_cuda_op_mul_mat_vec_q_id( void ggml_cuda_op_fused_mul_mat_vec_q_id(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + const ggml_tensor * bias_u, const ggml_tensor * bias_g, const char * src0_dd_u, const char * src0_dd_g, const float * src1_ddf_i, const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, ggml_unary_op unary_op, cudaStream_t stream) { - GGML_ASSERT(unary_op == GGML_UNARY_OP_SILU || - unary_op == GGML_UNARY_OP_RELU || - unary_op == GGML_UNARY_OP_GELU); + if (!bias_u && !bias_g) { + GGML_ASSERT(unary_op == GGML_UNARY_OP_SILU || + unary_op == GGML_UNARY_OP_RELU || + unary_op == GGML_UNARY_OP_GELU); + } else { + GGML_ASSERT(unary_op == GGML_UNARY_OP_SWIGLU_OAI); + GGML_ASSERT(bias_u && bias_g); + GGML_ASSERT(bias_u->data && bias_g->data); + GGML_ASSERT(bias_u->nb[1] == bias_g->nb[1]); + GGML_ASSERT(bias_u->ne[0] == dst->ne[0]); + GGML_ASSERT(bias_g->ne[0] == dst->ne[0]); + } GGML_ASSERT(src0_dd_u && src0_dd_g && ids && ids->data); const int64_t ne00 = src0->ne[0]; @@ -733,8 +761,9 @@ void ggml_cuda_op_fused_mul_mat_vec_q_id(ggml_backend_cuda_context & ctx, ggml_cuda_op_mul_mat_vec_q_impl(ctx, src0->type, ne00, ne0, dst->ne[2], - src0->nb[2], src1->nb[2], dst->nb[2], ids->nb[0], + src0->nb[2], src1->nb[2], dst->nb[2], ids->nb[0], bias_u ? bias_u->nb[1] : 0, src0_dd_u, src0_dd_g, src1_ddq_i, dst_dd_i, (const char *)ids->data, + bias_u ? bias_u->data : nullptr, bias_g ? bias_g->data : nullptr, row_low, row_high, src1_ncols, src1_padded_row_size, unary_op, stream); diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 94542f8a..09c4962c 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -28,6 +28,7 @@ void ggml_cuda_op_mul_mat_vec_q_id(ggml_backend_cuda_context & ctx, void ggml_cuda_op_fused_mul_mat_vec_q_id(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + const ggml_tensor * bias_u, const ggml_tensor * bias_g, const char * src0_dd_u, const char * src0_dd_g, const float * src1_ddf_i, const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, ggml_unary_op unary_op, cudaStream_t stream);