CUDA: fuse ffn_gate + glu into the mul_mat_q epilogue

ggml/src/ggml-cuda/common.cuh

@@ -1531,6 +1531,8 @@ struct ggml_cuda_mm_fusion_args_host {
     const ggml_tensor * gate_bias = nullptr;
     const ggml_tensor * x_scale = nullptr;
     const ggml_tensor * gate_scale = nullptr;
+    // Already-computed other half of a GLU pair, folded into this matmul's epilogue by MMQ.
+    const ggml_tensor * glu_up = nullptr;
     ggml_glu_op glu_op;
 };
 struct ggml_cuda_mm_fusion_args_device {

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -1780,6 +1780,24 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) {
     return use_mul_mat_vec_f;
 }
 
+// Counterpart of ggml_cuda_should_fuse_mul_mat_vec_q for the prefill path
+static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * tensor) {
+    const ggml_tensor * src0 = tensor->src[0];
+    const ggml_tensor * src1 = tensor->src[1];
+    const ggml_tensor * dst  = tensor;
+
+    const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
+                                   ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) &&
+                                   src0->view_src;
+
+    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+
+    return ggml_is_quantized(src0->type) && !bad_padding_clear &&
+           src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 &&
+           !ggml_cuda_should_use_mmvq(src0->type, cc, src1->ne[1]) &&
+           ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts =*/ 0);
+}
+
 static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
     ggml_tensor *       src0 = tensor->src[0];
     ggml_tensor *       src1 = tensor->src[1];
@@ -3787,6 +3805,27 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                 fused_node_count  = 3;
                 break;
             }
+
+            // Run the up matmul normally, then fuse the GLU into the gate matmul's epilogue,
+            // which reads the already-written up output
+            {
+                const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+                const ggml_glu_op glu_op = ggml_get_glu_op(glu);
+                const bool glu_ok = (glu_op == GGML_GLU_OP_SWIGLU || glu_op == GGML_GLU_OP_GEGLU);
+
+                // Each matmul is dispatched separately, so both have to be eligible.
+                if (op == GGML_OP_MUL_MAT && glu->ne[1] > 1 && glu_ok &&
+                        ggml_cuda_should_fuse_mul_mat_q(gate) && ggml_cuda_should_fuse_mul_mat_q(up) &&
+                        ggml_cuda_mmq_can_fuse_glu(gate->src[0]->type, cc)) {
+                    ggml_cuda_mm_fusion_args_host fusion_data{};
+                    fusion_data.glu_up = up;
+                    fusion_data.glu_op = glu_op;
+
+                    ggml_cuda_mul_mat_q(*cuda_ctx, up->src[0],   up->src[1],   up->src[2],   up);
+                    ggml_cuda_mul_mat_q(*cuda_ctx, gate->src[0], gate->src[1], gate->src[2], glu, &fusion_data);
+                    return 2; // skip next 2 as 3 nodes are handled
+                }
+            }
         }
     }
 

ggml/src/ggml-cuda/mmq.cu

@@ -83,7 +83,15 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
 }
 
 void ggml_cuda_mul_mat_q(
-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) {
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        const ggml_cuda_mm_fusion_args_host * fusion) {
+    // Requesting fusion where it was not compiled, or on the MUL_MAT_ID path below (which builds its
+    // mmq_args without glu_up), would drop the GLU without a diagnostic, so assert instead.
+    const ggml_tensor * glu_up = fusion ? fusion->glu_up : nullptr;
+    const ggml_glu_op   glu_op = fusion ? fusion->glu_op : GGML_GLU_OP_SWIGLU;
+    GGML_ASSERT(!glu_up || (!ids && ggml_cuda_mmq_can_fuse_glu(src0->type,
+                    ggml_cuda_info().devices[ggml_cuda_get_device()].cc)));
+    const float * glu_up_d = glu_up ? (const float *) glu_up->data : nullptr;
     GGML_ASSERT(        src1->type == GGML_TYPE_F32);
     GGML_ASSERT(        dst->type  == GGML_TYPE_F32);
     GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
@@ -171,7 +179,7 @@ void ggml_cuda_mul_mat_q(
             ne00, ne01, ne1, s01, ne11, s1,
             ne02, ne12, s02, s12, s2,
             ne03, ne13, s03, s13, s3,
-            ne1};
+            ne1, glu_up_d, glu_op};
         ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
         return;
     }