CUDA: only disable CUDA graphs when mul_mat_id actually needs a stream sync (llama/26802)

This commit is contained in:
Rafail Giavrimis 2026-08-11 18:50:03 +01:00 committed by Georgi Gerganov
parent 030656f26f
commit cdbe455102
1 changed files with 34 additions and 5 deletions

View File

@ -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