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

Code before:
            }
        }

        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");
            }

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

        // reserve with tg (token generation) graph to get the number of splits and nodes
        {
            auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get());
            if (!gf) {
                throw std::runtime_error("failed to allocate compute tg buffers");
            }

            n_splits_tg = ggml_backend_sched_get_n_splits(sched.get());
            n_nodes_tg  = ggml_graph_n_nodes(gf);
        }

        // reserve again with pp graph to avoid ggml-alloc reallocations during inference
        {
            // TODO: not sure if the following graph would be worster case for multi-stream KV caches:
            //
            // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get());
            //
            auto * gf = graph_reserve(n_tokens, n_seqs, n_tokens, mctx.get());
            if (!gf) {
                throw std::runtime_error("failed to allocate compute pp buffers");
            }
        }

        for (size_t i = 0; i < backend_ptrs.size(); ++i) {
            ggml_backend_t             backend = backend_ptrs[i];
            ggml_backend_buffer_type_t buft    = backend_buft[i];
            size_t size = ggml_backend_sched_get_buffer_size(sched.get(), backend);
            if (size > 1) {
                LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__,
                        ggml_backend_buft_name(buft),
                        size / 1024.0 / 1024.0);
            }
        }

        if (n_nodes_pp == n_nodes_tg) {

Review comment (suggestion):
I don't think this is a good way to do it. It is very fragile code that makes a lot of assumptions that are not guaranteed anywhere, and will break very easily and in a very difficult way to detect when making changes to other parts of the code.

Code after:
        int n_nodes_pp  = -1;

        int n_splits_tg = -1;
        int n_nodes_tg  = -1;

        llama_memory_context_ptr mctx;
        if (memory) {
            LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__);
            mctx = memory->init_full();
            if (!mctx) {
                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;