Repository: ggerganov/llama.cpp
File: src/llama-context.cpp (C++)

Code before:
                throw std::runtime_error("failed to initialize memory module");
            }
        }

        cross.v_embd.clear();

        // reserve pp (prompt processing) graph first so that buffers are only allocated once
        {
            auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
            if (!gf) {
                throw std::runtime_error("failed to allocate compute pp buffers");
            }

            if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
                ggml_backend_sched_alloc_graph(sched.get(), gf);

                bool fa_device_mismatch = false;
                for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
                    ggml_tensor * n = ggml_graph_node(gf, i);
                    if (n->op != GGML_OP_FLASH_ATTN_EXT) {
                        continue;
                    }
                    ggml_backend_dev_t device_fa = ggml_backend_get_device(
                        ggml_backend_sched_get_tensor_backend(sched.get(), n));

                    GGML_ASSERT(strncmp(n->name, "fattn-", 6) == 0);
                    const int il = std::stoi(n->name + 6);
                    ggml_backend_dev_t device_kv = model.dev_layer(il);
                    if (device_fa != device_kv) {
                        fa_device_mismatch = true;
                        break;
                    }
                }
                if (fa_device_mismatch) {
                    cparams.flash_attn = false;
                    LLAMA_LOG_INFO("%s: Flash Attention was auto, set to disabled\n", __func__);
                    if (ggml_is_quantized(params.type_v)) {
                        throw std::runtime_error("quantized V cache was requested, but this requires Flash Attention");
                    }
                    auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
                    if (!gf) {
                        throw std::runtime_error("failed to allocate compute pp buffers");
                    }
                } else {
                    cparams.flash_attn = true;
                    LLAMA_LOG_INFO("%s: Flash Attention was auto, set to enabled\n", __func__);
                }
            }

            n_splits_pp = ggml_backend_sched_get_n_splits(sched.get());
            n_nodes_pp  = ggml_graph_n_nodes(gf);

Review comment (suggestion):
Might be a bit more future-proof to use a global constant here and in `llama-graph.cpp` instead of a hardcoded string:

```c++
// llama-impl.h

#define LLAMA_TENSOR_NAME_FATTN "__fattn__"
```

With the convention being that underscores in tensor names are reserved for the application. This could avoid a problem where a developer names some other tensor `"fattn"`.

Code after:
                throw std::runtime_error("failed to initialize memory module");
            }
        }

        cross.v_embd.clear();

        // reserve pp (prompt processing) graph first so that buffers are only allocated once
        {
            auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
            if (!gf) {
                throw std::runtime_error("failed to allocate compute pp buffers");
            }

            if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
                ggml_backend_sched_alloc_graph(sched.get(), gf);

                const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FATTN) + 1;
                bool fa_device_mismatch = false;
                for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
                    ggml_tensor * n = ggml_graph_node(gf, i);
                    if (n->op != GGML_OP_FLASH_ATTN_EXT) {
                        continue;
                    }
                    ggml_backend_dev_t device_fa = ggml_backend_get_device(
                        ggml_backend_sched_get_tensor_backend(sched.get(), n));

                    // TODO: instead of the tensor names, use a map to keep track of which (FA) tensors belong to which layer
                    GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FATTN "-", prefix_len) == 0);
                    const int il = std::stoi(n->name + prefix_len);
                    ggml_backend_dev_t device_kv = model.dev_layer(il);
                    if (device_fa != device_kv) {
                        LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the Flash Attention tensor "
                            "is assigned to device %s (usually due to missing support)\n",
                            __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_fa));
                        // FIXME: fa_device_mismatch logic is wrong for --no-kv-offload, but this is broken anyways
                        fa_device_mismatch = true;
                        break;
                    }
                }
                if (fa_device_mismatch) {
                    cparams.flash_attn = false;
                    LLAMA_LOG_WARN("%s: Flash Attention was auto, set to disabled\n", __func__);
                    if (ggml_is_quantized(params.type_v)) {
                        throw std::runtime_error("quantized V cache was requested, but this requires Flash Attention");
                    }
                    auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
                    if (!gf) {
                        throw std::runtime_error("failed to allocate compute pp buffers");
                    }
                } else {
                    cparams.flash_attn = true;