diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index ac3da98627bc..1e581f3db5a7 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1516,6 +1516,10 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * x_scale = nullptr; const ggml_tensor * gate_scale = nullptr; ggml_glu_op glu_op; + // non-gated unary activation applied to the matmul output; COUNT = none + ggml_unary_op activation = GGML_UNARY_OP_COUNT; + // optional per-row operand multiplied after the activation (matmul -> act -> mul) + const ggml_tensor * activation_mul = nullptr; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; @@ -1524,6 +1528,8 @@ struct ggml_cuda_mm_fusion_args_device { const void * x_scale = nullptr; const void * gate_scale = nullptr; ggml_glu_op glu_op; + ggml_unary_op activation = GGML_UNARY_OP_COUNT; + const void * activation_mul = nullptr; }; struct ggml_cuda_kernel_launch_params { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 227487c27eda..168091b2a1d3 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3746,6 +3746,159 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return fused_node_count - 1; } + auto tensors_overlap = [](const ggml_tensor * a, const ggml_tensor * b) { + const int64_t a_start = (int64_t) a->data; + const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + const int64_t b_start = (int64_t) b->data; + const int64_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + return (b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end); + }; + + // mul_mat + reshape + add: same as mul_mat + add but tolerating a view between the + // projection and the residual add (e.g. GDN layers reshape the attn output first). + if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ADD }, { i + 2 })) { + ggml_tensor * mm_node = cgraph->nodes[i]; + ggml_tensor * reshape_node = cgraph->nodes[i + 1]; + ggml_tensor * add_node = cgraph->nodes[i + 2]; + + const ggml_tensor * bias_tensor = nullptr; + if (add_node->src[0] == reshape_node) { + bias_tensor = add_node->src[1]; + } else if (add_node->src[1] == reshape_node) { + bias_tensor = add_node->src[0]; + } + + if (bias_tensor && reshape_node->src[0] == mm_node && bias_tensor->type == GGML_TYPE_F32 && + bias_tensor->ne[0] == mm_node->ne[0] && + ggml_is_contiguous(mm_node) && ggml_is_contiguous(add_node) && + ggml_nelements(add_node) == ggml_nelements(mm_node) && + ggml_are_same_shape(add_node->src[0], add_node->src[1])) { + + const ggml_tensor * src0 = mm_node->src[0]; + const ggml_tensor * src1 = mm_node->src[1]; + const ggml_tensor * ids = mm_node->src[2]; + + // dst may alias the residual (in-place add, same-index) but not the matmul inputs + if (!tensors_overlap(add_node, src0) && !tensors_overlap(add_node, src1)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.x_bias = bias_tensor; + + ggml_tensor dst_tensor = *mm_node; + dst_tensor.data = add_node->data; + dst_tensor.buffer = add_node->buffer; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) { + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, &dst_tensor, &fusion_data); + return 2; + } + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, &dst_tensor, &fusion_data); + return 2; + } + } + } + } + + // mul_mat + optional reshape + non-gated unary activation (sigmoid/silu), + // optionally followed by an elementwise mul (act(mul_mat) * y). The variant that + // also folds the trailing mul is tried first so it wins over the plain unary_mul. + for (int with_mul = 1; with_mul >= 0; --with_mul) { + for (int with_reshape = 0; with_reshape < 2; ++with_reshape) { + const int unary_off = with_reshape ? 2 : 1; + const int mul_off = unary_off + 1; + const int n_ops = (with_mul ? mul_off : unary_off) + 1; + + ggml_op ops[4]; + int k = 0; + ops[k++] = GGML_OP_MUL_MAT; + if (with_reshape) { + ops[k++] = GGML_OP_RESHAPE; + } + ops[k++] = GGML_OP_UNARY; + if (with_mul) { + ops[k++] = GGML_OP_MUL; + } + + const int out_nodes[] = { i + n_ops - 1 }; + if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1)) { + continue; + } + + ggml_tensor * mm_node = cgraph->nodes[i]; + ggml_tensor * unary_node = cgraph->nodes[i + unary_off]; + ggml_tensor * out_node = cgraph->nodes[i + n_ops - 1]; + + const ggml_unary_op uop = ggml_get_unary_op(unary_node); + if (uop != GGML_UNARY_OP_SIGMOID && uop != GGML_UNARY_OP_SILU) { + continue; + } + + if (with_reshape) { + const ggml_tensor * reshape_node = cgraph->nodes[i + 1]; + if (reshape_node->src[0] != mm_node || unary_node->src[0] != reshape_node) { + continue; + } + } else if (unary_node->src[0] != mm_node) { + continue; + } + + // matmul output and the fused output must be contiguous and hold the same + // number of elements, so writing the result in matmul order is correct + if (!ggml_is_contiguous(mm_node) || !ggml_is_contiguous(out_node) || + ggml_nelements(out_node) != ggml_nelements(mm_node)) { + continue; + } + + const ggml_tensor * src0 = mm_node->src[0]; + const ggml_tensor * src1 = mm_node->src[1]; + const ggml_tensor * ids = mm_node->src[2]; + + const ggml_tensor * act_mul = nullptr; + if (with_mul) { + if (out_node->src[0] == unary_node) { + act_mul = out_node->src[1]; + } else if (out_node->src[1] == unary_node) { + act_mul = out_node->src[0]; + } else { + continue; + } + // no broadcast, and same per-row layout as the matmul output + if (!ggml_are_same_shape(out_node->src[0], out_node->src[1]) || + !ggml_is_contiguous(act_mul) || act_mul->type != GGML_TYPE_F32) { + continue; + } + // The fused kernel writes out_node while reading the matmul inputs across + // the full reduction; out_node must not alias them. Aliasing act_mul is + // safe (elementwise, same-index read before write, e.g. an in-place mul). + if (tensors_overlap(out_node, src0) || tensors_overlap(out_node, src1)) { + continue; + } + } else if (!ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) { + continue; + } + + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.activation = uop; + fusion_data.activation_mul = act_mul; + + // use the matmul geometry but write into the fused output buffer; the + // intermediate reshape/activation/mul nodes are elided + ggml_tensor dst_tensor = *mm_node; + dst_tensor.data = out_node->data; + dst_tensor.buffer = out_node->buffer; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) { + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, &dst_tensor, &fusion_data); + return n_ops - 1; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, &dst_tensor, &fusion_data); + return n_ops - 1; + } + } + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); return 2; diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index d7dbc8b99282..1f1498d25ad9 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -56,15 +56,20 @@ static __global__ void mul_mat_vec_f( bool use_bias = false; bool use_gate_bias = false; ggml_glu_op glu_op = ggml_glu_op::GGML_GLU_OP_SWIGLU; + ggml_unary_op activation = GGML_UNARY_OP_COUNT; + bool use_act_mul = false; const T * gate_x = nullptr; const float * x_bias = nullptr; const float * gate_bias = nullptr; + const float * act_mul = nullptr; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; use_bias = fusion.x_bias != nullptr; use_gate_bias = fusion.gate_bias != nullptr; glu_op = fusion.glu_op; + activation = fusion.activation; + use_act_mul = fusion.activation_mul != nullptr; if (use_gate) { gate_x = static_cast(fusion.gate); @@ -78,6 +83,9 @@ static __global__ void mul_mat_vec_f( } else { use_gate_bias = false; } + if (use_act_mul) { + act_mul = static_cast(fusion.activation_mul); + } } if (use_gate) { @@ -92,6 +100,9 @@ static __global__ void mul_mat_vec_f( if (use_gate_bias) { gate_bias += int64_t(sample_dst)*stride_sample_dst + channel_bias*stride_channel_dst; } + if (use_act_mul) { + act_mul += int64_t(sample_dst)*stride_sample_dst + channel_bias*stride_channel_dst; + } } const float2 * y2 = (const float2 *) y; @@ -368,13 +379,27 @@ static __global__ void mul_mat_vec_f( default: break; } + } else { + switch (activation) { + case GGML_UNARY_OP_SIGMOID: + value = ggml_cuda_op_sigmoid_single(value); + break; + case GGML_UNARY_OP_SILU: + value = ggml_cuda_op_silu_single(value); + break; + default: + break; + } + if (use_act_mul) { + value *= act_mul[tid*stride_col_dst + row]; + } } } dst[tid*stride_col_dst + row] = value; if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, gate_x, x_bias, gate_bias, sumf_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, activation, use_act_mul, gate_x, x_bias, gate_bias, act_mul, sumf_gate); } } @@ -389,7 +414,8 @@ static void mul_mat_vec_f_switch_fusion( const ggml_cuda_kernel_launch_params launch_params = {block_nums, block_dims, nbytes_shared, stream}; - const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || + fusion.activation != GGML_UNARY_OP_COUNT || fusion.activation_mul != nullptr; if constexpr (ncols_dst == 1) { if (has_fusion) { ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, @@ -444,7 +470,8 @@ void launch_mul_mat_vec_f_cuda( } } - const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || + fusion.activation != GGML_UNARY_OP_COUNT || fusion.activation_mul != nullptr; const int nbytes_shared = warp_size*sizeof(float) + (has_fusion ? warp_size*sizeof(float) : 0); const dim3 block_nums(nrows, nchannels_dst, nsamples_or_ntokens); @@ -674,7 +701,14 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]); fusion_local.gate_bias = fusion->gate_bias->data; } + if (fusion->activation_mul) { + GGML_ASSERT(fusion->activation_mul->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(fusion->activation_mul)); + GGML_ASSERT(ggml_nelements(fusion->activation_mul) == ggml_nelements(dst)); + fusion_local.activation_mul = fusion->activation_mul->data; + } fusion_local.glu_op = fusion->glu_op; + fusion_local.activation = fusion->activation; } const int64_t s01 = src0->nb[1] / ts_src0; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index e18ada5377d5..4306689fd0b2 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -529,6 +529,9 @@ static __global__ void mul_mat_vec_q( const float * x_scale = nullptr; const float * gate_scale = nullptr; ggml_glu_op active_glu; + ggml_unary_op active_activation = GGML_UNARY_OP_COUNT; + bool use_act_mul = false; + const float * act_mul = nullptr; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -538,6 +541,9 @@ static __global__ void mul_mat_vec_q( x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; + active_activation = fusion.activation; + use_act_mul = fusion.activation_mul != nullptr; + act_mul = (const float *) fusion.activation_mul; if constexpr (type == GGML_TYPE_NVFP4) { use_scale = fusion.x_scale != nullptr; use_gate_scale = fusion.gate_scale != nullptr && use_gate; @@ -549,6 +555,7 @@ static __global__ void mul_mat_vec_q( [[maybe_unused]] float x_biases[ncols_dst] = { 0.0f }; [[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f }; + [[maybe_unused]] float act_muls[ncols_dst] = { 0.0f }; [[maybe_unused]] float x_scales = 1.0f; [[maybe_unused]] float gate_scales = 1.0f; if constexpr (has_fusion) { @@ -564,6 +571,13 @@ static __global__ void mul_mat_vec_q( x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x]; } } + if (use_act_mul) { + act_mul = act_mul + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0; +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + act_muls[j] = act_mul[j * stride_col_dst + threadIdx.x]; + } + } if (use_gate_bias) { gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0; #pragma unroll @@ -683,6 +697,20 @@ static __global__ void mul_mat_vec_q( result = result * gate_value; break; } + } else { + switch (active_activation) { + case GGML_UNARY_OP_SIGMOID: + result = ggml_cuda_op_sigmoid_single(result); + break; + case GGML_UNARY_OP_SILU: + result = ggml_cuda_op_silu_single(result); + break; + default: + break; + } + if (use_act_mul) { + result *= act_muls[j]; + } } } dst[j*stride_col_dst + i] = result; @@ -691,7 +719,7 @@ static __global__ void mul_mat_vec_q( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, gate_bias, x_bias, x_scale, gate_scale, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, active_activation, use_act_mul, gate_bias, x_bias, x_scale, gate_scale, act_mul, act_muls, tmp_gate); } if constexpr (type != GGML_TYPE_NVFP4) { GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales); @@ -791,7 +819,8 @@ static void mul_mat_vec_q_switch_fusion( const uint32_t ids_stride, cudaStream_t stream) { const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || - fusion.x_scale != nullptr || fusion.gate_scale != nullptr; + fusion.x_scale != nullptr || fusion.gate_scale != nullptr || + fusion.activation != GGML_UNARY_OP_COUNT || fusion.activation_mul != nullptr; if constexpr (c_ncols_dst == 1) { if (has_fusion) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); @@ -1205,7 +1234,14 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); fusion_local.gate_scale = fusion->gate_scale->data; } + if (fusion->activation_mul) { + GGML_ASSERT(fusion->activation_mul->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(fusion->activation_mul)); + GGML_ASSERT(ggml_nelements(fusion->activation_mul) == ggml_nelements(dst)); + fusion_local.activation_mul = fusion->activation_mul->data; + } fusion_local.glu_op = fusion->glu_op; + fusion_local.activation = fusion->activation; } // If src0 is a temporary compute buffer, clear any potential padding. diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 81ed873ecc30..4e669f0d0028 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -97,6 +97,10 @@ __device__ __forceinline__ float ggml_cuda_op_silu_single(float x) { return x / (1.0f + expf(-x)); } +__device__ __forceinline__ float ggml_cuda_op_sigmoid_single(float x) { + return 1.0f / (1.0f + expf(-x)); +} + __device__ __forceinline__ float ggml_cuda_op_gelu_single(float x) { const float GELU_COEF_A = 0.044715f; const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;