Fuse compatible Q8 CUDA residual graph patterns (0xShug0#154)

Author: mirek190

external/ggml/src/ggml-cuda/common.cuh

@@ -1518,10 +1518,12 @@ struct ggml_cuda_mm_fusion_args_host {
     const ggml_tensor * gate = nullptr;
     const ggml_tensor * gate_bias = nullptr;
     ggml_glu_op glu_op;
+    bool residual_only = false;
 };
 struct ggml_cuda_mm_fusion_args_device {
     const void * x_bias = nullptr;
     const void * gate = nullptr;
     const void * gate_bias = nullptr;
     ggml_glu_op glu_op;
+    bool residual_only = false;
 };

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

@@ -3771,6 +3771,43 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
     return is_ok;
 }
 
+// Some model graphs reshape a matvec result before adding the residual. RESHAPE
+// is metadata-only and therefore cannot pass the generic compute-node fusion
+// validator. Validate this exact chain explicitly so the residual-only Q8_0
+// specialization can write the final result directly.
+static bool ggml_cuda_can_fuse_q8_0_mul_mat_reshape_add(
+        const struct ggml_cgraph * cgraph, int node_idx) {
+    if (node_idx + 2 >= cgraph->n_nodes) {
+        return false;
+    }
+
+    const ggml_tensor * mul_mat = cgraph->nodes[node_idx + 0];
+    const ggml_tensor * reshape = cgraph->nodes[node_idx + 1];
+    const ggml_tensor * add     = cgraph->nodes[node_idx + 2];
+
+    if (mul_mat->op != GGML_OP_MUL_MAT ||
+        !mul_mat->src[0] ||
+        mul_mat->src[0]->type != GGML_TYPE_Q8_0 ||
+        reshape->op != GGML_OP_RESHAPE ||
+        reshape->src[0] != mul_mat ||
+        add->op != GGML_OP_ADD ||
+        (add->src[0] != reshape && add->src[1] != reshape)) {
+        return false;
+    }
+
+    if (ggml_nelements(mul_mat) != ggml_nelements(reshape) ||
+        ggml_nelements(reshape) != ggml_nelements(add) ||
+        ggml_node_get_use_count(cgraph, node_idx + 0) != 1 ||
+        ggml_node_get_use_count(cgraph, node_idx + 1) != 1 ||
+        (mul_mat->flags & GGML_TENSOR_FLAG_OUTPUT) ||
+        (reshape->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return false;
+    }
+
+    const int out_nodes[] = { node_idx + 2 };
+    return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, 3, out_nodes, 1);
+}
+
 
 static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
                                int                                       node_idx,
@@ -4291,22 +4328,29 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
     fused_mul_mat_vec = false;
     fused_node_count  = 0;
 
-    // gate + add + glu + up + add
+    // mul_mat + optional metadata-only reshape + add
     for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
         const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
 
-        if (!ggml_can_fuse(cgraph, i, { op, bias_op })) {
+        const bool reshape_bridge =
+            op == GGML_OP_MUL_MAT &&
+            ggml_cuda_can_fuse_q8_0_mul_mat_reshape_add(cgraph, i);
+        if (!reshape_bridge && !ggml_can_fuse(cgraph, i, { op, bias_op })) {
             continue;
         }
 
         ggml_tensor * mm_node   = cgraph->nodes[i];
-        ggml_tensor * bias_node = cgraph->nodes[i + 1];
+        ggml_tensor * mm_output = reshape_bridge ? cgraph->nodes[i + 1] : mm_node;
+        ggml_tensor * bias_node = cgraph->nodes[i + (reshape_bridge ? 2 : 1)];
+        if (reshape_bridge && mm_output->src[0] != mm_node) {
+            continue;
+        }
 
         ggml_tensor * bias_tensor = nullptr;
         if (bias_op == GGML_OP_ADD) {
-            if (bias_node->src[0] == mm_node) {
+            if (bias_node->src[0] == mm_output) {
                 bias_tensor = bias_node->src[1];
-            } else if (bias_node->src[1] == mm_node) {
+            } else if (bias_node->src[1] == mm_output) {
                 bias_tensor = bias_node->src[0];
             } else {
                 continue;
@@ -4331,19 +4375,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
         }
 
         ggml_cuda_mm_fusion_args_host fusion_data{};
-        fusion_data.x_bias = bias_tensor;
+        fusion_data.x_bias        = bias_tensor;
+        fusion_data.residual_only = reshape_bridge;
 
         if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) {
             ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data);
             fused_mul_mat_vec = true;
-            fused_node_count  = 2;
+            fused_node_count  = reshape_bridge ? 3 : 2;
             break;
         }
 
         if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) {
             ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data);
             fused_mul_mat_vec = true;
-            fused_node_count  = 2;
+            fused_node_count  = reshape_bridge ? 3 : 2;
             break;
         }
     }

external/ggml/src/ggml-cuda/mmvq.cu

@@ -391,7 +391,7 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int
     return 1;
 }
 
-template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false>
+template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false, bool residual_only = false>
 __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mul_mat_vec_q(
         const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst,
@@ -401,6 +401,8 @@ static __global__ void mul_mat_vec_q(
         const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst,
         const uint32_t ids_stride) {
 
+    static_assert(!residual_only || has_fusion, "residual-only MMVQ requires fusion");
+
     constexpr int qk  = ggml_cuda_type_traits<type>::qk;
     constexpr int qi  = ggml_cuda_type_traits<type>::qi;
     constexpr int vdr = get_vdr_mmvq(type);
@@ -437,7 +439,10 @@ static __global__ void mul_mat_vec_q(
     const float * gate_bias = nullptr;
     ggml_glu_op active_glu;
 
-    if constexpr (has_fusion) {
+    if constexpr (residual_only) {
+        use_bias = true;
+        x_bias   = (const float *) fusion.x_bias;
+    } else if constexpr (has_fusion) {
         use_gate      = fusion.gate      != nullptr;
         use_bias      = fusion.x_bias    != nullptr;
         use_gate_bias = fusion.gate_bias != nullptr && use_gate;
@@ -478,7 +483,7 @@ static __global__ void mul_mat_vec_q(
 
     // partial sum for each thread
     float tmp[ncols_dst][rows_per_cuda_block] = {{0.0f}};
-    float tmp_gate[ncols_dst][rows_per_cuda_block] = {{0.0f}};
+    float tmp_gate[ncols_dst][residual_only ? 1 : rows_per_cuda_block] = {{0.0f}};
 
     const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y;
     const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row0*stride_row_x;
@@ -495,7 +500,7 @@ static __global__ void mul_mat_vec_q(
             for (int i = 0; i < rows_per_cuda_block; ++i) {
                 tmp[j][i] += vec_dot_q_cuda(
                     vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs);
-                if constexpr (has_fusion) {
+                if constexpr (has_fusion && !residual_only) {
                     if (use_gate) {
                         tmp_gate[j][i] += vec_dot_q_cuda(
                             vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs);
@@ -506,8 +511,8 @@ static __global__ void mul_mat_vec_q(
     }
 
     __shared__ float tmp_shared[nwarps-1 > 0 ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size];
-    __shared__ float tmp_shared_gate[(has_fusion && (nwarps-1 > 0)) ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size];
-    if constexpr (!has_fusion) {
+    __shared__ float tmp_shared_gate[(has_fusion && !residual_only && (nwarps-1 > 0)) ? nwarps-1 : 1][ncols_dst][residual_only ? 1 : rows_per_cuda_block][warp_size];
+    if constexpr (!has_fusion || residual_only) {
         (void) tmp_shared_gate;
     } else if (!use_gate) {
         (void) tmp_shared_gate;
@@ -519,7 +524,7 @@ static __global__ void mul_mat_vec_q(
 #pragma unroll
             for (int i = 0; i < rows_per_cuda_block; ++i) {
                 tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i];
-                if constexpr (has_fusion) {
+                if constexpr (has_fusion && !residual_only) {
                     if (use_gate) {
                         tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i];
                     }
@@ -542,14 +547,14 @@ static __global__ void mul_mat_vec_q(
 #pragma unroll
             for (int l = 0; l < nwarps-1; ++l) {
                 tmp[j][i] += tmp_shared[l][j][i][threadIdx.x];
-                if constexpr (has_fusion) {
+                if constexpr (has_fusion && !residual_only) {
                     if (use_gate) {
                         tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x];
                     }
                 }
             }
             tmp[j][i] = warp_reduce_sum<warp_size>(tmp[j][i]);
-            if constexpr (has_fusion) {
+            if constexpr (has_fusion && !residual_only) {
                 if (use_gate) {
                     tmp_gate[j][i] = warp_reduce_sum<warp_size>(tmp_gate[j][i]);
                 }
@@ -562,33 +567,35 @@ static __global__ void mul_mat_vec_q(
                 if (use_bias) {
                     result += x_biases[j];
                 }
-                if (use_gate) {
-                    float gate_value = tmp_gate[j][threadIdx.x];
-                    if (use_gate_bias) {
-                        gate_value += gate_biases[j];
-                    }
-                    switch (active_glu) {
-                        case GGML_GLU_OP_SWIGLU:
-                            result *= ggml_cuda_op_silu_single(gate_value);
-                            break;
-                        case GGML_GLU_OP_GEGLU:
-                            result *= ggml_cuda_op_gelu_single(gate_value);
-                            break;
-                        case GGML_GLU_OP_SWIGLU_OAI: {
-                            result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
-                            break;
+                if constexpr (!residual_only) {
+                    if (use_gate) {
+                        float gate_value = tmp_gate[j][threadIdx.x];
+                        if (use_gate_bias) {
+                            gate_value += gate_biases[j];
+                        }
+                        switch (active_glu) {
+                            case GGML_GLU_OP_SWIGLU:
+                                result *= ggml_cuda_op_silu_single(gate_value);
+                                break;
+                            case GGML_GLU_OP_GEGLU:
+                                result *= ggml_cuda_op_gelu_single(gate_value);
+                                break;
+                            case GGML_GLU_OP_SWIGLU_OAI: {
+                                result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
+                                break;
+                            }
+                            default:
+                                result = result * gate_value;
+                                break;
                         }
-                        default:
-                            result = result * gate_value;
-                            break;
                     }
                 }
             }
             dst[j*stride_col_dst + threadIdx.x] = result;
         }
     }
 
-    if constexpr (!has_fusion) {
+    if constexpr (!has_fusion || residual_only) {
         GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate);
     }
 }
@@ -681,6 +688,17 @@ static void mul_mat_vec_q_switch_fusion(
     const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
     if constexpr (c_ncols_dst == 1) {
         if (has_fusion) {
+            if constexpr (type == GGML_TYPE_Q8_0) {
+                if (fusion.residual_only) {
+                    mul_mat_vec_q<type, c_ncols_dst, true, small_k, true><<<block_nums, block_dims, nbytes_shared, stream>>>
+                        (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
+                         channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
+                         sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride);
+                    return;
+                }
+            }
+
+            GGML_ASSERT(!fusion.residual_only && "residual-only MMVQ is supported only for Q8_0");
             mul_mat_vec_q<type, c_ncols_dst, true, small_k><<<block_nums, block_dims, nbytes_shared, stream>>>
                 (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst,
                  channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
@@ -1078,6 +1096,7 @@ void ggml_cuda_mul_mat_vec_q(
             fusion_local.gate_bias = fusion->gate_bias->data;
         }
         fusion_local.glu_op = fusion->glu_op;
+        fusion_local.residual_only = fusion->residual_only;
     }
 
     // If src0 is a temporary compute buffer, clear any potential padding.