diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 13e1b8a2e737..f7b17eb15e5d 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; @@ -4124,6 +4271,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 a3cc50323155..cbb328e74d78 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 08d8e98dc2f7..2ae850b42401 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, diff --git a/ggml/src/ggml-cuda/moe-reduce.cu b/ggml/src/ggml-cuda/moe-reduce.cu new file mode 100644 index 000000000000..2a086fe20c52 --- /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 000000000000..c3e43b93b10d --- /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);