Skip to content

feat(cuda): fuse activations and residual add into mmv f/q epilogues #67

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
roberteg16 wants to merge 1 commit into
base: gfx11
Choose a base branch
from
Draft

feat(cuda): fuse activations and residual add into mmv f/q epilogues #67

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions ggml/src/ggml-cuda/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 {
Expand Down
153 changes: 153 additions & 0 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
40 changes: 37 additions & 3 deletions ggml/src/ggml-cuda/mmvf.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<const T *>(fusion.gate);
Expand All @@ -78,6 +83,9 @@ static __global__ void mul_mat_vec_f(
} else {
use_gate_bias = false;
}
if (use_act_mul) {
act_mul = static_cast<const float *>(fusion.activation_mul);
}
}

if (use_gate) {
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}
}

Expand All @@ -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<T, type_acc, ncols_dst, block_size, true, is_multi_token_id>, launch_params,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down
40 changes: 38 additions & 2 deletions ggml/src/ggml-cuda/mmvq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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) {
Expand All @@ -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
Expand Down Expand Up @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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.
Expand Down
4 changes: 4 additions & 0 deletions ggml/src/ggml-cuda/unary.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading