From ebb546b7e961bd46fd9ed0387ffd14ca86b6fe1b Mon Sep 17 00:00:00 2001 From: Rafail Giavrimis <47496212+grafail@users.noreply.github.com> Date: Tue, 11 Aug 2026 18:50:03 +0100 Subject: [PATCH] CUDA: only disable CUDA graphs when mul_mat_id actually needs a stream sync (#26802) --- ggml/src/ggml-cuda/ggml-cuda.cu | 39 ++++++++++++++++++++++++++++----- tests/test-backend-ops.cpp | 4 ++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1d4f4dfbd2..cb7e9330c8 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1865,6 +1865,37 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); } +// returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization +// [TAG_MUL_MAT_ID_CUDA_GRAPHS] +static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int cc) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + return true; + } + + if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) { + if (ggml_is_quantized(src0->type)) { + if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) { + return false; + } + } else if (GGML_CUDA_CC_IS_AMD(cc)) { + return false; + } + } + + if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) { + return false; + } + + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + return false; + } + + return true; +} + static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -1907,7 +1938,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } // note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization - // TODO: add asserts to verify this. should work with CUDA, HIP, etc. + GGML_ASSERT(ggml_cuda_mul_mat_id_needs_sync(dst, cc)); cudaStream_t stream = ctx.stream(); GGML_ASSERT(nb12 % nb11 == 0); @@ -2522,10 +2553,8 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (node->op == GGML_OP_MUL_MAT_ID) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc); - if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) { - // under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs - // TODO: figure out a way to enable for larger batch sizes, without hurting performance + if (ggml_cuda_mul_mat_id_needs_sync(node, cc)) { + // the mul_mat_id fallback path synchronizes the stream, so we cannot use CUDA graphs // ref: https://github.com/ggml-org/llama.cpp/pull/18958 use_cuda_graph = false; #ifndef NDEBUG diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 86dfabbf5b..c27f0b92e8 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9012,6 +9012,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 128, k, {12,1}, {1,1})); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 50, 200, k)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, true, 50, 200, k)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, false, 50, 200, k)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, true, 50, 200, k)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 16, 16, false, 50, 200, k)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 16, 16, true, 50, 200, k)); } @@ -9043,6 +9045,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 2, 2, b, 32, 8192, 64)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64)); } test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1));