CANN: Disable acl_graph for prefill stage (llama/15933)

Since the prefill length is not fixed, graphs constructed for the
prefill stage cannot be reused. For this reason, ACL graph
execution is disabled by default during prefill.
This commit is contained in:
hipudding 2025-09-11 15:59:37 +08:00 committed by Georgi Gerganov
parent f5ef0e25e2
commit dadf73665a
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
1 changed files with 15 additions and 0 deletions

View File

@ -2360,6 +2360,21 @@ static enum ggml_status ggml_backend_cann_graph_compute(
bool use_cann_graph = true;
bool cann_graph_update_required = false;
static bool prefill_use_graph = parse_bool(get_env("GGML_CANN_PREFILL_USE_GRAPH").value_or(""));
if (!prefill_use_graph) {
// Do not use acl_graph for prefill.
for (int i = 0; i < cgraph->n_nodes; i++) {
ggml_tensor * node = cgraph->nodes[i];
// TODO: Optimize here. Currently, we can only
// get seq_len by FA's input.
if (node->op == GGML_OP_FLASH_ATTN_EXT) {
// Q -> src[0], shape: [B, S, N, D]
use_cann_graph = (node->src[0]->ne[1] == 1);
break;
}
}
}
if (!cann_ctx->acl_graph_mode) {
use_cann_graph = false;
}