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

Code before:
        }
        if (is_one_copy) {
            one_copy.insert(i);
        }
    }

    kv_swap_bufs.resize(size);

    for (int i = 0; i < size; 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
        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());
    }

    // 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_swap_bufs[view[i]].k.data(), kv_swap_bufs[view[i]].k.size());
            memcpy(src[i].kv_self.v->data, kv_swap_bufs[view[i]].v.data(), kv_swap_bufs[view[i]].v.size());
        } else {
            // modify KV caches of decoder using data from correspond decoder KV caches directly
            memcpy(src[i].kv_self.k->data, src[view[i]].kv_self.k->data, ggml_nbytes(src[view[i]].kv_self.k));
            memcpy(src[i].kv_self.v->data, src[view[i]].kv_self.v->data, ggml_nbytes(src[view[i]].kv_self.v));
        }
    }

    // then modify one-copy decoder KV caches
    for (auto & i : one_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_swap_bufs[view[i]].k.data(), kv_swap_bufs[view[i]].k.size());
            memcpy(src[i].kv_self.v->data, kv_swap_bufs[view[i]].v.data(), kv_swap_bufs[view[i]].v.size());
        } else {
            // modify KV caches of decoder using data from correspond decoder KV caches directly
            memcpy(src[i].kv_self.k->data, src[view[i]].kv_self.k->data, ggml_nbytes(src[view[i]].kv_self.k));
            memcpy(src[i].kv_self.v->data, src[view[i]].kv_self.v->data, ggml_nbytes(src[view[i]].kv_self.v));
        }
    }


Review comment (performance):
```suggestion
            // modify KV caches of decoder using data from kv_swap_bufs
```

Code after:
                    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());
    }

    // since two-copy decoder KV caches are protected by kv_swap_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_swap_bufs
            WHISPER_PRINT_DEBUG("%s: two-copy decoder using   swap buffers: swap[%d] -> %d\n", __func__, view[i], i);
            memcpy(src[i].kv_self.k->data, kv_swap_bufs[view[i]].k.data(), kv_swap_bufs[view[i]].k.size());
            memcpy(src[i].kv_self.v->data, kv_swap_bufs[view[i]].v.data(), kv_swap_bufs[view[i]].v.size());
        } else {
            // modify KV caches of decoder using data from correspond decoder KV caches directly
            WHISPER_PRINT_DEBUG("%s: two-copy decoder without swap buffers:      %d  -> %d\n", __func__, view[i], i);
            memcpy(src[i].kv_self.k->data, src[view[i]].kv_self.k->data, ggml_nbytes(src[view[i]].kv_self.k));
            memcpy(src[i].kv_self.v->data, src[view[i]].kv_self.v->data, ggml_nbytes(src[view[i]].kv_self.v));
        }
    }

    // then modify one-copy decoder KV caches
    for (auto & i : one_copy) {
        // skip the decoder indices that require pointer swapping
        if (p_swap_set.find(i) != p_swap_set.end()) {