ggml: add SWIGLU_CLAMP (#27930)

* ggml: add SWIGLU_CLAMP

* add vulkan shader
This commit is contained in:
Aman Gupta
2026-08-30 23:00:02 +08:00
committed by GitHub
parent 2578138397
commit 0190529ec4
40 changed files with 768 additions and 35 deletions
+2 -1
View File
@@ -1539,6 +1539,7 @@ struct ggml_cuda_mm_fusion_args_host {
const ggml_tensor * x_scale = nullptr;
const ggml_tensor * gate_scale = nullptr;
ggml_glu_op glu_op;
float glu_limit = 0.0f;
};
struct ggml_cuda_mm_fusion_args_device {
const void * x_bias = nullptr;
@@ -1547,6 +1548,7 @@ struct ggml_cuda_mm_fusion_args_device {
const void * x_scale = nullptr;
const void * gate_scale = nullptr;
ggml_glu_op glu_op;
float glu_limit = 0.0f;
};
struct ggml_cuda_kernel_launch_params {
@@ -1673,4 +1675,3 @@ static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_ke
kernel<<<launch_params.block_nums, launch_params.block_dims, launch_params.shmem, launch_params.stream>>>(std::forward<Args>(args)... );
CUDA_CHECK(cudaGetLastError());
}
+15 -5
View File
@@ -1744,7 +1744,7 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up,
return false;
}
static constexpr std::array<ggml_glu_op, 3> valid_glu_ops = { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU_OAI };
static constexpr std::array<ggml_glu_op, 4> valid_glu_ops = { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU_OAI, GGML_GLU_OP_SWIGLU_CLAMP };
if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) {
return false;
@@ -2203,6 +2203,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_GLU_OP_GEGLU_QUICK:
ggml_cuda_op_geglu_quick(ctx, dst);
break;
case GGML_GLU_OP_SWIGLU_CLAMP:
ggml_cuda_op_swiglu_clamp(ctx, dst);
break;
default:
return false;
}
@@ -3595,6 +3598,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
fusion_data.x_scale = up_scale;
fusion_data.gate_scale = gate_scale;
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) {
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
@@ -3688,6 +3692,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
fusion_data.x_scale = up_scale;
fusion_data.gate_scale = gate_scale;
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) {
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
@@ -3744,6 +3749,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
fusion_data.x_bias = up_bias_tensor;
fusion_data.gate_bias = gate_bias_tensor;
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
fused_mul_mat_vec = true;
@@ -3757,6 +3763,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
fusion_data.x_bias = up_bias_tensor;
fusion_data.gate_bias = gate_bias_tensor;
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
fused_mul_mat_vec = true;
@@ -3781,8 +3788,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) {
ggml_cuda_mm_fusion_args_host fusion_data{};
fusion_data.gate = gate->src[0];
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.gate = gate->src[0];
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
fused_mul_mat_vec = true;
@@ -3792,8 +3800,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) {
ggml_cuda_mm_fusion_args_host fusion_data{};
fusion_data.gate = gate->src[0];
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.gate = gate->src[0];
fusion_data.glu_op = ggml_get_glu_op(glu);
fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
fused_mul_mat_vec = true;
@@ -4919,6 +4928,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
case GGML_GLU_OP_SWIGLU_OAI:
case GGML_GLU_OP_GEGLU_ERF:
case GGML_GLU_OP_GEGLU_QUICK:
case GGML_GLU_OP_SWIGLU_CLAMP:
return ggml_is_contiguous_1(op->src[0]);
default:
return false;
+7 -1
View File
@@ -56,6 +56,7 @@ 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;
float glu_limit = 0.0f;
const T * gate_x = nullptr;
const float * x_bias = nullptr;
const float * gate_bias = nullptr;
@@ -65,6 +66,7 @@ static __global__ void mul_mat_vec_f(
use_bias = fusion.x_bias != nullptr;
use_gate_bias = fusion.gate_bias != nullptr;
glu_op = fusion.glu_op;
glu_limit = fusion.glu_limit;
if (use_gate) {
gate_x = static_cast<const T *>(fusion.gate);
@@ -365,6 +367,9 @@ static __global__ void mul_mat_vec_f(
value = ggml_cuda_op_swiglu_oai_single(gate_value, value);
break;
}
case GGML_GLU_OP_SWIGLU_CLAMP:
value = ggml_cuda_op_swiglu_clamp_single(gate_value, value, glu_limit);
break;
default:
break;
}
@@ -374,7 +379,7 @@ static __global__ void mul_mat_vec_f(
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, glu_limit, gate_x, x_bias, gate_bias, sumf_gate);
}
}
@@ -675,6 +680,7 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor
fusion_local.gate_bias = fusion->gate_bias->data;
}
fusion_local.glu_op = fusion->glu_op;
fusion_local.glu_limit = fusion->glu_limit;
}
const int64_t s01 = src0->nb[1] / ts_src0;
+7 -1
View File
@@ -595,6 +595,7 @@ static __global__ void mul_mat_vec_q(
const float * x_scale = nullptr;
const float * gate_scale = nullptr;
ggml_glu_op active_glu;
float glu_limit = 0.0f;
if constexpr (has_fusion) {
use_gate = fusion.gate != nullptr;
@@ -604,6 +605,7 @@ 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;
glu_limit = fusion.glu_limit;
if constexpr (type == GGML_TYPE_NVFP4) {
use_scale = fusion.x_scale != nullptr;
use_gate_scale = fusion.gate_scale != nullptr && use_gate;
@@ -745,6 +747,9 @@ static __global__ void mul_mat_vec_q(
case GGML_GLU_OP_SWIGLU_OAI:
result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
break;
case GGML_GLU_OP_SWIGLU_CLAMP:
result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit);
break;
default:
result = result * gate_value;
break;
@@ -757,7 +762,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, glu_limit, gate_bias, x_bias, x_scale, gate_scale, tmp_gate);
}
if constexpr (type != GGML_TYPE_NVFP4) {
GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales);
@@ -1310,6 +1315,7 @@ void ggml_cuda_mul_mat_vec_q(
fusion_local.gate_scale = fusion->gate_scale->data;
}
fusion_local.glu_op = fusion->glu_op;
fusion_local.glu_limit = fusion->glu_limit;
}
// If src0 is a temporary compute buffer, clear any potential padding.
+75
View File
@@ -427,6 +427,81 @@ void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
swiglu_oai_cuda(src0_p, src1_p, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), alpha, limit, stream);
}
// swiglu_clamp
template <typename T>
static __global__ void swiglu_clamp_kernel(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, float limit) {
const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x;
if (i >= k) {
return;
}
const int64_t j0 = (i / n) * o0 + (i % n);
const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
dst[i] = (T) ggml_cuda_op_swiglu_clamp_single((float) gate[j0], (float) up[j1], limit);
}
template <typename T>
static void swiglu_clamp_cuda(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, const float limit, cudaStream_t stream) {
const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE;
swiglu_clamp_kernel<<<num_blocks, CUDA_GLU_BLOCK_SIZE, 0, stream>>>(gate, up, dst, k, n, o0, o1, limit);
}
void ggml_cuda_op_swiglu_clamp(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
void * src0_d = src0->data;
void * src1_d = src1 ? src1->data : src0->data;
const int64_t src0_o = src0->nb[1];
const int64_t src1_o = src1 ? src1->nb[1] : src0->nb[1];
void * dst_d = dst->data;
const int64_t nc = src1 ? src0->ne[0] : src0->ne[0] / 2;
cudaStream_t stream = ctx.stream();
GGML_ASSERT(ggml_is_contiguous_1(src0));
GGML_ASSERT(src0->nb[0] == ggml_element_size(src0));
GGML_ASSERT(ggml_is_contiguous(dst));
GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
GGML_ASSERT(src0->type == dst->type);
GGML_ASSERT(dst->ne[0] == nc);
GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0));
if (src1) {
GGML_ASSERT(ggml_is_contiguous_1(src1));
GGML_ASSERT(src1->nb[0] == ggml_element_size(src1));
GGML_ASSERT(src1->ne[0] == nc);
GGML_ASSERT(src0->type == src1->type);
}
const int32_t swapped = ggml_get_op_params_i32(dst, 1);
const float limit = ggml_get_op_params_f32(dst, 3);
if (src0->type == GGML_TYPE_F16) {
half * src0_p = (half *) src0_d;
half * src1_p = (half *) src1_d;
if (!src1) {
src0_p += swapped ? nc : 0;
src1_p += swapped ? 0 : nc;
}
swiglu_clamp_cuda(src0_p, src1_p, (half *) dst_d, ggml_nelements(dst), nc, src0_o / sizeof(half), src1_o / sizeof(half), limit, stream);
} else {
float * src0_p = (float *) src0_d;
float * src1_p = (float *) src1_d;
if (!src1) {
src0_p += swapped ? nc : 0;
src1_p += swapped ? 0 : nc;
}
swiglu_clamp_cuda(src0_p, src1_p, (float *) dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), limit, stream);
}
}
/* CUDA kernel + launcher for xIELU */
template <typename T>
+9
View File
@@ -83,6 +83,8 @@ void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_swiglu_clamp(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_geglu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
@@ -112,3 +114,10 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g
out_glu = out_glu * (1.0f + g);
return out_glu;
}
__device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, float up, float limit) {
gate = fminf(gate, limit);
up = fmaxf(fminf(up, limit), -limit);
return ggml_cuda_op_silu_single(gate) * up;
}