From 38d83db9f974a0696783fcd80b4e797bcf7deccc Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Mon, 3 Aug 2026 07:39:56 -1000 Subject: [PATCH 1/2] cuda: fuse the SwiGLU epilogue into the routed MoE MMQ kernel The merged gate/up result was written to global memory and read back by a separate SwiGLU op. This writes the activated result directly from the matmul, removing a round trip through the intermediate tensor and one dispatch per layer. Output is bit-identical. Prefill improves at every sequence length measured; decode is unaffected. Disable with GGML_CUDA_MMQ_GLU_FUSION_DISABLE. --- ggml/src/ggml-cuda/ggml-cuda.cu | 9 ++ ggml/src/ggml-cuda/mmq.cu | 195 +++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/mmq.cuh | 201 ++++++++++++++++++++++++++++---- 3 files changed, 383 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 13e1b8a2e73..f4523bd035b 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4124,6 +4124,15 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fused_node_count = 3; break; } + + // Prefill / large-batch routed experts: fold the SwiGLU into the MMQ epilogue so the + // gate and up products are never written to global memory. + if (ggml_cuda_can_fuse_mmq_glu(gate, up, glu)) { + ggml_cuda_mul_mat_q_fused_glu(*cuda_ctx, gate->src[0], up->src[0], src1, ids, glu); + fused_mul_mat_vec = true; + fused_node_count = 3; + break; + } } } diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index a3cc5032315..cbb328e74d7 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -2,6 +2,7 @@ #include "mmq.cuh" #include "quantize.cuh" #include "mmid.cuh" +#include "mmvq.cuh" static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { switch (args.type_x) { @@ -225,6 +226,200 @@ void ggml_cuda_mul_mat_q( ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); } +bool ggml_cuda_can_fuse_mmq_glu(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu) { + static const bool disabled = getenv("GGML_CUDA_MMQ_GLU_FUSION_DISABLE") != nullptr; + if (disabled) { + return false; + } + + if (gate->op != GGML_OP_MUL_MAT_ID || up->op != GGML_OP_MUL_MAT_ID || glu->op != GGML_OP_GLU) { + return false; + } + + if (ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU) { + return false; + } + + if (((const int32_t *) glu->op_params)[1] /* swapped */) { + return false; + } + + if (glu->src[0] != gate || glu->src[1] != up) { + return false; + } + + const ggml_tensor * gate_w = gate->src[0]; + const ggml_tensor * up_w = up->src[0]; + const ggml_tensor * src1 = gate->src[1]; + const ggml_tensor * ids = gate->src[2]; + + if (!ids || up->src[1] != src1 || up->src[2] != ids) { + return false; + } + + if (gate_w->type != up_w->type || !ggml_are_same_shape(gate_w, up_w) || !ggml_are_same_stride(gate_w, up_w)) { + return false; + } + + if (src1->type != GGML_TYPE_F32 || gate->type != GGML_TYPE_F32 || up->type != GGML_TYPE_F32 || + glu->type != GGML_TYPE_F32 || ids->type != GGML_TYPE_I32) { + return false; + } + + // The fused kernel writes dst with the same (row, expert slot, token) addressing the unfused + // mul_mat_id would use, so the GLU output has to be the plain contiguous result. + if (!ggml_is_contiguous(glu) || !ggml_are_same_shape(glu, gate)) { + return false; + } + + if (gate_w->ne[3] != 1 || src1->ne[3] != 1) { + return false; + } + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || !turing_mma_available(cc)) { + return false; + } + + if (!mmq_glu_fusion_supported(gate_w->type)) { + return false; + } + + if (!ggml_cuda_should_use_mmq(gate_w->type, cc, src1->ne[2], gate_w->ne[2])) { + return false; + } + + // Leave the batch-1 / small-batch dispatch (mul_mat_vec_q and its own fusion) untouched. + if (src1->ne[2] <= MMVQ_MAX_BATCH_SIZE) { + return false; + } + + // A fused row tile holds mmq_y/2 gate rows and mmq_y/2 up rows; require an exact fit so the + // row bounds check can stay off. + const int mmq_y = get_mmq_y_host(cc); + if (gate_w->ne[1] % (mmq_y/2) != 0) { + return false; + } + + if (getenv("GGML_CUDA_MMQ_GLU_FUSION_VERBOSE")) { + static bool logged = false; + if (!logged) { + logged = true; + fprintf(stderr, "%s: fusing mul_mat_id + SwiGLU into MMQ (type=%s, n_ff=%lld, n_expert=%lld, n_tokens=%lld)\n", + __func__, ggml_type_name(gate_w->type), (long long) gate_w->ne[1], + (long long) gate_w->ne[2], (long long) src1->ne[2]); + } + } + + return true; +} + +void ggml_cuda_mul_mat_q_fused_glu( + ggml_backend_cuda_context & ctx, const ggml_tensor * gate_w, const ggml_tensor * up_w, + const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ids && ids->type == GGML_TYPE_I32); + GGML_ASSERT(gate_w->type == up_w->type); + + cudaStream_t stream = ctx.stream(); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + + const int64_t ne00 = gate_w->ne[0]; + const int64_t ne01 = gate_w->ne[1]; + const int64_t ne02 = gate_w->ne[2]; + + const int64_t ne10 = src1->ne[0]; + const int64_t ne11 = src1->ne[1]; + const int64_t ne12 = src1->ne[2]; + const int64_t ne13 = src1->ne[3]; + + GGML_ASSERT(ne13 == 1); + GGML_ASSERT(gate_w->ne[3] == 1); + GGML_ASSERT(dst->ne[0] == ne01); + + const size_t ts_src0 = ggml_type_size(gate_w->type); + const size_t ts_src1 = ggml_type_size(src1->type); + const size_t ts_dst = ggml_type_size(dst->type); + + GGML_ASSERT(gate_w->nb[0] == ts_src0); + GGML_ASSERT(src1->nb[0] == ts_src1); + GGML_ASSERT(dst->nb[0] == ts_dst); + GGML_ASSERT(ids->nb[0] == ggml_type_size(ids->type)); + GGML_ASSERT(src1->nb[2] % src1->nb[1] == 0); + GGML_ASSERT(dst->nb[2] % dst->nb[1] == 0); + + const int64_t s01 = gate_w->nb[1] / ts_src0; + const int64_t s02 = gate_w->nb[2] / ts_src0; + const int64_t s1 = dst->nb[1] / ts_dst; + const int64_t s2 = dst->nb[2] / ts_dst; + + GGML_ASSERT((int64_t) (up_w->nb[1] / ts_src0) == s01); + GGML_ASSERT((int64_t) (up_w->nb[2] / ts_src0) == s02); + + // If a weight lives in a temporary compute buffer, clear any potential padding. + for (const ggml_tensor * w : { gate_w, up_w }) { + if (ggml_backend_buffer_get_usage(w->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size_data = ggml_nbytes(w); + const size_t size_alloc = ggml_backend_buffer_get_alloc_size(w->buffer, w); + if (size_alloc > size_data) { + GGML_ASSERT(ggml_is_contiguously_allocated(w)); + GGML_ASSERT(!w->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) w->data + size_data, 0, size_alloc - size_data, stream)); + } + } + } + + const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); + + const int64_t n_expert_used = ids->ne[0]; + const int64_t ne_get_rows = ne12 * n_expert_used; + GGML_ASSERT(dst->ne[1] == n_expert_used); + + ggml_cuda_pool_alloc ids_src1(ctx.pool(), ne_get_rows); + ggml_cuda_pool_alloc ids_dst(ctx.pool(), ne_get_rows); + ggml_cuda_pool_alloc expert_bounds(ctx.pool(), ne02 + 1); + + { + const int si1 = ids->nb[1] / ggml_element_size(ids); + const int sis1 = src1->nb[2] / src1->nb[1]; + + ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), + ne02, ne12, n_expert_used, ne11, si1, sis1, stream); + CUDA_CHECK(cudaGetLastError()); + } + + const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * sizeof(block_q8_1)/QK8_1 + + get_mmq_x_max_host(cc)*sizeof(block_q8_1_mmq); + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + + { + const int64_t s11 = src1->nb[1] / ts_src1; + const int64_t s12 = src1->nb[2] / ts_src1; + const int64_t s13 = src1->nb[3] / ts_src1; + + quantize_mmq_q8_1_cuda((const float *) src1->data, ids_src1.get(), src1_q8_1.get(), gate_w->type, ne10, + s11, s12, s13, ne10_padded, ne12*n_expert_used, 1, 1, stream); + CUDA_CHECK(cudaGetLastError()); + } + + const int64_t s12_q = ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); + const int64_t s13_q = ne12*s12_q; + + // nrows_x counts the gate and the up rows: a fused row tile draws half of its rows from each. + const mmq_args args = { + (const char *) gate_w->data, gate_w->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), + (float *) dst->data, + ne00, 2*ne01, ne_get_rows, s01, ne_get_rows, s1, + ne02, ne02, s02, s12_q, s2, + 1, 1, 0, s13_q, 0, + /*use_stream_k =*/ true, ne12, + (const char *) up_w->data}; + + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); +} + void ggml_cuda_op_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 08d8e98dc2f..2ae850b4240 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -3,6 +3,7 @@ #include "common.cuh" #include "vecdotq.cuh" #include "mma.cuh" +#include "unary.cuh" #include #include @@ -3352,6 +3353,98 @@ static __device__ __forceinline__ void mmq_write_back_mma( } } +// Types for which the fused SwiGLU epilogue is compiled. Every type listed here doubles the +// number of MMQ kernels compiled for it, so keep this to the types where the fusion pays off. +static constexpr bool mmq_glu_fusion_supported(ggml_type type) { + return type == GGML_TYPE_Q1_0 || type == GGML_TYPE_Q2_0; +} + +// Write-back for the fused SwiGLU epilogue. +// +// The row tile holds mmq_y/2 rows of the gate matrix in rows [0, mmq_y/2) and the mmq_y/2 matching +// rows of the up matrix in rows [mmq_y/2, mmq_y). Both halves therefore live in the registers of a +// single block, but not of a single warp, so the up half is staged through shared memory before the +// gate half consumes it. Only mmq_y/2 rows are written to dst, which is the SwiGLU output. +template +static __device__ __forceinline__ void mmq_write_back_mma_glu( + const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst, + const int stride, const int j_max, float * __restrict__ scratch) { +#if defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int granularity = mmq_get_granularity_device(mmq_x); + +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int tileC_IJ = mmq_get_granularity_device(0); + typedef tile tile_C; + constexpr int rows_per_warp = granularity; +#else + typedef tile<16, 8, int> tile_C; + constexpr int rows_per_warp = 2 * granularity; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + constexpr int nrows_glu = mmq_y/2; + static_assert(mmq_y % 2 == 0, "mmq_y must be even for the fused GLU epilogue"); + // The scratch buffer aliases the x tile, which holds mmq_y*mmq_get_mma_tile_x_k(type) ints. + // MMQ_MMA_TILE_X_K_Q8_0 is the smallest of those, so this bound is conservative. + static_assert(nrows_glu*mmq_x <= mmq_y*MMQ_MMA_TILE_X_K_Q8_0, "GLU scratch does not fit in the x tile"); + static_assert(mmq_get_nwarps_device()*tile_C::I == mmq_y, "nwarps*tile_C::I != mmq_y"); + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I); + + // Stage the up half. +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + + if (i < nrows_glu) { + continue; + } + + const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); + + scratch[(i - nrows_glu)*mmq_x + j] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + } + } + } + + __syncthreads(); + + // Apply the epilogue on the gate half and write the only output of the fused op. +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + + if (i >= nrows_glu) { + continue; + } + + const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); + + if (j > j_max) { + continue; + } + + const float gate = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + const float up = scratch[i*mmq_x + j]; + + dst[ids_dst[j]*stride + i] = ggml_cuda_op_silu_single(gate) * up; + } + } + } +#else + NO_DEVICE_CODE; + GGML_UNUSED_VARS(sum, ids_dst, dst, stride, j_max, scratch); +#endif // defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + // ------------------------------------------------------------------------------------------------------------------------------------- template @@ -3543,9 +3636,9 @@ struct mmq_type_traits { static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; }; -template +template static __device__ __forceinline__ void mul_mat_q_process_tile( - const char * __restrict__ x, const int offset_x, const int * __restrict__ y, + const char * __restrict__ x, const char * __restrict__ x_up, const int offset_x, const int * __restrict__ y, const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, const int stride_row_x, const int ncols_y, const int stride_col_dst, const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { @@ -3583,7 +3676,19 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { - load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + if constexpr (fuse_glu) { + // Load mmq_y/2 gate rows followed by the mmq_y/2 matching up rows into one x tile. + // need_check is forced on for the half loads because a single loader pass can cover + // more rows than mmq_y/2; the clamp keeps those writes inside the half tile. + constexpr int glu_nrows = mmq_y/2; + constexpr load_tiles_mmq_t load_tiles_glu = mmq_type_traits::load_tiles; + constexpr int glu_tile_stride = mmq_get_mma_tile_x_k(type); + + load_tiles_glu(x, tile_x, offset_x + kb0, glu_nrows - 1, stride_row_x); + load_tiles_glu(x_up, tile_x + glu_nrows*glu_tile_stride, offset_x + kb0, glu_nrows - 1, stride_row_x); + } else { + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + } { const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz; #pragma unroll @@ -3617,7 +3722,11 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( __syncthreads(); } - if (fixup) { + if constexpr (fuse_glu) { + __syncthreads(); // the x tile is reused as scratch space by the epilogue + mmq_write_back_mma_glu(sum, ids_dst, dst, stride_col_dst, tile_y_max_j, (float *) tile_x); + GGML_UNUSED(write_back); + } else if (fixup) { write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(mmq_x*mmq_y), mmq_y, mmq_y, mmq_x); } else { write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j); @@ -3627,7 +3736,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 -template +template #if defined(GGML_USE_HIP) #if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) @@ -3640,7 +3749,8 @@ template #endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA #endif // defined(GGML_USE_HIP) static __global__ void mul_mat_q( - const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, + const char * __restrict__ x, const char * __restrict__ x_up, + const int * __restrict__ y, const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst, const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, @@ -3653,12 +3763,23 @@ static __global__ void mul_mat_q( return; } +#if !(defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)) + // The fused epilogue relies on the mma x-tile layout, whose row stride does not depend on mmq_y. + if constexpr (fuse_glu) { + NO_DEVICE_CODE; + return; + } +#endif // !(defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)) + constexpr int nwarps = mmq_get_nwarps_device(); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int qk = ggml_cuda_type_traits::qk; constexpr int mmq_y = get_mmq_y_device(); + // With the fused epilogue a row tile consumes mmq_y weight rows but produces mmq_y/2 output rows. + constexpr int mmq_y_dst = fuse_glu ? mmq_y/2 : mmq_y; + const uint32_t nty = (nrows_x + mmq_y - 1) / mmq_y; // Number of tiles y // Initialize the ids for writing back data with just the index. @@ -3720,16 +3841,16 @@ static __global__ void mul_mat_q( } offset_y += (col_low + jt*mmq_x)*(sizeof(block_q8_1_mmq)/sizeof(int)); - offset_dst += it*mmq_y; + offset_dst += it*mmq_y_dst; const int tile_x_max_i = nrows_x - it*mmq_y - 1; const int tile_y_max_j = col_diff - jt*mmq_x - 1; - const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x; + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y_dst*stride_row_x; constexpr bool fixup = false; - mul_mat_q_process_tile - (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + mul_mat_q_process_tile + (x, x_up, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); return; } @@ -3800,16 +3921,16 @@ static __global__ void mul_mat_q( } offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int)); - offset_dst += it*mmq_y; + offset_dst += it*mmq_y_dst; const int tile_x_max_i = nrows_x - it*mmq_y - 1; const int tile_y_max_j = col_diff - jt*mmq_x - 1; - const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x; + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y_dst*stride_row_x; constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. - mul_mat_q_process_tile - (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + mul_mat_q_process_tile + (x, x_up, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); kbc += blocks_per_ne00.z; @@ -3819,6 +3940,12 @@ static __global__ void mul_mat_q( kb0_stop = min(blocks_per_ne00.z, uint32_t(kbc_stop - kbc)); } + if constexpr (fuse_glu) { + // The epilogue needs the complete k range of a tile inside one block. The launcher uses one + // block per output tile for the fused path, so the stream-k fixup tail is unreachable. + return; + } + if (kbc >= kbc_stop) { return; } @@ -3869,16 +3996,16 @@ static __global__ void mul_mat_q( } offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int)); - offset_dst += it*mmq_y; + offset_dst += it*mmq_y_dst; const int tile_x_max_i = nrows_x - it*mmq_y - 1; const int tile_y_max_j = col_diff - jt*mmq_x - 1; - const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x; + const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*mmq_y_dst*stride_row_x; constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. - mul_mat_q_process_tile - (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + mul_mat_q_process_tile + (x, x_up, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); } @@ -4027,6 +4154,9 @@ struct mmq_args { int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst; int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst; bool use_stream_k; int64_t ncols_max; + // Fused SwiGLU epilogue: when non-null, dst = silu(x @ y) * (x_up @ y) and nrows_x is the sum of + // the row counts of x and x_up. Only mmq_glu_fusion_supported() types honour this. + const char * x_up = nullptr; }; template @@ -4072,11 +4202,30 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); + if constexpr (mmq_glu_fusion_supported(type)) { + if (args.x_up) { + // One block per output tile: every tile is completed by a single block, so the fused + // epilogue always sees the full k range and the stream-k fixup is never needed. + GGML_ASSERT(args.nrows_x % mmq_y == 0); + const int ntiles_dst_glu = nty * ntx * ntzw; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + mul_mat_q<<>> + (args.x, args.x_up, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + ntx_fd); + return; + } + } + if (!args.use_stream_k) { if (args.nrows_x % mmq_y == 0) { constexpr bool need_check = false; mul_mat_q<<>> - (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + (args.x, nullptr, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, @@ -4084,7 +4233,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a } else { constexpr bool need_check = true; mul_mat_q<<>> - (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + (args.x, nullptr, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, @@ -4116,7 +4265,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a if (args.nrows_x % mmq_y == 0) { constexpr bool need_check = false; mul_mat_q<<>> - (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, + (args.x, nullptr, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, @@ -4134,7 +4283,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a } else { constexpr bool need_check = true; mul_mat_q<<>> - (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, + (args.x, nullptr, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, @@ -4266,6 +4415,14 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); void ggml_cuda_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); +// Fused routed gate/up mul_mat_id + SwiGLU. Writes only the SwiGLU result; the gate and up products +// are never materialized in global memory. +void ggml_cuda_mul_mat_q_fused_glu( + ggml_backend_cuda_context & ctx, const ggml_tensor * gate_w, const ggml_tensor * up_w, + const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); + +bool ggml_cuda_can_fuse_mmq_glu(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu); + void ggml_cuda_op_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, From 06789550f239f228fc75d9c25570826cce7fe890 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Mon, 3 Aug 2026 07:40:19 -1000 Subject: [PATCH 2/2] cuda: fold the MoE router weighting and expert reduction into one kernel The down projection wrote one slice per selected expert, which was then scaled by the router weight and reduced by an add chain. This applies the weight and accumulates across the selected experts before the final store, removing a full-size intermediate tensor and the round trip the add chain needed. Output is bit-identical. Prefill improves; decode is unaffected. Disable with GGML_CUDA_MOE_REDUCE_DISABLE. --- ggml/src/ggml-cuda/ggml-cuda.cu | 147 +++++++++++++++++++++++ ggml/src/ggml-cuda/moe-reduce.cu | 191 ++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/moe-reduce.cuh | 19 +++ 3 files changed, 357 insertions(+) create mode 100644 ggml/src/ggml-cuda/moe-reduce.cu create mode 100644 ggml/src/ggml-cuda/moe-reduce.cuh diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f4523bd035b..f7b17eb15e5 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -31,6 +31,7 @@ #include "ggml-cuda/mmq.cuh" #include "ggml-cuda/mmvf.cuh" #include "ggml-cuda/mmvq.cuh" +#include "ggml-cuda/moe-reduce.cuh" #include "ggml-cuda/norm.cuh" #include "ggml-cuda/opt-step-adamw.cuh" #include "ggml-cuda/opt-step-sgd.cuh" @@ -3835,6 +3836,117 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } +static bool ggml_cuda_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b) { + const int64_t a_start = (int64_t) a->data; + const int64_t a_end = a_start + (int64_t) ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + + const int64_t b_start = (int64_t) b->data; + const int64_t b_end = b_start + (int64_t) ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + + return (b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end); +} + +// MoE weighted-expert reduce, as emitted by llm_graph_context::build_moe_ffn: +// MUL(down_out, router_weights) -> n_exp x VIEW -> (n_exp - 1) x ADD +// The MUL writes a second full-size [n_embd, n_expert_used, n_tokens] F32 tensor that only +// exists so the add chain can read it back, which at large ubatch is tens of MB written and +// read again per layer. Folding the weighting into the reduction removes that round trip and +// is bit-identical to the unfused sequence. +static bool ggml_cuda_should_fuse_moe_weighted_reduce(const ggml_cgraph * cgraph, + int node_idx, + int n_exp, + bool & stage_weights) { + const ggml_tensor * mul = cgraph->nodes[node_idx]; + const ggml_tensor * x = mul->src[0]; + const ggml_tensor * w = mul->src[1]; + + stage_weights = false; + + if (!x || !w) { + return false; + } + + if (mul->type != GGML_TYPE_F32 || x->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_is_contiguous(mul) || !ggml_is_contiguous(x) || !ggml_are_same_shape(mul, x)) { + return false; + } + + const int64_t ne0 = mul->ne[0]; + const int64_t n_tok = mul->ne[2]; + + // the router weights must be a per (expert, token) scalar, broadcast along ne0 + if (w->ne[0] != 1 || w->ne[1] != n_exp || w->ne[2] != n_tok || w->ne[3] != 1) { + return false; + } + + // one block row per token + if (n_tok > 65535) { + return false; + } + + // n_exp views of the weighted tensor, one per expert slot, in slot order + for (int e = 0; e < n_exp; ++e) { + const ggml_tensor * v = cgraph->nodes[node_idx + 1 + e]; + + if (v->view_src != mul || v->type != GGML_TYPE_F32) { + return false; + } + if (v->ne[0] != ne0 || v->ne[1] != n_tok || v->ne[2] != 1 || v->ne[3] != 1) { + return false; + } + if (v->nb[0] != mul->nb[0] || v->nb[1] != mul->nb[2]) { + return false; + } + if ((const char *) v->data != (const char *) mul->data + e*mul->nb[1]) { + return false; + } + } + + // left-to-right add chain over those views + for (int k = 0; k < n_exp - 1; ++k) { + const ggml_tensor * add = cgraph->nodes[node_idx + 1 + n_exp + k]; + + const ggml_tensor * expected_a = (k == 0) ? cgraph->nodes[node_idx + 1] : cgraph->nodes[node_idx + n_exp + k]; + const ggml_tensor * expected_b = cgraph->nodes[node_idx + 2 + k]; + + if (add->src[0] != expected_a || add->src[1] != expected_b) { + return false; + } + if (add->type != GGML_TYPE_F32 || !ggml_is_contiguous(add)) { + return false; + } + if (add->ne[0] != ne0 || add->ne[1] != n_tok || add->ne[2] != 1 || add->ne[3] != 1) { + return false; + } + } + + // Aliasing. The fused kernel reads x and w while it is writing dst, so anything the + // graph allocator has placed under dst is a cross-block race. ggml_cuda_check_fusion_ + // memory_ranges is not used here because it ignores sources whose op is GGML_OP_NONE + // and because it can only reject, while the weights case has a cheap repair. + const ggml_tensor * dst = cgraph->nodes[node_idx + 2*n_exp - 1]; + + if (ggml_cuda_tensors_overlap(dst, x)) { + return false; + } + + if (ggml_cuda_tensors_overlap(dst, w)) { + // The router weights die at the MUL, so the allocator routinely packs the add-chain + // output over them. They are only n_expert_used*n_tokens floats, so the op snapshots + // them into pool scratch rather than giving up the fusion. + if (!ggml_is_contiguous(w)) { + return false; + } + + stage_weights = true; + } + + return true; +} + // try and fuse nodes and return the number of nodes to skip static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { @@ -3968,6 +4080,41 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // MoE router weighting + expert reduction [kill switch: GGML_CUDA_MOE_REDUCE_DISABLE=1] + static bool disable_moe_reduce = getenv("GGML_CUDA_MOE_REDUCE_DISABLE") != nullptr && + std::atoi(getenv("GGML_CUDA_MOE_REDUCE_DISABLE")); + + if (!disable_moe_reduce && node->op == GGML_OP_MUL && node->ne[3] == 1 && + node->ne[1] >= 2 && node->ne[1] <= GGML_CUDA_MOE_REDUCE_MAX_EXPERTS) { + + const int n_exp = (int) node->ne[1]; + const int n_nodes = 2*n_exp; // MUL + n_exp views + (n_exp - 1) adds + + std::vector ops; + ops.reserve(n_nodes); + ops.push_back(GGML_OP_MUL); + ops.insert(ops.end(), n_exp, GGML_OP_VIEW); + ops.insert(ops.end(), n_exp - 1, GGML_OP_ADD); + + const int out_node = i + n_nodes - 1; + + bool stage_weights = false; + + if (ggml_can_fuse_subgraph(cgraph, i, n_nodes, ops.data(), &out_node, 1) && + ggml_cuda_should_fuse_moe_weighted_reduce(cgraph, i, n_exp, stage_weights)) { + + static bool debug_moe_reduce = getenv("GGML_CUDA_MOE_REDUCE_DEBUG") != nullptr; + if (debug_moe_reduce) { + GGML_LOG_INFO("%s: fused moe weighted reduce: n_expert_used=%d ne0=%d n_tokens=%d staged=%d\n", + __func__, n_exp, (int) node->ne[0], (int) node->ne[2], (int) stage_weights); + } + + ggml_cuda_op_moe_weighted_reduce(*cuda_ctx, node->src[0], node->src[1], cgraph->nodes[out_node], + n_exp, stage_weights); + return n_nodes - 1; + } + } + // multi-(add or mul) if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { int n_fuse = 0; diff --git a/ggml/src/ggml-cuda/moe-reduce.cu b/ggml/src/ggml-cuda/moe-reduce.cu new file mode 100644 index 00000000000..2a086fe20c5 --- /dev/null +++ b/ggml/src/ggml-cuda/moe-reduce.cu @@ -0,0 +1,191 @@ +#include "moe-reduce.cuh" + +// Fused MoE router-weighting + expert reduction. +// +// build_moe_ffn emits, after the down projection: +// experts = ggml_mul(down_out, weights) // [n_embd, n_expert_used, n_tokens] +// view_e = ggml_view_2d(experts, ..., e) // n_expert_used views +// moe_out = (((view_0 + view_1) + view_2) + ...) // n_expert_used - 1 adds +// +// The ggml_mul materialises a second full-size [n_embd, n_expert_used, n_tokens] F32 +// tensor (tens of MB per layer at large ubatch) purely so that the add chain +// can read it back. This kernel folds the weighting into the reduction, so the weighted +// intermediate is never written or read: only down_out is streamed in and the +// [n_embd, n_tokens] result is streamed out. +// +// Numerics are bit-identical to the unfused pair. GGML_OP_MUL rounds each product to F32 +// before storing it, and the fused multi-add folds the slices left to right, so the fused +// kernel must round every product before accumulating and must accumulate in increasing +// expert-slot order. __fmul_rn / __fadd_rn keep nvcc from contracting a product and its +// accumulation into an FMA, which would be more accurate but would not match the baseline. + +template +static __global__ void k_moe_weighted_reduce(const float * __restrict__ x, + const float * __restrict__ w, + float * __restrict__ dst, + const int ne0, + const int64_t sx1, + const int64_t sx2, + const int64_t sw1, + const int64_t sw2, + const int64_t sd1) { + const int64_t i2 = blockIdx.y; + + const float * x_t = x + i2 * sx2; + const float * w_t = w + i2 * sw2; + float * d_t = dst + i2 * sd1; + + float wv[n_exp]; +#pragma unroll + for (int e = 0; e < n_exp; ++e) { + wv[e] = w_t[e * sw1]; + } + + for (int i0 = blockIdx.x * blockDim.x + threadIdx.x; i0 < ne0; i0 += blockDim.x * gridDim.x) { + float acc = __fmul_rn(x_t[i0], wv[0]); +#pragma unroll + for (int e = 1; e < n_exp; ++e) { + acc = __fadd_rn(acc, __fmul_rn(x_t[i0 + e * sx1], wv[e])); + } + d_t[i0] = acc; + } +} + +template +static __global__ void k_moe_weighted_reduce_v4(const float4 * __restrict__ x, + const float * __restrict__ w, + float4 * __restrict__ dst, + const int ne0v, + const int64_t sx1v, + const int64_t sx2v, + const int64_t sw1, + const int64_t sw2, + const int64_t sd1v) { + const int64_t i2 = blockIdx.y; + + const float4 * x_t = x + i2 * sx2v; + const float * w_t = w + i2 * sw2; + float4 * d_t = dst + i2 * sd1v; + + float wv[n_exp]; +#pragma unroll + for (int e = 0; e < n_exp; ++e) { + wv[e] = w_t[e * sw1]; + } + + for (int i0 = blockIdx.x * blockDim.x + threadIdx.x; i0 < ne0v; i0 += blockDim.x * gridDim.x) { + const float4 x0 = x_t[i0]; + float4 acc = make_float4(__fmul_rn(x0.x, wv[0]), __fmul_rn(x0.y, wv[0]), + __fmul_rn(x0.z, wv[0]), __fmul_rn(x0.w, wv[0])); +#pragma unroll + for (int e = 1; e < n_exp; ++e) { + const float4 xe = x_t[i0 + e * sx1v]; + acc.x = __fadd_rn(acc.x, __fmul_rn(xe.x, wv[e])); + acc.y = __fadd_rn(acc.y, __fmul_rn(xe.y, wv[e])); + acc.z = __fadd_rn(acc.z, __fmul_rn(xe.z, wv[e])); + acc.w = __fadd_rn(acc.w, __fmul_rn(xe.w, wv[e])); + } + d_t[i0] = acc; + } +} + +static constexpr int MOE_REDUCE_BLOCK = 256; + +template +static void launch_moe_weighted_reduce(const float * x, + const float * w, + float * dst, + const int64_t ne0, + const int64_t n_tok, + const int64_t sx1, + const int64_t sx2, + const int64_t sw1, + const int64_t sw2, + const int64_t sd1, + cudaStream_t stream) { + const bool use_v4 = ne0 % 4 == 0 && sx1 % 4 == 0 && sx2 % 4 == 0 && sd1 % 4 == 0 && + ((uintptr_t) x) % sizeof(float4) == 0 && ((uintptr_t) dst) % sizeof(float4) == 0; + + if (use_v4) { + const int64_t ne0v = ne0 / 4; + const dim3 grid((ne0v + MOE_REDUCE_BLOCK - 1) / MOE_REDUCE_BLOCK, n_tok, 1); + k_moe_weighted_reduce_v4<<>>( + (const float4 *) x, w, (float4 *) dst, ne0v, sx1 / 4, sx2 / 4, sw1, sw2, sd1 / 4); + } else { + const dim3 grid((ne0 + MOE_REDUCE_BLOCK - 1) / MOE_REDUCE_BLOCK, n_tok, 1); + k_moe_weighted_reduce<<>>( + x, w, dst, ne0, sx1, sx2, sw1, sw2, sd1); + } +} + +void ggml_cuda_op_moe_weighted_reduce(ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + ggml_tensor * dst, + int n_exp, + bool stage_weights) { + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(n_exp >= 2 && n_exp <= GGML_CUDA_MOE_REDUCE_MAX_EXPERTS); + + const int64_t ne0 = src0->ne[0]; + const int64_t n_tok = src0->ne[2]; + + const int64_t sx1 = src0->nb[1] / sizeof(float); + const int64_t sx2 = src0->nb[2] / sizeof(float); + int64_t sw1 = src1->nb[1] / sizeof(float); + int64_t sw2 = src1->nb[2] / sizeof(float); + const int64_t sd1 = dst->nb[1] / sizeof(float); + + const float * x = (const float *) src0->data; + const float * w = (const float *) src1->data; + float * d = (float *) dst->data; + + cudaStream_t stream = ctx.stream(); + + // dst aliases the router weights, so snapshot them before the reduce starts writing. + // n_expert_used*n_tokens floats, a few KiB at typical shapes. + ggml_cuda_pool_alloc w_stage(ctx.pool()); + + if (stage_weights) { + GGML_ASSERT(ggml_is_contiguous(src1)); + + const size_t n_w = ggml_nelements(src1); + w_stage.alloc(n_w); + + CUDA_CHECK(cudaMemcpyAsync(w_stage.ptr, w, n_w * sizeof(float), cudaMemcpyDeviceToDevice, stream)); + + w = w_stage.ptr; + sw1 = 1; + sw2 = n_exp; + } + +#define MOE_REDUCE_CASE(N) \ + case N: \ + launch_moe_weighted_reduce(x, w, d, ne0, n_tok, sx1, sx2, sw1, sw2, sd1, stream); \ + break; + + switch (n_exp) { + MOE_REDUCE_CASE(2) + MOE_REDUCE_CASE(3) + MOE_REDUCE_CASE(4) + MOE_REDUCE_CASE(5) + MOE_REDUCE_CASE(6) + MOE_REDUCE_CASE(7) + MOE_REDUCE_CASE(8) + MOE_REDUCE_CASE(9) + MOE_REDUCE_CASE(10) + MOE_REDUCE_CASE(11) + MOE_REDUCE_CASE(12) + MOE_REDUCE_CASE(13) + MOE_REDUCE_CASE(14) + MOE_REDUCE_CASE(15) + default: + GGML_ABORT("fatal error"); + } + +#undef MOE_REDUCE_CASE + + CUDA_CHECK(cudaGetLastError()); +} diff --git a/ggml/src/ggml-cuda/moe-reduce.cuh b/ggml/src/ggml-cuda/moe-reduce.cuh new file mode 100644 index 00000000000..c3e43b93b10 --- /dev/null +++ b/ggml/src/ggml-cuda/moe-reduce.cuh @@ -0,0 +1,19 @@ +#pragma once + +#include "common.cuh" + +// maximum n_expert_used handled by the fused MoE weighted reduce +// (bounded by ggml_can_fuse_subgraph's 32-node limit: the subgraph is 2*n_exp nodes) +#define GGML_CUDA_MOE_REDUCE_MAX_EXPERTS 15 + +// dst[i0, i2] = sum_e src0[i0, e, i2] * src1[0, e, i2] +// +// stage_weights copies src1 into pool scratch before the reduce runs. The caller needs it +// when the graph allocator has placed the (tiny, already dead) router-weight tensor inside +// the destination buffer, which would otherwise be a cross-block read/write race. +void ggml_cuda_op_moe_weighted_reduce(ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + ggml_tensor * dst, + int n_exp, + bool stage_weights);