diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 0a9804073032..6e868e806db0 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4341,6 +4341,65 @@ static bool ggml_cuda_match_moe_weighted_reduction( } +// GGML_CUDA_MOE_EP_VERIFY=1: compare against the existing fused reduction (REPEAT/GET_ROWS result input); +// =2: compare against the fully unfused MUL, MUL, VIEW, ADD nodes +static int ggml_cuda_moe_ep_verify_mode() { + static const int m = getenv("GGML_CUDA_MOE_EP_VERIFY") ? atoi(getenv("GGML_CUDA_MOE_EP_VERIFY")) : 0; + return m; +} + +// GGML_CUDA_MOE_EP_FUSION bit mask: 1 = scale gather into the weighted reduction, 2 = router RMS_NORM+SCALE+MUL, 4 = router matvec + top-k +static int ggml_cuda_moe_ep_bits() { + static const int bits = getenv("GGML_CUDA_MOE_EP_FUSION") ? atoi(getenv("GGML_CUDA_MOE_EP_FUSION")) : 0; + return bits; +} +static bool ggml_cuda_moe_ep_fusion_enabled() { + return (ggml_cuda_moe_ep_bits() & 1) != 0; +} + +// REPEAT(w_s [1,n_expert,1] -> [1,n_expert,n_tokens]) -> GET_ROWS(.., ids) -> MUL(experts, .) -> weighted reduction. +// The gather and the repeat only feed the expert-scale MUL, so the reduction kernel can read w_s[ids] directly. +static bool ggml_cuda_match_moe_scale_gather( + const ggml_cgraph * cgraph, + int node_idx, + ggml_cuda_moe_weighted_reduction_match & match, + const ggml_tensor *& table, + const ggml_tensor *& ids) { + if (node_idx + 3 >= cgraph->n_nodes) { + return false; + } + const ggml_tensor * rep = cgraph->nodes[node_idx]; + const ggml_tensor * get = cgraph->nodes[node_idx + 1]; + if (rep->op != GGML_OP_REPEAT || get->op != GGML_OP_GET_ROWS || get->src[0] != rep || + rep->type != GGML_TYPE_F32 || get->type != GGML_TYPE_F32) { + return false; + } + const ggml_tensor * w_s = rep->src[0]; + ids = get->src[1]; + if (!w_s || !ids || w_s->type != GGML_TYPE_F32 || !ggml_is_contiguous(w_s) || w_s->ne[0] != 1 || w_s->ne[2] != 1 || + w_s->ne[3] != 1 || ids->type != GGML_TYPE_I32 || ids->ne[2] != 1 || ids->ne[3] != 1 || + ids->nb[0] != sizeof(int32_t) || ids->nb[1] % sizeof(int32_t) != 0 || + rep->ne[0] != 1 || rep->ne[1] != w_s->ne[1] || rep->ne[3] != 1 || + get->ne[0] != 1 || get->ne[1] != ids->ne[0] || get->ne[2] != ids->ne[1] || get->ne[3] != 1 || + rep->ne[2] != ids->ne[1] || !ggml_is_contiguous(rep) || !ggml_is_contiguous(get)) { + return false; + } + if (!ggml_cuda_match_moe_weighted_reduction(cgraph, node_idx + 2, match) || match.expert_scale != get) { + return false; + } + const int count = match.node_count + 2; + std::vector ops(count); + for (int j = 0; j < count; ++j) { + ops[j] = cgraph->nodes[node_idx + j]->op; + } + const int output_idx = node_idx + count - 1; + if (!ggml_can_fuse_subgraph(cgraph, node_idx, count, ops.data(), &output_idx, 1)) { + return false; + } + table = w_s; + return true; +} + static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops, @@ -5395,13 +5454,51 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } - if (node->op == GGML_OP_MUL) { + // GGML_CUDA_MOE_EP_FUSION=1: fold the per-expert scale gather (REPEAT + GET_ROWS of down_exps_s by the selected ids) + // into the weighted reduction kernel + if (node->op == GGML_OP_REPEAT && ggml_cuda_moe_ep_fusion_enabled()) { + ggml_cuda_moe_weighted_reduction_match match; + const ggml_tensor * table = nullptr; + const ggml_tensor * ids = nullptr; + if (ggml_cuda_match_moe_scale_gather(cgraph, i, match, table, ids)) { + const int count = match.node_count + 2; + const int output_idx = i + count - 1; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, count, &output_idx, 1)) { + // GGML_CUDA_MOE_EP_VERIFY=1: run the fused kernel into scratch, run the unfused nodes, compare bit-for-bit at the output + if (ggml_cuda_moe_ep_verify_mode() != 0) { + static float * scratch = nullptr; static int tag = 0; + if (!scratch) { CUDA_CHECK(cudaMalloc(&scratch, (size_t) 32*1024*1024)); } + const size_t need = ggml_nelements(match.dst)*sizeof(float); + if (need <= (size_t) 32*1024*1024 && tag < 1000000) { + ggml_tensor tmp = *match.dst; + tmp.data = scratch; + ggml_cuda_op_moe_weighted_reduction(*cuda_ctx, match.experts, nullptr, match.weights, &tmp, table, ids); + ggml_cuda_tail_verify_pending()[match.dst] = { scratch, tag++ }; + } + return 0; + } + ggml_cuda_op_moe_weighted_reduction( + *cuda_ctx, match.experts, nullptr, match.weights, match.dst, table, ids); + return count - 1; + } + } + } + + if (node->op == GGML_OP_MUL && !(ggml_cuda_moe_ep_verify_mode() == 2 && ggml_cuda_moe_ep_fusion_enabled())) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { const int output_idx = i + match.node_count - 1; if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, match.node_count, &output_idx, 1)) { ggml_cuda_op_moe_weighted_reduction( *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst); + if (!ggml_cuda_tail_verify_pending().empty()) { + auto it = ggml_cuda_tail_verify_pending().find(match.dst); + if (it != ggml_cuda_tail_verify_pending().end()) { + ggml_cuda_op_tail_verify_cmp(*cuda_ctx, match.dst, it->second.scratch, it->second.tag); + static int nl = 0; + if ((++nl & (nl - 1)) == 0) fprintf(stderr, "TAILVERIFY compare launches captured/issued: %d\n", nl); + } + } return match.node_count - 1; } } @@ -6162,6 +6259,25 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 1; } + const int ep_out2 = i + 2; + if ((ggml_cuda_moe_ep_bits() & 2) && i + 2 < cgraph->n_nodes && cgraph->nodes[i + 1]->op == GGML_OP_SCALE && + cgraph->nodes[i + 1]->src[0] == node && cgraph->nodes[i + 2]->op == GGML_OP_MUL && + (cgraph->nodes[i + 2]->src[0] == cgraph->nodes[i + 1] || cgraph->nodes[i + 2]->src[1] == cgraph->nodes[i + 1]) && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_SCALE, GGML_OP_MUL }, { i + 2 }) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, &ep_out2, 1)) { + if (ggml_cuda_moe_ep_verify_mode() != 0) { + static float * scratch = nullptr; static int tag = 0; + if (!scratch) { CUDA_CHECK(cudaMalloc(&scratch, (size_t) 32*1024*1024)); } + ggml_tensor * mul_n = cgraph->nodes[i + 2]; + if (ggml_nelements(mul_n)*sizeof(float) <= (size_t) 32*1024*1024 && tag < 1000000 && + ggml_cuda_op_rms_norm_scale_mul_fused(*cuda_ctx, node, cgraph->nodes[i + 1], mul_n, scratch)) { + ggml_cuda_tail_verify_pending()[mul_n] = { scratch, tag++ }; + } + } else if (ggml_cuda_op_rms_norm_scale_mul_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2])) { + return 2; + } + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_SCALE }, {})) { ggml_cuda_op_rms_norm_scale_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); return 1; @@ -6820,6 +6936,10 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph if (match.expert_scale != nullptr) { params->add_alloc_dep( params->user_data, const_cast(match.expert_scale), match.dst); + if (ggml_cuda_moe_ep_fusion_enabled() && match.expert_scale->op == GGML_OP_GET_ROWS && match.expert_scale->src[1]) { + params->add_alloc_dep( + params->user_data, const_cast(match.expert_scale->src[1]), match.dst); + } } i += match.node_count - 1; } diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cu b/ggml/src/ggml-cuda/moe-weighted-reduction.cu index 11ec58497f1e..febecd41066f 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cu +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cu @@ -5,7 +5,10 @@ static __global__ void moe_weighted_reduction_f32(const float * __restrict__ exp const float * __restrict__ weights, float * __restrict__ dst, const int64_t n_embd, - const int n_expert_used) { + const int n_expert_used, + const float * __restrict__ scale_table, + const int32_t * __restrict__ scale_ids, + const int64_t ids_s1) { const int64_t token = blockIdx.x; const int64_t col = (int64_t) blockIdx.y * blockDim.x + threadIdx.x; if (col >= n_embd) { @@ -13,12 +16,16 @@ static __global__ void moe_weighted_reduction_f32(const float * __restrict__ exp } const uint64_t first_row = (uint64_t) token * n_expert_used; - const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + // scale_table != nullptr: the per-expert scale is gathered here, scale_table[ids[token, expert]], instead of + // being read from the REPEAT + GET_ROWS result (same float values) + const float first_scale = scale_table != nullptr ? scale_table[scale_ids[token * ids_s1]] : + expert_scale != nullptr ? expert_scale[first_row] : 1.0f; float sum = (experts[first_row * n_embd + col] * first_scale) * weights[first_row]; for (int expert = 1; expert < n_expert_used; ++expert) { const uint64_t row = first_row + expert; - const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; + const float scale = scale_table != nullptr ? scale_table[scale_ids[token * ids_s1 + expert]] : + expert_scale != nullptr ? expert_scale[row] : 1.0f; sum += (experts[row * n_embd + col] * scale) * weights[row]; } dst[token * n_embd + col] = sum; @@ -31,18 +38,23 @@ static void launch_moe_weighted_reduction(const float * experts, int64_t n_embd, int64_t n_tokens, int n_expert_used, + const float * scale_table, + const int32_t * scale_ids, + int64_t ids_s1, cudaStream_t stream) { constexpr int threads = 256; const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1); moe_weighted_reduction_f32 - <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); + <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used, scale_table, scale_ids, ids_s1); } void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, const ggml_tensor * expert_scale, const ggml_tensor * weights, - ggml_tensor * dst) { + ggml_tensor * dst, + const ggml_tensor * scale_table, + const ggml_tensor * scale_ids) { GGML_ASSERT(experts->type == GGML_TYPE_F32); GGML_ASSERT(weights->type == GGML_TYPE_F32); GGML_ASSERT(expert_scale == nullptr || expert_scale->type == GGML_TYPE_F32); @@ -51,6 +63,8 @@ void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, GGML_ASSERT(ggml_is_contiguous(weights)); GGML_ASSERT(expert_scale == nullptr || ggml_is_contiguous(expert_scale)); GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(scale_ids == nullptr || (scale_ids->type == GGML_TYPE_I32 && scale_ids->nb[0] == sizeof(int32_t) && + scale_ids->nb[1] % sizeof(int32_t) == 0 && scale_table != nullptr && scale_table->type == GGML_TYPE_F32)); const int64_t n_embd = experts->ne[0]; const int64_t n_expert_used = experts->ne[1]; @@ -60,6 +74,9 @@ void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, launch_moe_weighted_reduction((const float *) experts->data, expert_scale ? (const float *) expert_scale->data : nullptr, (const float *) weights->data, - (float *) dst->data, n_embd, n_tokens, (int) n_expert_used, stream); + (float *) dst->data, n_embd, n_tokens, (int) n_expert_used, + scale_table ? (const float *) scale_table->data : nullptr, + scale_ids ? (const int32_t *) scale_ids->data : nullptr, + scale_ids ? (int64_t) (scale_ids->nb[1] / sizeof(int32_t)) : 0, stream); CUDA_CHECK(cudaGetLastError()); } diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh index b72f947ab398..6bf359982354 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh @@ -4,4 +4,6 @@ void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, const ggml_tensor * expert_scale, const ggml_tensor * weights, - ggml_tensor * dst); + ggml_tensor * dst, + const ggml_tensor * scale_table = nullptr, + const ggml_tensor * scale_ids = nullptr); diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c24a5c1bf661..eb0f585f8dd5 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -952,6 +952,85 @@ void ggml_cuda_op_add_rms_norm_scale_fused( } } +// RMS_NORM -> SCALE -> MUL(weight): same reduction and expression order as the rms_norm_f32 kernel the unfused RMS_NORM takes +// (the scalar path, or with vec the float4 path with one float4 per thread), followed by the elementwise scale and multiply, +// so the result is bit-identical to the three separate nodes. +template +static __global__ void rms_norm_scale_mul_f32(const float * x, const float * w, float * dst, const int ncols, + const int64_t stride_row, const float eps, const float scale_out) { + ggml_cuda_pdl_lc(); + extern __shared__ float s_sum[]; + const int row = blockIdx.x; + const int tid = threadIdx.x; + ggml_cuda_pdl_sync(); + const float * xc = x + row*stride_row; + float * dstc = dst + (int64_t) row*ncols; + if constexpr (vec) { + const int n4 = ncols / 4; + float4 xv = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float tmp = 0.0f; + if (tid < n4) { + xv = ((const float4 *) xc)[tid]; + tmp = __fadd_rn(tmp, rms_norm_sumsq4(xv)); + } + tmp = block_reduce(tmp, s_sum); + const float scale = rms_norm_row_scale(tmp, ncols, eps); + if (tid < n4) { + const float4 wv = ((const float4 *) w)[tid]; + float4 r; + r.x = __fmul_rn(__fmul_rn(scale_out, __fmul_rn(scale, xv.x)), wv.x); + r.y = __fmul_rn(__fmul_rn(scale_out, __fmul_rn(scale, xv.y)), wv.y); + r.z = __fmul_rn(__fmul_rn(scale_out, __fmul_rn(scale, xv.z)), wv.z); + r.w = __fmul_rn(__fmul_rn(scale_out, __fmul_rn(scale, xv.w)), wv.w); + ((float4 *) dstc)[tid] = r; + } + } else { + float tmp = 0.0f; + for (int col = tid; col < ncols; col += 1024) { + const float xi = xc[col]; + tmp += xi * xi; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + for (int col = tid; col < ncols; col += 1024) { + dstc[col] = (scale_out * (scale * xc[col])) * w[col]; + } + } +} + +bool ggml_cuda_op_rms_norm_scale_mul_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rms, ggml_tensor * scale_t, ggml_tensor * mul_t, float * verify_scratch) { + const ggml_tensor * src0 = rms->src[0]; + const ggml_tensor * w = mul_t->src[0] == scale_t ? mul_t->src[1] : mul_t->src[0]; + if (src0->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || mul_t->type != GGML_TYPE_F32 || scale_t->type != GGML_TYPE_F32 || + src0->ne[0] < 1024 || src0->ne[2] != 1 || src0->ne[3] != 1 || src0->nb[0] != sizeof(float) || + !ggml_is_contiguous(w) || ggml_nelements(w) != src0->ne[0] || !ggml_is_contiguous(mul_t) || + !ggml_are_same_shape(mul_t, src0)) { + return false; + } + float eps, scale; + memcpy(&eps, rms->op_params, sizeof(float)); + memcpy(&scale, (const float *) scale_t->op_params + 0, sizeof(float)); + if (!(eps >= 0.0f) || ((const float *) scale_t->op_params)[1] != 0.0f) { + return false; + } + // The unfused RMS_NORM takes the float4 path when its operands are 16-byte aligned and ncols/4 <= 1024 fits one float4 + // per thread; wider rows use a different float4 geometry that this kernel does not reproduce, so they are not fused. + const int ncols = (int) src0->ne[0]; + const bool vec = rms_norm_vector_enabled() && ncols % 4 == 0 && ggml_cuda_is_aligned(src0, 16) && + ((uintptr_t) rms->data & 15) == 0; + if (vec && ncols / 4 > 1024) { + return false; + } + if (vec && (!ggml_cuda_is_aligned(w, 16) || ((uintptr_t) scale_t->data & 15) != 0 || ((uintptr_t) mul_t->data & 15) != 0)) { + return false; + } + const ggml_cuda_kernel_launch_params lp{ dim3(src0->ne[1], 1, 1), dim3(1024, 1, 1), 32 * sizeof(float), ctx.stream() }; + ggml_cuda_kernel_launch(vec ? rms_norm_scale_mul_f32 : rms_norm_scale_mul_f32, lp, (const float *) src0->data, (const float *) w->data, verify_scratch ? verify_scratch : (float *) mul_t->data, + (int) src0->ne[0], (int64_t) (src0->nb[1] / sizeof(float)), eps, scale); + return true; +} + void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) { const ggml_tensor * src0 = dst->src[0]; float * dst_d = (float *) scale_tensor->data; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 7da8554fdc2e..d6e7ea9bd1de 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -21,6 +21,7 @@ void ggml_cuda_op_add_rms_norm_scale_fused( void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor); +bool ggml_cuda_op_rms_norm_scale_mul_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rms, ggml_tensor * scale_t, ggml_tensor * mul_t, float * verify_scratch = nullptr); void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor); void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,