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
This commit is contained in:
Aman Gupta
2026-10-06 10:38:18 +03:00
committed by Georgi Gerganov
parent d7682b5814
commit cf5d9e3014
3 changed files with 136 additions and 10 deletions
+7
View File
@@ -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 {
+89
View File
@@ -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)) {
+40 -10
View File
@@ -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]);