From cf5d9e3014fa93a0edce366696ae9c6ad87154cf Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sat, 3 Oct 2026 02:26:27 +0800 Subject: [PATCH] CUDA: fuse shared experts into MMVQ (llama/29184) * CUDA: fuse shared experts into MMVQ * check if buffer is null * move stride_col_dst to fusion args --- ggml/src/ggml-cuda/common.cuh | 7 +++ ggml/src/ggml-cuda/ggml-cuda.cu | 89 +++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/mmvq.cu | 50 ++++++++++++++---- 3 files changed, 136 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 50713c012..928607c41 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1571,6 +1571,9 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * gate_scale = nullptr; ggml_glu_op glu_op; float glu_limit = 0.0f; + const ggml_tensor * shared_up = nullptr; + const ggml_tensor * shared_gate = nullptr; + ggml_tensor * shared_dst = nullptr; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; @@ -1580,6 +1583,10 @@ struct ggml_cuda_mm_fusion_args_device { const void * gate_scale = nullptr; ggml_glu_op glu_op; float glu_limit = 0.0f; + const void * shared_up = nullptr; + const void * shared_gate = nullptr; + float * shared_dst = nullptr; + uint32_t shared_stride_col_dst = 0; }; struct ggml_cuda_kernel_launch_params { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 043511b72..303c9f463 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1823,6 +1823,55 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } +static bool ggml_cuda_match_shared_expert(const ggml_cgraph * graph, int routed_idx, int shared_idx) { + if (routed_idx + 2 >= graph->n_nodes || shared_idx + 2 >= graph->n_nodes || shared_idx < routed_idx + 3) { + return false; + } + const int nodes[] = { routed_idx, routed_idx + 1, routed_idx + 2, shared_idx, shared_idx + 1, shared_idx + 2 }; + const ggml_op ops[] = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU, + GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; + const int outputs[] = { routed_idx + 2, shared_idx + 2 }; + if (!ggml_can_fuse_subgraph_ext(graph, nodes, 6, ops, outputs, 2)) { + return false; + } + + const ggml_tensor * routed = graph->nodes[routed_idx + 2]; + const ggml_tensor * shared = graph->nodes[shared_idx + 2]; + const ggml_tensor * gate = routed->src[0]; + const ggml_tensor * up = routed->src[1]; + const ggml_tensor * shared_gate = shared->src[0]; + const ggml_tensor * shared_up = shared->src[1]; + const auto is_pair = [&](const ggml_tensor * a, const ggml_tensor * b, int idx) { + return (a == graph->nodes[idx] && b == graph->nodes[idx + 1]) || + (b == graph->nodes[idx] && a == graph->nodes[idx + 1]); + }; + if (!is_pair(gate, up, routed_idx) || !is_pair(shared_gate, shared_up, shared_idx) || + !ggml_cuda_should_fuse_mul_mat(up, gate, routed) || + !ggml_cuda_should_fuse_mul_mat(shared_up, shared_gate, shared) || + !up->src[0]->buffer || + !ggml_cuda_should_fuse_mul_mat_vec_q(up)) { + return false; + } + const ggml_tensor * input = up->src[1]; + const ggml_tensor * weight = up->src[0]; + const ggml_tensor * shared_weight = shared_up->src[0]; + if (input->op != GGML_OP_RESHAPE || input->src[0] != shared_up->src[1] || + input->ne[1] != 1 || input->ne[3] != 1 || !ggml_is_contiguous(input) || + !ggml_is_contiguous(shared_up->src[1]) || !ggml_is_matrix(shared_up->src[1]) || + weight->type != shared_weight->type || weight->ne[0] != shared_weight->ne[0] || + weight->ne[1] != shared_weight->ne[1] || weight->nb[1] != shared_weight->nb[1] || weight->ne[3] != 1 || + !ggml_is_matrix(shared_weight) || !ggml_is_contiguous(shared_weight) || + !ggml_is_contiguous(shared_gate->src[0]) || !ggml_is_contiguous(routed) || !ggml_is_contiguous(shared)) { + return false; + } + if (shared_weight->op != GGML_OP_NONE || shared_gate->src[0]->op != GGML_OP_NONE || + ggml_get_glu_op(routed) != ggml_get_glu_op(shared) || + ggml_get_op_params_f32(routed, 3) != ggml_get_op_params_f32(shared, 3)) { + return false; + } + return true; +} + static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_TENSOR_BINARY_OP_LOCALS @@ -3459,6 +3508,25 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; + if (node->op == GGML_OP_MUL_MAT_ID && cuda_ctx->stream_context().concurrent_events.empty() && + ggml_cuda_match_shared_expert(cgraph, i, i + 3)) { + const int outputs[] = { i + 2, i + 5 }; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 6, outputs, 2)) { + ggml_tensor * routed = cgraph->nodes[i + 2]; + ggml_tensor * shared = cgraph->nodes[i + 5]; + const ggml_tensor * up = routed->src[1]; + ggml_cuda_mm_fusion_args_host fusion{}; + fusion.gate = routed->src[0]->src[0]; + fusion.glu_op = ggml_get_glu_op(routed); + fusion.glu_limit = ggml_get_op_params_f32(routed, 3); + fusion.shared_up = shared->src[1]->src[0]; + fusion.shared_gate = shared->src[0]->src[0]; + fusion.shared_dst = shared; + ggml_cuda_mul_mat_vec_q(*cuda_ctx, up->src[0], up->src[1], up->src[2], routed, &fusion); + return 5; + } + } + if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { @@ -4549,6 +4617,27 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph if (!disable_fusion) { // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute + ggml_cuda_set_device(cuda_ctx->device); + for (int i = 0; i + 5 < cgraph->n_nodes; ++i) { + if (cgraph->nodes[i]->op != GGML_OP_MUL_MAT_ID) { + continue; + } + for (int j = i + 3; j + 2 < cgraph->n_nodes; ++j) { + if (cgraph->nodes[j]->op == GGML_OP_MUL_MAT_ID && cgraph->nodes[j + 1]->op == GGML_OP_MUL_MAT_ID) { + break; + } + if (cgraph->nodes[j]->op != GGML_OP_MUL_MAT || !ggml_cuda_match_shared_expert(cgraph, i, j)) { + continue; + } + // Group both outputs before allocation so the shared result cannot alias intervening nodes. + std::rotate(cgraph->nodes + i + 3, cgraph->nodes + j, cgraph->nodes + j + 3); + ggml_tensor * up = cgraph->nodes[i + 2]->src[1]; + params->add_alloc_dep(params->user_data, up->src[1], cgraph->nodes[i + 5]); + params->add_alloc_dep(params->user_data, up->src[2], cgraph->nodes[i + 5]); + i += 5; + break; + } + } for (int i = 0; i < cgraph->n_nodes; ++i) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index a35dc9370..251577d2f 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -601,7 +601,7 @@ __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, h static __global__ void mul_mat_vec_q( const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, - const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, const uint32_t ids_stride) { @@ -625,14 +625,20 @@ static __global__ void mul_mat_vec_q( const int blocks_per_row_x = ncols_x / qk; constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi; - const uint32_t channel_dst = blockIdx.y; + const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1; + const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y; + if (shared_expert) { + vx = fusion.shared_up; + dst = fusion.shared_dst; + stride_col_dst = fusion.shared_stride_col_dst; + } uint32_t channel_x; uint32_t channel_y; uint32_t sample_dst; ggml_cuda_pdl_sync(); - channel_x = ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); + channel_x = shared_expert ? 0 : ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; sample_dst = blockIdx.z; @@ -656,7 +662,7 @@ static __global__ void mul_mat_vec_q( use_gate = fusion.gate != nullptr; use_bias = fusion.x_bias != nullptr; use_gate_bias = fusion.gate_bias != nullptr && use_gate; - vgate = fusion.gate; + vgate = shared_expert ? fusion.shared_gate : fusion.gate; x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; @@ -854,7 +860,7 @@ static __global__ void mul_mat_vec_q_moe( const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, - const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_row_x, const uint32_t stride_col_y, uint32_t stride_col_dst, const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint32_t ncols_dst, const uint32_t ids_stride) { const void * GGML_CUDA_RESTRICT vx = vx_ptr; @@ -869,6 +875,13 @@ static __global__ void mul_mat_vec_q_moe( constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); + const bool shared_expert = has_fusion && fusion.shared_up && blockIdx.y == gridDim.y - 1; + if (shared_expert) { + vx = fusion.shared_up; + dst = fusion.shared_dst; + stride_col_dst = fusion.shared_stride_col_dst; + } + // fuse gate, bias, scales, and glu_op into the up projection bool use_gate = false; const void * vgate = nullptr; @@ -881,7 +894,7 @@ static __global__ void mul_mat_vec_q_moe( if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; - vgate = fusion.gate; + vgate = shared_expert ? fusion.shared_gate : fusion.gate; x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; @@ -897,14 +910,14 @@ static __global__ void mul_mat_vec_q_moe( const int blocks_per_row_x = ncols_x / qk; constexpr int blocks_per_iter = vdr * warp_size / qi; - const uint32_t channel_dst = blockIdx.y; + const uint32_t channel_dst = shared_expert ? 0 : blockIdx.y; if (token_idx >= ncols_dst) { return; } ggml_cuda_pdl_sync(); - const uint32_t channel_x = ids[channel_dst + token_idx * ids_stride]; + const uint32_t channel_x = shared_expert ? 0 : ids[channel_dst + token_idx * ids_stride]; const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y; @@ -1050,7 +1063,7 @@ static void mul_mat_vec_q_moe_launch( constexpr int rows_per_block = 2; // 2 gives best perf based on tuning const int64_t nblocks_rows = (nrows_x + rows_per_block - 1) / rows_per_block; - const dim3 block_nums(nblocks_rows, nchannels_dst); + const dim3 block_nums(nblocks_rows, nchannels_dst + (fusion.shared_up != nullptr)); const dim3 block_dims(warp_size, ncols_dst); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); @@ -1187,7 +1200,7 @@ static void mul_mat_vec_q_switch_ncols_dst( constexpr bool c_halve_iters = decltype(halve_iters_tag)::value && c_promoted; - const std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, + const std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst + (fusion.shared_up != nullptr), nsamples_dst, warp_size, table_id, c_small_k, c_halve_iters); mul_mat_vec_q_switch_fusion( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, @@ -1454,6 +1467,23 @@ void ggml_cuda_mul_mat_vec_q( // non-negligible for some models such as gpt-oss-20b GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4); + if (fusion->shared_up) { + GGML_ASSERT(ids && fusion->gate && fusion->shared_gate && fusion->shared_dst); + GGML_ASSERT(!fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale); + GGML_ASSERT(ne11 == 1 && ne03 == 1 && ne13 == 1); + GGML_ASSERT(fusion->shared_up->type == src0->type && fusion->shared_gate->type == src0->type); + GGML_ASSERT(ggml_are_same_shape(fusion->shared_up, fusion->shared_gate)); + GGML_ASSERT(ggml_is_contiguous(fusion->shared_up) && ggml_is_contiguous(fusion->shared_gate)); + GGML_ASSERT(fusion->shared_up->ne[0] == ne00 && fusion->shared_up->ne[1] == ne01); + GGML_ASSERT(fusion->shared_up->nb[1] == nb01 && ggml_is_matrix(fusion->shared_up)); + GGML_ASSERT(fusion->shared_dst->type == GGML_TYPE_F32 && ggml_is_contiguous(fusion->shared_dst)); + GGML_ASSERT(fusion->shared_dst->ne[0] == ne0 && fusion->shared_dst->ne[1] == ne2); + fusion_local.shared_up = fusion->shared_up->data; + fusion_local.shared_gate = fusion->shared_gate->data; + fusion_local.shared_dst = (float *) fusion->shared_dst->data; + fusion_local.shared_stride_col_dst = fusion->shared_dst->nb[1] / ts_dst; + } + if (fusion->x_bias) { GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); GGML_ASSERT(fusion->x_bias->ne[0] == dst->ne[0]);