cuda: MoE prefill — fused SwiGLU epilogue and fused weighted reduction

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<ggml_op> 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;
+            }
         }
     }
 

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<int32_t> ids_src1(ctx.pool(), ne_get_rows);
+    ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows);
+    ggml_cuda_pool_alloc<int32_t> 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<char> 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,