mirror of
https://github.com/ggml-org/whisper.cpp.git
synced 2026-10-12 09:15:38 +02:00
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:
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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]);
|
||||
|
||||
Reference in New Issue
Block a user