Repository: ggerganov/whisper.cpp
File: whisper.cpp (C++)

Code before:
        for (int j = i + 1; j < size; j++) {
            if (i == view[j]) {
                two_copy.insert(i);
                is_one_copy = false;
                // detect symmetric diagram
                if (j == view[i]) {
                    p_swap_set.insert(i);
                    p_swap_set.insert(j);
                    p_swap_vec.emplace_back(i, j);
                }
                break;
            }
        }
        if (is_one_copy) {
            one_copy.insert(i);
        }
    }

    // beam-search helpers
    struct kv_buf {
        std::vector<uint8_t> k;
        std::vector<uint8_t> v;
    };

    static std::vector<kv_buf> kv_bufs(size);
    static bool is_resized = false;

    if (!is_resized) {
        for (int i = 0; i < size; i++) {
            kv_bufs[i].k.resize(ggml_nbytes(src[i].kv_self.k));
            kv_bufs[i].v.resize(ggml_nbytes(src[i].kv_self.v));
        }
        is_resized = true;
    }

    for (auto & i : two_copy) {
        // make a copy of KV caches
        memcpy(kv_bufs[i].k.data(), src[i].kv_self.k->data, kv_bufs[i].k.size());
        memcpy(kv_bufs[i].v.data(), src[i].kv_self.v->data, kv_bufs[i].v.size());
    }

    // since two-copy decoder KV caches are protected by kv_bufs, modify them first
    for (auto & i : two_copy) {
        // skip the decoder indices that require pointer swapping
        if (p_swap_set.find(i) != p_swap_set.end()) {continue;}

        if (two_copy.find(view[i]) != two_copy.end()) {
            // modify KV caches of decoder using data from kv_bufs
            memcpy(src[i].kv_self.k->data, kv_bufs[view[i]].k.data(), kv_bufs[view[i]].k.size());
            memcpy(src[i].kv_self.v->data, kv_bufs[view[i]].v.data(), kv_bufs[view[i]].v.size());
        } else {

Review comment (suggestion):
This is not OK - different whisper context will now share the same KV buffers, so this breaks parallel usage.
These buffers have to become part of `whisper_state` instead.

Also, no need for `is_resized` flag. Just always resize. The resize will be noop if the new size is less or equal to the old size.

Code after:
    std::set<int> two_copy; // decoder indices require two copies to safely modify KV caches

    // (buffer->decoder or decoder->decoder)
    std::set<int> one_copy; // decoder indices require one copy to safely modify KV caches

    // (decoder<->decoder)
    std::set<int> p_swap_set; // decoder indices able to swap KV-cache pointers
    std::vector<whisper_pair<int, int>> p_swap_vec;
    p_swap_vec.reserve(n_decoders);

    // see https://github.com/ggerganov/whisper.cpp/wiki
    for (int i = 0; i < n_decoders; i++) {
        // zero-copy (no modification)
        if (i == view[i] || view[i] < 0) {
            continue;
        }

        bool is_one_copy = true;
        // since we modify data sequentially, we only consider decoder indices after current index
        for (int j = i + 1; j < n_decoders; j++) {
            if (i == view[j]) {
                // detect symmetric diagram
                if (j == view[i]) {
                    p_swap_set.insert(i);
                    p_swap_set.insert(j);
                    p_swap_vec.emplace_back(i, j);
                } else {
                    two_copy.insert(i);
                    is_one_copy = false;
                }
                break;
            }
        }
        if (is_one_copy) {
            one_copy.insert(i);
        }
    }

    kv_swap_bufs.resize(n_decoders);

    for (int i = 0; i < n_decoders; i++) {
        kv_swap_bufs[i].k.resize(ggml_nbytes(src[i].kv_self.k));
        kv_swap_bufs[i].v.resize(ggml_nbytes(src[i].kv_self.v));
    }

    for (auto & i : two_copy) {
        // make a copy of KV caches
        WHISPER_PRINT_DEBUG("%s: store KV cache into swap: idx %d\n", __func__, i);
        memcpy(kv_swap_bufs[i].k.data(), src[i].kv_self.k->data, kv_swap_bufs[i].k.size());
        memcpy(kv_swap_bufs[i].v.data(), src[i].kv_self.v->data, kv_swap_bufs[i].v.size());
    }