Commit bed0a8566 for llama.cpp
commit bed0a856606ee4a24a164066f73d2379447033f5
Author: Aman Gupta <amangupta052@gmail.com>
Date: Sat Oct 3 02:26:27 2026 +0800
CUDA: fuse shared experts into MMVQ (#29184)
* CUDA: fuse shared experts into MMVQ
* check if buffer is null
* move stride_col_dst to fusion args
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<dim3, dim3> dims = calc_launch_params<type>(c_ncols_dst, nrows_x, nchannels_dst,
+ const std::pair<dim3, dim3> dims = calc_launch_params<type>(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<type, c_ncols_dst, c_small_k, c_halve_iters>(
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]);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index d725191b7..98ce061c5 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -5423,6 +5423,39 @@ struct test_mul_mat_id_fusion : public test_case {
}
};
+struct test_mul_mat_id_shared : public test_case {
+ const ggml_type type;
+ const int64_t n;
+ ggml_tensor * routed = nullptr;
+ ggml_tensor * shared = nullptr;
+
+ test_mul_mat_id_shared(ggml_type type, int64_t n) : type(type), n(n) {}
+
+ std::string vars() override { return VARS_TO_STR2(type, n); }
+ std::string op_desc(ggml_tensor *) override { return "MUL_MAT_ID_SHARED"; }
+ bool run_whole_graph() override { return true; }
+ double max_nmse_err() override { return 5e-4; }
+ std::vector<ggml_tensor *> fusion_test_nodes() override { return { routed, shared }; }
+
+ ggml_tensor * build_graph(ggml_context * ctx) override {
+ ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 512, n);
+ ggml_tensor * input_ids = ggml_reshape_3d(ctx, input, 512, 1, n);
+ ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 3, n);
+ ggml_set_name(ids, "ids");
+ ggml_tensor * up = ggml_new_tensor_3d(ctx, type, 512, 128, 8);
+ ggml_tensor * gate = ggml_new_tensor_3d(ctx, type, 512, 128, 8);
+ ggml_tensor * shared_up = ggml_new_tensor_2d(ctx, type, 512, 128);
+ ggml_tensor * shared_gate = ggml_new_tensor_2d(ctx, type, 512, 128);
+ routed = ggml_swiglu_split(ctx, ggml_mul_mat_id(ctx, gate, input_ids, ids), ggml_mul_mat_id(ctx, up, input_ids, ids));
+ shared = ggml_swiglu_split(ctx, ggml_mul_mat(ctx, shared_gate, input), ggml_mul_mat(ctx, shared_up, input));
+ return ggml_add(ctx, routed, ggml_reshape_3d(ctx, shared, 128, 1, n));
+ }
+
+ void initialize_tensors(ggml_context * ctx) override {
+ init_mul_mat_id_tensors(ctx, 8);
+ }
+};
+
// GGML_OP_OUT_PROD
struct test_out_prod : public test_case {
const ggml_type type_a;
@@ -10467,6 +10500,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, k));
}
test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3));
+ for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q8_0 }) {
+ for (int64_t n : { 1, 2, 3, 4, 8, 9 }) {
+ test_cases.emplace_back(new test_mul_mat_id_shared(type, n));
+ }
+ }
// gpt-oss issue with Vulkan mmq_id
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880));