From ff5a23376725e5a889fa0bfb987b7541dc304da0 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sun, 18 Oct 2026 23:18:47 +0530 Subject: [PATCH] cuda : fuse qwen4exp's SCALE -> SIGMOID -> SCALE -> hc_post chain (#30602) Port of the Vulkan HC_POST_GATE fusion (#29520). hc_combine builds the hc_post weights as 2*sigmoid(inject/hc); dsv4_hc_post_f32 now takes a gated template flag and computes scale_out*sigmoid(scale_in*p) from the raw weights, so the chain runs as one kernel instead of four. --- ggml/src/ggml-cuda/dsv4-hc.cu | 31 ++++++++++++++++++++++++------- ggml/src/ggml-cuda/dsv4-hc.cuh | 4 +++- ggml/src/ggml-cuda/ggml-cuda.cu | 25 +++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/dsv4-hc.cu b/ggml/src/ggml-cuda/dsv4-hc.cu index ca1d2dc8a..2d0b3cfea 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cu +++ b/ggml/src/ggml-cuda/dsv4-hc.cu @@ -145,7 +145,8 @@ static __global__ void dsv4_hc_pre_f32( dst[i0*sd0 + it*sd1] = scale * sum; } -template +// gated: post = gate_scale_out*sigmoid(gate_scale_in*post), the fused SCALE -> SIGMOID -> SCALE chain of qwen4exp +template static __global__ void dsv4_hc_post_f32( const float * x, const float * residual, @@ -167,7 +168,9 @@ static __global__ void dsv4_hc_post_f32( int64_t sc2, int64_t sd0, int64_t sd1, - int64_t sd2) { + int64_t sd2, + float gate_scale_in, + float gate_scale_out) { ggml_cuda_pdl_lc(); const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; const int64_t nr = n_embd * hc * n_tokens; @@ -182,7 +185,12 @@ static __global__ void dsv4_hc_post_f32( const int64_t idst = (ir / n_embd) % hc; const int64_t it = ir / (n_embd * hc); - float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1]; + float p = post[idst*sp0 + it*sp1]; + if constexpr (gated) { + p = gate_scale_out * (1.0f / (1.0f + expf(-(gate_scale_in * p)))); + } + + float sum = x[i0*sx0 + it*sx1] * p; if constexpr (has_comb) { for (int64_t isrc = 0; isrc < hc; ++isrc) { sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2]; @@ -272,12 +280,19 @@ void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst scale); } -void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst, + const ggml_tensor * gate_scale_in, const ggml_tensor * gate_scale_out) { const ggml_tensor * x = dst->src[0]; const ggml_tensor * residual = dst->src[1]; - const ggml_tensor * post = dst->src[2]; const ggml_tensor * comb = dst->src[3]; + // when fused, read the raw post weights from before the SCALE -> SIGMOID -> SCALE chain + const bool gated = gate_scale_in != nullptr; + const ggml_tensor * post = gated ? gate_scale_in->src[0] : dst->src[2]; + + const float gate_in = gated ? ggml_get_op_params_f32(gate_scale_in, 0) : 1.0f; + const float gate_out = gated ? ggml_get_op_params_f32(gate_scale_out, 0) : 1.0f; + GGML_ASSERT(x->type == GGML_TYPE_F32); GGML_ASSERT(residual->type == GGML_TYPE_F32); GGML_ASSERT(post->type == GGML_TYPE_F32); @@ -303,7 +318,8 @@ void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * ds const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream()); - auto kernel = comb ? dsv4_hc_post_f32 : dsv4_hc_post_f32; + auto kernel = comb ? (gated ? dsv4_hc_post_f32 : dsv4_hc_post_f32) + : (gated ? dsv4_hc_post_f32 : dsv4_hc_post_f32); ggml_cuda_kernel_launch(kernel, launch_params, (const float *) x->data, (const float *) residual->data, (const float *) post->data, comb ? (const float *) comb->data : nullptr, (float *) dst->data, @@ -312,5 +328,6 @@ void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * ds nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float), nbp0 / sizeof(float), nbp1 / sizeof(float), nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float), - nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float)); + nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float), + gate_in, gate_out); } diff --git a/ggml/src/ggml-cuda/dsv4-hc.cuh b/ggml/src/ggml-cuda/dsv4-hc.cuh index 2379aaefb..1885e7665 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cuh +++ b/ggml/src/ggml-cuda/dsv4-hc.cuh @@ -3,4 +3,6 @@ void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// gate_scale_in/gate_scale_out: the SCALE nodes of a fused SCALE -> SIGMOID -> SCALE chain feeding src[2] +void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst, + const ggml_tensor * gate_scale_in = nullptr, const ggml_tensor * gate_scale_out = nullptr); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index dc30492fa..2eb5e203e 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3550,6 +3550,26 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + // scale_out*sigmoid(scale_in*post) as the hc_post weights (qwen4exp hc_combine) + std::initializer_list hc_post_gate_ops = { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE, GGML_OP_DSV4_HC_POST }; + + if (is_equal(hc_post_gate_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 3 })) { + const ggml_tensor * scale_in = cgraph->nodes[node_idx]; + const ggml_tensor * sigmoid = cgraph->nodes[node_idx + 1]; + const ggml_tensor * scale_out = cgraph->nodes[node_idx + 2]; + + // the chain must feed hc_post->src[2]; a bias on either scale is not handled + if (ggml_check_edges(cgraph, node_idx, {{1, 0, 0}, {2, 0, 1}, {3, 2, 2}}) && + ggml_get_unary_op(sigmoid) == GGML_UNARY_OP_SIGMOID && + ggml_get_op_params_f32(scale_in, 1) == 0.0f && ggml_get_op_params_f32(scale_out, 1) == 0.0f && + scale_in->src[0]->type == GGML_TYPE_F32 && scale_out->type == GGML_TYPE_F32 && + ggml_are_same_shape(scale_in->src[0], scale_out)) { + int out_nodes[] = { node_idx + 3 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); + } + return false; + } + if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; } @@ -4513,6 +4533,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 1; } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE, GGML_OP_DSV4_HC_POST }, { GGML_UNARY_OP_SIGMOID })) { + ggml_cuda_op_dsv4_hc_post(*cuda_ctx, cgraph->nodes[i + 3], node, cgraph->nodes[i + 2]); + return 3; + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE }, { GGML_UNARY_OP_TANH })) { ggml_cuda_op_softcap(*cuda_ctx, cgraph->nodes[i + 2], node); return 2;