diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 02d8e1ca597..0f660694a44 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1493,6 +1493,7 @@ struct ggml_backend_cuda_context { int curr_stream_no = 0; ggml_cuda_mmb_context * mmb = nullptr; + bool mmb_opt_in = false; // set by ggml_backend_cuda_set_mmb_enabled, before the first graph bool mmb_after_compute = true; const void * mmb_first_split = nullptr; std::vector mmb_graph_sigs; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f3a2bf7d881..158bb676409 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -791,7 +791,9 @@ static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer if (padded_size > original_size) { ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size)); + // not cudaMemset: on HIP a synchronous memset invalidates a graph capture running on any other thread + CUDA_CHECK(cudaMemsetAsync((char *) tensor->data + original_size, 0, padded_size - original_size, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); } } return GGML_STATUS_SUCCESS; @@ -1894,7 +1896,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); return; } - if (ggml_cuda_mmb_supported_mm(src0, src1, dst)) { + if (ggml_cuda_mmb_supported_mm(ctx, src0, src1, dst)) { ggml_cuda_mul_mat_mmb(ctx, src0, src1, dst); return; } @@ -1966,7 +1968,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } } - if (ggml_cuda_mmb_supported_mmid(src0, src1, ids, dst)) { + if (ggml_cuda_mmb_supported_mmid(ctx, src0, src1, ids, dst)) { ggml_cuda_mul_mat_id_mmb(ctx, src0, src1, ids, dst); return; } @@ -4210,14 +4212,14 @@ static bool ggml_cuda_qsa_expand_alias_ok(const ggml_cuda_qsa_expand_args & args } // true when the BF16 WMMA GEMM would take this MUL_MAT / MUL_MAT_ID, so MMQ-based fusions leave it alone -static bool ggml_cuda_mmb_claims(const ggml_tensor * t) { +static bool ggml_cuda_mmb_claims(ggml_backend_cuda_context & ctx, const ggml_tensor * t) { if (!t || !t->src[0] || !t->src[1]) { return false; } if (t->op == GGML_OP_MUL_MAT_ID) { - return t->src[2] && ggml_cuda_mmb_supported_mmid(t->src[0], t->src[1], t->src[2], t); + return t->src[2] && ggml_cuda_mmb_supported_mmid(ctx, t->src[0], t->src[1], t->src[2], t); } - return t->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_supported_mm(t->src[0], t->src[1], t); + return t->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_supported_mm(ctx, t->src[0], t->src[1], t); } // ggml_cuda_try_fuse: node i was computed on its own and no later node is skipped @@ -4236,7 +4238,7 @@ static uint16_t * ggml_cuda_bf16_for_mmb_reader(ggml_backend_cuda_context & ctx, } const ggml_tensor * x = n->src[1]; if (n->op == GGML_OP_MUL_MAT && x && (x == t || (x->view_src == t && x->view_offs == 0)) && ggml_is_contiguous(x) && - ggml_nelements(x) == ggml_nelements(t) && ggml_cuda_mmb_supported_mm(n->src[0], x, n)) { + ggml_nelements(x) == ggml_nelements(t) && ggml_cuda_mmb_supported_mm(ctx, n->src[0], x, n)) { return ggml_cuda_mmb_cache_produce(ctx, x, (size_t) ggml_nelements(x)); } return nullptr; @@ -4268,7 +4270,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * gate = glu->src[0]; ggml_tensor * up = glu->src[1]; const bool paired = (gate == node && up == cgraph->nodes[i + 1]) || (gate == cgraph->nodes[i + 1] && up == node); - if (paired && ggml_cuda_mmb_supported_glu(gate->src[0], up->src[0], up->src[1], up->src[2], glu) && + if (paired && ggml_cuda_mmb_supported_glu(*cuda_ctx, gate->src[0], up->src[0], up->src[1], up->src[2], glu) && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }, { i + 2 })) { const int out_nodes[] = { i + 2 }; if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 1)) { @@ -4515,7 +4517,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } if (node->op == GGML_OP_GLU && i + 1 < cgraph->n_nodes && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU && node->src[1] && - !ggml_cuda_mmb_is_bf16_only(*cuda_ctx, node) && !ggml_cuda_mmb_claims(cgraph->nodes[i + 1])) { + !ggml_cuda_mmb_is_bf16_only(*cuda_ctx, node) && !ggml_cuda_mmb_claims(*cuda_ctx, cgraph->nodes[i + 1])) { ggml_tensor * next = cgraph->nodes[i + 1]; const bool has_ids = next->op == GGML_OP_MUL_MAT_ID; if (has_ids || next->op == GGML_OP_MUL_MAT) { @@ -4790,7 +4792,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph // [n_embd, n_used] intermediate and its separate reduction launch. if (getenv("GGML_CUDA_DISABLE_WEIGHTED_DOWN") == nullptr && node->op == GGML_OP_MUL_MAT_ID && i + 20 < cgraph->n_nodes && cgraph->nodes[i + 1]->op == GGML_OP_MUL && - !ggml_cuda_mmb_claims(node)) { + !ggml_cuda_mmb_claims(*cuda_ctx, node)) { constexpr int n_used = 10; constexpr int n_ops = 2 + n_used + (n_used - 1); ggml_tensor * mul = cgraph->nodes[i + 1]; @@ -4929,7 +4931,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph const bool use_mmq = shared_inputs && ggml_cuda_should_use_mmq(src0->type, cc, mmq_cols, n_experts) && ggml_cuda_should_use_mmq(src0_next->type, cc, mmq_cols, n_experts); - const bool compatible = use_mmq && !ggml_cuda_mmb_claims(node) && node->src[1]->type == GGML_TYPE_F32 && + const bool compatible = use_mmq && !ggml_cuda_mmb_claims(*cuda_ctx, node) && node->src[1]->type == GGML_TYPE_F32 && node->type == GGML_TYPE_F32 && next->type == GGML_TYPE_F32 && ggml_are_same_shape(src0, src0_next) && ggml_are_same_shape(node, next) && mmq_get_q8_1_ds_layout(src0->type) == mmq_get_q8_1_ds_layout(src0_next->type) && @@ -5263,7 +5265,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph const ggml_tensor * src1 = up->src[1]; const ggml_tensor * ids = up->src[2]; - if (op == GGML_OP_MUL_MAT_ID && ggml_cuda_mmb_supported_glu(gate->src[0], up->src[0], src1, ids, glu)) { + if (op == GGML_OP_MUL_MAT_ID && ggml_cuda_mmb_supported_glu(*cuda_ctx, gate->src[0], up->src[0], src1, ids, glu)) { ggml_cuda_mul_mat_id_mmb_glu(*cuda_ctx, gate->src[0], up->src[0], src1, ids, glu); fused_mul_mat_vec = true; fused_node_count = 3; @@ -5533,7 +5535,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph // HC gate GEMM [320 -> 10240] whose only consumer is the fused stream mix: GEMM + sigmoid + mix in one kernel if (node->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_gatemix() && i + 1 < cgraph->n_nodes && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { const ggml_tensor * w = node->src[0], * lo = node->src[1]; - if (ggml_is_quantized(w->type) && ggml_node_has_n_uses(cgraph, i, 1) && ggml_cuda_mmb_supported_mm(w, lo, node)) { + if (ggml_is_quantized(w->type) && ggml_node_has_n_uses(cgraph, i, 1) && ggml_cuda_mmb_supported_mm(*cuda_ctx, w, lo, node)) { ggml_cuda_hc_mix_args ma; const int count = ggml_cuda_hc_mix_closed(cgraph, i + 1, ma); if (count > 0 && ma.gate == node && ggml_cuda_hc_gate_mix(*cuda_ctx, w, lo, ma.xn, ma.dst, ma.hc, ma.scale, ma.bias)) return count; @@ -6169,7 +6171,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph ++nread; if (t == ca.out_inject) continue; if (t->op == GGML_OP_MUL_MAT && t->src[0]->type != GGML_TYPE_F32 && - ggml_cuda_mmb_supported_mm(t->src[0], t->src[1], t)) continue; + ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) continue; if (t->op == GGML_OP_MUL && n >= 1) { ggml_cuda_hc_mix_args ma; if (ggml_cuda_hc_mix_closed(cgraph, n - 1, ma) > 0 && (ma.xn == t->src[0] || ma.xn == t->src[1]) && (ma.xn == xn || ma.xn->view_src == xn)) continue; } if (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE) continue; ok = false; @@ -6186,7 +6188,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph for (const ggml_tensor * b : blks) { if (ggml_nrows(b) < 512 || b->ne[0] % 8 != 0) continue; bool producer_ok = false; // only producers that honour a BF16 mark - if (b->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_supported_mm(b->src[0], b->src[1], b)) producer_ok = true; + if (b->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_supported_mm(*cuda_ctx, b->src[0], b->src[1], b)) producer_ok = true; if (!producer_ok) { for (int k = 0; k < cgraph->n_nodes && !producer_ok; ++k) { ggml_cuda_moe_weighted_reduction_match mm; @@ -6257,14 +6259,14 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph ++nread; if (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE) continue; if (t->op == GGML_OP_MUL_MAT && t->src[0]->type != GGML_TYPE_F32 && - ggml_cuda_mmb_supported_mm(t->src[0], t->src[1], t)) continue; + ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) continue; ok = false; } if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, d)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, d); } for (int i = 0; i < cgraph->n_nodes; ++i) { // the HC gate GEMM feeding only the fused mix keeps its BF16 epilogue const ggml_tensor * t = cgraph->nodes[i]; - if (t->op != GGML_OP_MUL_MAT || t->src[0]->ne[0] != 320 || t->src[0]->ne[1] != 10240 || !ggml_cuda_mmb_supported_mm(t->src[0], t->src[1], t)) continue; + if (t->op != GGML_OP_MUL_MAT || t->src[0]->ne[0] != 320 || t->src[0]->ne[1] != 10240 || !ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) continue; bool ok = true; int nread = 0; for (int n = i + 1; n < cgraph->n_nodes && ok; ++n) { const ggml_tensor * u = cgraph->nodes[n]; @@ -6283,7 +6285,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph int xi = -1; for (int k = 0; k < i; ++k) { if (cgraph->nodes[k] == ex) { xi = k; break; } } if (xi < 0 || ex->op != GGML_OP_MUL_MAT_ID || ex->type != GGML_TYPE_F32) continue; - if (!ggml_cuda_mmb_supported_mmid(ex->src[0], ex->src[1], ex->src[2], const_cast(ex))) continue; + if (!ggml_cuda_mmb_supported_mmid(*cuda_ctx, ex->src[0], ex->src[1], ex->src[2], const_cast(ex))) continue; if (ex->src[0]->type != GGML_TYPE_IQ4_NL) continue; // other formats keep F32 expert outputs if (!ggml_node_has_n_uses(cgraph, xi, 1)) continue; if (ex->ne[0] % 8 != 0 || ggml_nrows(ex) < 512) continue; @@ -6306,12 +6308,12 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph auto reads = [](const ggml_tensor * t, const ggml_tensor * x) { for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) if (t->src[s] == x || t->src[s]->view_src == x) return true; return false; }; for (int i = 0; i < cgraph->n_nodes; ++i) { const ggml_tensor * t = cgraph->nodes[i]; - if (t->op == GGML_OP_MUL_MAT && (t->src[0]->type == GGML_TYPE_IQ4_NL || t->src[0]->type == GGML_TYPE_Q6_K) && ggml_cuda_mmb_supported_mm(t->src[0], t->src[1], t)) ggml_cuda_mmb_shadow_prepare(*cuda_ctx, t->src[0]); + if (t->op == GGML_OP_MUL_MAT && (t->src[0]->type == GGML_TYPE_IQ4_NL || t->src[0]->type == GGML_TYPE_Q6_K) && ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) ggml_cuda_mmb_shadow_prepare(*cuda_ctx, t->src[0]); } for (int i = 0; i < cgraph->n_nodes; ++i) { const ggml_tensor * glu = cgraph->nodes[i]; if (glu->op != GGML_OP_GLU || !glu->src[0] || !glu->src[1] || glu->src[0]->op != GGML_OP_MUL_MAT_ID) continue; - if (!ggml_cuda_mmb_supported_glu(glu->src[0]->src[0], glu->src[1]->src[0], glu->src[0]->src[1], glu->src[0]->src[2], glu)) continue; + if (!ggml_cuda_mmb_supported_glu(*cuda_ctx, glu->src[0]->src[0], glu->src[1]->src[0], glu->src[0]->src[1], glu->src[0]->src[2], glu)) continue; for (const auto * input : {glu->src[0]->src[0], glu->src[1]->src[0], glu->src[0]->src[1], glu->src[0]->src[2], glu->src[0], glu->src[1]}) { auto * root = const_cast(input->view_src ? input->view_src : input); params->add_alloc_dep(params->user_data, root, const_cast(glu)); @@ -6321,7 +6323,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph const ggml_tensor * t = cgraph->nodes[n]; if (!reads(t, glu)) continue; ++nread; - if (t->op == GGML_OP_MUL_MAT_ID && t->src[1] == glu && ggml_cuda_mmb_supported_mmid(t->src[0], t->src[1], t->src[2], t)) continue; + if (t->op == GGML_OP_MUL_MAT_ID && t->src[1] == glu && ggml_cuda_mmb_supported_mmid(*cuda_ctx, t->src[0], t->src[1], t->src[2], t)) continue; ok = false; } if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, glu)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, glu); @@ -7598,8 +7600,17 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +// opt this backend context in to the BF16 WMMA matmul path (mmb); call before the first graph +static void ggml_backend_cuda_set_mmb_enabled(ggml_backend_t backend, bool enable) { + GGML_ASSERT(ggml_backend_is_cuda(backend)); + ((ggml_backend_cuda_context *) backend->context)->mmb_opt_in = enable; +} + static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_cuda_set_mmb_enabled") == 0) { + return (void *)ggml_backend_cuda_set_mmb_enabled; + } if (strcmp(name, "ggml_backend_comm_init") == 0) { return (void *)ggml_backend_cuda_comm_init; } diff --git a/ggml/src/ggml-cuda/mmb.cu b/ggml/src/ggml-cuda/mmb.cu index 1a0ae824ae7..25a6b04b606 100644 --- a/ggml/src/ggml-cuda/mmb.cu +++ b/ggml/src/ggml-cuda/mmb.cu @@ -703,10 +703,10 @@ static const uint16_t * mmb_shadow_lookup(ggml_backend_cuda_context & ctx, const auto it = mmb_state(ctx).shadow.find(w->data); return it == mmb_state(ctx).shadow.end() ? nullptr : it->second; } -// the BF16 WMMA kernels are RDNA3.5 (gfx1151) work: other devices keep the MMQ/MMVQ paths -bool mmb_enabled() { - const int id = ggml_cuda_get_device(); - return GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[id].cc); +// RDNA3.5 (gfx1151) only, tuned for qwen4exp shapes. On gfx1151 (ROCm 7.2.1) it lost to MMQ on other archs: dense qwen35 prefill 3.4-3.9x slower, MoE 14-28% (PR #75). +// So each backend context opts in by model arch. Drop the opt-in when mmb matches MMQ on those archs. +bool mmb_enabled(const ggml_backend_cuda_context & ctx) { + return ctx.mmb_opt_in && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc); } int mmb_min_t() { return 512; } int mmb_f32split_mode(){ return 2; } @@ -761,16 +761,16 @@ void ggml_cuda_mmb_release_all(ggml_backend_cuda_context & ctx) { } // A producer writes the BF16 copy of t itself: the entry is found by the MMB GEMM that reads t next. uint16_t * ggml_cuda_mmb_cache_produce(ggml_backend_cuda_context & ctx, const ggml_tensor * t, size_t n) { - if (!mmb_enabled()) return nullptr; + if (!mmb_enabled(ctx)) return nullptr; return mmb_cache_insert(ctx, t, n); } uint16_t * ggml_cuda_mmb_cache_reserve(ggml_backend_cuda_context & ctx, const ggml_tensor * t, size_t n) { - if (!mmb_enabled() || ggml_nrows(t) < mmb_min_t()) return nullptr; + if (!mmb_enabled(ctx) || ggml_nrows(t) < mmb_min_t()) return nullptr; return ggml_cuda_mmb_slot_reserve(ctx, 0, t, n); } -bool ggml_cuda_mmb_supported_mm(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { - if (!mmb_enabled()) return false; +bool ggml_cuda_mmb_supported_mm(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + if (!mmb_enabled(ctx)) return false; const bool quant = mmb_quant_type(src0->type); const bool bf16w = src0->type == GGML_TYPE_BF16 && mmb_bf16w(); const bool f32w = src0->type == GGML_TYPE_F32 && mmb_f32split(); @@ -785,8 +785,8 @@ bool ggml_cuda_mmb_supported_mm(const ggml_tensor * src0, const ggml_tensor * sr return ggml_nrows(dst) == T; } -bool ggml_cuda_mmb_supported_mmid(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst) { - if (!mmb_enabled()) return false; +bool ggml_cuda_mmb_supported_mmid(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst) { + if (!mmb_enabled(ctx)) return false; if (!mmb_quant_type(src0->type) || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || ids->type != GGML_TYPE_I32) return false; if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) return false; const int64_t K = src0->ne[0], M = src0->ne[1], E = src0->ne[2]; @@ -920,8 +920,8 @@ void ggml_cuda_mul_mat_id_mmb(ggml_backend_cuda_context & ctx, const ggml_tensor CUDA_CHECK(cudaGetLastError()); } -bool ggml_cuda_mmb_supported_glu(const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * glu) { - if (!mmb_enabled() || !mmb_glu() || !gw || !uw || !src1 || !ids || !glu) return false; +bool ggml_cuda_mmb_supported_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * glu) { + if (!mmb_enabled(ctx) || !mmb_glu() || !gw || !uw || !src1 || !ids || !glu) return false; if (!mmb_quant_type(gw->type) || uw->type != gw->type) return false; if (!ggml_are_same_shape(gw, uw) || gw->nb[1] != uw->nb[1] || gw->nb[2] != uw->nb[2]) return false; if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || ggml_get_op_params_i32(glu, 1) != 0) return false; @@ -929,7 +929,7 @@ bool ggml_cuda_mmb_supported_glu(const ggml_tensor * gw, const ggml_tensor * uw, if (glu->src[0]->op != GGML_OP_MUL_MAT_ID || glu->src[1]->op != GGML_OP_MUL_MAT_ID) return false; if (glu->src[0]->src[0] != gw || glu->src[1]->src[0] != uw || glu->src[0]->src[1] != src1 || glu->src[1]->src[1] != src1 || glu->src[0]->src[2] != ids || glu->src[1]->src[2] != ids) return false; if (ggml_nelements(glu) != ggml_nelements(glu->src[0]) || glu->ne[0] != gw->ne[1]) return false; - return ggml_cuda_mmb_supported_mmid(gw, src1, ids, glu->src[0]) && ggml_cuda_mmb_supported_mmid(uw, src1, ids, glu->src[1]); + return ggml_cuda_mmb_supported_mmid(ctx, gw, src1, ids, glu->src[0]) && ggml_cuda_mmb_supported_mmid(ctx, uw, src1, ids, glu->src[1]); } void ggml_cuda_mul_mat_id_mmb_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * glu) { diff --git a/ggml/src/ggml-cuda/mmb.cuh b/ggml/src/ggml-cuda/mmb.cuh index b975fe044b6..5e547eef4cd 100644 --- a/ggml/src/ggml-cuda/mmb.cuh +++ b/ggml/src/ggml-cuda/mmb.cuh @@ -1,8 +1,9 @@ #pragma once #include "common.cuh" // Quantized-weight BF16 WMMA GEMM on gfx1151, from 512 tokens up. -bool ggml_cuda_mmb_supported_mm (const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst); -bool ggml_cuda_mmb_supported_mmid(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst); +// Off unless the backend context opts in (ggml_backend_cuda_set_mmb_enabled). +bool ggml_cuda_mmb_supported_mm (ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst); +bool ggml_cuda_mmb_supported_mmid(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst); void ggml_cuda_mul_mat_mmb (ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); void ggml_cuda_mul_mat_id_mmb(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); void ggml_cuda_mmb_begin_graph(ggml_backend_cuda_context & ctx); @@ -22,7 +23,7 @@ bool ggml_cuda_mmb_down16(); bool ggml_cuda_mmb_res16(); bool ggml_cuda_mmb_blk16(); bool ggml_cuda_hc_gate_mix(ggml_backend_cuda_context & ctx, const ggml_tensor * w, const ggml_tensor * lo, const ggml_tensor * xn, ggml_tensor * dst, int hc, float scale, float bias); -bool ggml_cuda_mmb_supported_glu(const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * glu); +bool ggml_cuda_mmb_supported_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * glu); void ggml_cuda_mul_mat_id_mmb_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * glu); void ggml_cuda_mmb_shadow_prepare(ggml_backend_cuda_context & ctx, const ggml_tensor * w); void ggml_cuda_mmb_release_all(ggml_backend_cuda_context & ctx); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ccfd5217574..abe70376cd3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -375,6 +375,16 @@ llama_context::llama_context( } } + // the CUDA/HIP BF16 WMMA matmul path (mmb) is tuned for qwen4exp: other archs keep MMQ + for (auto & backend : backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); + ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; + auto * set_mmb_fn = reg ? (void (*)(ggml_backend_t, bool)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_mmb_enabled") : nullptr; + if (set_mmb_fn) { + set_mmb_fn(backend.get(), model.arch == LLM_ARCH_QWEN4EXP); + } + } + llama_set_abort_callback(this, params.abort_callback, params.abort_callback_data); // graph outputs buffer diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 18ece03b53b..52438272d5d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -13810,6 +13810,16 @@ static bool run_fa_vec_slice(ggml_backend_t backend, ggml_backend_t backend_cpu, return n_fail == 0; } +// The CUDA/HIP BF16 WMMA matmul path (mmb) is off unless a context opts in (llama does it for qwen4exp). +// --mmb picks which path the >=512-token MUL_MAT/MUL_MAT_ID cases test: mmb, or MMQ as every other arch runs. +static void set_mmb(ggml_backend_t backend, bool enable) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + auto set_mmb_fn = (void (*)(ggml_backend_t, bool)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_mmb_enabled"); + if (set_mmb_fn) { + set_mmb_fn(backend, enable); + } +} + // Target and draft contexts must not share mutable MMB scratch, including during graph replay and teardown. static bool run_mmb_context_test(ggml_backend_dev_t dev, const char * op_names_filter, printer * output_printer) { if (!op_names_filter_selects(op_names_filter, "MMB_CONTEXT") || @@ -13829,6 +13839,7 @@ static bool run_mmb_context_test(ggml_backend_dev_t dev, const char * op_names_f context(ggml_backend_dev_t dev, int tokens) : backend(ggml_backend_dev_init(dev, nullptr)), cpu(ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr)), test(tokens, GGML_TYPE_Q4_0) { tensors.reset(ggml_init({ggml_tensor_overhead()*128 + ggml_graph_overhead(), nullptr, true})); GGML_ASSERT(backend && cpu && tensors); + set_mmb(backend.get(), true); // this test is about mmb scratch, so it ignores --mmb test.gf = ggml_new_graph(tensors.get()); output = test.build_graph(tensors.get()); for (ggml_tensor * t = ggml_get_first_tensor(tensors.get()); t; t = ggml_get_next_tensor(tensors.get(), t)) { @@ -13919,7 +13930,7 @@ static bool run_mmb_context_test(ggml_backend_dev_t dev, const char * op_names_f } static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mode mode, const char * op_names_filter, const char * params_filter, - printer * output_printer, const char * test_file_path, int parallel_workers) { + printer * output_printer, const char * test_file_path, int parallel_workers, bool mmb) { auto filter_test_cases = [](std::vector> & test_cases, const char * params_filter) { if (params_filter == nullptr) { return; @@ -14025,6 +14036,7 @@ static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mo if (b == NULL) { return; } + set_mmb(b.get(), mmb); ggml_backend_ptr b_cpu(ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, NULL)); if (b_cpu == NULL) { @@ -14206,7 +14218,7 @@ static void show_test_coverage() { static void usage(char ** argv) { printf("Usage: %s [mode] [-o ] [-b ] [-p ] [--output ] [--list-ops]", argv[0]); - printf(" [--show-coverage] [--test-file ] [-j ]\n"); + printf(" [--show-coverage] [--test-file ] [-j ] [--mmb ]\n"); printf(" valid modes:\n"); printf(" - test (default, compare with CPU backend for correctness)\n"); printf(" - grad (compare gradients from backpropagation with method of finite differences)\n"); @@ -14219,6 +14231,7 @@ static void usage(char ** argv) { printf(" --show-coverage shows test coverage\n"); printf(" --test-file reads test operators from a test file generated by test-export-graph-ops\n"); printf(" -j runs tests using parallel worker threads (default: 1, test mode only)\n"); + printf(" --mmb on|off: CUDA/HIP matmuls from 512 tokens use the BF16 WMMA path (on, default; qwen4exp) or MMQ (off; other archs)\n"); } int main(int argc, char ** argv) { @@ -14229,6 +14242,7 @@ int main(int argc, char ** argv) { const char * params_filter = nullptr; const char * test_file_path = nullptr; int parallel_workers = 1; + bool mmb = true; for (int i = 1; i < argc; i++) { if (strcmp(argv[i], "test") == 0) { @@ -14294,6 +14308,13 @@ int main(int argc, char ** argv) { usage(argv); return 1; } + } else if (strcmp(argv[i], "--mmb") == 0) { + if (i + 1 < argc && (strcmp(argv[i + 1], "on") == 0 || strcmp(argv[i + 1], "off") == 0)) { + mmb = strcmp(argv[++i], "on") == 0; + } else { + usage(argv); + return 1; + } } else { usage(argv); return 1; @@ -14332,6 +14353,7 @@ int main(int argc, char ** argv) { ggml_backend_ptr backend(ggml_backend_dev_init(dev, NULL)); GGML_ASSERT(backend != NULL); + set_mmb(backend.get(), mmb); ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads"); @@ -14346,7 +14368,7 @@ int main(int argc, char ** argv) { false, "", ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024, true)); - bool ok = test_backend(backend.get(), dev, mode, op_names_filter, params_filter, output_printer.get(), test_file_path, parallel_workers); + bool ok = test_backend(backend.get(), dev, mode, op_names_filter, params_filter, output_printer.get(), test_file_path, parallel_workers, mmb); if (ok) { n_ok++;